aiqbench analysis
This research introduces BEAST, a Bayesian Swin Transformer that leverages an innovative 4D parallelization strategy to train massive models for high-fidelity atmospheric forecasting. By efficiently scaling across thousands of GPUs, the system achieves significant speedups in generating uncertainty-quantified ensembles while maintaining predictive skill competitive with traditional numerical models. This advancement marks a critical step toward operationalizing probabilistic AI for climate science and extreme event prediction.
4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
arXiv:2609.12815v1 Announce Type: cross Abstract: We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$^\circ$ global resolution able to accurately quantify both aleatoric and epistemic uncertainty. To overcome the associated computational bottlenecks, we devise an orthogonal 4D-parallelization scheme that introduces a unique domain-tensor-parallelism strategy and a novel uncertainty parallel method, enabling us to fully leverage GPU capacity and efficiently scale model training. For a 2.4-billion-parameter model, we achieve a peak performance of 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. We train BEAST as a 700-million-parameter model with 96 random weight samples on 384 nodes on 40 years of data for nearly one million gradient updates. This model achieves predictive skill scores competitive with state-of-the-art probabilistic atmospheric AI models and numerical models, and can predict extreme events with exceptional skill, while generating large ensembles 3 to 4 times faster than the current-best AI model. Our contribution unlocks the potential of high-fidelity uncertainty quantification in atmospheric AI models, heralding a new era for AI-based models in climate and Earth system sciences.
Abstract:We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$^\circ$ global resolution able to accurately quantify both aleatoric and epistemic uncertainty. To overcome the associated computational bottlenecks, we devise an orthogonal 4D-parallelization scheme that introduces a unique domain-tensor-parallelism strategy and a novel uncertainty parallel method, enabling us to fully leverage GPU capacity and efficiently scale model training. For a 2.4-billion-parameter model, we achieve a peak performance of 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. We train BEAST as a 700-million-parameter model with 96 random weight samples on 384 nodes on 40 years of data for nearly one million gradient updates. This model achieves predictive skill scores competitive with state-of-the-art probabilistic atmospheric AI models and numerical models, and can predict extreme events with exceptional skill, while generating large ensembles 3 to 4 times faster than the current-best AI model. Our contribution unlocks the potential of high-fidelity uncertainty quantification in atmospheric AI models, heralding a new era for AI-based models in climate and Earth system sciences.
Atmospheric and Oceanic Physics (physics.ao-ph); Artificial Intelligence (cs.AI); Distributed, Parallel, and Cluster Computing (cs.DC); Machine Learning (cs.LG); Performance (cs.PF)
arXiv:2609.12815 [physics.ao-ph]
(or arXiv:2609.12815v1 [physics.ao-ph] for this version)
https://doi.org/10.48550/arXiv.2609.12815 arXiv-issued DOI via DataCite (pending registration)
Full text under the publisher's open licence · 205 words
Related coverage
Humanity’s Last Invention — Richard Socher of Recursive
Latent Space · Sep 14, 2026
Researchers Estimate Fidelity and Save Computation with New Quantum Scheduling System
Quantum Zeitgeist · Sep 10, 2026
Does preservation make sense before we know how to revive?
LessWrong · Jun 15, 2026
Generalization and infinite width
LessWrong · Aug 1, 2026