ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling

📄 arXiv: 2505.04802v2 📥 PDF

作者: Xiao Wang, Jong-Youl Choi, Takuya Kurihaya, Isaac Lyngaas, Hong-Jun Yoon, Xi Xiao, David Pugmire, Ming Fan, Nasik M. Nafi, Aristeidis Tsaris, Ashwin M. Aji, Maliha Hossain, Mohamed Wahib, Dali Wang, Peter Thornton, Prasanna Balaprakash, Moetasim Ashfaq, Dan Lu

分类: cs.LG, astro-ph.EP, cs.AI, cs.DC, physics.ao-ph

发布日期: 2025-05-07 (更新: 2025-09-02)

💡 一句话要点

ORBIT-2：通过可扩展的视觉基础模型实现天气和气候的降尺度

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 气候降尺度 视觉基础模型 Vision Transformer Reslim TILES 大规模并行 贝叶斯正则化

📋 核心要点

  1. 现有AI降尺度方法难以在不同变量和地理区域泛化，且受限于ViT自注意力的二次复杂度，阻碍了高分辨率气候建模。
  2. ORBIT-2通过Reslim轻量级架构和TILES分块序列缩放算法，实现了高效、鲁棒的全球气候降尺度，并显著降低了计算复杂度。
  3. ORBIT-2在65,536个GPU上实现了4.1 exaFLOPS的吞吐量，支持0.9公里分辨率的降尺度，并在7公里分辨率基准测试中取得了优异的精度。

📝 摘要（中文）

稀疏的观测数据和粗分辨率的气候模型限制了有效的区域决策，突显了对鲁棒降尺度的需求。然而，现有的AI方法在变量和地理位置上的泛化能力不足，并且受到Vision Transformer (ViT)自注意力二次复杂度的限制。我们提出了ORBIT-2，一个用于全球、超分辨率气候降尺度的可扩展基础模型。ORBIT-2包含两个关键创新：(1) Residual Slim ViT (Reslim)，一种具有残差学习和贝叶斯正则化的轻量级架构，用于高效、鲁棒的预测；(2) TILES，一种分块序列缩放算法，将自注意力复杂度从二次降低到线性，从而实现长序列处理和大规模并行。ORBIT-2可扩展到65,536个GPU上的100亿参数，实现高达4.1 exaFLOPS的持续吞吐量和74-98%的强扩展效率。它支持降尺度到0.9公里的全球分辨率，并处理高达42亿个token的序列。在7公里分辨率的基准测试中，ORBIT-2实现了高精度，相对于观测数据的R^2得分在0.98-0.99范围内。

🔬 方法详解

问题定义：论文旨在解决气候模型降尺度问题，即如何从粗分辨率的气候模型输出中推断出高分辨率的气候信息。现有AI方法，特别是基于Vision Transformer的方法，在处理长序列和高分辨率数据时面临计算复杂度高、泛化能力弱等问题。ViT的自注意力机制具有二次复杂度，限制了其在高分辨率气候数据上的应用。

核心思路：论文的核心思路是设计一种可扩展的视觉基础模型，该模型能够高效地处理长序列气候数据，并具有良好的泛化能力。通过引入轻量级架构Reslim和分块序列缩放算法TILES，降低计算复杂度，提高模型的可扩展性和训练效率。同时，利用残差学习和贝叶斯正则化增强模型的鲁棒性。

技术框架：ORBIT-2的整体框架包括数据预处理、模型训练和降尺度预测三个主要阶段。首先，对气候数据进行预处理，包括数据清洗、标准化等。然后，使用Reslim架构和TILES算法训练模型。最后，利用训练好的模型进行降尺度预测，将粗分辨率的气候数据转换为高分辨率的气候数据。该框架支持大规模并行计算，可以在多个GPU上进行高效训练。

关键创新：ORBIT-2的关键创新在于Reslim架构和TILES算法。Reslim是一种轻量级的ViT架构，通过减少参数数量和引入残差学习，提高了模型的训练效率和泛化能力。TILES算法通过将长序列数据分割成小块，并对每个小块进行自注意力计算，将自注意力复杂度从二次降低到线性，从而实现了长序列数据的处理。

关键设计：Reslim架构采用了深度可分离卷积和瓶颈结构，减少了参数数量。贝叶斯正则化用于防止过拟合。TILES算法的关键参数包括块大小和重叠区域大小，这些参数需要根据具体的气候数据集进行调整。损失函数采用均方误差（MSE）或其变体，用于衡量预测值与真实值之间的差异。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ORBIT-2在65,536个GPU上实现了4.1 exaFLOPS的持续吞吐量，强扩展效率达到74-98%。它支持降尺度到0.9公里的全球分辨率，并处理高达42亿个token的序列。在7公里分辨率的基准测试中，ORBIT-2实现了高精度，相对于观测数据的R^2得分在0.98-0.99范围内，表明其具有强大的降尺度能力。

🎯 应用场景

ORBIT-2可应用于气候变化研究、极端天气事件预测、农业规划、水资源管理等领域。通过提供高分辨率的气候信息，ORBIT-2可以帮助决策者更好地了解气候变化的影响，制定更有效的应对措施，并提高资源利用效率。未来，ORBIT-2有望成为气候变化研究的重要工具。

📄 摘要（原文）

Sparse observations and coarse-resolution climate models limit effective regional decision-making, underscoring the need for robust downscaling. However, existing AI methods struggle with generalization across variables and geographies and are constrained by the quadratic complexity of Vision Transformer (ViT) self-attention. We introduce ORBIT-2, a scalable foundation model for global, hyper-resolution climate downscaling. ORBIT-2 incorporates two key innovations: (1) Residual Slim ViT (Reslim), a lightweight architecture with residual learning and Bayesian regularization for efficient, robust prediction; and (2) TILES, a tile-wise sequence scaling algorithm that reduces self-attention complexity from quadratic to linear, enabling long-sequence processing and massive parallelism. ORBIT-2 scales to 10 billion parameters across 65,536 GPUs, achieving up to 4.1 exaFLOPS sustained throughput and 74--98% strong scaling efficiency. It supports downscaling to 0.9 km global resolution and processes sequences up to 4.2 billion tokens. On 7 km resolution benchmarks, ORBIT-2 achieves high accuracy with $R^2$ scores in the range of 0.98--0.99 against observational data.