数据科学

从 TB 级到一站式解决方案:AI 驱动的气候模型走向主流

在了解地球不断变化的气候的竞赛中,速度和准确性至关重要。但当今使用最广泛的气候模拟器往往难以满足需求:由于计算能力的限制,它们无法完全捕捉雷暴或高耸的热带云等关键的小规模过程。

为了捕捉这些特征,科学家们运行超高分辨率模拟,称为云解析模型 (CRM) 。这些模拟跟踪云层的形成和演变过程,但成本高昂,在十年的全球气候预测中运行一个模拟几乎是不可能的。

如果我们能够将这些详细模拟的智慧融入机器学习模型,在不放弃保真度的情况下将运行速度提高数十到数百倍,情况会怎样?

这正是 ClimSim-OnlineClimSim-Online 的承诺,ClimSim – Online 是一个可复制的框架,用于大规模开发和部署混合物理 – 机器学习气候模型。该框架由 NVIDIA Earth 2 以及由政府和学术界的国际气候建模者组成的联盟制作。该计划由位于哥伦比亚大学、由美国国家科学基金会资助的科技中心发起并提供支持,该中心正在探索 AI 驱动的气候模拟技术的未来。

从 TB 级到一站式:训练 AI 以模拟复杂的嵌套气候物理学

ClimSim-Online 基于 NeurIPS 2023 上推出的屡获殊荣的 ClimSimClimSim 数据集构建。该数据集提供给 ClimSim Hugging Face 存储库。此数据集使用 Energy Exascale 地球系统模型 – 多尺度建模框架 (E3SM-MMF) 创建。E3SM – MMF 是新一代气候模拟器,可在主机粗网格气候模型的每个大气列中嵌入数千个局部的计算密集型 CRM。这是一种生成气候预测的实验性方法,可减少精细尺度物理学通常必须作出的假设数量,但其计算成本极高,未用于主流国际预测。将嵌套物理特性外包给 AI 可能会改变这种状况。

主机气候模型以大约 1.5 度 (约 150 公里) 或更粗糙的水平分辨率运行,而每个嵌入式 CRM 以 2 公里的分辨率运行,以更精细的比例明确模拟云层和对流。

在模拟的 10 年时间里,E3SM-MMF 产生了 57 亿个惊人的样本,每个样本都描述了小规模物理过程如何改变大规模大气状态。这些过程包括流上升气流如何导致云形成,是什么导致微物理液滴形成,对流如何从单个云的尺度组织到大型有组织的云复合体,以及这些云系统如何与太阳和红外辐射相互作用,从而调节气候。

这个庞大的数据集可用作训练模拟子网格物理的机器学习模型的基础,并且可以取代昂贵的嵌入式 CRM,后者消耗约占总计算支出的 95%。这已催生了一场全球 Kaggle 竞赛Kaggle 竞赛,吸引了来自世界各地的 460 多个团队基于此高保真度气候数据集开发 ML 解决方案并对其进行基准测试,这有助于通过开放式协作创新加速进展。

A schematic diagram of the ClimSim dataset. The diagram shows the input variables on the left, which consist of a set of macro-scale state variables. On the right, the diagram displays the target variables, which primarily include the tendencies of those state variables due to unresolved processes.
图 1。ClimSim 数据集和基础机器学习问题的示意图。输入由一组宏观尺度状态变量组成;目标主要包括这些状态变量因未解决的进程而产生的趋势。

挑战?这些模型需要的不仅仅是离线的精确度。当集成到实时气候模拟器中 (年复一年地小时复一地运行) 时,它们必须保持稳定,而不会让虚拟大气漂入不切实际的状态。控制混合物理机器学习仿真的行为是一项巨大的挑战,尤其是在主机物理模型无法可微的情况下。一些简单的主机模型可以用可微代码重写,从而直接实现混合动力学的 ML 优化。但是,许多候选主机模型不容易以不同方式重写,或者是非常非线性,因此直接优化混合行为是不切实际的。涵盖数百万行源代码的功能齐全的气候模拟器就是一个很好的例子。

插入并模拟

ClimSim-Online 由 NVIDIA 率先推出,旨在为更广泛的 ML 社区提供混合气候建模。我们构建了一个可复制的容器化工作流,以避开运行功能齐全的气候模拟器的典型障碍 (例如对某些超级计算机和软件环境的依赖) ,这些障碍限制了社区与之交互的能力。用户只需一个 TorchScript 模型文件,即可将其经过训练的机器学习模型注入基于 Fortran 的 E3SM 气候模拟器,并启动混合仿真。无论是在本地工作站、HPC 集群还是云 VM 上,他们都可以插入标准化诊断来衡量成功与否。

A flowchart diagram of the E3SM-MMF simulation workflow. The diagram shows various modules represented as blocks, each executed during one model time step. In standard simulations, the 'physics-before-coupling' block includes the embedded Cloud Resolving Model (CRM). In hybrid mode, this block is replaced by a Machine Learning (ML) emulator. The entire workflow is containerized for easy deployment across different platforms.
图 2。E3SM-MMF 仿真工作流程图。每个块代表在一个模型时间步长期间执行的一个模块。在标准仿真中,“物理前合”模块包括嵌入式 CRM。在混合模式下,此模块将被 ML 模拟器取代。整个工作流已容器化,可轻松跨平台部署。

气候模拟,即插即用

整个系统在预加载了所有必需库和依赖项的容器中运行。只需加载、安装和模拟即可。用户可以在 ClimSim-Online 存储库中找到设置容器的说明。从访问数据、训练机器学习模型到运行和评估混合气候模拟,整个工作流程都可以在 ClimSim 存储库中找到。

突破性进展:稳定运行多年,逼真适应对流层顶

NVIDIA 技术研发组织的科学家现已使用这些新的 API 取得重要突破。在 7 月 10 日发表在《 Journal of Advances in Modeling Earth Systems》 (JAMES) 上的最新论文中,我们展示了使用 PhysicsNemo 在 ClimSim 数据集上训练的 U-Net 神经网络进行的多年稳定混合模拟。它在 ClimSim-Online 中建立了新的在线技能基准。PhysicsNemo 是一个开源深度学习框架,使用户能够探索、开发、验证和部署先进的科学和工程方法,将基于物理的知识与数据相结合。

真正的突破是什么呢?基于物理信息的机器学习

为了避免失控的模拟和不切实际的云行为,我们直接在神经网络架构中构建了微物理约束

  • 所有凝聚都遵循基于温度的相位划分,就像神经网络正在模拟的云解析模型一样。
  • 对流层顶上没有残留的冰云。

有了这些严格的限制,我们稳定了之前的漂移模拟,并大幅提高了云气候学的真实感,尤其是在热带地区,无约束的模型往往会过高估价高纬度的云层。

ClimSim-Online 从根本上加速了导致此解决方案的研究过程:能够快速迭代不断演变的下游混合模型病理学是找到最终为我们的科学侦探工作提供依据的线索的关键。

Two side-by-side line graphs show how different machine learning models perform over 12-month hybrid climate simulations. The left graph shows temperature error (in Kelvin), and the right graph shows moisture error (in grams per kilogram). Each line represents a different model. In both graphs, red lines represent U-Net models with physical constraints. These models have lower and more stable errors throughout the year. Cyan lines represent unconstrained U-Net models with higher and growing error. Blue lines show unconstrained MLP models, which become unstable early. Dashed black lines represent the internal unpredictability limit of the atmosphere—these are the lowest achievable errors, and the constrained models get close to them. The X-axis in both graphs shows months from one to 12.
图 3。一年模拟中每月全球平均的在线温度和湿度均方根误差。红线表示具有微物理学限制的混合模型,接近内部大气不可预测性的理论下限。蓝色和青色线条表示无限制 MLP 和 U-Net 基准的错误增长不稳定。

在我们的混合模拟中,我们观察到对流层内的温度偏差保持在 2 摄氏度以下,湿度偏差保持在每千克 1 克以下,这是多尺度建模框架下的一项新的先进结果。

我们看到了稳定 > 五年的模拟,其中包括显式云凝聚建模、真实地理和陆地 – 大气合,这是以前在此类混合模拟中没有展示过的里程碑。

准备起飞

ClimSim-Online 降低了 AI – 气候协作的门槛。您可以轻松完成以下操作:

  • 使用出色的模拟数据训练 ML 模型
  • 离线技能基准测试
  • 最重要的是,评估完整气候模拟器中的在线性能,这是对现实世界准备情况的终极测试。

无论您是渴望从事气候研究的 AI 研究人员,还是对混合建模的强大功能感到好奇的气候科学家,ClimSim-Online 都能为您提供加入下一波气候模拟浪潮的工具。

虽然我们已经展示了一种基于领域科学信息的方法来解决混合建模的一阶问题,但要将混合偏差降低到真正可以容忍的水平,仍有许多工作要做。我们需要新的想法。例如:强化学习社区能否找到与领域科学无关的更可靠的解决方案?现在,ClimSim-Online 可以轻松地对下游、不可微分的奖励信号进行采样,也许我们很快就会发现。混合物理 – 机器学习气候模拟的未来还在等着您。

 

标签