深度学习在地球系统预测中的实用化
地球是一个复杂系统。从温度波动等常规事件到干旱、冰雹和厄尔尼诺-南方振荡现象等极端事件，这些变异性会影响作物产量、导致航班延误，并引发洪水和森林火灾。对这些变异性进行精确及时的预测，可以帮助人们采取必要预防措施避免危机，或更好地利用风能和太阳能等自然资源。
基于Transformer的模型在其他AI领域的成功，促使研究人员尝试将其应用于地球系统预测。但这些努力遇到了几个主要挑战。其中最突出的是地球系统数据的高维性：直接应用Transformer的二次复杂度注意力机制计算成本过高。
大多数现有的基于机器学习的地球系统模型也输出单一的点预测，这通常是可能结果范围的平均值。然而，有时了解极端天气事件发生的10%概率可能比了解可能结果范围内的一般平均值更重要。此外，典型的机器学习模型没有物理定律或历史先例的防护栏，可能产生不太可能甚至不可能的预测结果。
Earthformer与立方体注意力
Transformer模型的核心是其"注意力机制"，使其在处理输出序列的每个元素时能够权衡输入序列不同部分的重要性。这种机制允许Transformer捕捉数据中时空长程依赖关系和关系，这是传统卷积神经网络或循环神经网络架构未能很好建模的。
然而，地球系统数据本质上是高维且时空复杂的。在我们NeurIPS 2022论文研究的SEVIR数据集中，每个数据序列包含25帧以5分钟间隔捕获的数据，每帧的空间分辨率为384 x 384像素。使用传统的Transformer注意力机制处理如此高维数据将极其昂贵。
在我们的NeurIPS 2022论文中，我们提出了一种名为立方体注意力的新型注意力机制，它将输入张量分解为立方体（或高维类似立方体的结构），并在每个立方体级别应用注意力。由于注意力的计算成本随张量大小呈二次方增长，在每个立方体内局部应用注意力比尝试一次性计算整个张量的注意力权重在计算上更容易处理。例如，沿时间轴分解可以使SEVIR数据集的成本降低384^2倍，因为每帧的空间分辨率为384 x 384像素。
当然，这种分解引入了一个限制：注意力在每个立方体内独立运行，立方体之间没有通信。为了解决这个问题，我们还计算了总结立方体注意力权重的全局向量。其他立方体可以将这些全局向量纳入它们自己的注意力权重计算中。
我们将采用立方体注意力的基于Transformer的模型称为Earthformer。Earthformer采用分层编码器-解码器架构，逐渐将输入序列编码为多级表示，并通过从粗到细的过程生成预测。每个层次包括一堆立方体注意力块。通过堆叠具有不同配置的多个立方体注意力层，我们能够有效探索有效的时空注意力。
我们实验了多种将输入张量分解为立方体的方法。我们的实证研究表明，“轴向"模式（沿时间、高度和宽度轴堆叠三个未平移的局部分解）既有效又高效。它在避免原始注意力指数计算成本的同时实现了最佳性能。
实验结果
为了评估Earthformer，我们在两个真实世界数据集上将其与六个最先进的时空预测模型进行比较：SEVIR，用于在不久的将来连续预测降水概率（“临近预报”）的任务；ICAR-ENSO，用于预测海表温度异常。
在SEVIR上，我们使用的评估指标是标准均方误差和临界成功指数，这是降水临近预报评估中的标准指标。CSI也称为交并比：在不同阈值下，它表示为CSI-thresh；它们的平均值表示为CSI-M。
在MSE和CSI上，Earthformer在所有六个基线模型中都表现优异。带有全局向量的Earthformer版本也一致优于没有全局向量的版本。
| 模型 | #参数(M) | GFLOPS | 指标 |
|---|---|---|---|
| Earthformer | 15.1 | 257 | CSI-M: 0.4419↑ |
在ICAR-ENSO上，我们报告了三个月移动平均Nino3.4指数的相关技能，该指数评估太平洋特定区域(170°-120°W, 5°S-5°N)海表温度异常预测的准确性。Earthformer在所有相关评估指标中一致优于基线，使用全局向量的版本进一步提高了性能。
| 模型 | #参数(M) | GFLOPS | 指标 |
|---|---|---|---|
| Earthformer | 7.6 | 23.9 | C-Nino3.4-M: 0.7329↑ |
PreDiff
扩散模型最近已成为许多AI任务的领先方法。扩散模型是生成模型，建立了一个向训练样本迭代添加高斯噪声的前向过程；然后模型学习在反向扩散过程中逐步移除添加的噪声，逐渐降低噪声水平，最终产生清晰高质量的生成结果。
在训练期间，模型学习其逐步执行的每个去噪步骤之间的一系列转移概率。因此，它本质上是一个概率模型，非常适合概率预测。
扩散模型的一个最新变体是潜在扩散模型：在传递到扩散模型之前，输入首先馈送到自动编码器，该编码器具有产生压缩嵌入（数据表示）的瓶颈层；然后扩散模型应用于压缩空间。
在我们即将发表的NeurIPS论文"PreDiff: Precipitation nowcasting with latent diffusion models"中，我们提出了PreDiff，这是一个使用Earthformer作为其核心神经网络架构的潜在扩散模型。
通过修改训练模型的转移概率，我们可以对模型输出施加约束，使其更可能符合某些先验知识。我们通过简单地移动学习分布的均值来实现这一点，直到它更好地符合我们希望施加的约束。
结果
我们在SEVIR上评估了PreDiff在预测近期降水强度（“临近预报”）的任务。我们使用预期的降水强度作为知识控制来模拟可能的极端天气事件，如暴雨和干旱。
我们发现，具有预期未来降水强度的知识控制有效地指导生成，同时保持保真度和对真实数据分布的遵守。例如，下图的第三行模拟了在极端情况（概率约为0.35%）下天气如何展开，其中未来平均强度超过μτ + 4στ。这种模拟对于估计极端暴雨情况下的潜在损害可能有价值。