Mosaic模型通过稀疏注意力提升气象预报的频谱保真度,实现高效精准预测。
(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models

- 采用网格对齐的块稀疏注意力,在原始分辨率上捕捉长程依赖。
- 1.5°分辨率下参数量214M,性能优于6倍更细分辨率的模型。
- 单卡H100可在12秒内完成10天24成员预报,适合实际业务应用。
我们提出Mosaic,一种概率性气象预报模型,解决机器学习气象预测中的三大频谱退化问题:统计性频谱衰减、架构导致的高频混叠及参数引起的高频泄漏。Mosaic通过学习的函数扰动生成集合成员,并在原生分辨率网格上运行,采用与硬件对齐的网格对齐块稀疏注意力机制,通过共享相邻查询的键值对,在线性计算成本下捕获长程依赖。在1.5°分辨率、214M参数条件下,Mosaic在关键变量上的表现达到或超越在6倍更细分辨率下训练的模型,成为1.5°模型中的最新最佳水平,生成的集合成员在所有可解析频率上均展现出近乎完美的频谱对齐。单张H100 GPU上,24成员10天预报耗时不足12秒。代码已开源于https://github.com/maxxxzdn/mosaic。
原文摘要 · Abstract (English)
We introduce Mosaic, a probabilistic weather forecasting model that addresses three failure modes of spectral degradation in ML-based weather prediction: spectral damping (statistical), high-frequency aliasing (architectural), and residual high-frequency leakage (parametric). Mosaic generates ensemble members through learned functional perturbations and operates on native-resolution grids via mesh-aligned block-sparse attention, a hardware-aligned mechanism that captures long-range dependencies at linear cost by sharing keys and values across spatially adjacent queries. At 1.5° resolution with 214M parameters, Mosaic matches or outperforms models trained on 6$\times$ finer resolution on key variables and achieves state-of-the-art results among 1.5° models, producing well-calibrated ensembles whose individual members exhibit near-perfect spectral alignment across all resolved frequencies. A 24-member, 10-day forecast takes under 12s on a single H100~GPU. Code is available at https://github.com/maxxxzdn/mosaic.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。