HEAL-ViT | 球形网格与Transformer的完美结合，引领机器学习预测新纪元！

集智书童公众号

发布于 2024-04-12 20:02:37

1070

发布于 2024-04-12 20:02:37

近年来，各种机器学习架构和技术在生成中期天气预报方面取得了成功。特别是基于视觉 Transformer （ViT）的模型（如Pangu-Weather、FuXi）表现出强大的性能，几乎可以“即开即用”，通过将天气数据视为直角网格上的多通道图像。尽管直角网格适合2D图像，但天气数据本质上是球形的，因此在直角网格上的极地会出现严重扭曲，导致用于模拟极地附近数据的计算不成比例。基于图的方法（如GraphCast）不会出现这个问题，因为它们将经纬网格映射到球面网格上，但通常对内存要求更高，且在训练和推理过程中往往需要更多的计算资源。虽然球面网格在空间上是同质的，但它并不容易被基于ViT的模型所建模，而这些模型隐式依赖于直角网格结构。作者提出了HEAL-ViT，一种新颖的架构，它使用基于球面网格的ViT模型，从而既受益于基于图模型所享有的空间同质性，又利用了 Transformer 的基于注意力的有效机制。 HEAL-ViT生成的天气预报在关键指标上超过了ECMWF IFS，并且在偏差累积和模糊度方面优于其他机器学习天气预报模型。此外，HEAL-ViT降低的计算占用使其也适用于业务使用，除了每6小时一次的预测模型，可能还需要其他模型来生成所需的全部业务预报集合。

1 Introduction

近年来，各种机器学习天气预测模型（MLWPs）在中期天气预报方面表现出了强大的性能，这被定义为从给定初始条件下生成10天预报的任务。MLWPs通常在ECMWF的ERA5数据集（Hersbach等人，2020年）上进行训练，并在关键指标上超过了通常被认为是数值天气预报（NWP）领域最先进技术的ECMWF IFS模型（Haiden等人，2018年）。多种模型结构都成功地生成了高质量的10天预报，其中突出的模型包括FourCastNet（Pathak等人，2022年）、Pangu-Weather（Bi等人，2023年）、GraphCast（Lam等人，2022年）和FuXi（Chen等人，2023年），这些模型在ERA5数据集（Hersbach等人，2020年）提供的原生0.25

{}^{\circ}

分辨率下表现出色。尽管所有MLWPs在IFS基准上均显示出均方根误差（RMSE）的降低，但它们往往更容易累积偏差，并且与IFS相比，生成的预报在区域能量谱测量下显得更加模糊。根据所选指标和变量的不同，不同的MLWPs表现各有优劣，例如，尽管GraphCast在大多数变量的均方根误差（RMSE）方面表现更好，但基于视觉 Transformer （ViT）Kolesnikov等人（2021年）的模型如Pangu-Weather，往往显示出减少的模糊性，并且不会像IFS那样强烈地累积偏差。

所有架构都采用“编码器-处理器-解码器”方案，其中（i）编码器将原始的0.25

{}^{\circ}

经纬度网格映射到较小的网格上，（ii）处理器在较小的网格上建模节点之间的相互作用，以及（iii）解码器将处理后的网格映射回经纬度网格。

FourCastNet、Pangu-Weather和FuXi在编码器和解码器中使用卷积层，将经纬度网格映射到直角网格。GraphCast使用图网络将网格映射到一个二十面体网格，该网格在球体上均匀分布节点。对于处理器，FourCastNet使用傅里叶神经算子，而Pangu-Weather和FuXi采用ViT架构，都在直角网格上进行操作。GraphCast在球形的二十面体网格上的节点间使用图网络进行消息传递。虽然基于ViT的方法可以利用注意力机制有效地学习网格中的长距离关系，但使用的直角网格过度表示了靠近极地的数据，因此需要不成比例的计算资源来模拟这种空间不均匀的表示。尽管基于图的方法在空间均匀表示上操作，但由于需要消息传递机制，它们往往更占用内存。

Our contribution

作者提出了HEAL-ViT，一种新颖的架构，它使得标准的视觉 Transformer 能够应用于球面网格，从而得益于球面网格提供的统一空间表示以及能够以内存高效的方式建模长距离关系的注意力机制。HEAL-ViT在内存和计算占用上更低，因为空间上均匀的球面网格在表示底层数据时更为高效，这使得它在需要组合多个模型以生成业务天气预报的操作使用中具有吸引力。在误差降低方面，HEAL-ViT与其他MLWPs相当，但在由纬向能量谱测量的偏差累积和模糊性方面表现出改进的性能。

Outline

在定义问题之后，作者简要概述了SWIN Transformer 模型，这是一种通用变体视觉 Transformer ，它在矩形网格上操作。然后作者描述了HEALPix网格以及如何将其适配以用于SWIN Transformer 。后续章节将介绍HEAL-ViT模型架构，网络参数化，训练细节，以及最终的结果。

2 Problem definition and notation

使用

Z^{t}

来表示时间

时的真实天气状态，天气预报问题可以定义为寻找转换函数

\Phi

以预测未来状态

Z^{t+\Delta t}=\Phi(Z^{t})

。由于

Z^{t}

包含了时间

的完整天气状态，它具有预测

Z^{t+\Delta t}

的足够信息。在实际应用中，作者只能获得真实天气

Z^{t}

的部分观测值

X^{t}

，作者的目标是学习真实转换函数

\Phi

的最佳可能近似

\phi

。由于

X^{t}

是真实状态

Z^{t}

的部分视角，作者提供额外的历史上下文形式

X^{t-\Delta t}

来学习准确的

\phi

：

X^{t+\Delta t}=\phi(X^{t},X^{t-\Delta t}) \tag{1}

在学到一个精确模型

\phi

之后，给定一个初始状态

X^{t_{0}}

，作者可以使用

\phi

以自回归的方式生成一个期望长度的预测：

遵循其他多层感知机模型（MLWPs），作者将时间间隔

\Delta t=6

小时固定，将作者的预测限制在最多10天（40步）内，并将符号

X^{t+k\Delta t}

简化为

X^{t+k}

。在作者的工作中，观察到的天气状态

X^{t}

包括表1中列出的地表和大气变量。除了观察到的天气状态

X^{t}

和

X^{t-1}

之外，作者还使用陆地-海洋 Mask （一个二元变量）和地表位势（捕捉地形）作为输入特征。总的来说，这意味着

\phi

总共消耗了110个特征作为输入：

X^{t}

和

X^{t-1}

各54个特征以及2个静态特征。

3 SWIN transformers on the HEALPix mesh

SWIN transformers

SWIN（移位窗口） Transformer Liu等人（2021年）是一种用于计算机视觉任务的通用型 Backbone 网络，已被广泛用于中程天气预报Chen等人（2023年）；Hu等人（2023年）；Bi等人（2023年）；Chen等人（2023年）。在SWIN Transformer 中，首先将图像“划分”成

M\times N

个块，有效地创建了一个矩形网格，网格中的每个节点表示输入图像的一个特定部分。块被组合成窗口，并在每个窗口内执行自注意力来模拟块之间的依赖关系。通过在交替层中移动窗口来学习跨窗口连接，如图1所示。

仅在窗口内执行自注意力限制了注意力机制的计算成本（该成本与窗口大小的平方成正比），而移动窗口则可以学习远距离关系。在所有行和列都具有相等节点数的矩形网格上移动窗口是微不足道的，但在球面网格上则具有挑战性，因为与赤道相比，极地区域的节点数较少。作者利用HEALPix像素化方案的有利特性，在球面网格上高效地移动窗口，并允许在不进行任何修改的情况下使用SWIN Transformer 。

The HEALPix mesh

HEALPix（Gorski等人，1999年，球面的层次化等面积纬度像素化）是一种像素化方案，它将球面划分为表示相等面积的四边形曲线路段的像素。这种技术最初是为了绘制和分析宇宙微波背景而开发的，HEALPix网格也已被用于球面CNN（Perraudin等人，2019年；Krachmalnicoff和Tomasi，2019年）。

在最粗的 Level 上，HEALPix网格有12个像素，代表12个四边形的“面”。网格通过将每个像素划分为4个子像素来细化，这些子像素代表原始四边形的四个象限。将这个过程重复

次可以将一个面划分为

(2^{n})^{2}

个像素，排列成大小为

2^{n}\times 2^{n}

像素的方形。在任何细化 Level 上的像素总数可以计算为

12*((2^{n})^{2})

。图2展示了0、1、2和3级细化程度的细化过程，表2显示了节点数量如何随细化 Level 的增加而增加。

HEALPix网格提供了一种对网格节点的规则排序方式，每个网格节点沿着东、西、北、南、东北、西北、东南和西南方向都有明确定义的邻居。因此，整个网格可以被定向为南北和东西方向，并且可以被“展平”，如图3所示。

3.2.1 Creating windows on the HEALPix mesh

每个HEALPix网格节点对应于球体的一个特定部分，因此它等同于SWIN Transformer 中使用的“图像块”。作者可以使用像素的分层细分来创建SWIN Transformer 进行局部自注意力所需的窗口。对于一个精细的细化 Level

和一个更粗糙的细化 Level

n-w

，作者可以看到每个粗糙像素“包含”

(2^{w})^{2}

个精细像素，并且每个精细像素恰好属于一个粗糙像素。因此，来自粗糙细化 Level 的像素可以用作非重叠窗口，在这些窗口内可以对精细像素执行自注意力。图4展示了在展平的HEALPix网格上窗口的布局。对于一个在细化 Level

的精细网格，窗口大小参数

可以调整，以使每个精细像素关注球体上更大的或更小的上下文。

3.2.2 Shifting windows on the HEALPix mesh

分层细分和等面积像素是许多球面网格中发现的特性，例如GraphCast所使用的二十面体网格。然而，HEALPix网格由于其以下特性，特别适合与SWIN Transformer 配合使用：

所有像素细分为4个子像素，每个子像素分别对应父像素的北、南、东、西四个象限。

像素在8个基本方向上具有明确定义的邻居：东北、西北、东南、西南、北、南、东和西。

给定一个细化 Level

，以及窗口大小参数

，可以通过以下步骤构建平移窗口：

将每个粗略像素在细化 Level

n-w

处分裂为细化 Level

n-w+1

的子像素。由于 Level

n-w

的像素定义了原始窗口，这些子像素代表了原始窗口的北、南、东、西四个象限。

将每个南向子像素与其西南、南方和东南邻居合并，以创建一个新的窗口。这3个邻居都保证来自不同的窗口，因为它们位于原始南向子像素的更南边。

同样地，将每个北方向子像素与其西北、北和东北邻居合并，东西象限的子像素也以此类推。

当这个分割-合并方案应用于细化层次为

n-w

的整个网格时，所有窗口都正好被向南和向西移动了一个象限，确保每个新窗口由四个独立原始窗口的象限组成。

生成的偏移窗口可以如图5所示进行可视化。在HEALPix网格中，每个偏移窗口由四个相邻原始窗口的四个象限组成，正如在直角网格中一样。这使得可以在HEALPix网格中以与直角网格相同的方式学习跨窗口的连接。在图5中用蓝色标记的极地区域保持无失真，因为靠近极地的斑块和窗口与靠近赤道的那些大小相同。

由于所有窗口包含相同数量的网格节点，且在HEALPix网格中所有网格节点代表相等的面积，因此所有像素在球面上执行相等"真实"面积窗口的局部自注意力。相比之下，在直角网格上，尽管所有窗口包含相同数量的网格节点，但靠近极地的窗口代表较小的面积，这减少了靠近极地自注意力的"真实"上下文。作者注意到有8个位置，一个中央HEALPix面与2个北部或2个南部HEALPix面相接，在这些交点处，移动的窗口每个只有3个象限，如图5中红色标记所示。

Efficiencies from the spherical mesh

作者可以通过比较在类似细化程度的HEALPix球面网格和直角网格中的节点数量来量化使用球面网格对ViT模型的益处。采用直角网格上的SWIN Transformer 的Pangu-Weather和FuXi，使用4x4大小的 Patch 将原始图像从721x1440像素减少到180x360大小的直角网格（即64,800个节点）。由于原始分辨率为0.25°，这个网格上的每个节点在赤道处代表1°x1°的区域，而在极地附近则代表更小的区域。Pangu-Weather使用形状为

6\times 12

的窗口来执行局部自注意力，从而构建了每个窗口72个节点的900个窗口。

在HEALPix网格上，使用细化 Level

n=2^{6}

，作者得到一个含有49,152个节点的网格，每个节点在球面上表示一个0.91

{}^{\circ}\times

0.91

{}^{\circ}

的区域。一个窗口参数

w=3

导致每个窗口有

(2^{3})^{2}=64

个节点，768个窗口覆盖整个网格。

与矩形网格相比，HEALPix网格节点减少了25%，同时每个节点代表的区域也更小。覆盖网格所需的窗口数量减少了15%，且这些窗口的大小与矩形网格中的窗口相似。由于所有HEALPix窗口具有相等的面积，因此所有网格节点都关注相同数量的输入上下文。换句话说，HEALPix网格需要更少的节点来表示更高分辨率的数据，这进而导致执行注意力机制所需的窗口减少，从而降低了内存和计算资源的占用。

4 HEALPix

在描述了HEALPix网格以及窗口是如何构建和移动之后，作者现在来描述整个模型的架构。HEALPix模型遵循编码器-处理器-解码器的结构，正如其他MLWPs（多尺度学习波前预测器）那样：

编码器使用一个简单的图网络将经纬度网格映射到一个HEALPix网格。

处理器使用SWIN Transformer 来学习网格节点之间的交互作用。

解码器最终使用一个简单的图网络将处理后的HEALPix网格映射回经纬度网格。

Encoder

类似于GraphCast，作者使用一个二分图

\mathcal{G}(\mathcal{V}^{G},\mathcal{V}^{M},\mathcal{E}^{G2M})

将经纬度网格嵌入到 HEALPix 网格中。

\mathcal{V}^{G}

是所有位于721x1440大小经纬度网格上的网格节点的集合。

\mathcal{V}^{M}

是在细化 Level

n=7

的 HEALPix 网格上，具有

196,608

个节点的所有网格节点的集合。由于每个 HEALPix 网格节点代表球面上的一个曲边四边形区域，每个网格节点都与位于该四边形内的所有网格节点相连，如图6所示。

\mathcal{E}^{G2M}

是连接网格节点到网格节点的有向边集合，用32维嵌入表示。由于每个网格节点恰好连接到一个网格节点，因此每个网格节点有一条边，共有1,038,240条边。

每个网格节点

\mathbf{v}^{G}_{i}

由特征

\mathbf{v}^{G.features}_{i}=[\mathbf{x}^{i-1}_{i},\mathbf{x}^{i}_{i},\mathbf{c} _{i}]

表示，其中

\mathbf{x}^{i}_{i}

是时间相关的天气状态，

\mathbf{c}_{i}

是网格节点的常数特征集。每个网格节点及其对应的边缘首先通过单个前馈层嵌入到一个维度为

的潜在空间中，然后使用 GeLU（Hendrycks 和 Gimpel，2016年）激活函数：

\mathbf{v}^{G^{\prime}}_{i} =[\mathbf{c}^{\text{G2M}}_{v^{G}_{s}\to v^{M}_{i}},\mathbf{v}^{G,features}_{i}] \tag{3}

\mathbf{v}^{G}_{i} =\text{GeLU}(\mathbf{W}^{\text{embeddore}}\mathbf{v}^{G^{\prime}} _{i}+\mathbf{b}^{\text{embeddore}})

网格节点是通过聚合来自入边的信息计算得出的，随后进行层归一化（Ba等人，2016年）。聚合的信息通过单个前馈层嵌入到一个维度为

的潜在空间中：

\mathbf{v}^{M^{\prime}}_{i} =\text{LayerNorm}(\sum_{\begin{subarray}{c}c\in\Omega\\ v^{G}_{s}\to v^{M}_{s}\end{subarray}}\mathbf{v}^{G}_{s}) \tag{4}

\mathbf{v}^{M}_{i} =\text{W}^{\text{G2M}}\mathbf{v}^{M^{\prime}}_{i}+\mathbf{b}^{ \text{G2M}}

总之，编码器消耗一个尺寸为

N_{\text{features}}\times 721\times 1440

的图像，并生成一个尺寸为

196608\times C

的网格。

Processor

一旦网格被嵌入到大小为

196608\times C

的网格中，处理器使用U-Net架构来学习网格节点之间的交互。首先，对编码后的网格通过2个SWIN块进行处理，然后使用下采样块将其粗化到一个具有细化 Level

n=2^{6}

的网格，该网格包含49,152个节点。大小为

49152\times 2C

的粗网格由48个SWIN块处理，然后上采样回原始的

196608\times C

大小的网格。使用跳过连接将下采样块的输入与上采样块的输出连接起来，然后由2个SWIN块处理。与典型的SWIN Transformer 一样，每隔一个SWIN块就会移动执行局部自注意的窗口。整个处理器堆栈如图7所示。

SWIN块 作者使用了如ViT-22BDehghani等人（2023年）架构中描述的ViT块，该块通过以下步骤将输入

转换为输出

：

遵循ViT-22B架构，作者在QKV投影中不使用偏置，并在计算注意力权重之前对 Query 和键向量进行层归一化：

下采样块 下采样块将网格从

196608\times C

粗化到大小为

49152\times 2C

的网格，有效地使网格细化 Level 变粗一个 Level 。下采样块在一个二分图网络

\mathcal{G}(\mathcal{V}^{M_{f}},\mathcal{V}^{M_{c}},\mathcal{E}^{M_{f}2M_{c}})

上使用简单的图网络，其中

\mathcal{V}^{M_{f}}

是细网格（细化 Level

2^{7}

）中的所有节点集合，

\mathcal{V}^{M_{c}}

是粗网格（细化 Level

2^{6}

）中的所有节点集合，而

\mathcal{E}^{M_{f}2M_{c}}

将细网格节点连接到粗网格节点。因为HEALPix网格是分层细化的，每个细网格节点都恰好连接到一个“父”粗网格节点。连接边由可学习的32维嵌入表示。每个粗网格节点首先从其子节点聚合传入信息，然后进行层归一化。接着，前馈层将聚合的信息嵌入到维度为

的潜在空间中：

每个粗节点包含4个细节点，下采样块将4个维度为

的细节点粗化为一个维度为

的粗节点，这与Pangu-Weather和FuXi中使用的下采样块相似。

上采样块 上采样块反转了下采样块的过程，并将网格从尺寸

49152\times 2C

精细化为

196608\times C

。与下采样块类似，上采样块在一个二分图

\mathcal{G}(\mathcal{V}^{M_{c}},\mathcal{V}^{M_{f}},\mathcal{E}^{M_{c}2M_{f}})

上使用一个简单的图网络，其中

\mathcal{V}^{M_{f}}

是细网格中所有节点的集合，

\mathcal{V}^{M_{c}}

是粗网格中所有节点的集合，而

\mathcal{E}^{M_{c}2M_{f}}

将粗网格节点连接到细网格节点。作者将细网格中的每个节点连接到粗网格中最近的四个节点，并且边使用一个可学习的32维嵌入表示。

Normalization and network parameterization

Input and output normalization

在输入到编码器之前，表1中的所有输入变量都通过在整个训练期间计算的均值和方差统计量进行归一化，以达到零均值和单位方差。目标也同样进行归一化处理，网络直接预测归一化的目标作为输出。对于每个变量和 Level ，分别计算均值和方差统计量。

Neural network parametrization

作者为编码器的输出选择了一个值

C=256

，即编码器将输入网格从

110\times 721\times 1440

映射为一个大小为

196608\times 256

的网格。因此，处理器中的下采样模块将网格从

1966608\times 256

粗化到一个大小为

49152\times 512

的网格。每个SWIN层内的MLP有一个隐藏层，其大小为4*D（其中

或

，取决于SWIN层作用于哪个网格）。SWIN块的自注意力模块使用一个头维度

d=32

，导致根据SWIN块处理的网格的不同，有8个或16个头。所有的边都由可学习的32维嵌入表示。

Training details

Training objective

作者优化HEAL-ViT以最小化经度加权的L1损失：

其中

是预测变量的数量，而

和

分别是输出网格中的纬度和经度数量。对于自回归训练（下面将描述），所有时间步的L1损失被平均以获得整体训练损失。

w(i)

表示第

个纬度上网格单元的权重，具体细节在6.2节中说明。

Training and testing split

作者使用1992年至2018年的ERA5数据进行训练，将2019年的数据作为开发集，2020年的数据作为测试集。开发集用于微调超参数，而测试集仅用于模型的最终评估。

Curriculum training schedule

作者首先对HEAL-ViT模型进行预训练，以预测单个步骤

X^{t+1}=\phi(X^{t},X^{t-1})

80,000步。一旦模型被训练以预测紧接的下一步，作者遵循GraphCast的方法，并使用自回归训练制度来提升模型在第一步之后的性能，如第2节所述。与GraphCast（训练多达12步自回归）和FuXi（训练20步自回归）相比，HEAL-ViT仅针对以下训练课程中的4步自回归进行微调：

2个自回归步骤，训练了4000个训练步骤。

3个自回归步骤，训练了4000个训练步骤。

4个自回归步骤，训练了4000个训练步骤。

Hyperparameters

作者使用了AdamW优化器Kinga等人（2015年）；Loshchilov和Hutter（2017年）（beta1 = 0.9, beta2 = 0.95），并设置了权重衰减为0.1。在预训练阶段，学习率初始设置为2.5e-4，并使用余弦调度进行衰减。自回归微调则采用固定的学习率1e-7进行。作者使用了8块A100 GPU，每个GPU放置1个示例，并采用激活预训练权重技术Chen等人（2016年）以减少训练模型所需的内存占用。数据处理和模型训练作者使用了PyTorch Paszke等人（2019年）和xarray Hoyer和Hamman（2016年）。预训练阶段大约需要2.5天，而微调过程还需要额外的2天。

6 Evaluation results

作者评估了HEAL-ViT在2020年的预测性能，并将其与ECMWF的IFS以及其他机器学习天气预测模型（MLWPs）进行了比较，使用WeatherBench2Rasp等人（2023年）的基准测试套件进行评估。与其他MLWPs一样，HEAL-ViT使用ERA5数据，每天从2020年的00UTC开始初始化。

Baselines

6.1.1 NWP Baselines

ERA5在00UTC使用了一个更长的同化窗口，这使得它能够以更高品质的状态来初始化天气模型。由于所有包括HEAL-ViT在内的MLWPs每天都是使用00UTC时的ERA5数据进行初始化的，因此HEAL-ViT相对于没有访问ERA5数据的业务IFS具有优势，这使得对MLWPs和IFS进行一对一的比较变得困难。WeatherBench2提供了“ERA5预报”，这是一组使用ERA5数据初始化的IFS预报，它们提供了一个与MLWPs进行比较的同等基准，因为现在所有模型都从相同的起始状态开始，并且可以与ERA5作为 GT 情况进行比较。因此，在下面的比较中，作者将使用ERA5数据初始化的IFS预报作为NWP的 Baseline ，并将其标记为ERA5-IFS。

图8展示了ERA5-IFS与业务IFS的性能差异，以及相对于基准模型

计算的均方根误差（RMSE）的归一化差异，计算方式为：

(\text{RMSE}_{A}-\text{RMSE}_{B})/\text{RMSE}_{B}

。作者注意到，在整个预报范围内，业务IFS在T2M预报上表现一致更好，而在WS10预报上表现较差。对于Z500和T850，与IFS相比，ERA5-IFS在较短的预报提前期表现更佳。对于所有变量，作者注意到任何性能差异在较长的预报提前期内会减少，但仍然存在。

6.1.2 MLWP baselines

作者使用GraphCast、Pangu-Weather和FuXi作为MLWP基准线，并从WeatherBench2获取评估统计数据。尽管所有MLWPs都是使用ERA5数据初始化并在2020年进行评估的，但GraphCast是在00UTC和12UTC的起始状态下进行初始化，而其他MLWPs仅在00UTC进行初始化。

Metrics

所有指标都是通过对整个网格使用面积加权平均值来计算的，以考虑到网格单元相对于赤道而言，在极地向代表较小区域的事实。每个纬度的权重按如下方式计算：

其中，

\theta_{i}^{u}

和

\theta_{i}^{l}

分别是纬度索引为

的网格单元的纬度上限和下限。在下面的定义中，

f_{t,l,i,j}

表示在时间

、层次

、以及网格单元

(i,j)

上一个变量的预测值。相同时间、层次和网格单元上的观测值由

o_{t,l,i,j}

表示。

6.2.1 Root Mean Squared Error (RMSE)

与WMO和ECMWF保持一致，WeatherBench2为每个变量和层次定义了RMSE（均方根误差）。

6.2.2 Anomaly Correlation Coefficient (ACC)

ACC表示预测异常与观测异常相对于气候学

c_{t,l,i,j}

之间的相关性：

f^{\prime}_{t,l,i,j} =f_{t,l,i,j}-c_{t,l,i,j} \tag{13}

o^{\prime}_{t,l,i,j} =o_{t,l,i,j}-c_{t,l,i,j}

\text{ACC}_{l} =\frac{1}{T}\sum_{t}^{T}\frac{\sum_{i}^{I}\sum_{j}^{J}w(i)f^{\prime }_{t,l,i,j}o^{\prime}_{t,l,i,j}}{\sqrt{\sum_{i}^{I}\sum_{j}^{J}w(i){f^{\prime}_ {t,l,i,j}}^{2}\sum_{i}^{I}\sum_{j}^{J}w(i){o^{\prime}_{t,l,i,j}}^{2}}}

气候学代表了特定一年中某一天和一天中某个时间段的平均天气情况，这是通过长时间的历史时期计算得出的。关于气候学的计算细节，作者参考了WeatherBench2Rasp等人（2023年）的研究。

6.2.3 Bias

偏置（Bias）的计算简单地表示为：

MLWPs倾向于强烈积累负偏差，这可能会影响其预测热带气旋强度能力的长期预报准确性（DeMaria和DeMaria 2024）。因此，在考虑其他质量指标时，也很有必要考虑偏差问题。

6.2.4 Energy Spectra

能量谱揭示了每个变量和层次在不同空间尺度上的变化程度。较小尺度上的能量越大，表明预报场的较小尺度结构越好，能量越小则预报越模糊。在WeatherBench2中，对于恒定纬度的线条，分别计算经向谱能量，并通过在

30^{\circ}<|lat|<60^{\circ}

之间平均经向谱来计算最终谱。作者对HEAL-ViT采用相同的处理过程。

7 Results

作者报告了2米温度（T2M）、10米风速（WS10）、500百帕高度处的位势（Z500）以及850百帕处的温度（T850）的结果。这些变量是中期天气预报中最常评估的一些变量。

Rmse

图9展示了HEAL-ViT和ERA5-IFS模型在T2M、WS10、Z500和T850上的均方根误差（RMSE），以及均方根误差的归一化差异。作者观察到，对于所有变量，在最初的18小时（3个预报步骤）内，ERA5-IFS的RMSE更低，之后HEAL-ViT模型则持续显示出更低的误差。作者回顾一下，在较短的预报提前期，ERA5-IFS在WS10、Z500和T850上的表现明显优于业务IFS，这表明IFS模型从以ERA5数据作为初始状态中强烈受益。

图10展示了HEAL-ViT与其他多种机器学习天气预测模型（MLWPs），即GraphCast、Pangu-Weather和FuXi的均方根误差（RMSE）。虽然HEAL-ViT的表现优于Pangu-Weather，但作者观察到GraphCast和FuXi的RMSE低于HEAL-ViT（除了WS10，在超过5天后，FuXi的错误最高）。

Acc

图11展示了HEAL-ViT和ERA5-IFS模型在检验变量上的准确率（ACC），以及准确率的归一化差异。作者观察到与均方根误差（RMSE）类似的变化趋势，即在前3到4个步骤中ERA5-IFS的准确率较高，之后HEAL-ViT模型则一致显示出更高的准确率。

图12展示了HEAL-ViT和其他MLWPs的准确率（ACC），再次显示出与均方根误差（RMSE）类似的趋势。HEAL-ViT的表现优于Pangu-Weather，而GraphCast和FuXi的性能更佳（除了在WS10上，FuXi在第5天后的准确率较差）。

Bias

图13展示了ERA5-IFS和所有多层感知器加权预测（MLWPs）的偏差。由于偏差衡量的是预测对变量的低估或高估程度，理想的预测偏差为零。虽然所有模型都积累了负偏差，但作者可以观察到ERA5-IFS在整个预测期间的偏差积累较慢。

相比之下，尽管MLWPs的RMSE低于ERA5-IFS，但它们往往更强烈地积累负偏差。例如，Pangu-Weather在T2M上显示出强烈的负偏差，而GraphCast和FuXi在WS10、Z500和T850上显示出强烈的负偏差。特别是，FuXi在偏差趋势中出现不连续性，这是由于在预报的第5天之后，FuXi采用了另一个模型。HEAL-ViT在偏差积累方面的表现往往优于其他MLWPs，尤其在T2M、WS10和Z500上表现明显，HEAL-ViT在这些变量上的偏差积累与ERA5-IFS相似。

Energy spectra

区域能量谱表明了一个预报的模糊程度，较小尺度上的能量越多，意味着预报在小尺度结构上表现越好。相较于NWPs，MLWPs的预报通常更模糊，这是由于选择了MAE或MSE作为损失函数所致。作者展示了6小时、1天和3天领先时间的区域能量谱，之后大多数MLWPs的能量谱相对恒定。同时展示了ERA5-IFS的能量谱，几乎在所有领先时间上都是恒定的。图14展示了MLWPs和ERA5-IFS的能量谱。作者看到对于T2M，所有MLWPs都有类似的谱（尽管在所有领先时间GraphCast倾向于更模糊），而ERA5-IFS在较小尺度上的能量明显更高。对于WS10，HEAL-ViT比FuXi和GraphCast拥有更高的能量，尽管仍然低于ERA5-IFS。对于Z500，MLWPs在较小尺度上的能量高于ERA5-IFS，并且随着领先时间的延长，差异增大。在这里，HEAL-ViT在较小尺度上的能量也比FuXi和GraphCast多，尽管到了第3天FuXi和HEAL-ViT的能量相当。作者还注意到，HEAL-ViT的能量谱在所有领先时间上都相对恒定，与ERA5-IFS相似。

光谱还揭示了编码器选择对预测场的影响。FuXi编码器使用常规的2D Patch ，在光谱中出现了尖峰，可能是由于 Patch 边界造成的。随着领先时间的增加，这些尖峰也变得更加明显，这表明任何 Patch 机制的影响在自动回归生成预测时会被加强。HEAL-ViT在较小尺度上也显示出一种“凸起”，这可能是由于其自身的编码器和网格到网格架构引起的。

8 Conclusions and future work

通过结合使用图网络构建球形网格，以及视觉 Transformer 学习网格节点间的交互作用，HEAL-ViT提供了一种架构，它结合了GraphCast（一种空间均匀网格）的独特优势与基于 Transformer 的方法（高效学习远程连接）的优点。这种方法并非特定于中程天气预报问题，它可以推广到任何使用经纬度网格上的地理空间数据的问题。在第3.3节中讨论的内存占用减少在操作环境中使用MLWPs时非常有价值，在这种环境中，需要同时运行多个MLWP以生成比当前MLWPs包含的更多变量的逐小时预报。与其他MLWPs一样，HEAL-ViT证明了与类似的NWP Baseline 相比，机器学习天气模型在天气预报质量的关键指标上能够超越NWP。此外，HEAL-ViT通过在能量谱中观察到的更好的偏差积累和减少模糊性，改进了现有的MLWPs。

HEAL-ViT的性能可能通过改进编码器、处理器和解码器组件来进一步优化。作者为编码器和解码器使用了基本的图结构，这些图结构没有包含关于网格和网格节点的绝对相对位置（如GraphCast所使用的图结构中那样）的任何先验知识。这种归纳偏差可能会提高编码器和解码器的质量，从而提升整个模型的性能。处理器由标准的视觉 Transformer 块组成，这些块没有融入在其他应用中已经显示出改进的变体，例如，可学习的相对位置偏差、层后归一化、缩放余弦注意力、计划性DropPath（Larsson等人，2017年）等。这些改进在FuXi中被整合，以显著增加模型的大小（相对于Pangu-Weather中使用的SWIN模型）。任何对ViT架构的改进都可能提高HEAL-ViT的性能。

由于HEAL-ViT网格在规则的纬度环上放置节点，其他学习遥相关性的选项实际上是微不足道的，例如，通过提供相邻的纬度环作为上下文窗口，而不是网格节点的窗口。

HEAL-ViT降低的内存和计算需求，以及能够使用标准的视觉 Transformer 架构的能力，使作者能够轻易地采用计算机视觉领域取得的任何进步。除了上述改进之外，作者还致力于探索将HEAL-ViT应用于与天气预报相关的领域，例如：超分辨率、概率性预报，以及学习可重用的高质量天气状态表示。

参考

[1].HEAL-ViT: Vision Transformers on a spherical mesh for medium-range weather forecasting.

本文参与腾讯云自媒体分享计划，分享自微信公众号。

原始发表：2024-04-02，如有侵权请联系 cloudcommunity@tencent.com 删除

机器学习

本文分享自集智书童微信公众号，前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体分享计划，欢迎热爱写作的你一起参与！