# Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction

arXiv: https://arxiv.org/abs/2604.16590 · HTML: https://mlsanity.com/paper/2604.16590/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 37 references that are in the map.
- First posted on arXiv 17 Apr 2026.

## Builds on (most cited first, 13 of 13)

- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Is Space-Time Attention All You Need for Video Understanding?](https://mlsanity.com/paper/2102.05095/) (2021)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [MaxViT: Multi-Axis Vision Transformer](https://mlsanity.com/paper/2204.01697/) (2022)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://mlsanity.com/paper/2312.15796/) (2023)
- [DiffDA: a Diffusion Model for Weather-scale Data Assimilation](https://mlsanity.com/paper/2401.05932/) (2024)
- [An Ensemble Score Filter for Tracking High-Dimensional Nonlinear Dynamical Systems](https://mlsanity.com/paper/2309.00983/) (2023)
- [Prithvi WxC: Foundation Model for Weather and Climate](https://mlsanity.com/paper/2409.13598/) (2024)
- [Generative Data Assimilation of Sparse Weather Station Observations at Kilometer Scales](https://mlsanity.com/paper/2406.16947/) (2024)
- [ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability](https://mlsanity.com/paper/2404.14712/) (2024)
- [ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling](https://mlsanity.com/paper/2505.04802/) (2025)

## Built on by (most cited first, 2 of 2)

- [Data-Assimilation-Assisted Reinforcement Learning for Power Grid Control under Load Uncertainty](https://mlsanity.com/paper/2609.27229/) (2026)
- [Stability of Measure-to-Measure Transformers on Sub-Gaussian Data](https://mlsanity.com/paper/2610.07717/) (2026)

## Abstract (arXiv)

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.