翻译：ClimaX: A foundation model for weather and climate

本文是论文ClimaX: A foundation model for weather and climate的翻译。

Abstract
1. Introduction
2. 背景与相关工作
3. Approach
3.3. 数据集
- 3.3.1. 预训练
- 3.3.2. 微调和评估
3.4 Training
- 3.4.1 Pretraining
- 3.4.2 Finetuning
4 实验
5 讨论与未来工作

Abstract

大多数最先进的天气和气候建模方法基于大气物理学的数值模型。这些方法旨在模拟多个变量之间的非线性动态和复杂相互作用，这在近似上是具有挑战性的。此外，许多这样的数值模型在对大气现象进行细粒度空间和时间分辨率建模时需要大量计算资源。最近基于机器学习的数据驱动方法则旨在通过使用深度神经网络学习数据驱动的功能映射，直接解决下游的预测或投影任务。然而，这些网络是使用精心策划和同质化的气候数据集进行训练的，用于特定的时空任务，因此缺乏数值模型的普适性。我们开发并展示了 ClimaX，这是一个灵活且具有通用性的深度学习模型，用于天气和气候科学，可以使用涵盖不同变量、时空覆盖范围和物理基础的异构数据集进行训练。ClimaX扩展了Transformer架构，引入了新的编码和聚合模块，可以有效利用可用的计算资源同时保持通用性。ClimaX在基于自监督学习目标的CMIP6气候数据集上进行预训练。然后，可以对预训练的ClimaX进行微调，以解决广泛的气候和天气任务，包括那些在预训练过程中未见过的大气变量和时空尺度。与现有的数据驱动基线相比，我们展示了ClimaX的这种通用性在天气预测和气候投影基准测试中取得了优越的性能，即使在较低分辨率和计算预算下进行预训练也是如此。源代码可在https://github.com/microsoft/ClimaX 上获取。

图1：ClimaX被构建为任何天气和气候建模任务的基础模型。在天气方面，这些任务包括各种领先时间跨度和分辨率下的标准预测任务，无论是全球还是区域性的。在气候方面，进行长期预测和从较低分辨率模型输出中获取降尺度结果是标准任务。

1. Introduction

天气和气候建模是科学和社会面临的普遍挑战。随着对极端天气事件和气候变化日益关注，改进天气预报以减灾和气候预测以制定长期政策和适应努力的需求日益增长。目前，全球天气和气候建模的数值方法通过各种一般环流模型（GCM）参数化。GCMs代表大气、陆地和海洋中能量和物质流动相关的微分方程系统，可以随时间积分以获取相关大气变量的预报。尽管在实践中非常有用，但GCMs也面临许多挑战，如准确表示细分辨率下的物理过程和初始条件，以及大规模数据同化(data assimilation)和计算模拟中的技术挑战。这些因素限制了它们在许多情况下的使用，特别是在非常短的时间尺度（例如几小时）内快速模拟大气变量，或者在长时间尺度（例如超过5-7天）内准确模拟（Zha+19）。

相比之下，针对大气变量预测的数据驱动方法不断增长，特别是在气象应用中。这里的关键思想是使用几十年的历史全球数据集（如ERA-5再分析数据集）来训练深度神经网络来预测目标大气变量。与GCMs不同，这些网络并没有明确基于物理，并且缺乏对地球系统科学的通用用途，因为它们是为特定的预测建模任务而训练的。然而，随着计算和数据集的增长，出现了越来越多的证据表明这些模型可以在许多情况下达到与最先进的数值模型竞争的准确性，例如现在降水预测(nowcasting of precipitation)和中期温度、风和湿度等变量的预报。尽管这些趋势令人鼓舞，但仍然存在对这些数据驱动方法在不同现实场景中的普遍性的担忧，例如对极端天气事件和长期气候预测的预报，特别是在有限的时空监督和计算预算下（Zha+19）。

上述挑战的变体广泛适用于机器学习（ML）各个学科。在自然语言处理和计算机视觉等学科中，人们普遍承认，使用监督学习训练的ML模型解决单一任务时在训练过程中需要大量标签，并且在其训练分布之外部署时容易出现脆弱性。最近的研究表明，通过预训练大型无监督“基础”模型（Bom+21）（如从互联网上抓取的文本和图像等巨大被动数据集）可以缓解监督瓶颈。预训练后，有许多方法可以在几乎没有或零额外监督的情况下对相同的模型进行微调。除了低目标监督外，这些模型还能更好地适应训练分布之外的变化（Hen+20a），提高了它们的可靠性。

受到以上成功的启发，本文研究了一个问题：我们如何设计和训练一个用于地球大气通用任务的基础模型？我们提出了ClimaX，一个用于天气和气候的基础模型。对于预训练任何基础模型，关键的方法是使用无监督的目标在大型数据集上训练深度结构。例如，许多语言和视觉基础模型使用生成建模在互联网规模数据集上训练大型Transformer。虽然在概念上简单，但这种扩展方法对于天气和气候领域来说充满挑战，我们将在下文中讨论并提出解决方案。

首先，目前尚不清楚什么构成了预训练ClimaX的互联网规模被动数据集。任何特定时间的历史天气和气候数据集的大小是固定的，并且每天以几乎恒定的速率增加，因为它对应于自然现象的处理传感器测量。我们的第一个关键建议是超越这些数据集，明确利用基于物理的气候模拟模型。今天许多这样的模型正在使用，例如CMIP6集合，其中包含来自49个不同气候建模组的约100个不同气候模型的运行。我们展示了这些模拟数据集的异构性作为预训练ClimaX的丰富和丰富的数据源。

第二，我们需要一个模型架构，能够恰当地接纳上述气候数据集的异构性。气候数据是高度多模态的，因为观测通常对应于许多不同的、无边界的(unbounded)变量，具有不同的数据类型（例如，压力、温度、湿度）。此外，许多观测数据集在时空覆盖上是不规则的，并且可能对应于不同的大气变量子集。我们通过重新利用视觉Transformer来解决ClimaX中的上述挑战。与早期工作不同，那里的输入数据被表示为一个图像，其中不同的大气变量被视为其通道（Pat+22; Bi+22），我们将它们视为不同的模态，以便即使在不规则数据集的情况下也能进行更灵活的训练。这带来了序列长度大幅增加的副作用，我们建议在自注意层之前使用交叉注意力样式的通道聚合方案来解决这个问题。

第三，我们需要一个预训练目标，可以学习大气变量之间的复杂关系，并允许有效地对下游任务进行微调。鉴于气候数据的时空特性，我们提出了一个用于预训练ClimaX的随机化预测目标。在这里，模型的目标是预测未来任意时间的任意一组输入变量。虽然简单直观，但我们展示了这样的预训练目标有助于对新任务和时间尺度进行微调，甚至超出预训练窗口，如次季节至季节的累积预测、气候预测和气候模型的降尺度。我们的实证研究表明，一个单一的预训练模型可以在许多任务（例如，多尺度天气预测、气候预测、降尺度）下进行微调，涉及不同的时空分辨率、地理区域和目标预测变量，包括训练过程中未见过的变量。值得注意的是，我们的基准结果在ClimateBench上是最先进的，与WeatherBench上的操作性综合预报系统（IFS）竞争性（即使我们的模型是在中等分辨率上使用最多80个NVIDIA V100 GPU进行训练）。

最后，我们展示了ClimaX的有希望的扩展定律，这些定律与更大数量的预训练数据集、更大的模型以及扩展到更高分辨率格网数据集的性能改进自然轴相一致。尤其是最后一个与最近和同时进行的数据驱动天气预报工作一致（Pat+22; Bi+22; Lam+22）。据我们所知，ClimaX是首个可以使用预训练期间异构气候数据有效扩展并在微调期间泛化到不同下游任务的数据驱动模型，为地球系统科学的新一代数据驱动模型铺平了道路。

2. 背景与相关工作

当今使用的当前天气和气候模型广泛依赖数值方法和计算模拟来预测和理解地球的天气和气候系统。这些任务包括各种数值天气预报（NWP）系统，它们使用计算机模拟来对天气条件进行短期预测，以及气候模型，它们使用类似的技术来模拟和预测地球气候的长期变化。尤其值得注意的是，在天气和气候模型的核心都有相同的原始方程组。

对于气候建模，地球系统模型（ESM）[Hur+13]，或“耦合模型”，将控制大气、冰冷层(cryosphere)、陆地和海洋过程的模拟耦合在一起被认为是最先进的技术。主要这些模拟基于广义环流模型（GCMs）[Sat04; Lyn08; Ado14; MD+21]，这些模型可以追溯到Phillips [Phi56]和Lorenz [Lor67]的工作，它们在旋转球上解决Navier-Stokes方程以模拟流体环流。这些模型经常被用来进行各种因素敏感性(factor sensitivity)研究，以检查某些强迫(forcing)因子的变化（例如温室气体浓度）如何影响全球或区域气候，并帮助进行气候预测(climate projections)以帮助理解未来的情况。

数值天气预报（NWP）模型与GCMs共享许多组件，尤其是大气组件[BTB15; Lyn08; Kal03]。然而，结合数据同化[LSZ15; Gro22]，这涉及将观测数据和各种大气和海洋测量数据与这些数值模型结合起来对于准确的预测和模拟是很重要的。天气和气候模型之间的另一个显著区别在于对基础方程的解的框架：天气模型为初值问题，而气候模型为边界值问题[BTB15]。这些解决方法的不同难度水平导致了气候模型往往是全球性的，通常在更粗糙的时空分辨率下，而天气模型可以从全球到局部和区域模型的非常高时空分辨率[War10]。

尽管它们取得了显著的成功，包括最近2021年的诺贝尔物理学奖[RRH22]，但围绕广义环流模型（GCMs）的限制存在着相当多的争论，特别是在模型之间的结构性错误以及当前的GCMs是设计用来复制观察到的气候[Bal+22]。气候科学界已经意识到这些挑战，导致了耦合模型比较项目（Coupled Model Intercomparison Project, CMIP）的创建，作为评估和比较不同气候模型性能的标准化协议[Mee+00]。正如我们将在以下部分看到的，CMIP不仅在推进我们对气候变化及其潜在影响的理解方面发挥着关键作用，其评估程序还产生了大量数据，使现代基于深度学习的方法对许多任务非常有吸引力。值得注意的是，将这些知识编码到一个具有更快推断和数据同化能力的“基础”机器学习模型中可以为更广泛的影响铺平道路。

2.1 数据来源

与计算机视觉或自然语言处理中的数据不同，天气和气候数据不仅基于传感器数据，而且还融合了来自多种来源的信息。例如，再分析天气数据通过数据同化[BTB15]将气象观测与过去的短期天气预报相结合。数据本身的测量是高度异质的，代表了各种具有不同数据类型的物理变量（例如压力、温度、湿度），这些变量以不同的相对稀疏的空间位置和不同的时间频率记录。这些测量可以与已知的物理学相结合，设计气候模拟，再次产生具有不同变量和不同尺度的数据。从机器学习的角度来看，可用数据的种类繁多，包括直接在陆地、海洋或大气中进行的天气测量数据，覆盖几十年的不同空间尺度的再分析天气数据，以及用于各种情景的物理学气候预测。值得注意的是，这些数据共享相同的基本方程组，但具有相当不同的特征。下面我们描述了天气和气候建模中最常用的两种数据来源。

2.1.1 CMIP6

耦合模型比较项目（CMIP）[Mee+00]是一个跨不同气候建模团体的国际努力，旨在比较和评估其全球气候模型。虽然CMIP的主要目标是提高对地球气候系统的理解并改善其模拟的准确性，但最近从他们的实验运行中获得的数据很容易在CMIP6[Eyr+16]存档中获得。在CMIP6中，“6”指的是该项目的最新阶段，有49个团体参与其中，他们的实验涵盖了温度、降水量、海平面等多种气候变量，来自数百个模型。这导致了从1850年以来各种气候情景的全球预测，所有这些都遵循相似的控制方程，但具有不同的强迫因子，例如影响气候的温室气体排放。

2.1.2 ERA5

欧洲中期天气预报中心（ECMWF）的ERA5再分析存档[Her+18; Her+20]是学习和基准天气预报系统的主要数据来源。一旦完成，ERA5再分析将包含从1950年开始的全球大气、陆地表面和海洋波浪的详细记录。目前可用的ERA5再分析数据将ECMWF的集成预测系统（IFS）[Wed+15]与可用观测相结合，以提供大气、海洋波浪和陆地表面量的最佳猜测。在原始形式下，可用的再分析数据是巨大的：1979年至2018年的40年数据，以地球球体的0.25°×0.25°全球纬度经度网格，以小时为间隔，包含37个不同高度水平的不同气候变量加上地表。该网格总共包含721×1440个纬度和经度的网格点。高度水平以压力水平表示。

2.2. Tasks

鉴于数据可用性的规模，在当前的数值方法需要越来越多的计算资源的情况下，尽管将真实观测数据整合到其中是困难的，机器学习在与天气和气候建模相关的许多任务中越来越多地找到了应用。在天气方面，感兴趣的主要任务是预测关键天气变量的未来值。这些任务根据感兴趣的时间和空间范围可以采取以下形式：

全球预测任务范围从几小时（即现在预测）到领先几天和几周（即短期和中期预测）。通常这些任务是在ERA5再分析数据集上进行评估的（见第2.1.2节），而欧洲中期天气预报中心（ECMWF）的运营IFS [Wed+15]是当前的全球气象预报基线。
区域预测任务可以涵盖从大陆北美或欧洲的天气预测到个别州、县或城市的范围。
季节内至季节间预测（S2S）[VR18; Vit+22]是指在2周至2个月之间的领先时间内预测天气的任务。S2S弥合了天气预测和季节性气候预测之间的差距，对于灾害减轻至关重要。通常在这么长的时间范围内，预测关键天气变量的瞬时值可能是一项困难的任务，因此重点通常放在某个时间范围内关键天气变量的平均值上，例如每周平均降水量。

虽然对于区域或S2S任务的深度学习方法比较少见，但最近和同时的工作大部分集中在全球预测任务上。Rasp和Thuerey [RT21]首次在气候模拟中使用预训练来实现基于数据的中期天气预测，并采用ResNet [He+16]，Weyn、Durran等人[WDC20]在立方球上使用CNN进行全球天气预测，Weyn、Durran等人[Wey+21]使用大量深度学习天气预测模型的集成来对亚季节天气进行预测，Keisler [Kei22]应用基于图神经网络的方法进行天气预测，Ravuri、Lenc等人[Rav+21]利用雷达的深度生成模型进行降水现在预测，Arcomano、Szunyogh等人[Arc+20]构建基于水库计算的低分辨率全球预测模型，而MetNet [Søn+20]则以雷达和卫星数据作为输入来预测概率性降水图。这些方法由用于流体动力学的一般机器学习模型补充[Li+20; Koc+21; Lu+21; Bra+22; BWW22]。最近的最先进的神经天气模型，如FourCastNet [Pat+22]、Pangu-weather [Bi+22]或GraphCast [Lam+22]，也执行全球预测任务，使用最高分辨率的0.25° ERA5数据，并在各自的硬件资源上进行了优化。

另一方面，气候任务需要处理更长时间范围。机器学习可以帮助的可能任务类别包括气候预测和气候模型降尺度化：

气候预测是生成在不同未来社会经济情景下气候变化估计的任务。通常，这意味着弄清楚气候系统对不同强迫因子（如温室气体和气溶胶排放）的反应。气候预测是理解和准备应对气候变化潜在影响的关键任务。虽然机器学习在这个领域的应用仍处于早期阶段，但最近已经有努力对这一领域的评估进行了标准化。其中一个例子是ClimateBench [WP+22]，它是一个基于CMIP6的基准数据集，旨在为旨在提高气候预测准确性的机器学习模型提供评估框架。该基准旨在为应用于气候预测的各种机器学习模型提供一种一致和可靠的评估方法。
机器学习思想在气候模型降尺度化方面的应用更加流行。全球气候模型通常具有粗糙的空间分辨率，这意味着它们只能在局部或区域尺度提供气候条件的粗略估计。此外，这些模拟通常反映出与观测数据趋势不符的系统偏差。气候模型降尺度化的目标是通过将全球气候预测与观测到的局部气候条件相关联，从而创建具有局部准确性的气候信息。这个过程提高了数据的空间和时间分辨率，使其更适合在局部和区域分析中使用。降尺度化方法可以分为动态方法和统计方法。动态方法将全球气候模型的输出与区域气候模型的输出相关联，而统计方法则利用数据驱动的方法推断所需的转换。动态方法在物理上是一致的，但可能较慢且存在较大的偏差，而统计方法需要大量数据来学习对目标输出场景有效的表达映射。与天气预测类似，深度学习在气候科学中也已经成为一种吸引人的替代方案。最近的方法包括用于模拟气候预测的代理模型，从现有数据集或模拟中提取语境线索，并执行气候模型降尺度化。气候模型降尺度化通常输入低分辨率的再分析数据和局部地形(orographic)信息，以获取高分辨率的局部信息。许多最近的方法基于卷积架构。

2.3. Foundation models

Bommasani、Hudson等人将“基础模型”一词用于在广泛数据上通过自监督训练可扩展的深度学习模型的新兴范式，然后可以通过微调等方式调整适应各种下游任务。目前值得注意的例子包括语言领域的BERT、GPT和PaLM，以及视觉语言领域的CLIP、Florence和BEiT。除了从网络抓取的数据外，这种范式还开始在蛋白质设计等各种科学领域取得成功。这些模型的关键意义在于模型能力的涌现(emergence)以及在不同任务、领域和模态之间方法的同质化(homogenization)，这是由规模化的迁移学习原则所实现的。尽管基础模型本身应该被视为不完整，但它可以提供一个共同基础，从而可以派生出各种特定任务的模型。尽管在天气和气候科学与机器学习交叉领域的当前研究主要集中于为每个感兴趣的任务设计单独的模型，但实际上存在相当多样化的大规模数据，这些数据在这些任务之间共享着相同的物理和地质基础。最近有一些工作提出了卫星图像和遥感的预训练技术，但到目前为止，它们尚未应用于天气和气候中的多感知数据和变量。

3. Approach

鉴于大规模数据源的可用性，以及各种天气和气候任务之间的共享物理和地质基础，我们的目标是构建一个可推广的深度学习基础模型。该模型需要能够输入不同变量的异构数据集，并根据物理基础提供时空覆盖。因此，我们首先仔细研究输入表示，然后设计一个模型来处理它们的异构性 - 包括本地、全局和跨变量。

3.1. 输入表示

我们对格网预测任务感兴趣，模型接受形状为 V × H × W 的输入并预测形状为 V’ × H’ × W’ 的输出。V 是指输入变量的数量，可以是气候条件如位势和温度，或者气候迫降因子如$CO_2$和$SO_2$。H 和 W 是指输入数据的空间分辨率，这取决于我们如何将地球网格化。这种一般表示捕捉了地球系统科学中广泛的下游任务。同样，V’、H’、W’ 是指预测输出的变量和空间分辨率。我们主要使用两种空间分辨率：5.625°（32 × 64 网格点）和 1.40625°（128 × 256 网格点）。从语义上讲，一个 H × W 的地图可以表示整个地球或特定区域如北美洲。

3.2. 模型架构

我们的目标是设计一个基础模型，可以在异构数据源上进行预训练，然后微调以解决各种下游天气和气候任务。从第 3.1 节可以看出，可以将任务视为具有 V 输入通道和 V’ 输出通道的图像到图像转换问题。这使得任何图像架构都是自然的选择，如UNet、ResNet或Vision Transformers（ViT）。然而，气候和天气任务的设置要广泛得多，我们可能需要对区域性或甚至空间不完整的数据进行预测，预测未知的气候变量，或者根据预训练的数据在不同分辨率上微调模型。目前基于 CNN 的架构在这些场景中不适用，因为它们要求输入完美网格化，包含一组固定的变量，并具有固定的空间分辨率。另一方面，基于 Transformers 的架构通过将类似图像的数据视为一组token，提供了更好的灵活性。因此，我们基于 Vision Transformers（ViT）构建了 ClimaX 架构，并提出了两个主要的架构变化，即变量标记化(variable tokenization)和变量聚合(variable aggregation)，以进一步提高灵活性和通用性，接下来我们将详细描述。

3.2.1. 变量标记化

对于形状为 V × H × W 的输入，ViT 将输入标记化为一个序列，其中每个序列由 (H/p) × (W/p) = h × w 个补丁(patch)组成，每个补丁的大小为 V × $p^2$，其中 p 是补丁大小。这种标记化方案对图像数据效果很好，因为 V 始终是 RGB 通道，对于所有数据集都是相同的。然而，对于气候和天气数据，这并不适用，因为不同数据集之间的物理变量数量可能不同。例如，在 CMIP6 项目中，每个数据集包含了不同气候模型的模拟数据，因此具有不同的基础变量集。因此，我们提出了变量标记化，这是一种新颖的标记化方案，它将输入中的每个变量分别标记化。具体而言，将输入的每个变量作为形状为 H × W 的空间地图标记化为一个 h × w 补丁的序列，这将总共产生 V × h × w 个补丁。最后，将每个输入补丁的大小为 $p^2$ 的线性嵌入到一个维度为 D 的向量中，其中 D 是选择的嵌入大小。因此，变量标记化模块的输出维度为 V × h × w × D。图 3 展示了我们提出的标记化方案。

图3：变量标记化。每个变量都被独立标记化。

3.2.2. 变量聚合

虽然变量标记化使 ClimaX 能够从具有不同输入变量数量的数据集中学习，但它存在两个固有问题。首先，它导致了一个长度为 V × h × w 的序列，随着变量数量的增加呈线性增长。由于我们使用注意力来建模序列，内存复杂度随变量数量的增加呈二次方增长。这在计算上是昂贵的，因为我们的实验中最多可以有 48 个输入变量。此外，由于我们将每个变量分别标记化，输入序列将包含具有非常不同物理基础的不同变量的标记，这可能会给注意力层学习带来困难。因此，我们提出了变量聚合来解决这两个问题。对于 h × w 地图中的每个空间位置，我们执行交叉注意力操作，其中查询是一个可学习的向量，而键和值是该位置上 V 个变量的 V 个嵌入向量。交叉注意力模块为每个空间位置输出一个单一向量，从而将序列长度减少到 h × w，大大降低了计算成本。此外，序列现在包含统一语义的统一标记，使注意力层的任务变得更容易。图 4 展示了我们提出的变量聚合。

图 4：基于位置的变量聚合将长度为 V × h × w 的序列减少到 h × w。

3.2.3. Transformer

在变量聚合之后，我们需要一个序列模型来生成输出的标记。虽然原则上可以使用任何通用的序列模型，但我们建议扩展标准的 Vision Transformer（ViT）。此外，由于标准的 ViT 将图像建模视为纯粹的序列到序列问题，它可以执行一些其他变体无法完成的任务[Liu+21; Liu+22]，比如学习来自空间不完整数据的任务，其中输入不一定形成完整的网格。这在我们在第 4.2.2 节中考虑的区域预测任务中非常有用。在实验中，我们报告了使用 8 个注意力层、嵌入大小为 1024 和隐藏维度为 1024 × 4 的结果。在注意力层之后，我们使用一个预测头，它接收一个标记并输出大小为 V’× $p^2$ 的向量。预测头是一个具有 1024 隐藏维度的 2 层 MLP。我们在附录 A 中提供更多细节。

3.3. 数据集

3.3.1. 预训练

我们认为，CMIP6 的多样性和规模为预训练大规模基础模型提供了有吸引力的机会。然而，处理不同数据源中不一致的变量集合可能是一个挑战。在这项工作中，我们仅使用来自五个不同数据源（MPI-ESM、TaiESM、AWI-ESM、HAMMOZ、CMCC）的变量子集，这些数据源包含从1850年到2015年的全球气候情景预测，时间间隔为6小时，如表8所述。由于变量的原始分辨率不同，我们选择通过重新网格化将它们简化为常用的分辨率[Ras+20；RT21]，即5.625°（32×64个网格点）和1.40625°（128×256个网格点）。【Regridding 是使用 xesmf Python 包[Zhu18]进行的，使用了双线性插值。】

3.3.2. 微调和评估

我们使用附录 C.2 中描述的 ERA5 重分析数据作为用于微调和评估各种与天气相关的下游任务的数据集来源。由于其体积庞大，通常将高分辨率数据重新网格化[Ras+20；RT21]到低分辨率，如5.625°（32×64个网格点）和1.40625°（128×256个网格点），以适应可用的计算约束。我们遵循 Rasp 和 Thuerey [RT21] 的评估程序，并使用这些数据来评估我们的机器学习模型在不同前瞻时段的预测性能。关于各个数据集的更多细节请参见它们各自的实验部分。

3.4 Training

3.4.1 Pretraining

图2：ClimaX 的预训练阶段。变量使用变量分隔的标记化进行编码，然后使用变量聚合进行聚合。这些与位置嵌入和先导时间嵌入一起馈送到 ViT 骨干。

我们对 CMIP6 数据进行 ClimaX 的预训练，以预测未来的天气条件，给定当前的条件。换句话说，对于特定时间 t 的形状为 V × H × W 的天气快照 $X_t$，ClimaX 学习预测相同形状的未来天气场景 $X_{t+\Delta t}$，在领先时间 $\Delta t$。为了获得一个在各种时间预测任务中通用的预训练模型，我们在预训练期间将领先时间从 6 小时到 168 小时（即 1 周）随机化。我们将领先时间嵌入添加到token中，以通知模型预测未来的时间长度。领先时间嵌入模块是一个单层 MLP，将标量映射到嵌入大小 D 的向量。图 2 描述了 ClimaX 在预训练中的前向传递。对于输入 $X_t$，我们从 U[6, 168] 中采样领先时间 $\Delta t$，并获得相应的真实值 $X_{t+\Delta t}$。输入变量使用变量分词分别进行标记，然后在每个空间位置聚合，得到 h × w 统一的token序列。我们将token与领先时间嵌入和位置嵌入相加，然后将序列馈送到 ViT 主干。最后一个注意力层的输出馈送到预测头部，将序列转换回原始形状 V × H × W。我们采用了纬度加权均方误差[Ras+20]作为目标函数。给定预测 $\tilde{X}_{t+\Delta t}$ 和真实值 $X_{t+\Delta t}$，损失计算为：

其中 L(i) 是纬度加权因子：

其中 lat(i) 是网格的第 i 行对应的纬度。纬度加权项考虑了当我们网格化圆形地球时的非均匀性。靠近赤道的网格单元具有更大的面积，因此应该分配更多的权重。

3.4.2 Finetuning

ClimaX具有四个可学习的组件，包括token嵌入层、变量聚合模块、注意力块和预测头部。我们对ClimaX在各种下游任务中的性能进行评估，将其分类为两种微调场景：一种是下游变量属于预训练变量集合的情况，另一种是预训练期间未见过的变量。在第一种情况下，我们微调整个模型，在后一种情况下，我们将嵌入层和预测头部替换为新初始化的网络，并微调或冻结其他两个组件。我们在第4节中提供了每个下游任务的更多细节。

4 实验

我们对ClimaX进行微调，针对各种不同的下游任务评估其性能和通用性。我们将这些任务分类为预测、气候预测和气候降尺度。实验旨在回答以下问题：

ClimaX在全球预测方面与当前最先进的NWP系统相比表现如何？
我们能否微调ClimaX以对特定区域或从预训练开始的不同时间范围进行预测？
ClimaX在与预训练完全不同的气候任务上表现如何？

除了主要实验外，我们还分析了ClimaX的扩展性，即ClimaX的性能随着数据规模、模型容量和数据分辨率的增加而提高的情况。最后，我们进行了全面的消融研究，以了解微调ClimaX时计算与性能之间的权衡关系。

4.1 神经网络基线

在全球预测方面，我们将ClimaX与IFS [Wed+15]进行比较，后者是天气预测中的当前黄金标准。对于没有基线的任务，我们与UNet [RFB15; GB22]和ResNet [He+16]进行比较，后者是视觉任务中常用的两个CNN基线之一。我们从Weatherbench [Ras+20]借用ResNet架构。这些基线的具体架构细节见附录A.2。

4.2 预测

4.2.1 全球预测

给定特定时间t的全球天气条件$X_t$，我们希望预测未来时间$X_{t+\Delta t}$的天气情况，其中$\Delta t$是预测时间。输入变量包括7个垂直层次的6个大气变量、3个表面变量和3个常数场，总共48个输入变量。变量的详细信息见表9。我们评估ClimaX在预测四个目标变量方面的表现：500hPa的位势高度（Z500）、850hPa的温度（T850）、地面2米处的温度（T2m）和地面10米处的经向风速（U10）。Z500和T850是大多数中期NWP模型的两个标准验证变量，并且通常用于以前深度学习工作的基准测试，而两个表面变量T2m和U10则与人类活动相关。我们考虑七种预测时间：6小时，{1、3、5、7}天，2周和1个月，涵盖了现在预测到短期和中期预测以及更长期的范围。我们将每个预测变量在每个预测时间的预测视为一个单独的任务，并为每个任务微调一个单独的模型。我们在第4.6节讨论了替代微调协议。

我们将ClimaX与IFS和两个CNN基线在5.625°和1.40625°分辨率的ERA5数据集上进行比较。按照[Ras+20]的做法，我们将数据分为三组，其中训练数据来自1979年至2015年，验证数据在2016年，测试数据在2017年和2018年。我们通过纬度加权的MSE损失（公式（1））微调ClimaX，并使用同样的损失训练其他深度学习基线模型。对于所有深度学习模型，我们根据验证损失进行提前停止，并在测试集上评估最佳检查点。对于IFS，我们从TIGGE存档[Bou+10]中下载了2018年的预测数据。我们在以前的工作中常用的两个指标：纬度加权的均方根误差（RMSE）和纬度加权的异常相关系数（ACC）上比较所有方法。这两个指标的计算公式见附录D.1。较低的RMSE和较高的ACC表明性能更好。

图5：在5.625°分辨率的ERA5上进行全球预测的表现

图6：在1.40625°分辨率的ERA5上进行全球预测的表现。

图5和图6分别显示了ClimaX和基线在5.625°和1.40625°上的表现。在低分辨率下，IFS在6小时到5天的预测任务上表现优于ClimaX。然而，在更长的预测时间上，尤其是在14天预测上，ClimaX的表现与IFS相当或略优。在更高的分辨率下，即使在短期内，ClimaX的表现也与IFS非常接近，并且在7天及更长时间的预测上表现更好。RMSE和ACC的趋势是相似的。两个CNN基线的表现相似，但在所有任务上都落后于ClimaX和IFS。我们在附录D.2中包含了其他额外的任务特定基线[Pat+22; Bi+22; Lam+22]。这些基线是在更高分辨率的ERA5（0.25°）上训练的，因此不能直接比较。

4.2.2. 区域预测

在某些情况下，我们无法进行全球预测，特别是当我们只能访问区域数据时。在本节中，我们评估ClimaX在北美地区的相关变量的区域预测能力，在这个任务中，我们的目标是根据北美地区当前的天气状况来预测未来的天气。我们从1.40625°的ERA5数据中创建了一个新的数据集，该数据集具有相同的变量集合，但重点放在北美地区。我们称这个数据集为ERA5-NA，并在附录C.2中详细介绍了如何构建它。训练、验证和测试的分割方式与第4.2.1节类似。图7说明了ClimaX在这个任务上的微调过程，唯一的区别是输入现在只包含属于北美地区的标记。

图7：北美区域预测的微调设置。

由于以前的工作中未考虑这个任务，我们将ClimaX与两个CNN基线ResNet和UNet以及从头开始训练的ClimaX版本（我们称之为Cli-ViT）进行了比较。此外，我们对两个ClimaX模型进行了微调，其中一个是在1.40625° CMIP6上进行了预训练，另一个是在5.625°数据上进行了预训练。为了将低分辨率模型微调到高分辨率数据上，我们遵循了常见的做法，即插值位置嵌入[Dos+20; Tou+21]。我们将这个模型表示为ClimaX-pos-interp。我们使用经纬度加权的RMSE作为评估指标。

图8：关键变量在区域（北美）预测中的表现。

图8比较了ClimaX和基线的性能。在不同的目标变量和预测时间方面，ClimaX是表现最好的方法。有趣的是，即使在较低分辨率数据上进行预训练，ClimaX-pos-interp在预测Z500和T850方面仍然表现出色，仅在3天的T2m预测中表现略逊于ResNet。这个结果表明，ClimaX可以在任务中获得强大的性能，这些任务的空间覆盖范围与预训练不同，甚至空间分辨率也不同。

4.2.3. 季节性累积预测

季节性累积（S2S）预测是在2周至2个月的时间范围内进行预测的任务，它连接了天气预报和气候预测之间的差距。与其他两个成熟的任务相比，S2S预测受到的关注要少得多，尽管在灾害缓解方面具有重要的社会经济价值。最近的研究提出了基于传统机器学习[Hwa+19；Pro+18；TL18]和深度学习[Wey+21；Zho+21；Ore+19]的数据驱动方法，但它们的表现常常落后于标准基准测试[Mou+23a]上的自适应偏差校正方法[Mou+23b]。在这里，我们遵循S2S竞赛（https://s2s-ai-challenge.github.io/），旨在预测第3-4周和第5-6周的双周平均统计数据，分别对应于2周和4周的先行时间。我们构建了ERA5-S2S，这是一个新的数据集，从5.625° ERA5中获取输入变量，但输出变量从先行时间平均到未来的2周。

我们在T850、T2m、U10和V10的S2S预测中比较了ClimaX与ResNet、UNet和ClimaViT。表1比较了ClimaX和基线模型的RMSE。ClimaX在所有变量上均取得最低的误差，并且随着先行时间的增加，与最佳基线模型UNet之间的性能差距更大。ClimaX在其零起点训练对应模型ClimaViT上也有显著的性能提升，显示了我们预训练过程在捕获对各种时间预测任务普遍有用的特征方面的有效性。

表1：ClimaX和基线模型在5.625° ERA5-S2S预测任务上的RMSE。

4.3. 气候预测

为了进一步测试ClimaX的通用性，我们在ClimateBench上评估了该模型，这是一个为气候预测测试机器学习模型而设计的最新基准。ClimateBench的目标是根据四个人为强迫因素（二氧化碳（$CO_2$）、二氧化硫（$SO_2$）、黑碳（BC）和甲烷（$CH_4$）），预测全球表面温度的年平均分布、昼夜温差、降水量以及降水量的90th百分位数。我们注意到这不是一个时间建模任务，因为我们不是根据过去预测未来。而是回答像在特定$CO_2$水平下年平均温度将是多少这样的问题？特别要注意的是，输入变量和任务本身与预训练时完全不同。

图9：ClimateBench的微调流程。不同的输入和输出变量需要不同的嵌入层和预测头部。注意力层可以冻结或微调。

图9展示了ClimaX在ClimateBench上的微调流程。由于在预训练期间看不到输入和输出变量，我们用新初始化的网络替换预训练的嵌入层和预测头部，同时保留注意力层和变量聚合模块。我们考虑了两种微调协议，一种是冻结（ClimaXfrozen），另一种是微调（ClimaX）注意力层。此外，我们在图2的流程中引入了两个组件。我们使用过去十年的强迫因素历史来预测特定年份，创建一个形状为T×V×H×W的输入。输入的每个时间片段按照变量分词、变量聚合和注意力层的常规流程进行，输出一个形状为T×h×w×D的特征张量，其中D是嵌入大小。然后，特征张量通过全局平均池化层进行降维，将维度降至T×D。最后，在馈送到预测头之前，通过一个交叉注意力层聚合了10年的历史数据，线性转换D维特征向量为H×W的地图。历史数据聚合和全局池化模块是原始ClimaX架构的两个新增设计，这些设计受到[WP+22]中神经网络基准的启发。

我们将ClimaX与ClimaX$_{\text{frozen}}$、Cli-ViT以及ClimateBench中最佳基线模型进行比较。按照[WP+22]，我们使用标准均方误差（不包含加权项的方程式（1））作为损失函数。我们对所有方法进行RMSE、NRMSE$_s$（空间）、NRMSE$_g$（全球）和Total = NRMSEs + 5 × NRMSE$_g$ [WP+22]的评估。度量细节见附录D.1。表2展示了结果。ClimaX$_{\text{frozen}}$在预测两个与温度相关的变量方面表现最佳，其次是ClimaX。这表明预训练的注意力层在看似不相关的任务中可以作为强大的特征提取器。当下游数据稀缺时（ClimateBench只有754个数据点），进一步微调注意力层可能会导致过拟合，从而轻微影响性能。在两个与降水相关的任务中，ClimaX$_{\text{frozen}}$在NRMSE$_s$和NRMSE$_g$方面略低于ClimateBench基线，但在RMSE上表现更好。我们推测这是因为ClimaX在预训练期间没有观察到降水变量，它与其他变量有非常不同的行为特征。

表2：ClimaX和ClimateBench基线模型在ClimateBench上的性能。Spatial和Global分别表示标准化均方根误差NRMSEs和全局平均NRMSEg，Total是Spatial和Global的加权组合。

4.4. 气候模型降尺度化

由于其高计算成本，气候模型通常在粗网格上运行。尽管这些预测有助于理解大尺度气候趋势，但它们无法提供足够的细节来分析区域和局部现象。降尺度化旨在获得更高分辨率的预测，并减少这些模型输出中的偏差。为了评估ClimaX在气候模型降尺度化任务中的适用性，我们基于CMIP6和ERA5数据源构建了一个新的数据集，用于粗输入和更高分辨率目标。具体来说，我们使用来自CMIP6的MPI-ESM数据集以及表8中列出的变量作为输入（分辨率为5.625°），并训练单独的模型将其降尺度到每个ERA5目标变量（分辨率为1.40625°）。我们将ClimaX与Cli-ViT以及两个CNN基线模型UNet和ResNet进行比较，因为最近的深度降尺度化方法[Van+17；Rod+18；Höh+20；VKG19；LGD20]都是基于卷积的。由于我们无法获取高分辨率辅助数据（如海拔和地形信息），因此无法与当前最佳的深度降尺度化方法YNet [LGD20]进行比较。对于所有方法，我们首先使用双线性插值将输入插值到所需输出的分辨率，然后将其馈送到模型中。我们使用RMSE、皮尔逊相关系数和平均偏差对所有方法进行评估，这些指标在现有的深度降尺度化作品中经常使用[Van+17；LGD20]。度量的细节见附录D.1。

表3：ClimaX和基线模型在从MPI-ESM（5.625°）降尺度到ERA5（1.40625°）任务中的性能。

表3定量比较了ClimaX和基线模型。ClimaX在所有三个目标变量上均实现了最低的RMSE和最接近0的平均偏差，并且在皮尔逊相关系数方面与基线模型表现类似。虽然ClimaX预训练用于进行预测，但它成功捕捉了气象数据的空间结构，这有助于降尺度化等下游任务。图10展示了ClimaX对三个目标变量进行降尺度化预测的可视化结果。输入数据的分辨率要低得多，并且与实际情况存在很多偏差。虽然预测结果缺少一些细节，但它成功地捕捉了ERA5数据的一般结构并消除了输入的偏差。

图10：ClimaX对关键变量进行降尺度化预测的示例可视化。

4.5. 规模定律分析

Transformer模型已经展示了在语言[Kap+20；Hof+22]、视觉[Zha+22a]甚至多模态任务[Hen+20b；Hen+21；Ree+22b]方面具有有利的规模性质。也就是说，它们的性能会随着数据量和模型容量的增加而提高，前提是有足够的计算资源支持。在本节中，我们研究了ClimaX在气象预测中的规模定律。图11展示了ClimaX的性能随着数据量和模型容量的变化情况。x轴表示预训练数据量（以Gigabytes为单位），对应于1到5个CMIP6数据集，y轴显示了ClimaX在3天预测任务上的RMSE。我们通过将嵌入维度从128变化到1024来比较四个不同容量的ClimaX模型。所有实验都在5.625°数据上进行。两个最大模型的误差率随着数据量和模型大小的增加而持续下降。这突显了ClimaX从多样和异构数据源中学习的独特能力，使我们能够通过简单地在更多数据上进行预训练来进一步提高性能。然而，两个较小的模型并没有像较大的那样良好地扩展，增加数据量并没有带来太多的改进，有时甚至会影响性能。这个结果表明，较大的模型不仅表现更好，而且更具数据效率。

图11：与预训练期间观察到的CMIP6 5.625°数据相关的不同变量的ERA5 3天预测误差。较大的模型具有更高的样本效率。

除了数据量和模型容量外，数据分辨率在天气和气候背景下是另一个重要的扩展维度。在许多视觉任务中，如分类，了解图像的一般高级结构就足以进行准确预测。然而，为了建模控制天气和气候的基础复杂物理过程，模型需要查看输入的细节，以了解数据的空间和时间结构以及不同变量之间的相互作用。高分辨率数据包含更细致的细节和天气条件的局部过程，这些在低分辨率数据中不存在，并且为训练深度学习模型提供了更强的信号。图12比较了ClimaX在全球预测上对5.625°和1.40625°数据进行预训练和微调后的性能。除了1天和3天的T2m预测外，ClimaX（1.40625°）一直表现出比低分辨率模型更低的RMSE和更高的ACC。我们注意到，由于内存问题，对于高分辨率数据，我们必须使用更大的块大小（与低分辨率数据相比增加到4），我们可以通过减小块大小进一步提高ClimaX在1.40625°数据上的性能，因为模型能够更好地捕捉细节。

图12：与数据分辨率相关的性能缩放。尽管使用了较大的补丁尺寸，但ClimaX（1.40625°）在几乎所有任务中始终比低分辨率模型表现更好，除了在1天和3天lead时间的T2m预测方面。

4.6. 消融研究

在主要的预测结果中，我们为每个目标变量和每个lead时间分别微调一个单独的ClimaX模型，因为我们发现这种协议导致了最佳的性能。然而，这可能会消耗大量计算资源，因为微调成本随着目标变量和lead时间的数量呈线性增长。在本节中，我们考虑了不同的微调替代方案，以研究计算和性能之间的权衡。

4.6.1. 我们应该为每个变量单独微调ClimaX还是一次性全部微调？

与为每个目标变量单独微调ClimaX不同，我们可以选择一次性微调，同时预测输入中的所有变量，我们将其称为ClimaX-all-vars。图13显示，ClimaX-all-vars在大多数任务中实现了与ClimaX可比较的性能，仅在预测T2m方面表现较差。这表明，在有限的预算下，可以一次性微调ClimaX以同时预测所有目标变量，而不会损失太多性能。

4.6.2. 我们应该进行迭代预测还是直接预测？

为了避免为每个lead时间微调不同的模型，我们可以微调ClimaX以在短时间范围（如6小时）内进行预测，并在推断过程中展开预测，以在较长时间范围内进行预测。我们将这种模型称为ClimaX-iter，其中iter代表迭代预测[Ras+20]。我们注意到，为了展开多个步骤，ClimaX-iter必须预测所有输入变量，或者换句话说。这提供了一个好处，即微调一个单一模型，该模型可以预测任何目标变量的任何lead时间。图13显示，ClimaX-iter在1天内的预测上表现良好，但在较长lead时间下性能显著下降。这并不奇怪，因为ClimaX-iter没有经过微调来预测未来的多个步骤，导致错误迅速积累。我们可以采用多步目标来微调，如Pathak，Subramanian等人[PAT+22]，以获得更好的结果。

图13：ClimaX及其变体在天气预测中的性能。ClimaX-cont是一个lead时间条件模型，我们对其进行微调以在6小时到7天的时间范围内进行预测。ClimaX-iter在6小时lead时间上进行预测，并展开预测以在更长的时间范围内进行预测。ClimaX-all-vars预测了输入中所有变量在特定lead时间的未来状态。

4.6.3. 我们可以微调ClimaX以适用于所有lead时间吗？

避免为每个lead时间单独微调的另一种方法是微调一个lead时间条件模型。具体来说，在微调过程中，我们将lead时间从6小时随机化到7天，类似于预训练设置。请注意，与ClimaX-iter不同，我们仍然必须为每个目标变量微调一个单独的模型。我们将这种模型称为ClimaX-cont，其中cont代表连续，在先前的作品中使用了这个标准术语[Ras+20]。图13显示，ClimaX-cont在6小时到7天的预测上表现竞争力十足，但无法外推到训练中未见过的2周和1个月的lead时间。我们也可以将lead时间从6小时随机化到1个月，但这意味着模型看到了更少的数据点，可能会影响性能。

每组权重的微调成本是一个常数C，对于8×V100s约为15小时。在不同的微调协议中，ClimaX是最昂贵的，其总成本为C ×#variables×#lead_times，与目标变量和lead时间的数量成线性关系。接下来是ClimaX-all-vars和ClimaX-cont，它们的总成本分别为C × #lead_times和C × #variables。最后，ClimaX-iter是最便宜的微调协议，我们只需要微调一个单一模型，该模型适用于所有目标变量和所有lead时间。性能与计算成本成正比，因为ClimaX是表现最佳的模型，而ClimaX-iter是表现最差的模型。

5 讨论与未来工作

数据集、模型架构和计算的扩展在人工智能的各个子领域中产生了革命性影响，从自然语言和语音处理到计算机视觉，以及在生物学和化学等科学应用中。特别是，它导致了通用基础模型的出现，这些模型在大型数据集和计算集群上进行训练，并且可以高效地适应各种下游任务，无论是在计算还是数据监督方面。我们的工作代表了在天气和气候数据驱动模型中实现这种广泛扩展和普适性的开创性努力。这种方法超越了传统数值建模和现有数据驱动预测方法的局限性。与ClimaX不同，数值模型只在计算方面扩展，而不是在数据集大小方面扩展，而现有的数据驱动模型通常局限于特定任务，缺乏在各种任务上的通用适用性。

除了语言和视觉方面的传统考虑之外，像ClimaX这样的基础模型通过使用模拟数据集和网格分辨率为扩展提供了新的机会。为了简化我们的方法，我们选择使用包含在先前关于数据驱动预测的研究中进行基准测试的标准变量的预训练数据集。此外，我们避免使用模拟未来情景的数据集，以防止对气候预测任务的潜在泄漏。未来的研究可以探索将观测和模拟数据集结合起来，其中包括更广泛范围的气候变量、更高的时空分辨率，甚至扩展到未来情景。

此外，我们展示了分辨率在ClimaX扩展中发挥了关键作用。由于计算限制，我们在低到中等分辨率上训练了ClimaX。尽管如此，我们的经验趋势表明，扩展到更高分辨率（0.25°）很可能会带来更好的结果。

未来的扩展工作可以从更好的序列建模架构中受益，特别是那些为多模式时空输入设计的架构。正如我们在ClimaX中看到的，气候数据集的通道数量远远大于标准多模式设置（例如，音频-视频、视觉-语言模型）。此外，在实践中，不同气候数据集的分辨率范围也相当大。这种异构性大大增加了标准架构（如ViT）对于输入序列的原始长度。未来，我们相信研究单一的多尺度架构（例如，[Fan+21]）可以通过学习推断与大气现象在逐渐增加的空间分辨率上相关的特征，从而有助于扩展到这种多样化的多分辨率和多模式数据集。

总之，我们相信我们方法的普适性具有超越本工作考虑的任务的潜力。有趣的是，探索预训练的ClimaX骨干在其他地球系统科学任务中的泛化性能，例如预测极端天气事件和评估人为对气候变化的贡献，以及与天气和气候条件密切相关的更广泛的领域，例如农业、人口统计学和精算学。