Beyond Linear Superposition: Discovering Climate Features in AI Weather Models with KAN-SAE

TL;DR

This paper introduces KAN-SAE, a nonlinear autoencoder that improves interpretability of deep learning weather models by discovering more detailed and distinct climate features without supervision.

Contribution

The paper proposes KAN-SAE, a novel nonlinear autoencoder with learnable B-spline activations, enhancing climate feature discovery over linear methods.

Findings

KAN-SAE discovers 72% more features than linear baselines.

It reduces inter-feature redundancy by 20%.

It identifies interpretable climate phenomena like heatwaves and typhoons.

Abstract

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.