Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866908975333638144 |
|---|---|
| author | Wang, Xiao Zhang, Zezhong Lyngaas, Isaac Yoon, Hong-Jun Choi, Jong-Youl Liang, Siming Wang, Janet Chipilski, Hristo G. Aji, Ashwin M. Bao, Feng van Leeuwen, Peter Jan Lu, Dan Zhang, Guannan |
| author_facet | Wang, Xiao Zhang, Zezhong Lyngaas, Isaac Yoon, Hong-Jun Choi, Jong-Youl Liang, Siming Wang, Janet Chipilski, Hristo G. Aji, Ashwin M. Bao, Feng van Leeuwen, Peter Jan Lu, Dan Zhang, Guannan |
| contents | Accurate weather and climate prediction relies on data assimilation (DA), which estimates the Earth system state by integrating observations with models. While exascale computing has significantly advanced earth simulation, scalable and accurate inference of the Earth system state remains a fundamental bottleneck, limiting uncertainty quantification and prediction of extreme events. We introduce a unified one-stage generative DA framework that reformulates assimilation as Bayesian posterior sampling, replacing the conventional forecast-update cycle with compute-dense, GPU-efficient inference. At the core is STORM, a novel spatiotemporal transformer with a global attention linear-complexity scaling algorithm that breaks the quadratic attention barrier. On 32,768 GPUs of the Frontier supercomputer, our method achieves 63% strong scaling efficiency and 1.6 ExaFLOP sustained performance. We further scale to 20 billion spatiotemporal tokens, enabling km-scale global modeling over 177k temporal frames, regimes previously unreachable, establishing a new paradigm for Earth system prediction. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2604_16590 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction Wang, Xiao Zhang, Zezhong Lyngaas, Isaac Yoon, Hong-Jun Choi, Jong-Youl Liang, Siming Wang, Janet Chipilski, Hristo G. Aji, Ashwin M. Bao, Feng van Leeuwen, Peter Jan Lu, Dan Zhang, Guannan Machine Learning Artificial Intelligence Accurate weather and climate prediction relies on data assimilation (DA), which estimates the Earth system state by integrating observations with models. While exascale computing has significantly advanced earth simulation, scalable and accurate inference of the Earth system state remains a fundamental bottleneck, limiting uncertainty quantification and prediction of extreme events. We introduce a unified one-stage generative DA framework that reformulates assimilation as Bayesian posterior sampling, replacing the conventional forecast-update cycle with compute-dense, GPU-efficient inference. At the core is STORM, a novel spatiotemporal transformer with a global attention linear-complexity scaling algorithm that breaks the quadratic attention barrier. On 32,768 GPUs of the Frontier supercomputer, our method achieves 63% strong scaling efficiency and 1.6 ExaFLOP sustained performance. We further scale to 20 billion spatiotemporal tokens, enabling km-scale global modeling over 177k temporal frames, regimes previously unreachable, establishing a new paradigm for Earth system prediction. |
| title | Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2604.16590 |