微软团队发布第一个基于AI的天气和气候基础模型 ClimaX

大多数最先进的天气和气候建模方法都是基于大气的物理学数值模型。这些方法旨在模拟非线性动力学和多个变量之间的复杂相互作用，这些变量很难近似。此外，许多此类数值模型的计算量很大，尤其是在以细粒度的空间和时间分辨率对大气现象进行建模时。

近期基于机器学习的数据驱动方法，旨在通过使用深度神经网络学习数据驱动的函数映射，来直接解决下游预测或投影任务。然而，这些网络是使用针对特定时空任务的精选和同质气候数据集进行训练的，因此缺乏数值模型的通用性。

微软自主系统与机器人研究小组以及微软研究院科学智能中心开发了 ClimaX，这是一种灵活且可推广的天气和气候科学深度学习模型，可以使用跨越不同变量、时空覆盖和物理基础的异构数据集进行训练。

ClimaX 使用新颖的编码和聚合块扩展了 Transformer 架构，这些块允许有效使用可用计算，同时保持通用性。ClimaX 在源自 CMIP6 的气候数据集上使用自我监督学习目标进行了预训练。然后可以对预训练的 ClimaX 进行微调，以解决广泛的气候和天气任务，包括那些涉及预训练期间看不到的大气变量和时空尺度的任务。

研究表明，与现有的数据驱动基线相比，ClimaX 中的这种普遍性导致天气预报和气候预测基准的卓越性能，即使在较低分辨率和计算预算下进行预训练也是如此。

该研究以「ClimaX: A foundation model for weather and climate」为题，于 2023 年 1月 24 日发布在 arXiv 预印平台。

当前一个关键的见解是，认识到天气和气候科学中的所有预测和建模任务都基于物理现象及其与当地和全球地理的相互作用。因此，应当设计一个以许多不同尺度模拟大量天气和气候变量的基础模型，从而对这些物理定律和相关的地理相互作用进行编码。

当前最先进的数值天气和气候模型基于对大型微分方程系统的模拟，这些微分方程基于不同地球系统的已知物理学将能量和物质的流动联系起来。因此，通常需要最先进的数值天气和气候模型才能在高分辨率的大型超级计算机上运行。尽管这些方法很成功，但众所周知，这些模型在长期和短期内都存在弱点和局限性。

另一方面，技术的进步带来了来自卫星、雷达和其他气象传感器的大量数据。这些数据还可以为天气和气候建模提供有价值的信息，特别是在更精细的时间和空间分辨率下，同时可能解释不太了解的复杂物理学。然而，当前的大规模数值天气和气候模型很难吸收这种规模的数据。

机器学习（ML）模型可以提供替代权衡，从数据和计算的规模中获益。近期，为中短期天气预报扩展深度学习系统的尝试已经取得了巨大的成功，通常已经在感兴趣的关键变量上匹配了当前最先进的数值天气模型。然而，由于大多数 ML 模型都是针对特定数据集的特定预测任务进行训练的，它们缺乏地球系统科学的通用实用性，因此不完全基于物理学。

从机器学习的角度来看，大量可用数据——从陆地、沙子或大气的直接天气测量，数十年来在不同空间尺度上重新分析的天气数据，到各种场景的物理信息气候预测——是为天气和气候建模建立完全基于物理的基础模型的卓有成效的基础。尤其如此，因为天气和气候数据通常共享同一组方程（尽管具有相当不同的特征）。

ClimaX 架构和框架

在自然语言处理或计算机视觉等学科中，众所周知，经过训练以使用监督学习解决单个任务的 ML 模型在训练期间需要标签，并且在训练分布之外部署时很脆弱。近年来，预训练大型无监督「基础」模型因此成为一种新范式，缓解了监督瓶颈。预训练后，有很多方法可以在几乎没有或没有（即零样本）额外监督的情况下，在任意跨度的任务上微调同一模型。

ClimaX 遵循预训练-微调范式。对于预训练 ClimaX，研究人员的第一个关键建议是超越标准的均质天气数据集，而是利用基于物理的气候模拟数据集，由于来自多个小组的各种气候模拟，这些数据集非常丰富。通过仅使用可用数据集的一小部分，研究表明这些数据集中的异质性已经足以作为一个丰富而丰富的预训练数据集。

但要做到这一点，研究人员需要一个模型架构能够恰当地包含那些高度多模态的气候数据集的异质性，因为观测通常对应于许多不同的、无限的变量。此外，许多观测数据集是不规则的，因为它们的时空覆盖范围不同，对应于大气变量的不同子集。

ClimaX 的核心是一个基于 Vision Transformers（ViT）的多维图像到图像转换架构。基于 ViT 的架构特别适合对天气和气候现象进行建模，因为它们自然地标记了类似于不同时空输入的多尺度数据的空间性质，并且还提供了将标记化扩展到广泛的多通道特征的机会。然而，要将 ViT 架构重新用于 ClimaX，需要进行两个根本性的改变：变量标记化和变量聚合。

变量标记化：图像数据的标准 ViT 标记化方案将输入分成大小相等的块，并将这些块在宽度、高度和通道维度上展平为一个向量。然而，这对于气候和天气数据来说并不是那么简单，因为不同数据集之间的物理变量数量可能会有所不同。具体来说，在该研究的例子中，每个气候预训练数据子集都包含不同模型的模拟数据，因此具有不同的基础变量。因此，研究人员提出变量标记化，将变量视为单独的模式，即使在不规则数据集的情况下也能实现更灵活的训练。

图示：变量标记化，将变量视为单独的模式，以实现更灵活的训练。（来源：论文）

变量聚合：变量标记化有两个固有的问题。首先，它产生的序列随输入变量的数量线性增加，这在计算上作为 ViT 的自注意层的输入是不可行的。其次，输入很容易包含具有不同物理基础的不同变量的标记。因此，研究人员提出变量聚合，这是一种交叉注意操作，可为每个空间位置输出大小相等的嵌入向量。

图示：变量聚合，一种交叉注意操作，为每个空间位置输出大小相同的嵌入向量。（来源：论文）

对各种下游任务进行微调

论文重点介绍了 ClimaX 在各种与天气和气候相关的下游任务中的表现，研究人员将这些任务分为天气预报（全球、区域、次季节和季节）、气候预测和气候降尺度。ClimaX 由于其四个可学习的组件而非常灵活：令牌嵌入层、变量聚合模块、注意力块和预测头。如果下游变量与预训练变量重叠，还可以微调整个模型。如果在预训练期间看不到变量，研究人员会用新初始化的网络替换嵌入层和预测头，并对其他两个组件进行微调或冻结。

图示：用于气候预测任务的示例微调管道。一组不同的输入和输出变量需要不同的嵌入层和预测头。注意层可以被冻结或微调。（来源：论文）

亮点

全球天气预报

图示：使用 ClimaX（未来 6 小时至 1 个月）的关键天气变量（温度：T2m、T850、风：U10、V10）的预测结果可视化。（来源：论文）

ClimaX 在对相同的 ERA5 数据进行微调时，即使在中等分辨率 (1.40625˚) 上，即使在短期和中期预测方面不比 IFS 表现得更好，但在长期预测方面的表现也相当出色。

图示：ClimaX 对关键天气变量（温度：T2m、T850、风：U10、地势：Z500）的全球预报性能与不同前置时间范围内使用的最先进的数值天气预报系统 IFS 的比较。ClimaX 已经接近中短期预测，同时在更长的提前期范围内变得更好。（来源：论文）

气候预测

气候预测帮助气候科学家了解温室气体浓度或气溶胶排放等各种强迫因素对长期气候状况的影响。最近引入了 ClimateBench，以持续评估机器学习方法，以提高气候预测的准确性。该任务明显不同于预训练机制，其输入和输出与预训练期间完全不同。尽管如此，将 ClimaX 注意力层转移到此任务，仍会产生与 ClimateBench 中当前最先进的基线相当或更好的性能。

气候模型降尺度

由于空间分辨率较粗糙，气候模型通常无法提供足够的细节来分析区域和当地现象。降尺度有助于提供更高分辨率的气候预测，并通过将这些模型与更高分辨率的当地气候条件相关联来减少这些模型输出的偏差。研究人员通过使用较低分辨率气候模型的预测作为输入，并将再分析天气数据中的相应值作为较高分辨率的目标来评估 ClimaX 在此任务上的表现。研究人员发现 ClimaX 在所有关键指标上再次优于其他基于深度学习的基线。

表：ClimaX 在从 MPI-ESM (5.625˚) 到 ERA5 (1.40625˚) 的降尺度方面比其他深度学习基线表现更好。（来源：论文）

图示：使用 ClimaX 对关键气候变量（温度：T2m、T850）进行降尺度预测的可视化。（来源：论文）

缩放分析

当给定更多计算、数据或参数时，基于 Transformer 的机器学习架构已经发现了有利且可预测的缩放特性。ClimaX 模型也是如此。研究人员发现这些趋势很有希望，因为与其他领域目前流行的具有数十亿参数的架构相比，这里只扩展到相当小的模型。此外，还有大量公开可用的天气和气候数据，尚未利用这些数据来预训练更大的模型。

图示：ClimaX 的比例规律分析。更大的模型和更多的数据不断提高关键任务的性能，更大的模型也具有更高的样本效率。（来源：论文）

使用数据驱动方法推进天气和气候建模

该团队发布 ClimaX，旨在进一步推动数据驱动的天气和气候建模。该团队的目标是让任何人都能轻松使用最新的机器学习方法来解决大量问题，从局部范围内的近期预测到涉及天气和气候变量的长期过程建模。ClimaX 朝着为各种此类任务提供单一起点的想法迈出了一大步。研究人员迫不及待地想看看这个新兴领域的未来。

论文链接：https://arxiv.org/abs/2301.10343

相关报道：https://microsoft.com/en-us/research/group/autonomous-systems-group-robotics/articles/introducing-climax-the-first-foundation-model-for-weather-and-climate/