Tropical Cyclone Attributes in Sparse Autoencoder Features of an AI Weather Model
Authors: Jiapei Liu, Ruidong Li, Ting Sun, Guangheng Ni
Organizations: State Key Laboratory of Hydro-science and Engineering, Department of Hydraulic Engineering, Tsinghua University, Beijing 100084, China · Department of Risk and Disaster Reduction, University College London, London WC1E 6BT, UK
Artificial intelligence (AI) weather models now forecast weather events skillfully, but how they represent these events internally remains unknown. Sparse autoencoders (SAEs) provide sparse features for probing hidden states, but how information about multiple physical attributes of an event is organized across them remains unclear. Using tropical cyclones in GraphCast as a test case, we analyze 2,115 six-hourly timesteps from 205 cyclones, relating SAE activations to position, intensity, and pressure structure across five Layer-8 decompositions and 16 processor layers. Each attribute relates to multiple features; position and pressure-structure relations overlap in a smaller subset. This organization recurs across independent SAEs and layers without feature alignment. In deeper layers, the leading position feature aligns better with position and pressure structure six hours ahead than at present. These results suggest that the organization of physical-attribute relations provides a reproducible level of description for interpreting cyclone-related information in GraphCast.
Figures & tables
Figure 1: Temporal distributions of cyclone-attribute relations for GraphCast Layer-8 feature f613 in 2018–2019: (a) position F1, (b) signed intensity deviation D , and (c) pressure-structure correlation.
Figure 2: Dorian lifecycle in GraphCast Layer-8 feature f613 in 2019. (a–j) Storm-centred maps of MSLP anomaly (a–e) and corresponding feature activation (f–j) at the five dates and UTC times indicated. (k–m) Six-hourly trajectories of position F1, signed intensity deviation D , and pressure-structure correlation from 24 August at 18 UTC to 7 September at 12 UTC; open circles mark the displayed map times. Outer map ticks show absolute longitude and latitude.
Figure 3: Layer-8 organization of cyclone-attribute relations. (a) Dictionary-wide landscape for the 3,589 features with finite summaries for all three measures, with the F1-selected Top 4 highlighted. Colour represents the temporal median pressure-structure correlation ρ . Feature labels (e.g., f613 ) denote SAE feature IDs. (b–e) Spatial patterns of the Top 4 on the common 1 September 2019, 18 UTC Dorian–Juliette field. (f–h) Complete 2018–2019 distributions of position F1, signed intensity deviation D , and pressure-structure correlation.
Figure 4: Organization and physical-time alignment across GraphCast processor depth. (a–c) Layer-wise temporal medians of position F1, signed intensity deviation D , and finite-only pressure-structure correlation for the F1-selected Top 4 of SAE seed 50108. (d–f) Five-seed fixed-feature t+6 -minus- t differences in the same three measures for contemporaneously selected F1-leading features. Grey boxes summarize the temporal distribution of the median difference across five seeds at each matched time. Boxes span the interquartile range; whiskers extend to the outermost observed values within 1.5 interquartile ranges. Red lines connect temporal medians of these differences, and blue points show each seed’s temporal median. For correlation boxes and connected medians, only matched times with at least three finite seed values are used.
While AI weather models now rival operational forecasts, how they represent the atmosphere internally remains an open question: feature attribution reveals which input patterns matter, not what the model computes or how it combines information internally. We train sparse autoencoders (SAEs) on GraphCast to uncover its learned concepts, using atmospheric rivers as our phenomenon of focus. Both standard and Matryoshka SAEs show GraphCast computes atmospheric river intensity, measured by integrated vapor transport (IVT), as a stable internal variable, despite IVT being neither an input nor a target. In contrast to the unstructured concept retrieval of the standard SAE, the Matryoshka SAE orders concepts by importance and exposes their relations. Atmospheric river concepts persist across depth and direct interventions confirm causality. This method offers a way to find internal variables and determine which of them the model actually relies on, which is a prerequisite for asking whether those variables remain meaningful as the phenomenon changes under a warming climate.
Madelyn Mathai, Timothy B. Higgins, Kevin M. Grise +2
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.
Minjong Cheon
Department of Computer Science and Engineering Sejong University
Time-series forecasting informs critical decisions in energy dispatch, industrial operations, and environmental monitoring; understanding the patterns models rely on is essential for assessing reliability and identifying failures. Input attribution identifies important variables and time segments but offers limited insight into internal features, while standard sparse autoencoder (SAE) objectives do not directly constrain cross-variable structure or temporal continuity. We introduce TSAE, a structured sparse autoencoder for forecasting representations that decomposes hidden states into individually inspectable features. TSAE organizes cross-variable structure through shared and variable-routed private dictionaries, separates feature detection from magnitude estimation with gated encoding, and constrains neighboring sparse-code changes according to raw-segment similarity. These mechanisms support analysis of variable context, activation strength, and temporal evolution. Forecast-consistency fine-tuning further improves preservation of the frozen forecaster's outputs. The accompanying TSEVAL protocol separately audits fidelity, feature coherence, and physical calibration to ground feature interpretation. In three-seed experiments with frozen PatchTST on ETTh1, ETTh2, and ETTm1, TSAE achieves the lowest hidden-state reconstruction error, normalized forecast-reconstruction error (NFRE), and feature transition rate among five SAEs at comparable per-token activity. NFRE decreases by 4.3-26.4% relative to the next-best mean. Dataset-dependent tradeoffs in selectivity, physical correlation, and calibration show that fidelity and temporal-stability gains require independent semantic validation to support feature interpretation.