面向地球系统预测的全局注意力线性复杂度超规模生成数据同化
机器学习
2026-04-21 v1 人工智能
摘要
准确的天气和气候预测依赖于数据同化（DA），该方法通过将观测与模型整合来估计地球系统状态。尽管超规模计算显著推进了地球模拟，但可扩展且准确的地球系统状态推断仍然是一个根本性瓶颈，这限制了不确定性定量化和极端事件预测。我们引入了一个统一的单阶段生成DA框架，将同化重新表述为贝叶斯后验抽样，取代了传统的预报-更新循环，采用计算密集、GPU高效的推理。核心是STORM（Spatiotemporal Transformer with Optimized Reduced Model），一种具有全局注意力线性复杂度缩放算法的新型时空转换器，突破了注意力的二次方块。我们在Frontier超级计算机的32,768个GPU上实现了63%的强度扩展效率和1.6 ExaFLOP的持续性能。我们进一步扩展到200亿时空token，使其能够在177k时间帧上进行km尺度的全球建模，先前无法到达的规模，为地球系统预测树立了新的范式。
引用
@article{arxiv.2604.16590,
title = {Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction},
author = {Xiao Wang and Zezhong Zhang and Isaac Lyngaas and Hong-Jun Yoon and Jong-Youl Choi and Siming Liang and Janet Wang and Hristo G. Chipilski and Ashwin M. Aji and Feng Bao and Peter Jan van Leeuwen and Dan Lu and Guannan Zhang},
journal= {arXiv preprint arXiv:2604.16590},
year = {2026}
}