Beyond Linear Superposition: Discovering Climate Features in AI Weather Models with KAN-SAE
On the map: Mechanistic interpretability → Interpreting time-series and weather models
Abstract (arXiv)
Deep learning weather prediction models achieve remarkable predictive skill yet remain largely opaque: we know little about how they represent physical climate phenomena internally. Mechanistic interpretability through Sparse Autoencoders (SAEs) offers a principled route to decomposing these representations, but existing SAEs assume strictly linear feature superposition - a constraint ill-suited for the highly nonlinear atmospheric dynamics encoded in modern transformers. We introduce KAN-SAE, a sparse autoencoder whose encoder replaces the standard ReLU with learnable per-feature B-spline activations drawn from Kolmogorov-Arnold Networks (KANs), allowing each latent dimension to develop its own nonlinear gating profile. Applied to Sonny, KAN-SAE discovers 975 alive features (vs. 566 for a linear baseline, a 72% improvement) with 20% lower inter-feature redundancy and comparable reconstruction fidelity. Without any climate supervision, KAN-SAE identifies an interpretable European heatwave feature spatially concentrated over western Europe, and a western Pacific typhoon tracker confirmed by causal steering experiments. Our results demonstrate that nonlinear activations are essential for mechanistic interpretability of deep learning weather prediction models, recovering climate features that remain invisible to linear baselines.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Interpreting time-series and weather models, one of 13 directions in Mechanistic interpretability. Builds on 13 of its 20 references that are in the map. First posted on arXiv 17 May 2026.
Builds on
13 of 20 references.
References · 13
- 2022Scalable Diffusion Models with Transformerscited by 6,054
- 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,420
- 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,165
- 2022Toy Models of Superpositioncited by 971
- 2024KAN: Kolmogorov-Arnold Networkscited by 649
- 2024Scaling and evaluating sparse autoencoderscited by 553
- 2024Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2cited by 394
- 2023ClimaX: A foundation model for weather and climatecited by 337
- 2019Tackling Climate Change with Machine Learningcited by 265
- 2023GenCast: Diffusion-based ensemble forecasting for medium-range weathercited by 177
- 2025Towards mechanistic understanding in a data-driven weather model: internal activations reveal interpretable physical featurescited by 12
- 2025Step by Step Networkcited by 3
1 more
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Mechanistic Interpretability of Atmospheric Rivers in GraphCast
- 2026SAE-Xplainers: Rule-Based Feature Interpretation for Extreme Earth Events
- 2024CRA5: Extreme Compression of ERA5 for Portable Global Climate and Weather Research via an Efficient Variational Transformer
- 2026Tracing the Unlabeled Storm: Cross-Variable Transfer in a Lagrangian Atmospheric JEPA Framework
- 2025GAIA: A Foundation Model for Operational Atmospheric Dynamics
- 2026TSAE: Structured Sparse Autoencoders for Interpreting Time-Series Forecasting Models
- 2026Marchuk: Efficient Global Weather Forecasting from Mid-Range to Sub-Seasonal Scales via Flow Matching
- 2026Mechanistic Interpretability Tool for AI Weather Models