论文
地球系统预测中百亿亿次生成数据同化的线性复杂度全局关注
Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
摘要
准确的地球系统预测需要从不完整的观测中进行状态推断,但传统的两阶段数据同化(DA)在计算上是令人望而却步的,因为重复的基于偏微分方程的集合预测、观测更新和中间数据移动限制了高分辨率下的集合大小。我们引入了 STORM,这是一种单阶段生成 AI 框架,它将 DA 重新表述为基于扩散的贝叶斯后验采样,用可扩展的 AI 推理取代在线 PDE 集成预测。它进一步将时空 Transformer 与全局注意力算法相结合,通过可扩展的梯度传播将复杂性从二次降低到线性,从而实现高分辨率、长上下文地球建模。 STORM 在 Frontier 上可扩展到 74,400 个 GPU,具有 96--99\% 的强大扩展效率和高达 6 ExaFLOP 的持续 BF16 吞吐量,同时支持 32,768 个成员的集成在 34 秒内在 4,096 个 GPU 上进行不确定性量化。它可扩展至 200 亿个时空标记和 177,000 个时间帧。飓风跟踪和长期气候再分析证明了准确性的提高,包括更长的时间背景带来的好处以及仅预测所错过的极端温度的恢复。