[论文解读] Prithvi WxC: Foundation Model for Weather and Climate

Johannes Schmude, Sujit Roy+27 more|arXiv (Cornell University)|Sep 20, 2024

Geophysics and Gravity Measurements被引用 7

一句话总结

Prithvi WxC 是一个 2.3B-parameter transformer-based foundation model，在 160 MERRA-2 大气变量上训练，设计用于预测、下尺度化、重力波参数化和极端事件估计，在 Hugging Face 开源发布。

ABSTRACT

Triggered by the realization that AI emulators can rival the performance of traditional numerical weather prediction models running on HPC systems, there is now an increasing number of large AI models that address use cases such as forecasting, downscaling, or nowcasting. While the parallel developments in the AI literature focus on foundation models -- models that can be effectively tuned to address multiple, different use cases -- the developments on the weather and climate side largely focus on single-use cases with particular emphasis on mid-range forecasting. We close this gap by introducing Prithvi WxC, a 2.3 billion parameter foundation model developed using 160 variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2). Prithvi WxC employs an encoder-decoder-based architecture, incorporating concepts from various recent transformer models to effectively capture both regional and global dependencies in the input data. The model has been designed to accommodate large token counts to model weather phenomena in different topologies at fine resolutions. Furthermore, it is trained with a mixed objective that combines the paradigms of masked reconstruction with forecasting. We test the model on a set of challenging downstream tasks namely: Autoregressive rollout forecasting, Downscaling, Gravity wave flux parameterization, and Extreme events estimation. The pretrained model with 2.3 billion parameters, along with the associated fine-tuning workflows, has been publicly released as an open-source contribution via Hugging Face.

研究动机与目标

• 在全球与区域任务中凸显基础模型方法在天气与气候中的需求与必要性。
• 开发一个可扩展的编码器-解码器变换器架构，能够处理大量的标记和多样的拓扑结构。
• 采用混合目标进行训练，将屏蔽重建与基于气候学的预测结合起来。
• 评估零样本的重建和预测性能，以及下游任务如下尺度化和重力波参数化。
• 公开发布预训练权重和微调工作流，以便在社区中广泛使用。

提出的方法

• 采用局部和全局注意力的编码器-解码器变换器，在窗口内实现长程依赖。
• 2D 视觉变换器，输入掩蔽率为 50%，并使用灵活的标记化方案（窗口和每个窗口的标记）。
• 预训练目标将屏蔽重建与基于气候学的预测相结合，包括归一化残差形式（等式 2）。
• 在 160 MERRA-2 变量（1980-2019）上进行训练，包含静态输入、气候项 C_t，以及前导时间参数 δt 和 δτ。
• 两阶段预训练：初始阶段采用掩蔽和随机前导时间，随后进行针对预测的微调，使用调整后的损失和自回归回放。
• 用于下游的下尺度化和重力波参数化的微调工作流。

实验结果

研究问题

• RQ1 一个单一的基础模型是否能够以零样本能力同时处理全球与区域的天气-气候任务？
• RQ2 哪一个预训练目标最能同时支持重建和随变量前导时间变化的预测？
• RQ3 在自回归预测、下尺度化和飓风路径预测方面，大型 WxC 模型相较于专门任务模型的表现如何？
• RQ4 将气候学和静态特征整合到统一的预训练目标中的有效策略是什么？
• RQ5 对下游任务（下尺度化、重力波参数化、极端事件）的实际微调考虑因素是什么？

主要发现

• 在局部或全球掩蔽下，使用高掩蔽（剩余数据低至 5-25%）时，零样本重建可行。
• 零样本预测在短期前导时间（6-12 小时）表现最强，合理地延伸到 ~66 小时，相较于某些基线在此之后表现下降。
• 在飓风路径预测中，Prithvi WxC 相较于在 MERRA-2 或 ERA5 数据上训练的 FourCastNet 变体，具有更优的轨迹精度和登陆时刻（示例： Ida，2021；平均轨迹误差 63.9 km vs 201.9–262.3 km）。
• 下尺度化实验显著优于插值基线（MERRA-2 T2m：空间 RMSE 3.22→0.73；时间 RMSE 2.46→0.64；相关性 0.89→0.98）。
• CORDEX 下尺度化到 12.5 km 分辨率也取得类似改进（空间 RMSE 1.89→0.44；时间 RMSE 1.14→0.37；相关性 0.99→1.00）。
• 该模型已在 Hugging Face 上公开发布，包含用于下游任务的预训练权重和微调工作流。