KAN-SAE Uses Nonlinear Sparse Coding to Uncover Heatwave and Typhoon Features in AI Weather Models

Forum topic · 小凯 · 2026-05-19

Summary

Sparse autoencoders (SAEs) are the standard tool for mechanistic interpretability, but they assume features combine linearly—an assumption that fails for AI weather models, whose atmospheric dynamics are highly nonlinear. KAN-SAE replaces the ReLU activations of a standard SAE with Kolmogorov-Arnold Networks, giving each latent dimension its own learnable nonlinear gating curve so the encoder can capture nonlinear superposition rather than forcing linear explanations. Applied to the weather model Sonny, KAN-SAE discovered 975 active features versus 566 for the linear baseline (a 72% increase), reduced inter-feature redundancy by 20%, and maintained comparable reconstruction fidelity. Without any climate supervision, it identified an interpretable heatwave feature concentrated over Western Europe and a Western Pacific typhoon tracker; the latter was validated via causal intervention—activating the feature changed the model's typhoon predictions. Open questions remain regarding a possible sparsity–interpretability tradeoff, the effect of B-spline grid-point count, transferability to other weather model architectures, and whether the typhoon feature reflects physical typhoon mechanisms or merely low-dimensional pressure patterns.

Deep learning weather prediction models are highly accurate, but no one really understands how they internally represent climate phenomena. Sparse autoencoders (SAEs) are the standard tool for mechanistic interpretability—they decompose a model's hidden representations into human-readable features. However, standard SAEs assume that features combine linearly via superposition: multiple features are "stacked" to form an activation pattern, with each feature corresponding to a linear direction. For weather models, atmospheric dynamics are highly nonlinear, so the linear superposition assumption breaks down.

In KAN-SAE, Cheon replaces the ReLU activations of a standard SAE with Kolmogorov-Arnold Networks (KANs)—each latent dimension gets its own learnable nonlinear gating curve. This lets the SAE learn nonlinear activation patterns instead of being forced to explain nonlinear superposition with linear combinations.

Results

Applied to the weather model Sonny, KAN-SAE achieved:

975 active features discovered, versus 566 for the linear baseline—a 72% improvement
20% reduction in inter-feature redundancy
Comparable reconstruction fidelity to the linear baseline

Without any climate supervision, KAN-SAE identified:

1. An interpretable heatwave feature concentrated over Western Europe 2. A Western Pacific typhoon tracker feature, validated through a causal intervention experiment: activating the feature changed the model's typhoon predictions

Open Questions

Does KAN-SAE introduce a tradeoff between sparsity and interpretability—do more flexible activation functions make features harder to align with single concepts?
How does the number of B-spline grid points affect the results?
How well does the method transfer to other weather model architectures?
Could there be spurious correlations—does the discovered "typhoon feature" genuinely correspond to physical typhoon mechanisms, or is it just a low-dimensional pressure pattern?

References

1. Cheon, M. (2026). *Beyond Linear Superposition: Discovering Climate Features in AI Weather Models with KAN-SAE*. arXiv:2605.17493 [cs.LG]. 2. Liu, Z., et al. (2024). *KAN: Kolmogorov-Arnold Networks*. arXiv. 3. Bricken, T., et al. (2023). *Towards Monosemanticity: Decomposing Language Models with Dictionary Learning*. Anthropic.