Does Aurora Encode Atmospheric Structure? Latent Regime Analysis and Attribution

Technical Summary: Does Aurora Encode Atmospheric Structure? Latent Regime Analysis and Attribution

Problem Statement
While deep learning (DL) foundation models like Aurora, FourCastNet, and GraphCast have demonstrated superior accuracy and computational efficiency in weather forecasting compared to traditional Numerical Weather Prediction (NWP), they operate as opaque "black boxes." Unlike NWP, which relies on explicit physical equations, DL models encode atmospheric states into high-dimensional, non-interpretable latent representations. This opacity hinders operational trust, as meteorologists cannot verify if a model's predictions are physically consistent or merely the result of spurious correlations. Although Explainable AI (XAI) has been applied to meteorological CNNs, the internal dynamics of billion-parameter transformer-based foundation models remain largely unexplored. Specifically, it is unknown whether these models organize their latent spaces by distinct meteorological regimes (e.g., seasons vs. storms) or if they can attribute forecasts to the correct 3D vertical atmospheric structures.

Methodology
The authors investigate the internal representations of Aurora, a foundation model based on a 3D Swin Transformer V2 U-Net architecture pre-trained on heterogeneous datasets (ERA5, HRES, CMIP6). The study addresses two core research questions (RQs) using the AuroraSmallPretrained checkpoint:

1. Latent Regime Analysis (RQ1): To determine if the latent space is organized by meteorological regimes, the authors applied Spatially Pooled Principal Component Analysis (PCA) to the model's bottleneck latent feature map ($Z$).

   • Data: 2005 hurricane season data over the US, filtered for extreme storm events ($v_{max} \ge 24.0$ m/s) and calm conditions ($v_{max} \le 15.0$ m/s), alongside seasonal segments.
   • Technique: Statistical descriptor vectors (mean, std, RMS) were computed for spatially pooled regions. Singular Value Decomposition (SVD) was used to identify dominant modes of variation.
   • Validation: Bootstrap resampling ($N=1000$) was employed to quantify the stability of eigenvectors and ensure the observed patterns were not artifacts of sample size.

2. Local Attribution and Vertical Structure (RQ2): To determine if the model isolates specific storm systems and reflects their 3D structure, the authors applied Layer-wise Relevance Propagation (LRP) to the "Great Storm of 1987."

   • Adaptation: The LRP method was adapted to the Swin V2 backbone by wrapping the shifted-window self-attention mechanism to preserve the computational graph during cyclic spatial shifts (torch.roll), ensuring relevance conservation.
   • Process: A binary spatial mask corresponding to the target region (Europe) was applied to the latent activations. Relevance was propagated backward to the input space to generate heatmaps for variables like temperature, wind, and Mean Sea Level (MSL) pressure.
   • Quantitative Validation: A single-shot perturbation test was conducted on the North Atlantic sector. The top 1% of pixels with the highest relevance scores were masked, and the resulting forecast distortion ratio ($D$) was compared against random masking to verify the faithfulness of the attribution.

Key Results

• Latent Organization (RQ1): The analysis reveals that Aurora's latent space is primarily organized by seasonal cycles rather than specific dynamic storm events.

  • The first principal component (PC1) captures 24.1% of total variance, functioning as a seasonal axis that linearly separates winter and summer regimes with high stability (cosine similarity of 0.998 in bootstrap validation).
  • Conversely, extreme storm events do not form a linearly separable cluster. Storm and calm regimes exhibit overlapping distributions, and the storm-regime PC3 shows significant instability (cosine similarity of 0.650), likely due to the spatial heterogeneity of individual cyclone tracks and limited sample size.

• Attribution and Vertical Structure (RQ2): The LRP analysis demonstrates that the model attends to features consistent with the 3D vertical structure of the Great Storm of 1987.

  • Spatial Focus: Unlike a baseline where relevance is geographically locked to static features (e.g., topography), the model's relevance for the storm is tightly localized to the cyclone's head and frontal boundaries.
  • Vertical Coherence: Vertical analysis shows that surface wind activations contribute to latent features encoding the upper troposphere (approx. 150 hPa). This suggests the model has learned to link surface anomalies with their necessary vertical drivers without explicit physical instruction.
  • Perturbation Impact: Masking the top 1% of LRP-identified pixels degraded the forecast (MSE) by a factor of 3.31× more than random masking ($MSE_{LRP} = 5.30 \times 10^4$ vs. $MSE_{rand} = 1.60 \times 10^4$). This confirms that the identified features are critical drivers of the model's inference.

Significance and Claims
The paper claims that Aurora moves beyond simple pattern matching. Despite being a purely data-driven model, it:

1. Disentangles seasonal modes in its latent space, organizing atmospheric states by the annual cycle.
2. Encodes vertical atmospheric structure, successfully linking surface conditions to upper-atmosphere dynamics.
3. Prioritizes physically relevant features over static geographical shortcuts, as evidenced by the perturbation tests.

Limitations
The authors acknowledge several constraints:

• The analysis was performed on the small version of Aurora due to computational limits.
• Storm definitions relied solely on scalar wind magnitude, lacking vector metrics like vorticity.
• The latent analysis used a relatively small sample size, and the overlap in seasonal regimes may reflect the limitations of linear PCA in resolving cyclic topologies.
• The attribution study focused on a single storm instance, suggesting a need for broader evaluation across diverse meteorological phenomena.