Source-linked AI summary
Mechanistic Interpretability of Atmospheric Rivers in GraphCast
Madelyn Mathai, Timothy B. Higgins, Kevin M. Grise, Chirag Agarwal, Antonios Mamalakis
TL;DR
The paper asks what AI weather models compute internally, beyond which inputs influence their predictions. It uses sparse autoencoders to study atmospheric-river concepts in GraphCast and finds a stable internal representation of AR intensity that can be traced and causally probed.
Problem
Feature attribution identifies input patterns that affect predictions but does not reveal what a model computes internally or how it combines information.
Source: §1 · p. 1Method
Results
GraphCast represents AR intensity, measured by IVT, as a stable internal variable despite IVT being neither an input nor a target, and interventions on CONCEPT 99 alter AR properties.
Source: §1 · p. 1, §3.2 · pp. 3–4, §4 · p. 4Takeaways & Limitations
The approach moves weather-model interpretability beyond input attribution toward examining internal computation, concept organization, and forecast dependence.
Source: §1 · p. 2The study characterizes only one parent–child pair, while many more were returned by the sweep.
Source: §4 · p. 4
Abstract
from arXiv · showhide
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.
1 Introduction
The paper moves beyond input attribution to examine what GraphCast computes internally, using sparse autoencoders to study its atmospheric-river concepts. It finds that GraphCast represents AR intensity as a stable internal variable and that intervention on a concept affects AR behavior.
1 Introduction
- Post-hoc feature attribution shows which input patterns affect predictions, but not what weather models learned internally.Earlier climate-model SAE work applied a standard SAE at a single depth, serving as an implementation rather than a framework.Source: §1 · p. 1
- The study trains Standard and Matryoshka sparse autoencoders on GraphCast layers 0, 8, and 15 to identify, organize, and trace concepts.The Matryoshka dictionary makes the AR concept traceable in relation to other concepts.Source: §1 · p. 1
- GraphCast computes AR intensity as a stable internal variable, although it is neither an input nor a training target.The analysis measures intensity with integrated vapor transport (IVT).Source: §1 · p. 1
- Intervening on the AR concept weakens real atmospheric rivers when removed and intensifies their precipitation when written in.The authors describe the concept as causal and frame the approach as extending weather-model interpretability beyond input attribution.Source: §1 · p. 2
2 Methods
The methods apply dictionary-learning SAEs to GraphCast and use the Matryoshka SAE’s ordering to examine relationships among concepts. Spatial activation footprints provide a physical basis for testing whether later concepts are nested within earlier ones.
2 Methods
- GraphCast is a 16-layer GNN trained on ERA5 that forecasts 227 atmospheric fields at six-hour intervals on a 0.25° mesh.It predicts six atmospheric variables at 37 pressure levels and five surface variables, including accumulated precipitation.Source: §2 · p. 2
- Matryoshka SAE prefixes reconstruct the input independently, ordering latents from frequently active concepts in the core to rarer, more specific concepts in outer groups.The ordered groups are G0=[0, 255], G1=[256, 511], G2=[512, 1023], G3=[1024, 2047], and G4=[2048, 4095].Source: §2 · p. 2
- The method measures nesting as spatial containment: the fraction of a child’s node-timestep firings at which its parent also fires.Mesh activation footprints let the authors assess containment geometrically rather than relying on activation co-occurrence alone.Source: §2 · p. 2
3 Results
Matryoshka concepts reveal both intensity-based nesting among atmospheric-river features and a causal role for CONCEPT 99. Interventions show that changing this concept alters AR moisture and precipitation, while its child CONCEPT 3153 does not drive the same effects.
3.1 Matryoshka SAE shows nesting structure for global ARs
- CONCEPT 99 is the strongest global AR concept, with a full-period mean correlation of r = 0.51 between activation and regional-maximum IVT across four regions.Figure 2A compares monthly regional-mean concept activation with regional-maximum IVT during AR timesteps.Source: §3.1 · pp. 2–3, Figure 2 · p. 3
- CONCEPT 3153 is contained in CONCEPT 99 at 0.953, and no other G0 or G1 latent contains it, making 99 its unambiguous single parent.CONCEPT 99 is active at essentially every timestep and touches all but roughly 10% of mesh nodes over the record.Source: §3.1 · p. 3
- CONCEPT 99 nests CONCEPT 3153 by intensity: the child co-fires at the same sites but more often on strong- and extreme-IVT events.The child’s firing sites average 570 versus 460 kg m−1s−1 IVT, and it fires above 750 kg m−1s−1 nearly twice as often as its parent (26% versus 14%).Source: §3.1 · p. 3, Figure 2 · p. 3
3.2 Steering shows causality of CONCEPT 99
- Removing CONCEPT 99 weakens ARs, while amplifying it initially raises IVT and column water vapor before enhanced precipitation depletes the column within a day.Dialing the concept changes within-AR precipitation strongly and monotonically while leaving wind nearly unchanged, consistent with a moisture rather than dynamical channel.Source: §3.2 · pp. 3–4, Figure 3 · p. 4
4 Conclusion and climate change impact
GraphCast computes atmospheric-river intensity as a causal variable even though it was never given that quantity. Matryoshka dictionaries keep the AR concept addressable across depth, while the paper characterizes only one parent–child pair.
4 Conclusion and climate change impact
- AR intensity, which warming is expected to change, is a causal variable GraphCast computes despite never being given it.The paper connects atmospheric rivers to extreme precipitation and flooding.Source: §4 · p. 4
- Across depth, Matryoshka keeps the AR concept in the broad core, while the Standard SAE scatters the same function across arbitrary indices.Stable addressing makes it possible to trace a single concept across depth.Source: §4 · p. 4
- The authors characterized only one parent–child pair, so their account locates the full AR computation in a circuit of related concepts rather than a single concept.Figure 4 plots concept breadth against correlation with AR intensity across layers 0, 8, and 15 for both SAEs.Source: §4 · p. 4, Figure 4 · p. 4
A SAE Training
Six SAEs share a 4,096-latent dictionary width and fixed top-32 sparsity, but differ in normalization, reconstruction objective, and optimization schedule. Their training losses are shown in architecture-specific operating spaces, so only within-panel trends are comparable.
A SAE Training
- All six SAEs map 512-dimensional activations to 4,096 latents and retain the 32 largest pre-activations per sample, fixing sparsity without an ℓ1 penalty.They also use an auxiliary Top-K term to revive dead latents and Adam with a learning rate of 2 × 10−4.Source: A SAE Training · p. 7
- The Standard SAE trains for roughly 2.85 × 10^6 updates over 10 epochs, while the Matryoshka SAE trains for 3 × 10^5 updates over about half an epoch.Schedules were not matched; both training losses plateaued well before their schedules ended.Source: A SAE Training · p. 7
- Figure 5 reports reconstruction-loss trends within each SAE’s own operating space; absolute loss values across panels are not comparable.The Matryoshka curve is smoothed over its raw values because it logs every 500 steps.Source: Figure 5 · p. 7, A SAE Training · p. 7
B Matryoshka Nesting
The paper defines nesting through spatial containment: whether a parent fires at the same node and timestep as a child. This measured relation is a partial order, and strong containment is rare in the broad sweep of candidate pairs.
B Matryoshka Nesting
- Spatial containment measures how often a parent fires at the same node and timestep as a child, among the child’s firing events.A latent fires at a node when its activation exceeds zero; timestep co-occurrence alone would not distinguish nesting for broadly active latents.Source: B Matryoshka Nesting · p. 8
- Containment is a measured partial order, not an architectural binding: a latent may be nested under zero, one, or several parents.Dividing by the child’s firing count makes the relation one-directional, and the reported proportion uses a 95% Wilson interval.Source: B Matryoshka Nesting · p. 8
- Of 1,835,008 candidate pairs, 7,995 (0.44%) exceed containment 0.5 and 1,436 (0.078%) reach 0.9.The overwhelming majority have essentially zero containment; the sweep serves as its own negative control.Source: B Matryoshka Nesting · p. 8
- A later-group latent is nested when its 95% Wilson lower bound exceeds 0.5, strongly nested above 0.9, and exclusively nested when no other early latent meets the strong criterion.Strong nesting need not be exclusive because a narrow latent can lie in the overlap of multiple broad concepts.Source: B Matryoshka Nesting · p. 8
C Activation Steering
Concept edits change an SAE code and add its decoded activation-space pattern back to the model. The edit is rescaled to raw activation units using the factor that undoes the architecture’s normalization.
C Activation Steering
- Standard SAE normalization centers and rescales each activation vector separately, so its edit is multiplied by that vector’s own magnitude inside the forward pass.The factor is the length of the mean-centered activation and is known only during the forecast.Source: C Activation Steering · pp. 8–9
- Matryoshka normalization uses fixed per-dimension training ranges, allowing the edit’s conversion factor to be precomputed and applied during forecasting.The per-dimension minima and maxima map each dimension to −1 and +1.Source: C Activation Steering · pp. 8–9
- A concept edit changes code z_c by δz_c and adds the decoded pattern δz_cW_dec[c] back to the activation.The code can be scaled by β or set to a chosen value b when silent.Source: C Activation Steering · p. 8
- The decoded change is converted from normalized to raw activation units using the same factor that undoes the SAE’s normalization.This conversion preserves the model’s activation scale as the edited forward pass continues.Source: C Activation Steering · p. 9