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

arXiv: https://arxiv.org/abs/2505.04802 · HTML: https://mlsanity.com/paper/2505.04802/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 55 references that are in the map.
- First posted on arXiv 7 May 2025.

## Builds on (most cited first, 23 of 23)

- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [MaxViT: Multi-Axis Vision Transformer](https://mlsanity.com/paper/2204.01697/) (2022)
- [SRDiff: Single Image Super-Resolution with Diffusion Probabilistic Models](https://mlsanity.com/paper/2104.14951/) (2021)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [FLatten Transformer: Vision Transformer using Focused Linear Attention](https://mlsanity.com/paper/2308.00442/) (2023)
- [Transformer for Single Image Super-Resolution](https://mlsanity.com/paper/2108.11084/) (2021)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [ClimateLearn: Benchmarking Machine Learning for Weather and Climate Modeling](https://mlsanity.com/paper/2307.01909/) (2023)
- [Prithvi WxC: Foundation Model for Weather and Climate](https://mlsanity.com/paper/2409.13598/) (2024)
- [torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models](https://mlsanity.com/paper/2004.09910/) (2020)
- [PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers](https://mlsanity.com/paper/2102.03161/) (2021)
- [Conditional diffusion models for downscaling and bias correction of Earth system model precipitation](https://mlsanity.com/paper/2404.14416/) (2024)
- [ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability](https://mlsanity.com/paper/2404.14712/) (2024)
- [Algorithmic Hallucinations of Near-Surface Winds: Statistical Downscaling with Generative Adversarial Networks to Convection-Permitting Scales](https://mlsanity.com/paper/2302.08720/) (2023)
- [Deep-learning based down-scaling of summer monsoon rainfall data over Indian region](https://mlsanity.com/paper/2011.11313/) (2020)
- [Adaptive Patching for High-resolution Image Segmentation with Transformers](https://mlsanity.com/paper/2404.09707/) (2024)
- [Ultra-Long Sequence Distributed Transformer](https://mlsanity.com/paper/2311.02382/) (2023)
- [Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier](https://mlsanity.com/paper/2405.15780/) (2024)

## Built on by (most cited first, 10 of 10)

- [On Neural Scaling Laws for Weather Emulation through Continual Training](https://mlsanity.com/paper/2603.25687/) (2026)
- [Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction](https://mlsanity.com/paper/2604.16590/) (2026)
- [Towards Scaling Law Analysis For Spatiotemporal Weather Data](https://mlsanity.com/paper/2604.05068/) (2026)
- [Transforming the Use of Earth Observation Data: Exascale Training of a Generative Compression Model with Historical Priors for up to 10,000x Data Reduction](https://mlsanity.com/paper/2605.08633/) (2026)
- [Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction](https://mlsanity.com/paper/2607.02829/) (2026)
- [TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling](https://mlsanity.com/paper/2608.15211/) (2026)
- [Precipitation Downscaling Using Foundation Model-Conditioned Diffusion](https://mlsanity.com/paper/2608.25858/) (2026)
- [4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling](https://mlsanity.com/paper/2609.12815/) (2026)
- [Structure-Guided Masked Autoencoders for Ultra-High Resolution Scientific Image Understanding](https://mlsanity.com/paper/2609.30682/) (2026)
- [Scalable In-Domain Self-Supervised Foundation Model for Dense Representation Transfer in High-Resolution Plant Imaging](https://mlsanity.com/paper/2609.32183/) (2026)

## Abstract (arXiv)

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 R2 scores in the range of 0.98--0.99 against observational data.