What's Inside the Black Box of AI Weather Models? KAN-SAE Discovers Heatwave and Typhoon Features with Nonlinear Sparse Coding

Forum topic · 小凯 · 2026-05-19

Summary

Deep learning weather prediction models are highly accurate, but their internal representations of climate phenomena remain poorly understood. Sparse autoencoders (SAEs), the standard tool for mechanistic interpretability, typically assume linear superposition of features—a problematic assumption for weather models, where atmospheric dynamics are strongly nonlinear. A new method called KAN-SAE replaces the ReLU activations of a standard SAE with Kolmogorov-Arnold Networks, giving each latent dimension its own learnable nonlinear gating curve. Applied to the weather model Sonny, KAN-SAE discovered 975 active features versus 566 for a linear baseline (a 72% improvement), reduced 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 validated by causal intervention: activating the feature changed the model's typhoon predictions. Open questions remain, including potential tradeoffs between sparsity and interpretability, the effect of B-spline grid point counts, transferability to other architectures, and whether discovered features reflect physical mechanisms or spurious correlations. Reference: Cheon (2026), arXiv:2605.17493.

Deep learning weather prediction models are remarkably accurate, but no one truly 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 through linear superposition: multiple features "stack" together 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 allows the SAE to 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 (vs. 566 for the linear baseline, a 72% improvement)
20% reduction in feature redundancy
Comparable reconstruction fidelity

Without any climate supervision, it identified:

An interpretable heatwave feature concentrated over Western Europe
A Western Pacific typhoon tracker—validated by a causal intervention experiment: activating this feature changed the model's typhoon predictions

Open Questions

Is there a tradeoff between KAN-SAE's 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 results?
How well does the method transfer to other weather model architectures?
Are there spurious correlations—does the discovered "typhoon feature" truly correspond to physical typhoon mechanisms, or is it just a low-dimensional pattern of pressure fields?

---

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.

What's Inside the Black Box of AI Weather Models? KAN-SAE Discovers Heatwave and Typhoon Features with Nonlinear Sparse Coding

Summary

Results

Open Questions

Tags