MetMamba:基于时空Mamba模式的区域天气预报

MetMamba: Regional Weather Forecasting with Spatial-Temporal Mamba Model

https://arxiv.org/pdf/2408.06400

ABSTRACT

Deep Learning based Weather Prediction (DLWP) models have been improving rapidly over the last few years, surpassing state of the art numerical weather forecasts by significant margins. While much of the optimization effort is focused on training curriculum to extend forecast range in the global context, two aspects remains less explored: limited area modeling and better backbones for weather forecasting.

We show in this paper that MetMamba, a DLWP model built on a state-of-the-art state-space model, Mamba, offers notable performance gains and unique advantages over other popular backbones using traditional attention mechanisms and neural operators. We also demonstrate the feasibility of deep learning based limited area modeling via coupled training with a global host model.

基于深度学习的天气预报（DLWP）模型近年来取得了迅速的发展，显著超越了最先进的数值天气预报模型。尽管大部分优化工作都集中在训练课程上，以在全球范围内延长预报范围，但有两个方面仍然探索较少：有限区域建模和更适用于天气预报的骨干网络。

本文展示了MetMamba，一种基于最先进的状态空间模型Mamba构建的DLWP模型，该模型在使用传统注意力机制和神经算子的其他流行骨干网络上具有显著的性能提升和独特优势。本文还通过全局宿主模型的耦合训练，证明了基于深度学习的有限区域建模的可行性。

Introduction

随着气候变化导致的极端天气事件增多，天气预报备受关注。研究者通过复杂数值方案和方程提高预报准确性，但计算复杂度高，需大规模计算集群，成本高昂且耗时长，使集合预报变得昂贵或不可行。

基于深度学习的天气预报模型（DLWP），通过训练ECMWF的ERA5数据集，表现出色。FourCastNet是首个与ECMWF IFS竞争的数据驱动模型，随后其他模型也在各项指标上超越IFS。这些模型能以较低计算量捕捉全球天气趋势。

尽管粗网格全球预报对长期预测有重要意义，但极端天气多为对流形式，细网格预报也难以反映，需本地化调整。因此，有限区域模型（LAM）在多领域至关重要。本研究在LAM背景下进行，旨在探索深度学习在LAM中的可行性。

Mamba模型在连续数据上表现卓越，计算复杂度低，其视觉变体还具有类似卷积网络的有效感受野，是Transformer主导DLWP模型的有力替代品，可构建更深模型和更高分辨率，捕捉更复杂动态。

模型

Result

[5] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven highresolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214, 2022.

[6] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al. Graphcast: Learning skillful medium-range global weather forecasting. arXiv preprint arXiv:2212.12794, 2022.

[7] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Pangu-weather: A 3d highresolution model for fast and accurate global weather forecast. arXiv preprint arXiv:2211.02556, 2022.

[8] Lei Chen, Xiaohui Zhong, Feng Zhang, Yuan Cheng, Yinghui Xu, Yuan Qi, and Hao Li. Fuxi: a cascade machine learning forecasting system for 15-day global weather forecast. npj Climate and Atmospheric Science, 6(1):190, 2023.

[9] Kang Chen, Tao Han, Junchao Gong, Lei Bai, Fenghua Ling, Jing-Jia Luo, Xi Chen, Leiming Ma, Tianning Zhang, Rui Su, et al. Fengwu: Pushing the skillful global medium-range weather forecast beyond 10 days lead. arXiv preprint arXiv:2304.02948, 2023.

[10] Stephan Rasp, Stephan Hoyer, Alexander Merose, Ian Langmore, Peter Battaglia, Tyler Russel, Alvaro SanchezGonzalez, Vivian Yang, Rob Carver, Shreya Agrawal, Matthew Chantry, Zied Ben Bouallegue, Peter Dueben, Carla Bromberg, Jared Sisk, Luke Barrington, Aaron Bell, and Fei Sha. Weatherbench 2: A benchmark for the next generation of data-driven global weather models, 2023.