Prithvi WxC: Foundation Model for Weather and Climate
作者: Johannes Schmude, Sujit Roy, Will Trojak, Johannes Jakubik, Daniel Salles Civitarese, Shraddha Singh, Julian Kuehnert, Kumar Ankur, Aman Gupta, Christopher E Phillips, Romeo Kienzler, Daniela Szwarcman, Vishal Gaur, Rajat Shinde, Rohit Lal, Arlindo Da Silva, Jorge Luis Guevara Diaz, Anne Jones, Simon Pfreundschuh, Amy Lin, Aditi Sheshadri, Udaysankar Nair, Valentine Anantharaj, Hendrik Hamann, Campbell Watson, Manil Maskey, Tsengdar J Lee, Juan Bernabe Moreno, Rahul Ramachandran
分类: cs.LG, physics.ao-ph
发布日期: 2024-09-20
💡 一句话要点
Prithvi WxC：用于天气和气候的23亿参数开源基础模型，支持多种下游任务。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 天气预报 气候预测 基础模型 Transformer 深度学习 降尺度 极端天气事件
📋 核心要点
- 现有AI天气模型多针对单一用例，缺乏通用性，无法有效适应多种下游任务。
- Prithvi WxC通过编码器-解码器架构和混合训练目标，构建了一个可适应多种天气气候任务的基础模型。
- 实验表明，Prithvi WxC在自回归预测、降尺度等任务上表现出色，并已开源模型和微调流程。
📝 摘要（中文）
本文介绍了Prithvi WxC，一个拥有23亿参数的基础模型，它利用来自现代回顾性研究和应用分析第2版（MERRA-2）的160个变量进行训练。该模型旨在弥补天气和气候领域中专注于特定用例（特别是中期预报）与人工智能领域中通用基础模型之间的差距。Prithvi WxC采用基于编码器-解码器的架构，并融合了Transformer模型的概念，以有效捕捉输入数据中的区域和全局依赖关系。该模型设计用于处理大量tokens，从而在高分辨率下对不同拓扑中的天气现象进行建模。此外，它采用混合目标进行训练，结合了掩码重建和预测的范例。该模型在一系列具有挑战性的下游任务上进行了测试，包括自回归展开预测、降尺度、重力波通量参数化和极端事件估计。预训练模型及其微调工作流程已通过Hugging Face公开发布。
🔬 方法详解
问题定义：现有天气和气候领域的AI模型开发主要集中在特定用例上，例如中期天气预报，缺乏像AI领域那样能够有效适应多种不同任务的基础模型。这限制了模型的通用性和可复用性，需要针对每个新任务重新训练或进行大量修改。
核心思路：Prithvi WxC的核心思路是构建一个通用的天气和气候基础模型，通过在大规模数据集上进行预训练，使其能够学习到天气和气候系统的基本规律和表示。然后，可以通过微调的方式，将该模型应用于各种下游任务，例如天气预报、气候预测、降尺度等。这样可以大大减少针对每个任务单独训练模型所需的时间和资源。
技术框架：Prithvi WxC采用基于编码器-解码器的Transformer架构。编码器负责将输入的天气和气候数据编码成高维表示，解码器则负责根据该表示生成预测结果。该模型使用来自MERRA-2的160个变量作为输入，这些变量涵盖了大气、海洋和陆地等多个方面的信息。模型设计允许处理大量的tokens，以捕捉不同分辨率和拓扑结构下的天气现象。
关键创新：Prithvi WxC的关键创新在于其混合训练目标，它结合了掩码重建和预测两种范例。掩码重建是指随机掩盖输入数据的一部分，然后让模型预测被掩盖的部分。这种方法可以帮助模型学习到输入数据中的内在结构和依赖关系。预测是指让模型根据过去的天气和气候数据预测未来的情况。这种方法可以帮助模型学习到天气和气候系统的演变规律。
关键设计：Prithvi WxC拥有23亿参数，这使得它能够学习到复杂的天气和气候模式。模型使用AdamW优化器进行训练，学习率设置为一个合适的值。损失函数是掩码重建损失和预测损失的加权和。编码器和解码器都由多个Transformer层组成，每个Transformer层都包含自注意力机制和前馈神经网络。模型还使用了位置编码来表示输入数据中的空间信息。
🖼️ 关键图片
📊 实验亮点
Prithvi WxC在多个下游任务上表现出竞争力。例如，在自回归展开预测任务中，该模型能够生成较长时间范围内的天气预报。在降尺度任务中，该模型能够将低分辨率的气候数据转换为高分辨率的数据，从而更好地模拟局地天气现象。此外，该模型还能够用于重力波通量参数化和极端事件估计。
🎯 应用场景
Prithvi WxC可广泛应用于天气预报、气候预测、极端天气事件分析、农业气象等领域。它能够帮助提高天气预报的准确性和可靠性，为气候变化研究提供更精确的数据支持，并为农业生产提供更科学的指导。该模型的开源发布将促进天气和气候领域的AI研究和应用。
📄 摘要（原文）
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.