4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling

• URL: http://arxiv.org/abs/2609.12815v1
• Date: Fri, 11 Sep 2026 13:08:17 GMT
• Title: 4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
• Abstract summary: We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$circ$ global resolution.<n>We numerically 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.<n>This model achieves predictive skill scores competitive with state-of-the-art atmospheric AI models and numerical models, and can predict extreme events with exceptional skill.
• Score: 1.0277586860424235
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Explorative Modeling: Unlocking a Third Pretraining Axis and End-to-End Generation [89.51408170909548]
  We introduce Explorative Modeling, a new paradigm that factors the training loop, exploring K candidate matches between model generations and data, and training on the best.<n>First, increasing exploration adds a third pretraining axis beyond parameters and data for existing generative models.<n>Second, XMs enable end-to-end reconstructive generative modeling, matching diffusion on control tasks with 16-256x fewer inference steps.
  arXiv  Detail & Related papers  (2026-07-29T18:25:17Z)
• Perforated Neural Networks for Keyword Spotting [0.0]
  This paper presents the application of Perforated Backpropagation to keyword spotting on the Edge Impulse platform.<n>We demonstrate that dendritic models outperform traditional architectures at every level of parameter count.<n>The best dendritic model achieved a test accuracy of 0.933 with only 1,500 parameters, versus the baseline accuracy of 0.921 requiring approximately 4,000 parameters.
  arXiv  Detail & Related papers  (2026-05-15T06:02:19Z)
• Energy Scaling Laws for Diffusion Models: Quantifying Compute and Carbon Emissions in Image Generation [50.21021246855702]
  We propose an adaptation of Kaplan scaling laws to predict GPU energy consumption for diffusion models based on computational complexity (FLOPs)<n>Our approach decomposes diffusion model inference into text encoding, iterative denoising, and decoding components, with the hypothesis that denoising operations dominate energy consumption due to their repeated execution across multiple inference steps.<n>Our results validate the compute-bound nature of diffusion inference and provide a foundation for sustainable AI deployment planning and carbon footprint estimation.
  arXiv  Detail & Related papers  (2025-11-21T08:12:47Z)
• FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale [91.84761739154366]
  FourCastNet 3 advances global weather modeling by implementing a scalable, geometric machine learning (ML) approach to probabilistic ensemble forecasting.<n>FourCastNet 3 delivers forecasting accuracy that surpasses leading conventional ensemble models and rivals the best diffusion-based methods.<n>Its computational efficiency, medium-range probabilistic skill, spectral fidelity, and rollout stability at subseasonal timescales make it a strong candidate for improving meteorological forecasting and early warning systems through large ensemble predictions.
  arXiv  Detail & Related papers  (2025-07-16T11:22:18Z)
• Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism [0.08388591755871733]
  Accurate modeling of complex atmospheric dynamics requires large neural networks and gigabyte-sized data samples.<n>We introduce WeatherMixer, a multi-layer-perceptron-based architecture whose workload scales linearly with input size.<n>To cope with the computational demand, we propose Jigsaw, a novel model parallelization scheme.
  arXiv  Detail & Related papers  (2025-07-08T07:57:08Z)
• ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability [10.88886669820126]
  We introduce the Oak Ridge Base Foundation Model for Earth System Predictability (ORBIT) ORBIT is the largest model of its kind and surpasses the current climate AI foundation model size by a thousandfold. Performance scaling tests on the Frontier supercomputer have demonstrated that ORBIT achieves 684 petaFLOPS to 1.6 exaFLOPS sustained throughput.
  arXiv  Detail & Related papers  (2024-04-23T03:39:57Z)
• STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning [82.03481509373037]
  Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy in visual input environments. We introduce Transformer-based wORld Model (STORM), an efficient world model architecture that combines strong modeling and generation capabilities. Storm achieves a mean human performance of $126.7%$ on the Atari $100$k benchmark, setting a new record among state-of-the-art methods.
  arXiv  Detail & Related papers  (2023-10-14T16:42:02Z)
• Multi-fidelity Hierarchical Neural Processes [79.0284780825048]
  Multi-fidelity surrogate modeling reduces the computational cost by fusing different simulation outputs. We propose Multi-fidelity Hierarchical Neural Processes (MF-HNP), a unified neural latent variable model for multi-fidelity surrogate modeling. We evaluate MF-HNP on epidemiology and climate modeling tasks, achieving competitive performance in terms of accuracy and uncertainty estimation.
  arXiv  Detail & Related papers  (2022-06-10T04:54:13Z)
• ANNETTE: Accurate Neural Network Execution Time Estimation with Stacked Models [56.21470608621633]
  We propose a time estimation framework to decouple the architectural search from the target hardware. The proposed methodology extracts a set of models from micro- kernel and multi-layer benchmarks and generates a stacked model for mapping and network execution time estimation. We compare estimation accuracy and fidelity of the generated mixed models, statistical models with the roofline model, and a refined roofline model for evaluation.
  arXiv  Detail & Related papers  (2021-05-07T11:39:05Z)
• Using Bayesian deep learning approaches for uncertainty-aware building energy surrogate models [0.0]
  Machine learning surrogate models are trained to emulate slow, high-fidelity engineering simulation models. Deep learning models exist that follow the Bayesian paradigm. We show that errors can be reduced by up to 30% when the 10% of samples with the highest uncertainty are transferred to the high-fidelity model.
  arXiv  Detail & Related papers  (2020-10-05T15:04:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.