[论文解读] ClimaX: A foundation model for weather and climate

Tung Nguyen, J. Brandstetter+3 more|arXiv (Cornell University)|Jan 24, 2023

Meteorological Phenomena and Simulations被引用 168

一句话总结

ClimaX 是一个灵活的基于 transformer 的基础模型，在异质的 CMIP6 气候数据上进行预训练，从而能够针对多样的天气和气候任务进行微调，具备更好的泛化能力和具有竞争力的基准。

ABSTRACT

Most state-of-the-art approaches for weather and climate modeling are based on physics-informed numerical models of the atmosphere. These approaches aim to model the non-linear dynamics and complex interactions between multiple variables, which are challenging to approximate. Additionally, many such numerical models are computationally intensive, especially when modeling the atmospheric phenomenon at a fine-grained spatial and temporal resolution. Recent data-driven approaches based on machine learning instead aim to directly solve a downstream forecasting or projection task by learning a data-driven functional mapping using deep neural networks. However, these networks are trained using curated and homogeneous climate datasets for specific spatiotemporal tasks, and thus lack the generality of numerical models. We develop and demonstrate ClimaX, a flexible and generalizable deep learning model for weather and climate science that can be trained using heterogeneous datasets spanning different variables, spatio-temporal coverage, and physical groundings. ClimaX extends the Transformer architecture with novel encoding and aggregation blocks that allow effective use of available compute while maintaining general utility. ClimaX is pre-trained with a self-supervised learning objective on climate datasets derived from CMIP6. The pre-trained ClimaX can then be fine-tuned to address a breadth of climate and weather tasks, including those that involve atmospheric variables and spatio-temporal scales unseen during pretraining. Compared to existing data-driven baselines, we show that this generality in ClimaX results in superior performance on benchmarks for weather forecasting and climate projections, even when pretrained at lower resolutions and compute budgets. The source code is available at https://github.com/microsoft/ClimaX.

研究动机与目标

• 推动创建一个面向天气与气候的通用、数据驱动的基础模型，以克服特定任务的局限性。
• 利用异质的、具有物理信息的气候数据集（CMIP6）进行预训练，以提升通用性。
• 开发使变量模态和不规则时空覆盖成为可能的架构创新。
• 展示微调在预测、气候投影和下尺度任务上的多样性适用性。
• 评估数据、模型规模和分辨率对扩展性的影响，为未来研究提供指南。

提出的方法

• 在 Vision Transformer (ViT) 上扩展两个关键模块：变量分词（对每个输入变量单独分词）和变量聚合（在每个空间位置进行跨注意力以融合变量）。
• 添加位置和领先时间嵌入，指导跨任意时间范围的时序预测。
• 在 CMIP6 数据上对 ClimaX 进行随机化预测目标的预训练，以在 lead times 6h–168h 内从当前状态预测未来状态。
• 使用纬度加权均方误差来考虑全球网格在球面面积上的差异。
• 对 ClimaX 进行下游任务的微调，包括在预训练中看到的变量以及未见过的变量，采用模块化微调策略。
• 在预测、气候投影和下尺度基准上进行评估，并与神经基线和数值天气预报（NWP）基线进行比较。

实验结果

研究问题

• RQ1 一个经过预训练的 ClimaX 模型能否适应具有异质数据的多种天气和气候任务？
• RQ2 与专门基线相比，ClimaX 在全球/区域预测、亚季-至季节性预测、气候投影和下尺度方面的表现如何？
• RQ3 数据集大小、模型容量和分辨率对 ClimaX 性能的影响（扩展定律）是什么？
• RQ4 哪些微调策略最能将 ClimaX 转移到未见变量和变化的时空分辨率？

主要发现

• 单一的预训练 ClimaX 能针对不同的领先时间、分辨率和区域，对广泛的任务进行微调。
• ClimaX 在 ClimateBench 上达到最先进的性能，并且在 WeatherBench 上与运营的 IFS 相竞争，即使计算资源中等。
• 扩展分析表明，更多的预训练数据、更大的模型和更高分辨率的输入都带来性能提升。
• 消融研究支持有效的微调策略，包括处理未见变量和具有共享组件的多任务。
• 由于变量分词和聚合，架构对不规则数据覆盖和变量集合保持鲁棒性。