Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
In one sentence. STORM is a one-stage diffusion-based data assimilation framework with linear-complexity attention, scaling to 74,400 GPUs and 32,768-member ensembles.
Abstract
Accurate Earth system prediction requires state inference from incomplete observations, but conventional two-stage data assimilation (DA) is computationally prohibitive because repeated PDE-based ensemble forecasts, observation updates, and intermediate data movement limit ensemble size at high resolution. We introduce STORM, a one-stage generative AI framework that reformulates DA as diffusion-based Bayesian posterior sampling, replacing online PDE ensemble forecasts with scalable AI inference. It further combines a spatiotemporal transformer with a global-attention algorithm that reduces complexity from quadratic to linear through scalable gradient propagation, enabling high-resolution, long-context Earth modeling. STORM scales to 74,400 GPUs on Frontier with 96--99\% strong-scaling efficiency and up to 6 ExaFLOPs sustained BF16 throughput, while enabling 32,768-member ensembles for uncertainty quantification in 34 seconds on 4,096 GPUs. It scales to 20 billion spatiotemporal tokens and 177,000 temporal frames. Hurricane tracking and long-term climate reanalysis demonstrate improved accuracy, including benefits from longer temporal context and recovery of temperature extremes missed by forecast-only predictions.
Labels
- Topic
- Data assimilation, Software, scaling & compression
- Type
- New model or method
- Method
- Transformer / ViT, Diffusion / flow matching / score-based
- Domain
- Global
- Data
- ERA5 / reanalysis
- Status
- Research
- Identifiers
- arXiv:2604.16590 OpenAlex:W7155033104
- Added
- 2026-10-22
Labels and summary are generated by Claude and reviewed weekly. Wrong? Edit this entry.