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

This repository contains the code for training ViT and ClimaX backbones, Sparse Autoencoders (SAEs), and SAE-Xplainers, as described in the paper SAE-Xplainers: Rule-Based Feature Interpretation for Extreme Earth Events.

• 📝📊 The paper is available on ArXiv.
• 🔨🔬 We provide a tool to investigate ClimaX features on SeasFire and ClimateNet available at this Demo Link.

Overview Framework

📑 Table of Contents

• Installation
• Repository Structure
• Data
• Models Finetuning
• Sparse Autoencoders Training
• SAE-Xplainers Training
• SAE-Xplainers Evaluation
• Citation

🛠️ Installation

Install the required Python libraries:

pip install -r requirements/requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121

Requirements: Python 3.10+, CUDA 12.1 (for GPU training).

⚠️ See .env_template for the environment variables required to run this repository. Copy it to .env and fill in the values for your setup.

📂 Repository Structure

├── ClimaX/              # ClimaX backbone (git submodule)
├── televit/             # TeleViT backbone (git submodule)
├── configs/             # Hydra configs
├── data/                # Datasets (SeasFire, ClimateNet) — populated after download
├── logs/                # Training logs, checkpoints, WandB runs — generated after running scripts
├── models/              # Pretrained checkpoints (e.g. ClimaX pretrained weights)
├── requirements/        # Python dependency files
├── scripts/
│   ├── downstream/      # ClimaX/ViT/ProtoSeg finetuning
│   ├── sae/             # TopK + GeoTopK training
│   ├── xplainer/        # SAE-Xplainers training
│   └── absorption/      # Absorption score evaluation
├── src/                 # Core source code
└── stats/               # Precomputed statistics (e.g. normalization stats, dataset stats)

Note: ClimaX/ and televit/ are git submodules. Clone with git clone --recurse-submodules <repo-url>, or run git submodule update --init --recursive after a regular clone.

💽 Data

⬇️ Downloading Data

SeasFire — Download version 0.3 from Zenodo and decompress it into your target data folder (default: ./data).

ClimateNet — First, move climatenet-splits/train.txt and climatenet-splits/test.txt extracted from the NERSC portal into your target ClimateNet data folder (default: ./data/climatenet). Then run:

python src/data/download_clim.py

to extract all tiles.

🔨 Models Finetuning

⬇️ ClimaX Pretrained Download

Download the pretrained checkpoint and place it in your target model folder (default: ./models).

⏳ ClimaX and ViT Finetuning

SeasFire

bash scripts/downstream/climax_training.sh
bash scripts/downstream/vit_training.sh

ClimateNet

bash scripts/downstream/climax_training_clim.sh
bash scripts/downstream/vit_training_clim.sh

🤖 Prototype Training on ClimaX

Warm-up stage (SeasFire and ClimateNet)

bash scripts/downstream/proto/climax_experiments_proto.sh
bash scripts/downstream/proto/climax_experiments_clim_proto.sh

Finetuning stage (SeasFire and ClimateNet)

Before running: set proto_path (checkpoint from the last epoch of the warm-up stage) in the YAML config.

bash scripts/downstream/proto/climax_experiments_proto_finetune.sh
bash scripts/downstream/proto/climax_experiments_clim_proto_finetune.sh

📊 Downstream Results

Performance of ClimaX, ViT, and ProtoSeg finetuned on SeasFire and ClimateNet. * denotes results reported from the original papers.

SeasFire

Metric	ClimaX	ViT	ProtoSeg	U-Net++*
PRAUC ↑	0.6205	0.6150	0.5161	0.6100

ClimateNet

Metric	ClimaX	ViT	ProtoSeg	DeepLabV3+*
IoU ↑	0.5385	0.4972	0.4254	0.5247

🔍 Sparse Autoencoders Training

Scripts to reproduce Table 1 (ClimaX & ViT × SeasFire & ClimateNet × TopK & GeoTopK, 8 configurations total) are in ./scripts/sae. Run each individually, e.g.:

bash scripts/sae/sae_training.sh

Before running: set model_path (ViT and ClimaX checkpoint) in the YAML config.

📊 SAE Results

Comparison of TopK and GeoTopK on the last residual stream of ClimaX and ViT, finetuned on SeasFire and ClimateNet respectively. Bold indicates the best score; underline indicates a near-tie (difference ≤ 0.005).

SeasFire

Metric	ClimaX (TopK)	ClimaX (GeoTopK)	ViT (TopK)	ViT (GeoTopK)
$\text{R}^2$ ↑	0.927	0.937	0.825	0.840
$\text{R}^2_{\text{worst}}$ ↑	0.881	0.901	0.767	0.786
$\text{R}^2_{\text{event}}$ ↑	0.801	0.807	0.651	0.633
MSE ↓	1.84e-3	1.59e-3	1.77e-1	1.63e-1
Dead Feat. ↓	0.159	0.002	0.021	0.000

ClimateNet

Metric	ClimaX (TopK)	ClimaX (GeoTopK)	ViT (TopK)	ViT (GeoTopK)
$\text{R}^2$ ↑	0.943	0.953	0.987	0.993
$\text{R}^2_{\text{worst}}$ ↑	0.910	0.924	0.983	0.988
$\text{R}^2_{\text{event}}$ ↑	0.880	0.893	0.974	0.980
MSE ↓	1.01e-3	0.84e-3	12.85e-3	7.12e-3
Dead Feat. ↓	0.513	0.391	0.811	0.442

📈 Scaling Experiments

To reproduce the scaling experiments, edit nb_k and nb_concepts in sae_training.sh and sae_training_scale.sh to one of the following pairs (nb_k, nb_concepts): (8, 8192) or (16, 16384).

🔬 GeoTopK Ablation

Ablations for the two geographic modulation components: localization-aware adaptive normalization (LOAN) and per-latitude SAE band bias (LatBias), can be run via:

bash scripts/sae/ablation/loan_sae_training_scale.sh
bash scripts/sae/ablation/latbias_sae_training_scale.sh

For the scalar-only, shift-only, and scale-only variants, modify the plSAE.forward function directly in src/model/sae_module.py.

📏 SAE-Xplainers Training

The following scripts directly compute the SAE-Xplainers across the three Skope-Rules parametrizations and with or without spatial context.

Before running: set sae_ckpt_path (SAE checkpoint), and output_act_dir (output model activations) in the YAML config.

GeoTopK Interpretations

bash scripts/xplainer/white_box_training_multi_adapt_spa.sh
bash scripts/xplainer/white_box_training_multi_adapt_spa_clim.sh

ClimaX Neurons Interpretations

bash scripts/xplainer/white_box_training_multi_adapt_spa_neurons.sh
bash scripts/xplainer/white_box_training_multi_adapt_spa_clim_neurons.sh

ProtoSeg Prototypes Interpretations

bash scripts/xplainer/white_box_training_multi_adapt_spa_proto.sh
bash scripts/xplainer/white_box_training_multi_adapt_spa_clim_proto.sh

🏛️ SAE-Xplainers Evaluation

🎯 Evaluation Accuracy

First, the SAE-Xplainers are aggregated via this python script that also outputs the aggregated evaluation accuracy:

python src/eval/merge_rules.py

Before running: set runs (Skope-Rules runs) in the YAML config.

💧 Absorption Scores

The absorption scores across both datasets are directly computed on the set of SAE-Xplainers trained on GeoTopK via the following scripts:

bash scripts/absorption/absorption.sh
bash scripts/absorption/absorption_clim.sh

Before running: set rules_csv (SAE-Xplainers CSV files) in the bash script.

SAE-Xplainers Quantitative Analysis

🔥 Fire Conditions

The high fire risk True Conditions and Inverse Conditions scores for the paper Table 2, are computed via the python script:

python src/eval/fire_conds.py

Before running: set path_rules (target SeasFire SAE-Xplainers CSV file), and test_data (SAE codes and SeasFire labels) in the YAML config.

Ratio of rule sets containing true (Tr) and inverse (Inv) conditions for fire-specific dimensions.

Metric	Neurons	Prototypes	GeoTopK
$\text{Tr}_{\text{cond}}$ ↑	0.33	0.50	0.46
$\text{Inv}_{\text{cond}}$ ↓	0.16	0.20	0.02

⚖️ Sampling Comparison

The scores in Table 7 are computed via the following python script:

python src/eval/sampling_comp.py

Before running: set runs (multiple seed runs) in the YAML config.

Comparison of sampling strategies for training SAE-Xplainers on ClimaX finetuned on SeasFire. Results are averaged across three runs, reported as mean ± standard deviation.

Sampling	Stability ↓	$\text{Weight}_{\text{loc}}$ ↓	$\text{Tr}_{\text{cond}}$ ↑
Random	0.688 ± 0.023	0.633 ± 0.059	0.490 ± 0.064
Adversarial	0.705 ± 0.017	0.574 ± 0.071	0.497 ± 0.074

🧑‍🔬 Investigation Tool

We also provide for this paper a global model interpretations framework available at this Demo Link, to investigate the SAE's features identify in the ClimaX models fine-tuned on SeasFire and CliamteNet.

📄 Citation

The paper is currently under-review to cite please refer to our ArXiv version:

@misc{porta2026saexplainersrulebasedfeatureinterpretation,
      title={SAE-Xplainers: Rule-Based Feature Interpretation for Extreme Earth Events}, 
      author={Hugo Porta and Emanuele Dalsasso and Chang Xu and Theo Gnassounou and Devis Tuia},
      year={2026},
      eprint={2608.20117},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2608.20117}, 
}