Interpretability of AI weather models via intermediate decoding

Matthias Beylich; Kirsten I. Tempest; Tobias Selz; George C. Craig

physics.ao-phOct 16, 2026

Interpretability of AI weather models via intermediate decoding

Authors: Matthias Beylich, Kirsten I. Tempest, Tobias Selz, George C. Craig

Organizations: Meteorological Institute, LMU Munich

Abstract

AI weather prediction models are often regarded as black boxes. Despite their impressive skill, there is little understanding of how the input data is processed during inference to produce the forecast. This limitation arises from the use of high-dimensional latent representations in which individual dimensions are no longer directly interpretable in physical terms. To gain insight into how forecasts are constructed, we propose the method of intermediate decoding for AI weather models and apply it to GraphCast, Aurora and ArchesWeather. In this approach, latent states from intermediate depths are decoded back into meteorological fields. To do this, we use the decoder that has been trained to output the forecast at the end of processing. Since latent states at different depths do not share a consistent basis, we train affine transformations that align intermediate latent states with the decoder's expected input space. Applying this framework to the three models yields the following results. First, the skill of forecasts at intermediate stages of processing improves gradually through the processor layers. Second, bypassing all processor blocks shows that the encoder-decoder pair is not an identity mapping returning the input, but instead plays a significant role in the forecast construction. Third, by examining increments of individual processor blocks across model depth, we show that all three models progressively refine the forecast from coarse to finer spatial scales. Finally, we assess the physical consistency of the intermediate decoding forecasts and find that large scales remain physically balanced throughout processing, while small scales become more balanced with each additional processor block.

Figures & tables

Figure 1: Residual network architecture. The encoder maps the input to the latent state

\tilde{X}_{0}

, which is updated by

L

processor blocks, each adding an increment

F_{i}(\tilde{X}_{i})

, creating the intermediate states

\tilde{X}_{i}

. The decoder maps the final latent state

\tilde{X}_{L}

to the forecast.

Figure 2: Architectural overview of GraphCast, Aurora, and ArchesWeather. Gray boxes denote the encoder and decoder, blue boxes groups of processor blocks, with the corresponding latent states

\tilde{X}_{i}

indicated above each group and the dimensions of the latent state below (for Aurora and ArchesWeather: vertical levels

\times

latitude

\times

longitude

\times

channels; for GraphCast: mesh nodes

\times

channels). GraphCast operates at constant resolution, whereas Aurora and ArchesWeather downscale and upscale between stages, changing both spatial resolution and channel dimension. Skip connections pass the output of the first stage to later stages, either by addition (

+

) or by concatenation along the channel dimension (conc). For Aurora, the skip connections are shown as implemented in model code version v1.8.0 (see text).

Model	1% convergence (epochs)	0.1% convergence (epochs)
GraphCast	4–7 (median 5)	6–13 (median 10)
Aurora	2–18 (median 2)	2–26 (median 4)
ArchesWeather	2–7 (median 3.5)	2–14 (median 7)

Table 1: Epochs required for translators to converge to within 1% and 0.1% relative improvement in validation loss (range and median).

Figure 3: Progression of extracted latent space channels across model depth for a forecast valid at 12:00 UTC on 10 March 2020. Shown is the progression of a representative channel for each model on the respective internal model grid, without and with the use of translators. For GraphCast, every second latent state on the icosahedral mesh is shown. For Aurora and ArchesWeather, latent states from the first and last layer are shown on the lat/lon grid on a single vertical level. Colorbar limits are set to the 99th percentile of the absolute raw latent values over all shown blocks of each model.

Figure 4: Comparison of covariance structures across model depth using Frobenius cosine similarity, for untranslated (raw) and translated latent states. Covariance matrices are computed from one year of data (2020) for GraphCast, Aurora, and ArchesWeather, and compared across all pairs of latent states; high similarity indicates shared covariance structure. Top panels show covariance similarity for untranslated latent states; bottom panels show covariance similarity for latent states after applying the respective translators.

Figure 5: Model loss and relative MSE across latent states for intermediate decoding, averaged over 48 runs (the 10th of each month in 2020 at 00, 06, 12, and 18 UTC). Panels (a–c) show each model’s native loss function for GraphCast, Aurora, and ArchesWeather, respectively, for direct (dotted) and tuned (solid) decodings. Panels (d–f) show the relative MSE of the tuned decodings per variable for the surface variables ( u_10m , v_10m , t_2m , pres_msl , and, for GraphCast only, prec_6h ) and the variables on the 500 hPa level ( geopot , temp , qv , u , v , and w ). For Aurora and ArchesWeather, dashed lines mark the first latent state of a new stage.

Figure 6: Comparison of encoder–decoder output and full model prediction for GraphCast, Aurora, and ArchesWeather, for a forecast valid at 12:00 UTC on 10 March 2020. All panels show increments relative to the most recent input state: (left) the difference between the two input states, $X^{t}-X^{t-1}$ ; (middle) the encoder–decoder-only increment, $\hat{X}^{t+1}_{0}-X^{t}$ ; and (right) the increment of the full model prediction, $\hat{X}^{t+1}_{L}-X^{t}$ . Shown variables are geopotential at 500 hPa and 2 m temperature. Colorbar limits are set to 60 % of the maximum absolute value for each row.

Figure 7: Increments $\Delta\hat{X}^{t+1}_{i}$ of consecutive intermediate decodings of GraphCast for processor blocks 1, 2, 14, and 15, for the forecast valid at 12:00 UTC on 10 March 2020. Top row: geopotential at 500 hPa ( $\mathrm{m^{2}\,s^{-2}}$ ); bottom row: 10 m zonal wind ( $\mathrm{m\,s^{-1}}$ ). Note the different color scales.

Figure 8: Power spectra $\ell\,P_{i}(\ell)$ of the increments $\Delta\hat{X}^{t+1}_{i}$ as a function of total wavenumber $\ell$ and layer depth $i$ , for geopotential at 500 hPa (top) and 10 m zonal wind (bottom), for GraphCast, Aurora, and ArchesWeather, averaged over the 48 evaluation runs (the 10th of each month in 2020 at 00, 06, 12, and 18 UTC). Black lines show the spectral centroid of each increment (Eq. 16 ). The first increment ( $i=0$ ) is that of the encoder–decoder-only forecast, $\hat{X}^{t+1}_{0}-X^{t}$ . Spectra of ArchesWeather are truncated at $\ell=120$ , corresponding to its $1.5^{\circ}$ resolution. Colorbar limits are set to the 0.2 and 99.8 percentiles of all values in each row.

Figure 9: Vertical profiles of physical balances for the full model runs, averaged over the 48 evaluation runs (the 10th of each month in 2020 at 00, 06, 12, and 18 UTC). Top: ratio of ageostrophic to geostrophic wind, $r_{\mathrm{geo}}$ ; bottom: ratio of divergence to vorticity, $r_{\mathrm{div}}$ , for all three models and the ERA5 reference. Dotted lines show ERA5 scaled in spectral space to match each model’s spectral energy at each total wavenumber $\ell$ .

Figure 10: Ratio of ageostrophic to geostrophic wind, $r_{\mathrm{geo}}$ , and ratio of divergence to vorticity, $r_{\mathrm{div}}$ , at 500 hPa for the intermediate decodings of GraphCast, Aurora, and ArchesWeather as a function of model depth, averaged over the 48 evaluation runs (the 10th of each month in 2020 at 00, 06, 12, and 18 UTC). Dotted lines show large scales ( $\ell\leq 20$ ), solid lines small scales ( $\ell>20$ ). For reference, horizontal lines show the corresponding values for the two scale ranges in ERA5.

Interpretable ML

Interpretable machine learning develops models and explanations that make predictions, evidence, and decision processes understandable to people.

See all 310 papers

The physics of AI weather models
May 22, 2026George Craig, Tobias Selz, Matthias Beylich +1AI for Science Representational Similarity Analysis
Meteorological Institute, LMU Munich · Karlsruhe Institute of Technology
Mechanistic Interpretability Tool for AI Weather Models
Apr 22, 2026Kirsten I. Tempest, Matthias Beylich, George C. CraigInterpretable ML Mechanistic Interpretability
Meteorological Institute Munich, Ludwig-Maximilians-Universität, 80333 Munich, Germany · Deutsches Zentrum für Luft- und Raumfahrt, Oberpfaffenhofen, Germany
Missing the Butterfly and Predicting the Past: Features or Bugs of Accurate AI Weather Models?
Aug 26, 2026Pedram Hassanzadeh, Weidong Li, Y. Qiang Sun +4Weather Forecasting Chaotic Time-Series Prediction
Department of Geophysical Sciences, University of Chicago, Chicago, USA · Committee on Computational and Applied Mathematics, University of Chicago, Chicago, USA. · Data Science Institute, University of Chicago, Chicago, USA. +3

May 22, 2026J/K move · Enter open · S save

The physics of AI weather models

George Craig, Tobias Selz, Matthias Beylich, Kirsten I. Tempest

Meteorological Institute, LMU Munich · Karlsruhe Institute of Technology

AI for Science Representational Similarity Analysis Weather Forecasting

Open paper AskPDF arXiv

Could it be that AI weather models are solving physical equations, although they may not be the equations used by conventional NWP models? We compute correlations of forecast skill and Centered Kernel Alignment, providing evidence that different AI weather models represent the atmosphere in similar ways, despite differences in architecture and capacity. We argue that the architecture and training of the AI models constrains the form of the physical laws that they might simulate. In particular, we propose that the models implement a particle description of the atmosphere, where the latent variables at each mesh point correspond to the position of a particle in the high dimensional latent space. We hypothesize that the movement of the particles follows a gradient flow in the latent space towards a minimum of a learned free energy functional. Analysis of the GraphCast and Aurora models show that they make changes on large spatial scales in the early processor layers and move to smaller scale with increasing layer depth, consistent with the gradient flow hypothesis.

Abstract

Interpretable ML

Explore similar work

The physics of AI weather models

Mechanistic Interpretability Tool for AI Weather Models

Missing the Butterfly and Predicting the Past: Features or Bugs of Accurate AI Weather Models?