- The paper demonstrates that Aurora's latent space is primarily structured by seasonal variations rather than extreme storm regimes, with PCA revealing a prominent annual cycle separation.
- It employs Spatially Pooled PCA and Layer-wise Relevance Propagation to quantitatively show that upper-tropospheric features are crucial for accurate storm forecasting.
- The study underscores the need for nonlinear regime discovery to further unravel complex atmospheric dynamics and enhance operational weather forecasting trust.
Aurora's Internal Representations: Latent Regime Organization and Attribution
Introduction
The development of Transformer-based foundation models for weather forecasting—exemplified by Aurora—marks a significant shift from physics-based numerical weather prediction (NWP) toward direct data-driven emulation of atmospheric dynamics. These models, while outperforming NWP baselines both in efficiency and in several skill metrics, present an interpretability barrier: their high-dimensional latent states and complex attention mechanisms obscure the physical basis for their predictions, limiting operational trust. This paper systematically investigates the internal representations of Aurora, addressing (1) the structure of its latent space with respect to meteorological regimes and (2) the model's capacity to encode the physical drivers of extreme events, specifically vertical storm structure.
Model Architecture and Methodology
Aurora utilizes a 3D Swin Transformer V2 U-Net backbone, combined with a Perceiver IO frontend for flexible ingestion and vertical compression of heterogeneous atmospheric inputs. The latent representations at the bottleneck aggregate physical variables and vertical structure across spatial grids, making them a natural target for regime analysis. For interpretability, the study introduces two primary techniques:
- Spatially Pooled PCA: Applied to the latent feature maps to assess whether the internal representations admit linear separability of meteorological regimes (seasonal cycles, storm events).
- Layer-wise Relevance Propagation (LRP): Adapted for the Swin backbone to attribute Aurora's predictions to specific input locations and variables, resolving the information flow both horizontally and vertically.
The study utilizes ERA5 reanalysis data for evaluation, focusing on the small pretrained Aurora variant.
Latent Regime Analysis
PCA performed on spatially pooled descriptors of the latent activation tensors reveals that Aurora's internal geometry is primarily structured by the annual cycle. The first principal component (PC1) explains 24.1% of variance and separates winter from summer regimes with minimal overlap (cosine stability 0.998 across bootstrap samples), while fall and spring remain less distinguished, particularly in higher components. Contrastive projections reinforce the distinction between winter/summer and partial overlap between transitional seasons.
Conversely, extreme storm events, identified through maximum surface wind magnitude, do not manifest as distinct or stably separable clusters in the latent space. Storm and calm regimes have overlapping distributions even when employing linear contrastive axes, with storm-regime PCs exhibiting substantially lower stability (mean PC3 similarity 0.650, SD 0.269). This suggests that, at the level of global linear geometry, storm-specific signatures are dominated by other sources of variance, and their spatial heterogeneity undermines global, regime-specific structure in the latent manifold.
These findings imply that while Aurora efficiently encodes seasonality—a volumetric, predictable pattern—it does not develop global, linearly separable features for rare or spatially heterogeneous regimes such as cyclonic storms. The organization of the latent manifold, therefore, reflects dominant, slowly varying drivers of atmospheric variability rather than isolated, high-impact events.
Attribution and Vertical Structure Encoding
Local attribution via LRP, targeted at the 1987 Great Storm, offers evidence that Aurora's representations are physically grounded. By adapting relevance propagation to the Swin Transformer's shifted window attention, the study produces input-space attribution heatmaps which consistently localize storm-relevant activations to the cyclone's frontal structures and vertical core across multiple variables. Notably, surface wind anomalies are linked to latent features corresponding to upper-tropospheric levels (~150 hPa), indicating that the model learns vertical coupling—an essential feature for correct synoptic evolution—without explicit vertical structure supervision.
Quantitative evaluation via perturbation testing further supports the faithfulness of Aurora's internal attributions. Masking the top 1% of input-space pixels by LRP relevance damages local forecast quality by a factor of 3.31 relative to random masking. This selectivity demonstrates that the model prioritizes physically relevant features within its attention and does not primarily rely on static or geographical artefacts, a crucial criterion for operational acceptability.
Implications and Future Directions
The results establish that Aurora, despite its nonlinear, high-capacity data-driven design, internalizes coherent atmospheric structure at the scale of both seasonal variability and localized vertical dynamics during extreme events. This contributes to the growing body of work suggesting that foundation models, when trained on sufficiently heterogeneous datasets, are capable of learning implicit physical constraints and coupling relationships.
However, the absence of linearly segregated storm representations in the latent space highlights the limitations of linear analysis and points to the need for nonlinear or topology-aware regime discovery techniques. Furthermore, the perturbation and attribution analysis, while promising, is limited in scope—focused on a single storm and a narrow input mask. Comprehensive evaluation across a range of event types and attribution approaches (e.g., LeGrad or attention-centric XAI methods) is necessary for robust trust in operational settings.
Scaling these analyses to larger model variants, employing richer regime definitions (beyond scalar wind thresholds), and mapping the implications of Aurora's representations for downstream tasks such as uncertainty quantification and rapid event detection, are compelling directions for ongoing research.
Conclusion
The study systematically audits the internal representations of Aurora with respect to meteorological regime structure and attribution, demonstrating strong organization along seasonal axes and physically meaningful vertical coupling in extreme event encoding. These findings, verified via robust statistical validation and perturbation analysis, underscore both the promise and the interpretative challenges of Transformer-based weather models. Continued advances in XAI methodologies and high-dimensional regime characterization will be critical for unlocking the operational reliability and scientific insight offered by next-generation Earth-system foundation models.