用百亿参数扩散模型提升天气气候预测精度与稳定性
AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
- 基于Swin结构的像素级扩散模型，结合窗口并行技术实现高效扩展
- 在0.25°分辨率下达到95.5%弱扩展效率，90天预测仍保持稳定
- 适合高性能计算与高分辨率气候模拟研究者使用
生成式机器学习为理解复杂地球系统动态提供了新机遇。近期基于扩散的方法相比确定性方法能缓解谱偏差并改善集合校准，但在高分辨率下难以稳定扩展。本文提出AERIS，一个参数量1.3至800亿的像素级Swin扩散变压器，用于填补该空白；并引入SWiPe，一种将窗口并行与序列、流水线并行结合的技术，可在不增加通信开销或全局批次大小的前提下拆分基于窗口的Transformer。在Aurora超算（10,080节点）上，AERIS以1×1图像块尺寸在0.25° ERA5数据集上实现10.21 ExaFLOPS（混合精度）持续性能和11.21 ExaFLOPS峰值性能，弱扩展效率达95.5%，强扩展效率为81.6%。AERIS超越IFS ENS，在季节尺度至90天内保持稳定，彰显百亿参数扩散模型在气象与气候预测中的潜力。
原文摘要 · Abstract (English)
Generative machine learning offers new opportunities to better understand complex Earth system dynamics. Recent diffusion-based methods address spectral biases and improve ensemble calibration in weather forecasting compared to deterministic methods, yet have so far proven difficult to scale stably at high resolutions. We introduce AERIS, a 1.3 to 80B parameter pixel-level Swin diffusion transformer to address this gap, and SWiPe, a generalizable technique that composes window parallelism with sequence and pipeline parallelism to shard window-based transformers without added communication cost or increased global batch size. On Aurora (10,080 nodes), AERIS sustains 10.21 ExaFLOPS (mixed precision) and a peak performance of 11.21 ExaFLOPS with $1 \times 1$ patch size on the 0.25° ERA5 dataset, achieving 95.5% weak scaling efficiency, and 81.6% strong scaling efficiency. AERIS outperforms the IFS ENS and remains stable on seasonal scales to 90 days, highlighting the potential of billion-parameter diffusion models for weather and climate prediction.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。