[论文解读] ArchesWeather: An efficient AI weather forecasting model at 1.5° resolution

Guillaume Couairon, Christian Lessig+2 more|arXiv (Cornell University)|May 23, 2024

Meteorological Phenomena and SimulationsEarth and Planetary Sciences被引用 3

一句话总结

ArchesWeather 是一个在 1.5° 分辨率下计算效率极高的 AI 天气预报模型，它用一种新颖的跨层级注意力（CLA）机制替代了 3D 局部注意力，仅用数个 GPU 天的训练预算便实现了最先进性能。其表现优于 1.4° 的 50 人份 NeuralGCM 集成模型，并在关键指标上与 Pangu-Weather 和 Stormer 比较持平或更优，同时可在标准硬件上实现低成本推理。

ABSTRACT

One of the guiding principles for designing AI-based weather forecasting systems is to embed physical constraints as inductive priors in the neural network architecture. A popular prior is locality, where the atmospheric data is processed with local neural interactions, like 3D convolutions or 3D local attention windows as in Pangu-Weather. On the other hand, some works have shown great success in weather forecasting without this locality principle, at the cost of a much higher parameter count. In this paper, we show that the 3D local processing in Pangu-Weather is computationally sub-optimal. We design ArchesWeather, a transformer model that combines 2D attention with a column-wise attention-based feature interaction module, and demonstrate that this design improves forecasting skill. ArchesWeather is trained at 1.5° resolution and 24h lead time, with a training budget of a few GPU-days and a lower inference cost than competing methods. An ensemble of four of our models shows better RMSE scores than the IFS HRES and is competitive with the 1.4° 50-members NeuralGCM ensemble for one to three days ahead forecasting. Our code and models are publicly available at https://github.com/gcouairon/ArchesWeather.

研究动机与目标

• 设计一种计算效率更高的 AI 天气预报模型，在 1.5° 分辨率下保持高精度。
• 通过证明 3D 局部注意力在天气建模中存在计算上的次优性，挑战其为最优的假设。
• 通过一种新型跨层级注意力（CLA）机制实现不同气压层之间的全局特征交互，从而提升预报能力。
• 在不增加模型规模或训练成本的前提下，降低训练成本和推理时间，同时匹配或超越更大、更昂贵的模型性能。
• 研究 ERA5 数据中分布偏移的影响，并通过在近期数据上微调来提升性能。

提出的方法

• 该模型采用 3D Swin U-Net 变压器架构，并引入地球特异性位置偏差，受 Pangu-Weather 启发，但将 3D 局部注意力替换为非局部跨层级注意力（CLA）层，以实现层间特征交互。
• CLA 层实现了跨气压层的全局注意力，克服了 3D 局部注意力因仅限于相邻层间交互而造成的计算低效问题。
• 模型训练目标为预测差值 $X_{t+ au} - X_t$ 而非 $X_{t+ au}$，以提升训练稳定性和泛化能力。
• 风速分量（U、V、W）作为输入特征，且模型通过条件输入日期和月份以改善时间泛化能力。
• 在 2007–2018 年 ERA5 数据上进行微调阶段，以缓解分布偏移问题，并提升对近期大气条件的性能。
• 通过平均多个独立训练模型（如 ArchesWeather-M4、ArchesWeather-L2）的输出生成集合预测，以提升准确性。

实验结果

研究问题

• RQ1 尽管具有物理合理性，Pangu-Weather 等模型中的 3D 局部注意力在 AI 天气预报中是否计算上次优？
• RQ2 跨气压层的非局部注意力机制是否能在不增加模型规模或训练成本的前提下提升预报精度？
• RQ3 在近期 ERA5 数据（2007–2018）上微调是否能减少由历史数据分布偏移引起的误差？
• RQ4 能否仅用数个 GPU 天的训练预算，在 1.5° 分辨率下训练出具备竞争力的模型？
• RQ5 ArchesWeather 在 24 小时预报时延下的性能与 NeuralGCM、Pangu-Weather 和 Stormer 等最先进模型相比如何？

主要发现

• ArchesWeather-M 在 2 块 A100 GPU 上训练约 2 天，相对于 IFS HRES 的平均相对均方根误差改进（RRH）达到 -5.0%，优于 1.4° 的 50 人份 NeuralGCM 集成模型。
• ArchesWeather-L2 集成模型（两个 L 模型）在关键高层大气变量（Z500、T850、Q700、U850、V850）上的 RMSE 低于 1.4° 的 50 人份 NeuralGCM 集成模型。
• 与无 CLA 的基线模型相比，CLA 模块将 Z500 RMSE 降低了 10.5%，使其更接近 IFS HRES 的性能。
• 在 2007–2018 年 ERA5 数据上微调，相比在 1979–2018 全周期数据上训练，Z500 RMSE 减少 1.2%，T2m RMSE 减少 0.001°C。
• 在单块 A100 GPU 上，ArchesWeather-M 模型的推理时间仅为每 24 小时预报约 0.25 秒，支持低成本部署。
• 该模型以约 10 V100 天的训练预算实现具有竞争力的性能，较 SphericalCNN 低 40 倍，较 Stormer 低 256 倍，同时在精度上匹配或超越二者。