深度学习在地球系统预测中的实践突破

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

基于Transformer的模型在其他AI领域的成功，促使研究人员尝试将其应用于地球系统预测。但这些努力面临几个主要挑战：首先是地球系统数据的高维性——直接应用Transformer的二次复杂度注意力机制计算成本过高；其次，现有机器学习模型多输出单点预测，无法提供概率分布；最后，典型机器学习模型缺乏物理定律或历史先例的约束，可能产生不合理输出。

Earthformer与立方体注意力

在2022年NeurIPS会议上发表的论文《Earthformer: Exploring space-time transformers for Earth system forecasting》中，研究团队提出名为"立方体注意力"的新型机制。该机制将输入张量分解为立方体（高维立方体类比），并在每个立方体级别应用注意力计算。由于注意力计算成本随张量大小呈平方增长，局部立方体计算大幅降低了计算复杂度。例如在SEVIR数据集上，沿时间轴分解可使计算成本降低384²倍。

为解决立方体间缺乏通信的问题，模型还计算全局向量来汇总各立方体的注意力权重，其他立方体可将全局向量纳入自身注意力计算。

Earthformer采用分层编码器-解码器架构，逐步将输入序列编码为多级表示，并通过粗到细的过程生成预测。每个层级包含堆叠的立方体注意力块，通过不同配置的多层堆叠实现高效时空注意力建模。

实验显示，“轴向"分解模式（沿时间、高度和宽度轴堆叠三个未平移的局部分解）在效果和效率上表现最佳。

PreDiff：概率预测与物理约束

在即将发表于NeurIPS 2023的《PreDiff: Precipitation nowcasting with latent diffusion models》中，团队提出PreDiff模型。该模型基于潜在扩散模型，使用Earthformer作为核心架构。

扩散模型通过前向过程逐步添加高斯噪声，再通过反向扩散过程逐步去噪实现数据生成。这种内在概率特性非常适合概率预测。潜在扩散模型首先通过自编码器将输入压缩为潜在表示，然后在压缩空间中进行扩散过程。

通过修改训练模型的转移概率，可以对输出施加约束，使其更符合先验知识。该方法通过简单平移学习分布的均值来实现约束合规。

实验结果

在SEVIR数据集上的降水临近预报任务中，Earthformer在均方误差（MSE）和临界成功指数（CSI）上全面优于六种基线模型。带全局向量的版本性能进一步提升。

在ICAR-ENSO数据集的海表温度异常预测任务中，Earthformer在所有评估指标上均优于基线模型，使用全局向量后性能进一步改善。

PreDiff在SEVIR数据集上的测试表明，基于预期降水强度的知识控制能有效指导生成，同时保持对真实数据分布的忠实度。该能力对极端暴雨情况下的潜在损害评估具有重要价值。