SAE-Xplainers: Rule-Based Feature Interpretation for Extreme Earth Events

AI summary

Overview

Research area: Explainable AI (XAI) for deep learning models applied to weather and climate (W&C) data, specifically extreme Earth events (ExEE). Combines sparse autoencoders (SAEs), geographic conditioning, and rule-based surrogate models.

Technical level: Advanced. Familiarity with transformer architectures, sparse autoencoders, FiLM conditioning, and interpretability methods (SHAP, LIME, Anchors, concept-based XAI) is assumed.

Scope: The paper proposes a location-aware sparse autoencoder (GeoTopK) and a rule-based interpretation framework (SAE-Xplainers) to make deep learning predictions of wildfires, tropical cyclones, and atmospheric rivers human-understandable.

What This Paper Is About

Deep learning models can forecast and detect extreme weather and climate events, but their opacity blocks adoption in operational settings where transparency is essential. Sparse autoencoders—effective at extracting human-readable concepts from language and vision models—fail on W&C data for two reasons: environmental drivers are inherently location-dependent (similar conditions produce different outcomes in different regions), and SAE feature activations cannot be semantically interpreted by inspecting input-space maps, since no Weather-Language models exist and the spatio-temporal patterns are too complex. This paper builds both a location-aware SAE and a rule-based interpreter to close that gap.

Key Contributions

1. GeoTopK, a geographic location-aware sparse autoencoder that is, to the authors' knowledge, the first domain-conditioned SAE training strategy. It uses a learned location encoder as a FiLM-style adapter to modulate token activations before the TopK sparsity projection.

2. SAE-Xplainers, an ensemble of rule-based surrogate models (trained via Skope-Rules) that faithfully approximate the binarized activations of individual SAE features and unfold them into human-readable if-then rules grounded in input variables.

3. Two new evaluation metrics tailored to ExEE data: an event-conditioned reconstruction score (R²_event, macro-averaged across imbalanced event classes) and a worst-case latitude-band score (R²_worst), which captures reconstruction fidelity under geographic domain shift.

4. A scalable method for detecting feature absorption in SAEs — grouping features by rule similarity (Jaccard distance plus Hungarian assignment), measuring co-occurrence ratios, and validating against SAEBench absorption scores — plus adversarial geographic negative sampling to prevent interpretations from degenerating into location-based shortcuts.

Main Findings

• GeoTopK improves reconstruction across the board. It outperforms the standard TopK SAE on R², MSE, R²_worst, and R²_event for both ClimaX and ViT encoders on both datasets. Representative gains: R² of 0.937 versus 0.927 (ClimaX/SeasFire) and 0.993 versus 0.987 (ViT/ClimateNet).

• Dead features drop dramatically. The fraction of unused features on ClimateNet falls from 0.513 to 0.391 (ClimaX) and from 0.811 to 0.442 (ViT), indicating that geographic modulation prevents neural superposition from producing polysemantic features. On SeasFire the reduction on ClimaX is even starker: 0.159 to 0.002.

• SAE-Xplainers are highly faithful. Aggregate rule-set accuracy reaches roughly 93% on SeasFire (22 points above neuron-level rules, 9 points above prototype-based rules) and roughly 85% on ClimateNet (9 and 5 points above the respective baselines). This confirms that monosemantic SAE features yield more interpretable surrogates than polysemantic model neurons.

• Spatial context matters far more for cyclone and river detection than for fire forecasting. Adding spatial statistics to the rules boosts ClimateNet accuracy by 13 points versus only 3 points on SeasFire, reflecting that tropical cyclones and atmospheric rivers are large, long-range phenomena whereas fires are highly localized.

• Feature absorption is detectable and consistent with rule quality. On both datasets the absorption score rises with rule accuracy (except for an edge case at low complexity on ClimateNet). Fires, being localized, produce higher absorption only at higher rule complexity; cyclones and rivers co-occur strongly even at low complexity.

• Rules align with the scientific literature. Among the top-100 fire-specific GeoTopK features, 46% of rule sets contain established fire-risk conditions (high vapor pressure deficit, low total precipitation, low relative humidity), while only 2% contain inverse (low-risk) conditions — fewer than the 16% and 20% seen for neurons and prototypes respectively. Of the top-50 explanations, 92% reference at least one of seven literature-supported fire drivers.

• Sample-level interpretations stay tractable. With k = 4 sparsity, at most four features activate per patch, keeping per-prediction explanations concise for human review.

Methodology in Plain English

The researchers first fine-tune two encoders — ClimaX (a climate-specific foundation model) and ViT (a general vision transformer) — on two global datasets: SeasFire, which forecasts fire occurrence eight days ahead from 14 environmental variables, and ClimateNet, which detects tropical cyclones and atmospheric rivers from six variables spanning the whole globe.

From each encoder they extract the activations of the last transformer layer for every image patch. Before feeding these activations into a sparse autoencoder, they pass the patch's latitude and longitude through a small neural network that outputs scaling and shifting vectors, bounded by a tanh function and applied to the activations as a feature-wise affine transformation. This modulation is inverted before reconstruction, so the SAE must not simply memorize location — it must use location to disentangle genuinely different environmental regimes. The modified autoencoder (GeoTopK) selects the top k features per patch and is trained with a reconstruction loss plus an auxiliary loss to revive unused features.

For interpretation, each of the 4,096 SAE features gets its own rule-based surrogate model, trained to predict whether that feature activates on a given patch. Negative examples are drawn adversarially from patches geographically close to where the feature fires, which prevents the rules from simply encoding "this feature is about Africa." Rules can be extended with statistics (mean, max, min) computed over the patches that the transformer attended to most, capturing spatial context. Feature similarity across rule sets then exposes groups of co-occurring, absorbed features.

Why This Matters

Impact on research. The paper reframes XAI for climate from "where does the model look" (attribution maps) to "what does the model represent" (semantic concepts), and does so without needing a Weather-Language model to auto-label features. It also provides the first domain-conditioned SAE training strategy, a transferable idea for any geospatial or otherwise location-dependent modality.

Real-world applications:

• Wildfire risk forecasting and early-warning systems, where forecasters need to know which environmental drivers triggered an alert.
• Tropical cyclone detection and tracking, supporting evacuation and infrastructure decisions.
• Atmospheric river detection for flood prediction and water-resource management.
• Climate adaptation and insurance risk modeling, where regulators increasingly demand auditable model reasoning.

Industry relevance. Weather agencies, reinsurers, energy utilities, and agricultural firms all consume climate predictions but face governance requirements for explainability. A framework that produces literature-consistent rules from a trained model could accelerate the operational certification of deep learning systems in these sectors, and the released open-source toolkit lowers the barrier to auditing existing climate models.

Future Directions

• Rule robustness. The authors note that extracted rules could be made more stable through constrained SAE training or knowledge-grounded surrogate models, since rule boundaries currently depend on the sampling strategy and Skope-Rules parametrization.

• Richer surrogate architectures. The current rule sets handle spatial context only through aggregate statistics over attended patches; more expressive models could capture full spatio-temporal dependencies, particularly important for cyclone and river detection where dead-feature rates remain high.

• Cross-domain transfer. The framework could extend to species distribution modeling, air quality forecasting, and other Earth system tasks where local context drives outcomes.

• Reducing residual superposition. Even with GeoTopK, ClimateNet retains substantial dead-feature rates, suggesting that further architectural or loss-based interventions are needed before interpretations are uniformly reliable across event types.

Target Audience

Interpretability researchers working on sparse autoencoders and concept-based XAI beyond language and vision; climate and Earth system scientists who want to audit deep learning models rather than treat them as black boxes; machine learning engineers building operational forecasting or detection systems for weather agencies and risk industries; and graduate students entering the intersection of geospatial deep learning and explainable AI.

Authors’ abstract

The emergence of large-scale Weather and Climate (W&C) datasets offers new opportunities for modeling extreme Earth events (ExEE) and their impacts using deep learning. However, their adoption in operational settings remains limited by the lack of models' interpretability. While for conventional text and image modalities, tools such as Sparse Autoencoders (SAEs) have proven effective for extracting human-understandable concepts, their use for the analysis of ExEE remains challenging due to the nature of W&C data. To address this, we introduce (i) a geographic location-based modulation of the inputs of SAE to capture the local semantic meaning of environmental patterns, and (ii) an ensemble of rule-based SAE-Xplainers to interpret the resulting high-dimensional features derived from complex, multi-modal environmental predictors. We evaluate our method on three ExEE types: the prediction of fires, and the detection of tropical cyclones and atmospheric rivers. We show that SAE input modulation improves both reconstruction performance and feature utilization, and that our SAE-Xplainers enable faithful interpretation of complex climatic patterns by unfolding them into human-understandable rules that are consistent with the scientific literature, while also supporting the identification of feature absorption.