Towards mechanistic understanding in a data-driven weather model: internal activations reveal interpretable physical features

TL;DR

This paper tackles the opacity of data-driven weather models by seeking mechanistic interpretations of GraphCast's internal representations. It applies sparse autoencoders to hidden activations, uncovering interpretable features across diurnal, seasonal, and large-scale weather phenomena, including sea-ice extent, tropical cyclones, and atmospheric rivers. The authors demonstrate causal control by perturbing a tropical cyclone feature and showing physically consistent changes in predictions, supported by hydrostatic balance and mass-conservation diagnostics. The work provides a path toward trustworthy, scientifically valuable use of data-driven weather models and suggests avenues for discovering novel physical pathways.

Abstract

Large data-driven physics models like DeepMind's weather model GraphCast have empirically succeeded in parameterizing time operators for complex dynamical systems with an accuracy reaching or in some cases exceeding that of traditional physics-based solvers. Unfortunately, how these data-driven models perform computations is largely unknown and whether their internal representations are interpretable or physically consistent is an open question. Here, we adapt tools from interpretability research in Large Language Models to analyze intermediate computational layers in GraphCast, leveraging sparse autoencoders to discover interpretable features in the neuron space of the model. We uncover distinct features on a wide range of length and time scales that correspond to tropical cyclones, atmospheric rivers, diurnal and seasonal behavior, large-scale precipitation patterns, specific geographical coding, and sea-ice extent, among others. We further demonstrate how the precise abstraction of these features can be probed via interventions on the prediction steps of the model. As a case study, we sparsely modify a feature corresponding to tropical cyclones in GraphCast and observe interpretable and physically consistent modifications to evolving hurricanes. Such methods offer a window into the black-box behavior of data-driven physics models and are a step towards realizing their potential as trustworthy predictors and scientifically valuable tools for discovery.

Figures (13)

• Figure 1: GraphCast interpretability pipeline. (A) Standard encode-process-decode architecture of GraphCast. Atmospheric variables are encoded onto an internal Graph Neural Network (GNN) where processing occurs. (B) Capturing the activations at some intermediate layer of the model, we display dense, uninterpretable nodal embeddings as global maps. By learning a transformation such that these fields can be written as a sparse linear combination of feature vectors, we uncover interpretable abstractions in intermediate GraphCast processing layers.
• Figure 2: Features on many timescales. (A) Power spectrum of global mean activation over time. Distinctive peaks indicate the existence of features with diurnal, seasonal, and annual oscillations. (B) Two example seasonal feature time series, one activating in the northern hemisphere winter and the other in the southern hemisphere winter. (C) Feature 1710 shows strong seasonality and its activation tracks northern sea-ice extent (overlaid in red), even though no ice extent information is processed by GraphCast. Feature 1437 shows similar strong seasonality but instead tracks southern sea-ice extent (overlaid in red). Another annual feature, feature 655, corresponds to surface heating in desert regions and seasonally migrates. (D) Some example strongly diurnal features. From top to bottom: daytime activation in especially arid regions; ocean basin activation in early morning; precipitation patterns resembling the ITCZ; corresponding negative of precipitation patterns, or especially dry ocean regions; rain forests, activating primarily in the Amazon during the day but also strongly in Indonesia and Africa during their respective sunlight hours.
• Figure 3: Grid-locked features: spurious features activate on the grid representation of GraphCast.
• Figure 4: TC feature interpretability. (A) Comparison of three year average of the ground truth TC dataset Kim2025AWeather; the highest F1 score feature, feature 3243 from our $l=8$, $k=32$, $n_l=4096$ SAE; and one of the most informative neurons at layer 8, neuron 19. The average activation of feature 3243 aligns strongly with the ground truth dataset, while the single neuron is mostly uninformative. (B) One example of localized feature activation. Top: feature 3243 activation for the duration of Hurricane Ida (2021) as it makes landfall. Bottom: 10 meter wind magnitudes from the same times. (C) A second example of localized feature activation. Top: feature 3243 activation for Typhoon Hagibis (2019) in the Pacific Basin. Bottom: 10 meter wind magnitudes from the same times.
• Figure 5: Artificially increasing (decreasing) a hurricane feature (feature 3243) during the GraphCast forward increases (decreases) hurricane strength prediction in a physically plausible manner (A,B) Maximum 10 meter wind speeds and minimum mean sea level pressure for all steered GraphCast hurricane forecasts, along with comparison to ERA5. (C) Paths of modified hurricanes. (D,E) Mass continuity and hydrostatic residual show the range of modified $\gamma$ parameters for which physical consistency is maintained. (F) Gradient wind balance demonstrates that feature modification increases pressure gradient and wind speed in a manner that still maintains force balance around eye of hurricane. Solid lines show measured azimuthal wind speed in a frame relative to the hurricane. Dashed lines show theoretical wind speeds calculated from measured pressure gradients.