用AI替代传统计算密集型气象预测流程，实现超大规模不确定性量化。
Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
- 将数据同化转为扩散模型后验采样，跳过耗时的物理模拟步骤。
- 4096块GPU上32768个成员集合仅需34秒，达6 ExaFLOPs算力利用率。
- 适合超大规模气候建模与灾害预警，尤其需长时序上下文的场景。
精准地球系统预测需从不完整观测中推断状态，但传统两阶段数据同化因重复的偏微分方程（PDE）集合预报、观测更新和中间数据传输，在高分辨率下严重受限于集合规模。本文提出STORM，一种单阶段生成式AI框架，将数据同化重构为基于扩散的贝叶斯后验采样，以可扩展的AI推理替代在线PDE集合预报。其结合时空变压器与全局注意力算法，通过可扩展梯度传播将复杂度从二次降低至线性，支持高分辨率、长上下文地球建模。STORM在Frontier系统上实现74,400块GPU的并行，强缩放效率达96%-99%，峰值持续BF16吞吐量达6 ExaFLOPs；在4,096块GPU上实现32,768成员集合的不确定性量化，耗时仅34秒。模型支持200亿时空标记与17.7万时间帧。飓风追踪与长期气候再分析实验表明，其精度优于传统预报，得益于更长的时间上下文及对温度极端事件的恢复能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。