深度学习革新地球系统预测：立方体注意力与扩散模型

让深度学习切实应用于地球系统预测

地球是一个复杂的系统。从温度波动等常规事件，到干旱、冰雹和厄尔尼诺-南方涛动等极端事件，这些变化都可能影响农作物产量、延误航班，并引发洪水和森林火灾。对这些变化进行精确、及时的预测，可以帮助人们采取必要的预防措施以避免危机，或更好地利用风能和太阳能等自然资源。

Transformer 模型在其他 AI 领域的成功，促使研究人员尝试将其应用于地球系统预测。但这些努力遇到了几个主要挑战。其中最重要的是地球系统数据的高维度性：朴素地应用 Transformer 的二次复杂度注意力机制，计算成本过于高昂。

大多数现有的基于机器学习的地球系统模型也输出单一的确定性预测，这些预测通常是大范围可能结果的平均值。然而，有时了解发生极端天气事件有10%的可能性，比了解一系列可能结果中的一般平均值更为重要。此外，典型的机器学习模型没有受到物理定律或历史先例的约束，可能产生不太可能甚至不可能的预测结果。

在近期的工作中，某中心的团队应对了所有这些挑战。在 NeurIPS 2022 上发表的研究《Earthformer: Exploring space-time transformers for Earth system forecasting》，提出了一种称为“立方体注意力”的新颖注意力机制，它使得 Transformer 能够更高效地处理大规模、多维数据。

而在即将发表于 NeurIPS 2023 的研究《PreDiff: Precipitation nowcasting with latent diffusion models》中，研究人员展示了扩散模型既能实现概率性预测，又能对模型输出施加约束，使其与历史记录和物理定律更加一致。

Earthformer 与立方体注意力

Transformer 模型的核心是其“注意力机制”，它能够在处理输出序列的每个元素时，衡量输入序列不同部分的重要性。这种机制使得 Transformer 能够捕捉数据中时空长程的依赖性和关系，这是传统的基于卷积神经网络或循环神经网络的架构所未能很好建模的。

然而，地球系统数据本质上是高维度且时空复杂的。例如，在 NeurIPS 2022 论文研究的 SEVIR 数据集中，每个数据序列包含 25 帧数据，捕获间隔为五分钟，每帧的空间分辨率为 384 x 384 像素。使用传统的 Transformer 注意力机制来处理如此高维度的数据将极其昂贵。

在这篇论文中，研究者提出了一种名为立方体注意力的新颖机制，它将输入张量分解为“立方体”（立方体的高维类比），并在每个立方体级别上应用注意力。由于注意力的计算成本随张量大小呈二次方增长，在每个立方体内局部应用注意力比一次性计算整个张量上的注意力权重在计算上要可行得多。例如，沿时间轴分解可以为 SEVIR 数据集带来 384² 倍的成本降低，因为每帧空间分辨率为 384 x 384 像素。

当然，这种分解引入了一个限制：注意力在每个立方体内独立运作，立方体之间没有通信。为了解决这个问题，研究还计算了总结立方体注意力权重的全局向量。其他立方体可以在计算自己的注意力权重时考虑这些全局向量。

研究者将采用立方体注意力的基于 Transformer 的模型命名为 Earthformer。Earthformer 采用分层编码器-解码器架构，该架构逐渐将输入序列编码为多级表示，并通过从粗到细的过程生成预测。每个层级都包含一堆立方体注意力块。通过堆叠多个具有不同配置的立方体注意力层，能够高效地探索有效的时空注意力。

研究者尝试了多种将输入张量分解为立方体的方法。实验研究表明，“轴向”模式（即沿着时间、高度和宽度轴堆叠三个未偏移的局部分解）既有效又高效。它在避免原始注意力指数级计算成本的同时，实现了最佳性能。

实验结果

为了评估 Earthformer，研究者在两个真实世界数据集上将其与六种最先进的时空预测模型进行了比较：SEVIR（用于持续预测近期降水概率的任务，即“临近预报”）和 ICAR-ENSO（用于预测海表温度异常）。

在 SEVIR 上，使用的评估指标是标准均方误差和临界成功指数。CSI 也称为交并比：在不同阈值下，它被表示为 CSI-thresh；它们的平均值表示为 CSI-M。在 MSE 和 CSI 上，Earthformer 在所有指标上均优于所有六个基线模型。带有全局向量的 Earthformer 也始终优于没有全局向量的版本。

模型	参数量（百万）	GFLOPS	CSI-M↑	CSI-219↑	CSI-181↑	MSE(10⁻³)↓
持续性基准	-	-	0.2613	0.0526	0.0969	11.5338
UNet	16.6	33	0.3593	0.0577	0.1580	4.1119
ConvLSTM	14.0	527	0.4185	0.1288	0.2482	3.7532
PredRNN	46.6	328	0.4080	0.1312	0.2324	3.9014
PhyDNet	13.7	701	0.3940	0.1288	0.2309	4.8165
E3D-LSTM	35.6	523	0.4038	0.1239	0.2270	4.1702
Rainformer	184.0	170	0.3661	0.0831	0.1670	4.0272
Earthformer (无全局向量)	13.1	257	0.4356	0.1572	0.2716	3.7002
Earthformer	15.1	257	0.4419	0.1791	0.2848	3.6957

在 ICAR-ENSO 上，报告了三个月移动平均的 Nino3.4 指数的相关技能，该指数评估太平洋特定区域海表温度异常预测的准确性。Earthformer 在所有关注的评估指标上始终优于基线模型，而使用全局向量的版本进一步提高了性能。

模型	参数量（百万）	GFLOPS	C-Nino3.4-M↑	C-Nino3.4-WM↑	MSE(10⁻⁴)↓
持续性基准	-	-	0.3221	0.447	4.581
UNet	12.1	0.4	0.6926	2.102	2.868
ConvLSTM	14.0	11.1	0.6955	2.107	2.657
PredRNN	23.8	85.8	0.6492	1.910	3.044
PhyDNet	3.1	5.7	0.6646	1.965	2.708
E3D-LSTM	12.9	99.8	0.7040	2.125	3.095
Rainformer	19.2	1.3	0.7106	2.153	3.043
Earthformer (无全局向量)	6.6	23.6	0.7239	2.214	2.550
Earthformer	7.6	23.9	0.7329	2.259	2.546

PreDiff

扩散模型最近已成为许多 AI 任务的主流方法。扩散模型是生成模型，它建立了一个向训练样本迭代添加高斯噪声的前向过程；然后模型学习在反向扩散过程中逐步去除添加的噪声，逐渐降低噪声水平，最终产生清晰高质量的生成结果。

在训练过程中，模型学习其逐步执行的每个去噪步骤之间的一系列转移概率。因此，它本质上是一个概率模型，非常适合于概率性预测。

扩散模型的一个近期变体是潜在扩散模型：在传递给扩散模型之前，输入首先被送入一个自编码器，该自编码器具有一个产生压缩嵌入的瓶颈层；然后扩散模型在这个压缩空间中应用。

在即将发表的 NeurIPS 论文《PreDiff: Precipitation nowcasting with latent diffusion models》中，研究者提出了 PreDiff，这是一个以 Earthformer 作为其核心神经网络架构的潜在扩散模型。

通过修改训练后模型的转移概率，可以对模型输出施加约束，使其更可能符合某些先验知识。这是通过简单地偏移学习分布的均值来实现的，直到它更好地符合希望施加的约束。

结果

研究在 SEVIR 数据集上评估了 PreDiff 的降水强度临近预报任务。使用预期的降水强度作为知识控制，来模拟可能发生的暴雨和干旱等极端天气事件。

研究发现，结合预期未来降水强度的知识控制能有效引导生成，同时保持对真实数据分布的真实性和一致性。例如，下图的第三行模拟了在极端情况（概率约为 0.35%）下天气如何演变，该情况下未来的平均强度超过 μτ + 4στ。这种模拟对于估计极端暴雨情况下的潜在损失可能很有价值。

更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）