中文

ORBIT-2:面向气象和气候下游的超规模视觉基础模型

机器学习 2025-09-03 v2 地球与行星天体物理 人工智能 分布式、并行与集群计算 大气与海洋物理

摘要

稀疏观测和粗糙分辨率的气候模型限制了有效的区域决策,凸显了鲁棒性下游分辨率的需求。然而,现有的AI方法在变量和地理空间上的泛化性方面存在挑战,并且受制于视觉Transformer(ViT)自注意力的二次复杂度。我们引入ORBIT-2,这是一个用于全球、超高分辨率气候下游的可扩展基础模型。ORBIT-2包含两个关键创新:(1)残差精简ViT(Reslim),一种具有残差学习和贝叶斯正则化的轻量级架构,用于高效、稳健的预测;(2)TILES,一种基于切片的序列缩放算法,将自注意力复杂度从二次降低到线性,从而实现长序列处理和大规模并行。ORBIT-2在65,536个GPU上扩展到100亿参数,实现了最高可达4.1 exaFLOPS的持续吞吐量和74%--98%的强扩展效率。它支持将下游分辨率降低到0.9公里全球分辨率,并处理最长达42亿个token的序列。在7公里分辨率的基准测试中,ORBIT-2以0.98--0.99的R²分数实现了高精度,与观测数据相符。

关键词

引用

@article{arxiv.2505.04802,
  title  = {ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling},
  author = {Xiao Wang and Jong-Youl Choi and Takuya Kurihaya and Isaac Lyngaas and Hong-Jun Yoon and Xi Xiao and David Pugmire and Ming Fan and Nasik M. Nafi and Aristeidis Tsaris and Ashwin M. Aji and Maliha Hossain and Mohamed Wahib and Dali Wang and Peter Thornton and Prasanna Balaprakash and Moetasim Ashfaq and Dan Lu},
  journal= {arXiv preprint arXiv:2505.04802},
  year   = {2025}
}