Paper deep dive

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

Deifilia Kieckhefen, Juan Pedro Gutiérrez Hermosillo Muriedas, Lars Helge Heyen, Mathis Bode, Iida Hakulinen, Andreas Herten, Chelsea Maria John, Thorsten Kurth, Anni Moisala, Asena Karolin Özdemir, Kaleb Phipps, Oskar Taubert, Arvid Weyrauch, Markus Götz, Charlotte Debus

Year: 2026Venue: arXiv preprintArea: physics.ao-phType: PreprintEmbeddings: 83

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%

Last extracted: 9/19/2026, 3:05:59 AM

Summary

The paper introduces BEAST, a Bayesian Swin Transformer for atmospheric forecasting at 0.25° global resolution, capable of quantifying both aleatoric and epistemic uncertainty. To address computational bottlenecks, the authors propose a 4D parallelization scheme combining Domain-Tensor Parallelism (DTP) and Uncertainty Parallelism (UP) alongside Data Parallelism (DP). This approach enables training a 2.4-billion-parameter model on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer, achieving 3.96 EFLOP/s. The model demonstrates competitive predictive skill and generates ensembles 3-4 times faster than current best AI models.

Entities (9)

BEAST · model · 98%JUPITER supercomputer · hardware · 95%NVIDIA GH200 · hardware · 95%Domain-Tensor Parallelism · method · 92%Uncertainty Parallelism · method · 92%Bayesian Swin Transformer · model-architecture · 90%Aleatoric uncertainty · concept · 85%Epistemic Uncertainty · concept · 85%ERA5 · dataset · 80%

Relation Signals (6)

BEAST → runson → JUPITER supercomputer

confidence 95% · we achieve a peak performance of 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer.

BEAST → uses → Domain-Tensor Parallelism

confidence 95% · We leverage distributed matrix multiplications to implement this domain–tensor parallelization (DTP) scheme... We apply this concept to all neural network layer types in BEAST.

BEAST → uses → Uncertainty Parallelism

confidence 95% · Next, we target the challenge of evaluating the m forward passes... by introducing sampling parallelism... We extend the GPU grid... by an orthogonal parallelization axis... That means that we have u DTP instances... This allows us to reduce the communication

BEAST → quantifies → Aleatoric uncertainty

confidence 90% · BEAST... able to accurately quantify both aleatoric and epistemic uncertainty.

BEAST → quantifies → Epistemic Uncertainty

confidence 90% · BEAST... able to accurately quantify both aleatoric and epistemic uncertainty.

BEAST → trainedon → ERA5

confidence 85% · Data parallelism (DP) processes separate, non-overlapping ERA5 data chunks in parallel.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

79,621 characters extracted from source content.

Expand or collapse full text

4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling Thanks: †These authors contributed equally. Corresponding author: C. Debus (charlotte.debus@kit.edu) Deifilia Kieckhefen11†, Juan Pedro Gutiérrez Hermosillo Muriedas11†, Lars Helge Heyen11, Mathis Bode22, Iida Hakulinen33, Andreas Herten22, Chelsea Maria John22, Thorsten Kurth44, Anni Moisala33, Asena Karolin Özdemir11, Kaleb Phipps11, Oskar Taubert33, Arvid Weyrauch11, Markus Götz11, and Charlotte Debus11 Affiliation: Affiliation: 11Karlsruhe Institute for Technology, Germany Affiliation: Affiliation: 22Forschungszentrum Jülich, Germany Affiliation: Affiliation: 33CSC – IT Center for Science Ltd., Finnland Affiliation: Affiliation: 44NVIDIA, Switzerland Affiliation: Abstract We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25° 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.42.4-billion-parameter model, we achieve a peak performance of 3.96 EFLOP/s3.96\,EFLOP/s on 20 48020\,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. We train BEAST as a 700700-million-parameter model with 9696 random weight samples on 384384 nodes on 4040 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 33 to 44 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. Index Terms: Atmospheric Forecasting, Uncertainty Quantification, Exascale Computing, Artificial Intelligence, Domain Parallelism, Sampling Parallelism I Overview Artificial intelligence (AI) models for atmospheric forecasting are revolutionizing weather and climate sciences, owing to their promise of generating predictions with skill scores competitive to numerical models at a fraction of the computational resources. Current atmospheric AI models learn a set of point-value (scalar) weights through iterative optimization over a finite dataset. However, this approach disregards the fact that the results of the optimization process can be ambiguous, meaning that the weights could also have different numeric values. From a modeling point of view, this corresponds to the epistemic uncertainty of the model – a systematic error originating from the incompleteness of the training datasets. The uncertainty in the model weights is dependent on the size and quality of the dataset: the more accurately the training dataset represents the true data distribution of the modeled problem, the more certain the weights are and the better the prediction becomes. For atmospheric AI models, this has two direct consequences. For one, the prediction of rare weather phenomena, such as extreme events, is challenging, as they are sparsely represented in the training data, or not at all. Second, the prediction quality is directly coupled to the size of the dataset [1], as more data items typically provide a better approximation of the true data distribution. Hence, modeling rare physical processes of Earth system components for which only small-sized datasets are available remains difficult. After the recent successes in AI-based weather forecasting with deterministic large-scale neural networks [2, 3, 4, 5], probabilistic models for quantifying the uncertainty of a prediction have become the next frontier. Current approaches to uncertainty quantification in AI models focus on the uncertainty of the atmospheric input state to generate ensemble forecasts, similar to ensemble generation via input perturbation in numerical models, as used for HENS [6, 7]. In particular, diffusion-based methods like GenCast [8], functional generative networks (FGN) [9], FourCastNet3 (FCN3) [10], and AIFS-CRPS [11] have achieved impressive results. However, these existing approaches capture only the so-called aleatoric uncertainty, which stems from the inherent chaotic nature of the physical processes in the atmosphere. In contrast, modeling epistemic uncertainty in atmospheric AI models has received little attention so far. FGN and HENS approximate Bayesian inference with deep ensembles [12, 13], i.e., training multiple models from independently randomly initialized parameters; however, a probabilistic model formulation that captures both epistemic and aleatoric uncertainty has not yet been realized. Bayesian neural networks (BNNs) provide a principled framework for modeling epistemic uncertainty. BNNs extend classical neural networks by replacing the scalar weights with probability distributions, thereby transforming point predictions into predictive distributions. Compared to alternative methods, like input perturbation or diffusion, BNNs offer a more comprehensive and theoretically grounded approach to uncertainty estimation, leading to better-calibrated and more informative ensemble predictions [14]. The distributional parameterization of BNN weights allows for distribution tails to be better captured, leading to more accurate predictions of extreme events. Moreover, BNNs can be trained to competitive prediction accuracy on much less data compared to their deterministic (non-Bayesian) counterparts, owing to their capability of accounting for the under-representation of the data distribution in small training datasets [15]. Hence, they have great potential for AI-based modeling of other components of the Earth system, where training data is much scarcer. However, training BNNs requires significantly more computational resources than training standard neural networks. For one, the probabilistic parameterization of network weights effectively doubles the number of parameters that must be learned, increasing the GPU memory requirements of model training. Second, each forward pass on a data batch has to be performed on a set of random weight samples drawn from the current weight distribution, and the resulting predictions, i.e., the ensemble, have to be aggregated for the backward pass. Evaluation of multiple samples per data item linearly increases computational cost and memory demand. For large-scale atmospheric models with hundreds of millions to billions of parameters, these costs become prohibitive. As a result, large-scale implementations of BNNs in this domain have not yet been realized to date. We set out to break this barrier and realize the first-ever Bayesian model for atmospheric forecasting at global 0.25° resolution. We strategically accelerate model training at all scales, from optimizing single-node performance by leveraging lower precision data types, just-in-time compilation, and fast GPU-to-GPU communication, to minimizing inter-node communication necessary for Bayesian model training. This allows training of a 2.42.4-billion-parameter Bayesian Shifted Window (Swin) Transformer for 12 h12\,h and 24 h24\,h forecasts at a peak performance of 3.96 EFLOP/s3.96\,EFLOP/s on 51205120 nodes on the JUPITER supercomputer. I State of the Art Regardless of the application, parallelism and scalability are paramount to enabling the training of large-scale models. The conventional way of accelerating training is data parallelism (DP) [16, 17], where each process holds a copy of the model weights, independently performs a forward–backward pass of its local subset of the dataset, and computes the corresponding gradients. The gradients are then averaged through an allreduce operation, and the weights are updated locally. The number of iterations required for an epoch can be reduced by increasing the number of data-parallel processes to increase the global batch size. While data parallelism scales well in a computational manner, having large global batch sizes can cause models to generalize poorly, ultimately limiting scalability [18]. Recent work on AI models in weather and climate science has focused extensively on so-called foundation models, i.e., very large models that are pre-trained in a self-supervised fashion on large heterogeneous datasets, and subsequently fine-tuned for downstream tasks [19, 20, 21]. In consequence, model sizes have been growing dramatically beyond the billion-parameter mark. The affirmation of neural scaling laws for AI weather models [1], stating that predictive skill improves with quadratically more data and weights, further fuels this trend. As growing model sizes quickly exceed the memory limitations of a single GPU, the field of model parallelism has emerged to distribute model weights and associated computations across multiple GPUs. There exist multiple forms of model parallelism, as there are different ways (axes) to parallelize the model architecture. The simplest form of model parallelism is pipeline parallelism, where sequential stages, i.e., blocks of layers, are distributed across GPUs [22]. While pipeline parallelism circumvents the GPU-memory bottleneck and entails simple communication schemes with only point-to-point operations, the inherent sequentiality of running through the network layers in the forward-backward pass causes GPU idle time, precluding efficient GPU utilization. Another drawback of pipeline parallelism is the necessity to perform asynchronous gradient updates, which degrades predictive performance [23]. An orthogonal approach to pipeline parallelism is tensor parallelism, i.e., distributing the weight matrices of each layer across multiple GPUs. Megatron-LM [24] was among the first works to utilize tensor parallelism, efficiently enabling the training of multi-billion-parameter Transformer models by parallelizing across attention heads and feed-forward weight matrices. In Transformer models specifically, the quadratic complexity in sequence length of the attention mechanism provides a particular computational burden and memory bottleneck for long sequence lengths. This is oftentimes addressed with sequence parallelism [25]. Next to the memory demand of the large weight matrices themselves, the associated gradients and optimizer states can also lead to a severe memory bottleneck. To deal with this, hybrid data–model-parallel approaches have been introduced, such as Zero Redundancy Optimizer (ZeRO) [26] or Fully Sharded Data Parallelism (FSDP) [27]. ZeRO partitions the parameters, activations, and optimizer states across GPUs, and uses allgather operations to reconstruct each layer upon demand. The approach achieves superlinear scaling due to reduced memory footprints. FSDP shards data batches and model parameters. In a given layer, the weights are gathered to compute the full activations. Afterward, the unnecessary activations are discarded to reduce the memory footprint. The main memory bottleneck for training large language models (LLMs) originates from large weight matrices; in contrast, Vision-Transformer (ViT)-like architectures with long sequence lengths (≈106≈ 10^6) face a different challenge: large activation memory footprints. While this can be mitigated by techniques such as gradient offloading to the CPU [28, 26] or gradient recomputation [29], these strategies can come at the cost of computational performance. Thus, novel parallelization strategies are required to address the significant bottlenecks arising from the size of the data items. As one of the first works in the field, FourCastNet v1 [4] employed tensor parallelism on its Adaptive Fourier Neural Operator (AFNO) architecture by applying feature parallelism, splitting the fast Fourier transforms across processes. It also distributes the multi-layer perceptron (MLP) at the end of an AFNO block, sharding model parameters in alternating row and column dimensions. The approach enables scaling up to 38083808 NVIDIA A100 GPUs, attaining 140.8 PFLOP/s140.8\,PFLOP/s in mixed precision, which amounts to 11.9 %11.9\,\% of peak performance [4]. ORBIT [30] used a ViT [31], trained on global 1.4° resolution with up to 91 channels, and scaled beyond the 1-billion-parameter model size. ORBIT introduces Hybrid-STOP, which fully shards the model’s weights as well as data batches, thus combining tensor parallelism and FSDP. Additionally, communication is optimized by using layer wrapping, i.e., overlapping computation and prefetching of subsequent model layers. Through these means, they train a model with up to 113 billion parameters across 49 15249\,152 GPUs. ORBIT-2 performs a spatial downscaling task, parallelized through the TILES approach [32]. In TILES, the spatial dimension is divided into overlapping spatial tiles, upon which each GPU calculates local attention. With their biggest model of 10B parameters, they achieve a sustained computational throughput of 368 PFLOP/s368\,PFLOP/s on 512512 nodes (20482048 GH200 GPUs) on the Alps system. AERIS [33] scales training of a Swin-Transformer-based [34] diffusion model with 1×11× 1 patch embedding up to 10 08010\,080 nodes and utilizes data, pipeline, and Sequence-Window (SWiPe) parallelism. SWiPe partitions the image across windows, and further parallelizes the sequence within each window with DeepSpeed Ulysses’s sequence parallelism [25]. Their biggest model (80B parameters) reaches 10.21 EFLOP/s10.21\,EFLOP/s on 10 08010\,080 nodes of the Aurora system, corresponding to roughly 1010 forecasts per second. Domain–Tensor-Parallelism (DTP)DTP instanceGPU 1GPU 2GPU 3GPU 4Spatial parallelChannel parallel Uncertainty-Parallelism (UP)UP instanceDTP 1DTP 2DTP 3DTP 4Aleatoric / noise injectionData-Parallelism (DP)DP instanceUP 1UP 2UP 3UP 4Data parallel ERA5 Hourly Data Time frame: 1979–2019 Format: zarr, chunked in time and space (1,1⋯1,m⋱n,1⋯n,m) pmatrixW_1,1&·s&·s&W_1,m\\ & && \\ &&&\\ && & \\ W_n,1&·s&·s&W_n,m pmatrixTensor parallel Epistemic / weight sampling Fig. 1: Schematic of BEAST’s 4D parallelization hierarchy from left to right. Domain‑Tensor‑Parallelism (DTP) decomposes inputs across spatial and channel dimensions. Uncertainty‑Parallelism (UP) runs independent aleatoric noise injections and epistemic weight samplings. Data‑Parallelism (DP) processes separate, non‑overlapping ERA5 data chunks in parallel. I Innovations Our main innovations overcome the challenges in training large Bayesian neural networks by introducing a 4-dimensional parallelization strategy to alleviate memory bottlenecks and to maximize per-evaluation throughput. The architecture for our Bayesian Exascale Atmospheric Swin Transformer (BEAST) features a 2D-Swin Transformer [34] with one down- and up-sampling layer. We train BEAST according to the Bayesian model formulation p⁡(y∗∣x∗,)=∫p⁡(y∗∣x∗,)p(∣) p(y^* x^*,D)= p(y^* x^*,w)\,p(w )\,dw where D is the training data, w are the model weights, x∗x^* is the input and y∗y^* is the output. The exact computation of the posterior p⁡(∣)p(w ) is generally intractable. To circumvent this, variational inference [35, 36] approximates the posterior by a known parameterized distributional family q. The mean-field assumption uses an independent Gaussian distribution for each weight, and assumes the parameters θ=μ,σθ=\μ,σ\ to be independent for each of the weights to make computation feasible. During training, the parameters θ are optimized towards minimizing the Kullback-Leibler divergence (KLKL) [37] between the variational distribution q and the ideal weight distribution p. This is encoded by the evidence lower bound (ELBOELBO) ELBO=W∼q[log⁡p⁡(Y|X,W)]⏟data fitting−λKL(q||p)⏟prior matching = E_W q [ p(Y|X,W) ]_data fitting- λ\,KL (q||p )_prior matching where p⁡(Y|X,W)p(Y|X,W) is the probability of the training labels according to the model given the training data and weights, which are sampled from the variational distribution for the expected value, and λ is an optional “cold posterior” hyperparameter factor [38]. Figure 1outlines how we take advantage of multiple parallelization axes within our implementation of BEAST to circumvent the complexity and memory requirements inherent to the model. I-A Domain–Tensor Parallelism Our first optimization target is to reduce per-GPU memory footprint. With a single data item, i.e., atmospheric input state, being several hundred megabytes large, current state-of-the-art AI models for atmospheric forecasting struggle to go beyond the 11-billion-parameter mark [30]. Accordingly, a significant proportion of the available GPU memory is consumed by data items and corresponding activations. In the Bayesian formulation, this problem becomes even more pronounced, as the memory demand of a BNN compared to a deterministic model is increased proportionally to 1. the number of parameters used to parameterize the distribution of each model weight (two in our case); and 2. the number of samples drawn—and thus, forward passes through the model and associated activations—per data item. Existing strategies for model parallelization are designed for training LLMs and thus prioritize maximizing the model size through tensor and pipeline parallelism. Long sequence lengths are handled by sharding the input data along the sequence dimension with context parallelism. Scientific image-based applications are much larger in terms of sequence lengths (order of 11 million) with the number of input variables in the order of 100100. The latent representations of the channels are projected to a larger embedding space; thus, sharding of the data along both the sequence and channel dimensions holds the biggest potential in alleviating memory bottlenecks. We pursue this approach of domain parallelism and combine it with tensor parallelism for all neural network layers. That means that, in contrast to existing parallelization schemes, our approach distributes not only the model parameter and optimizer states, but also shards each input data item along both the spatial and channel dimensions. We leverage distributed matrix multiplications to implement this domain–tensor parallelization (DTP) scheme [39]. By sharding both the data and weights and formulating a linear layer as a distributed matrix multiplication, we are able to parallelize individual and arbitrary layers without the need for allgather operations or re-sharding, simultaneously maximizing GPU memory usage and reducing the data I/O-workload. We define a DTP instance over a GPU grid of [c channels, s spatial] ranks, where each rank loads [1/c×1/s][1/c× 1/s] fraction of each data item and holds [1/c][1/c] of the model weights. Ranks within the same spatial group hold the same copy of the model weights, i.e., they compute on different data shards and are domain-parallel only. The parameters of the spatial groups are synchronized by allreduce operations of the gradients. Ranks within the same channel group are fully domain- and tensor-parallel. We apply this concept to all neural network layer types in BEAST. Convolutions are performed by distributing the filters across the channel groups. Each rank computes the local convolutions and exchanges the necessary input data with the remaining ranks in the channel to resolve the full convolution. A schematic of the parallelization strategy used in channel-parallel attention layers is shown in Figure 2. The query–key–value (QKV) map in each attention layer is computed by a distributed linear layer across the channels. Each rank then holds the local sequence and a fraction of the QKV embedding dimension. The attention calculation is then trivially parallelized across the heads, allowing us to leverage optimized single GPU attention kernels such as FlashAttention-2 [40] or cuDNNs fused dot product attention via NVIDIA’s Transformer Engine library [41] without any additional communication, while maintaining domain parallelism across channels on the input to the Transformer block. Head concatenation and the subsequent MLP are again performed by distributed layers, combining the information across the channel shards. The windows are shifted between every Transformer block by half a window, and halos are communicated between the ranks in spatial groups where necessary. Input at time ttInput at time t−Δt- t concat.Patch Embed.2×2 Conv2DSwinTransformerDownsample2×2 Conv2DSwinTransformerUpsample2×2 Conv2DTSwinTransformerPatch Recov.2×2 Conv2DT concat. along variablesOutput at time t+Δt+ tOutput at time t+2Δt+2 t Swin Transformer Block (windows in parallel)LocalinputLocalinputlocal channelsseq.GPU 1GPU 2Dist.LinearP2Pcomm.QKVEmbed. dim / cseq.QKVsoftmax()× pmatrix 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ pmatrix× QKTV× headsFlashAttention-2 / Transformer Enginesoftmax()× pmatrix 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ 16.38895pt& 16.38895pt& 16.38895pt\\ pmatrix× QKTVDist.LinearP2Pcomm.LocaloutputLocaloutputlocal channelsseq. Fig. 2: BEAST combines a downsampling–upsampling backbone of Swin‑Transformer blocks with convolutional patch embedding and recovery. Unlike typical AI models, it ingests two consecutive samples (t and t−Δt- t). The lower schematic illustrates the Swin blocks’ novel parallelization scheme using domain–tensor-parallelism. I-B Uncertainty Parallelism Next, we target the challenge of evaluating the m forward passes of the random weight samples for each data item by introducing sampling parallelism [42]. We extend the GPU grid of [c channels, s spatial] ranks by an orthogonal parallelization axis with u groups, along which we distribute the m random weight samples. That means that we have u DTP instances, each holding a full model on c×sc× s GPUs. Each of these model instances draws m/um/u independent random weight samples and performs the forward pass on them. To compute the mean and standard deviation of the ensemble, the sample predictions need to be aggregated: each rank computes the sum and squared sum of its local predictions, Su=∑i=1muxi(u)S_u= _i=1^m_ux_i^(u) and Qu=∑i=1mu(xi(u))2Q_u= _i=1^m_u (x_i^(u) )^2, respectively. SuS_u and QuQ_u are then aggregated across all u ranks in an allreduce operation to obtain the global sum and global square sum S=∑uSu,Q=∑uQuS= _uS_u, Q= _uQ_u. From these global sums, the global mean and global standard deviation of the predictions are computed. This allows us to reduce the communication in the prediction aggregation to two tensors instead of all the m predictions from the local weight samples. Through sampling parallelism, we can efficiently increase the number of random weight samples per forward–backward pass, thereby improving prediction quality without being bound to the memory capacities of a single accelerator. Sampling parallelism provides an additional major advantage: the parallel evaluations of the random weight samples are all carried out on the same input data batch. They can therefore be combined with data augmentation, which improves predictive generalization and training convergence [43]. While there are several data augmentation options from computer vision available, we chose a different route and used parallel data augmentation to include aleatoric uncertainty into our Bayesian model. We use a modification of conditional layer normalization [44], inspired by [9] (noise independent of spatial grid) and [11], which we coin noise input conditioning for ensembles (NICE). In each forward pass, a random vector ϵε of 3232 independent standard normal distributed values is drawn. The scale and shift parameters for each layer normalization are generated by shallow MLPs, with the noise ϵε as the input. This effectively injects noise into the model and can itself be used to produce ensemble outputs, thus making the model fully uncertainty aware. We call this combination of sampling parallelism with random noise injection uncertainty parallelism (UP). We handle the coordination of the data I/O within UP groups with synchronized random seeds. Uncertainty parallelism allows us to evaluate more random weight samples by scaling to higher node counts, thereby improving the approximation of the current weight distribution and, ultimately, model quality. I-C Data Parallelism Uncertainty parallelism can be combined orthogonally with data parallelism to accelerate and improve model training at the same time, adding an additional orthogonal scaling axis with d groups to our GPU grid. Training with d data parallel (DP) instances and u uncertainty parallel instances thus requires d×ud× u DTP instances in total. However, given the size of the input data items and the model, a realistic setup can only run with one or two local data items, i.e., local batch size, or random weight samples per GPU. Hence, data parallel scaling out also adds more data items to the effective global batch size. In contrast to sampling parallelism, in which having larger sample numbers can lead to better predictive performance, scaling out data parallelism will lead to the occurrence of large batch effects [18]: AI models tend to converge to sharper local minima when the number of data items in a gradient update step becomes too large. For atmospheric models, this effect starts to arise at relatively small global batch sizes of around 1616 to 3232 [45]. The degree to which model training can be accelerated through plain data parallelism is thus limited. With the additional usage of uncertainty parallelism, we circumvent this limitation. I-D Performance Optimization We further maximize per-GPU computational throughput with the following performance optimizations. Since each GPU in a DTP instance loads a shard of the input data sample, we pre-shard the dataset accordingly in the channel and spatial (longitude) domain, thereby facilitating file access for data staging. To further stabilize data I/O and alleviate the strain on the file system, we perform data access in waves of 128128 ranks in each wave and de-synchronize file system access every 5050 iterations. We fuse the MLPs of the NICE normalization and use grouped normalization to eliminate communication across the split channel domain. Our implementation of the domain–tensor-parallel layers requires us to overwrite the backward function of each layer type. This allows us to also combine gradient communication in the backward pass with respect to the input and weights. This step essentially halves communication in the backward pass, significantly improving performance. We handle communication for UP and DP groups via PyTorch’s DistributedDataParallel wrapper [16]. We also apply post-grad accumulation hooks on the weight synchronization between spatial ranks to better overlap communication with the backward pass. Additionally, we found that the computation of the ELBO loss takes significantly longer than the computation of a regular RMSE loss in a non-Bayesian neural network. To alleviate this bottleneck, we vectorize the computation in the prior matching term, which does not require communication, and overlap it with the computation for the data-fitting term, which does require communication of the locally obtained ensemble members. IV Performance Measurement We evaluate our 4D parallelization scheme, consisting of domain-, tensor-, uncertainty-, and data-parallelism, on the task of training BEAST with mean-field variational inference. TABLE I: Software stack on JUPITER and LUMI Component JUPITER LUMI Container Runtime Apptainer SingularityCE Operating System RHE Linux 9.9 SUSE 15.6 Acc Library CUDA 13.2 ROCM 6.2.4 XCCL NCCL 2.29.7 RCCL 2.22.5 PyTorch 2.12.0 2.7.1 Python 3.12 3.12 IV-A Model Architecture The model architecture of BEAST, visualized in Figure 2, is based on a 2D Swin Transformer [34]: The input data is embedded via a 2D non-overlapping convolution with a spatial patch size of 2×22× 2. Then, multi-headed self-attention is performed on non-overlapping windows in a Swin-Transformer-manner. In between each attention layer, the entire global domain is rolled by (⌊window size/2⌋ size/2 ) patches in both spatial dimensions. Periodic boundaries are used in the latitude dimension. The data is then downsampled by a 2D convolution with a patch size of 22, which also doubles the hidden dimension. Another Swin Transformer block is applied before the data is upsampled back with a transposed convolution and an MLP to mix across the channels, before another Swin Transformer block is applied. Finally, the target, with the original spatial dimensions, is recovered with a transpose convolution. IV-B Training objective We train BEAST to perform 1212- and 2424-hourly forecasts based on two previous time steps, i.e., atmospheric variables at times t and t−12ht-12h are used as input to predict atmospheric variables at times (t+12h,t+24h)(t+12h,t+24h). For training, we use the ERA5 dataset [46], obtained from WeatherBench2 [47] and stored in the zarr format. The data uses an equiangular grid at a global resolution of 0.25° (≈ 30 km30\,km). We consider 8484 input variables in total—six surface variables: 10 m10\,m u-velocity, 10 m10\,m v-velocity, mean sea level pressure, 2 m2\,m temperature, sea surface temperature, total precipitation, and six pressure variables: geopotential, specific humidity, temperature, and the u, v, and w components of wind velocity. Pressure variables are considered at 1313 pressure levels of 1000, 925, 850, 700, 600, 500, 400, 300, 250, 1000, 1000, 1000, and 1000 hPa10009258507006005004003002501000100010001000\,hPa. Constant inputs of soil type, topography, and land masks are also used. Variables from different time steps are concatenated along the channel dimension. Data is normalized with per-variable Z-score normalization. Hourly data from 1979–2019 is used for training, and data from 2020–2025 is used for evaluations. For minimizing the ELBO as loss function, we use an Adam optimizer with an initial learning rate of 10−410^-4, decaying to 10−510^-5 with a cosine decay. Linear warm-up is used for the first 10 00010\,000 gradient steps. We train the model using mixed precision format: for Transformer Engine and FlashAttention in the Transformer layers, we leverage Bfloat16 precision. All other computations are performed in TensorFloat32 or Float32 precision, to reflect the sensitivity of the Bayesian optimization objective to numerical rounding errors. IV-C System specifications We perform training of BEAST on two systems: JUPITER and LUMI. JUPITER is a EuroHPC JU supercomputer hosted at Forschungszentrum Jülich, Germany. It is the first exascale machine in Europe and ranks #5 in the TOP500 of June 2026, reaching 1 EFLOP/s1\,EFLOP/s in the HPL benchmark. JUPITER has been designed for energy-efficient, high-scaling numerical and AI applications [48], and is currently in the finalization of its build-up phase. The system utilizes 23 53623\,536 NVIDIA GH200 superchips, each combining a 7272-core NVIDIA Grace CPU (120 GB120\,GB LPDDR5 memory) with a Hopper GPU (96 GB96\,GB HBM3 memory) through a unique high-bandwidth bus called NVLink-C2C (900 GB/s900\,GB/s bandwidth). The full system features 58845884 nodes, of which about 5200 were accessible to us. Within a node, the GPUs are all-to-all connected with pair-wise bandwidths of 300 GB/s300\,GB/s. Similarly, all CPUs are all-to-all-connected with pair-wise bandwidths of 200 GB/s200\,GB/s. Each node features four NVIDIA InfiniBand NDR200 network adapters with each 200 Gbit/s200\,Gbit/s injection bandwidth, interconnected in a Dragonfly+ topology. Up to 240240 nodes are combined in a Dragonfly group, internally connected over two levels of switches in a fat tree (local) topology. Twenty-five of these groups form the JUPITER network, with 3030 direct links between each pair of groups; each inter-switch link provides 400 Gbit/s400\,Gbit/s bandwidth. Adaptive routing is enabled to counteract congestion due to the tapering of the network outside of the groups. LUMI is a EuroHPC JU supercomputer hosted by the CSC – IT Center for Science in Finland. It is a pre-exascale Cray EX supercomputer supplied by Hewlett Packard Enterprise (HPE). As of June 2026, LUMI ranks #11 in the TOP500 list, reaching 386 PFLOP/s386\,PFLOP/s in the HPL benchmark. The majority of the system’s compute power is found in the LUMI-G hardware partition, which consists of 29782978 nodes. Each node is equipped with four AMD MI250x GPUs and a single 6464-core AMD EPYC Trento CPU. The MI250X GPU is a multi-chip module containing two GPU dies named Graphics Compute Die (GCD) by AMD. Each die includes 110110 compute units and has access to a 64 GB64\,GB slice of HBM2e memory, resulting in a total of 220220 compute units and 128 GB128\,GB of total memory per MI250x module. LUMI uses a Dragonfly network topology in which compute nodes are organized into groups interconnected via switches. The compute nodes are linked through the HPE Cray Slingshot-11 200 Gbit/s200\,Gbit/s network interconnect (NIC). The HPE Cray Slingshot NIC features high-performance RDMA and hardware acceleration for MPI and SHMEM-based software. TABLE I: BEAST configurations used in scaling experiments and model training. Outer windows Inner windows Name Params (B) TFLOPTFLOP Dim. Win. size Heads Layers Dim. Win. size Heads Layers JUPITER-2.4B-64K 2.362.36 3557.9 30723072 [360360, 180180] 2424 22 61446144 [180180, 3030] 4848 22 LUMI-1.1B-64K 1.051.05 2238.2 20482048 [360360, 180180] 1616 22 40964096 [180180, 3030] 3232 22 JUPITER-0.7B-1.8K 0.680.68 423.8 10241024 [3030, 6060] 1616 44 20482048 [2020, 3030] 3232 66 IV-D Software The software stack employed in this work is based on containerized environments, leveraging Apptainer / Singularity Community Edition [49] to ensure portability and reproducibility across high-performance computing (HPC) systems. Apptainer enables unprivileged container execution and is widely adopted in scientific and high-performance computing contexts. On JUPITER, the base container is the NVIDIA NGC PyTorch image (version 26.04-py3 [50]), available on the NVIDIA NGC registry. For LUMI, we rely on a provided custom ROCm+PyTorch container [51]. Domain and tensor parallelism of the PyTorch layers was done by using the library jigsaw [39], and the Bayesian model was obtained by converting the deterministic torch model with torch_blue [52]. Profiles and traces for optimization were recorded using the built-in torch profiler and NVIDIA Nsight Systems 2026.2.1. Power draw and memory utilization measurements for the hardware accelerators were sampled during training using perun [53] for both systems. IV-E Performance metrics The most relevant performance metrics when training neural networks are predictive performance and overall training time. For a fixed-size dataset, training time is influenced by two factors: the time it takes to iterate through the entire dataset in batched forward-backward passes (time-per-epoch) and the number of epochs, i.e., iterations over the entire dataset, necessary for model convergence (convergence rate). While the latter is generally addressed algorithmically, the former can be directly tied to the computational throughput of performing one batch iteration (forward–backward pass and weight update). To assess the computational throughput of the batch iterations in our scaling experiments, we measure the time for 1. a forward–backward-update step, and 2. one full epoch using Python’s high-resolution time module with time.perf_counter_ns() and the torch.cuda.streams.Event() to set markers to monitor accelerator computations. The former includes the overhead associated with model training: data I/O, checkpointing, and evaluation; whereas the latter captures exclusively the computation that is offloaded to hardware accelerators. We further estimate the number of floating-point operations (FLOPFLOP) of a forward–backward pass on a DTP instance, using PyTorch’s torch.utils.flop_counter tools. Based on this, we report the number of FLOP/sFLOP/s for 1. a forward–backward-update-step as peak performance and 2. for a full epoch as sustained performance by dividing the number of FLOPFLOP by the measured time. We also report accelerator metrics related to ideal use during the forward–backward-pass, including percentage of idle time, memory utilization, and average power draw. In addition to these computational throughput numbers, we also report the commonly used metric images/sec from computer vision. In atmospheric modeling, this corresponds to the forecast throughput or forecast/sforecast/s, which is estimated based on the average batch processing time, and the total number of data items per batch (DP instances × local batch size). We also introduce the orthogonal metric of sample/ssample/s, determined via the product of forecast/sforecast/s × UP instances × local sampling size, for ensemble generation throughput. This metric indicates throughput regarding the evaluation of random weight samples of the BNN and is especially relevant for domain scientists during inference over an evaluation dataset, indicating how long it takes to generate a meaningful ensemble. Next to compute performance, we evaluate the probabilistic predictive skill scores of a fully trained BEAST, where we report the commonly used root mean square error (RMSE) of the ensemble average prediction, probabilistic spread skill ratio (SSR), and continuous ranked probability score (CRPS) metrics for different atmospheric variables. These metrics cover three different aspects: The RMSE of the ensemble mean measures how well the average forecast matches the target. The SSR quantifies whether the model’s uncertainty is well-calibrated. The CRPS, which we calculate using the fair CRPS estimator [54], combines these two aspects to measure overall probabilistic prediction quality. We compare our predictive skill scores against the current-best probabilistic AI model, FGN [9] vs. ERA5, and the IFS ensemble forecast [55] vs. ERA5, both obtained from [47]. We further assess the skill of BEAST on predicting extreme events, by evaluating three exemplary use-cases, using the recently published ExtremeWeatherBench (EWB) Python package [56]: a) a heat wave, b) a freeze, and c) a tropical cyclone (TC). The heat wave occurred from August 8, 2022, until August 16, 2022, and affected southern UK, France, and Spain. The freeze affected Western and Central Europe from December 10, 2022, until December 20, 2022. For the TC, we used Damien, a category 3 cyclone, which occurred between February 3, 2020, and February 9, 2020, over Western Australia. These cases probe both prolonged large-scale temperature anomalies and the emergence and subsequent trajectory of a compact, rapidly evolving severe-weather system. V Performance Results Table I describes the three model configurations used in our experiments. The JUPITER-2.4B-64K and LUMI-1.1B-64K model architectures are optimized to leverage the available hardware and train the maximal model size as efficiently as possible. Due to the size of these model configurations, the local batch size and number of local samples are restricted to 11. However, the long sequence length considered in the attention windows (64 00064\,000) causes the training time per sample to become prohibitively long for training a model until convergence. Therefore, we define JUPITER-0.7B-1.8K, a smaller but deeper model configuration designed to maximize predictive skill. For all three configurations, we distribute one DTP instance on a GPU mesh of [s,c]=[4,2][s,c]=[4,2], i.e., 88 GPUs (22 nodes) on JUPITER and 88 GCDs (11 node) on LUMI, when performing training steps (forward–backward-update passes). At inference, where the memory footprint is much smaller due to the absence of gradients and optimizer states, we distribute one DTP instance on a GPU mesh of [s,c]=[2,2][s,c]=[2,2], i.e., 44 GPUs. V-A Baseline performance We define baseline performance on the LUMI-1.1B-64K and JUPITER-2.4B-64K model configurations. The baseline configurations require 1616 GPUs on both systems: a domain–tensor group of [44, 22] × 22 UP groups, and a single data parallel group. This translates to 44 compute nodes on JUPITER, and 22 nodes on LUMI. Based on an average batch time (forward–backward-update) of 1.68 s1.68\,s on the JUPITER system, and an estimated 3.56 PFLOP3.56\,PFLOP operations given the model configuration, JUPITER-2.4B-64K reaches a peak performance of 1.06 PFLOP/(s Node)1.06\,PFLOP/(s\,Node), and a sustained performance of 0.92 PFLOP/(s Node)0.92\,PFLOP/(s\,Node). This translates to 0.59 forecast/s0.59\,forecast/s. We measured the amount of idle time and time spent waiting for communication by profiling the system on 2020 gradient updates, showing ∼ 25 %25\,\% of the time is spent waiting on communication, ∼ 1 %1\,\% device idle time, and ∼ 72 %72\,\% of the time is being used for computation. Within an epoch, memory measurements for the individual GPUs report 71.45 GB71.45\,GB consistently, utilizing 74.43 %74.43\,\% of the total accelerator memory. With the LUMI-1.1B-64K configuration, the average forward–backward-pass runs for 4.47 s4.47\,s. Based on an estimated 2.24 PFLOP2.24\,PFLOP, we reach single node peak performance of 0.5 PFLOP/(s Node)0.5\,PFLOP/(s\,Node), or 0.22 forecast/s0.22\,forecast/s. Memory measurements show each GCD reporting an average of 48.87 GB48.87\,GB, or 76 %76\,\% GPU memory utilization. Additionally, we measured the power draw of the hardware accelerators during training to provide a realistic estimate of the energy cost of training and running our Bayesian model [57]. We observe an average power draw of 455.6 W455.6\,W on each MI250X card, and 3.64 kW3.64\,kW for training LUMI-1.1B-64K in the baseline configuration. By comparison, we measured a consistent power draw of 514.36 W514.36\,W per GPU on the JUPITER system, resulting in a total GPU power draw of 8.22 kW8.22\,kW for training JUPITER-2.4B-64K over the four nodes in the baseline configuration. 3b 3c (a) Weak scaling computational throughputDP=1← =16→↑ higher is better16642561024409616 3845050100100150150200200250250300300350350GPUsThroughput (TFLOPS/s/GPU) (b) Strong scaling performanceUP=2← =16→ 5120 Nodes →↑ higher is better41664256102440961133101030301001003003001,0001,0003,0003,00010,00010,000NodesPeak performance (PFLOPS/s) (c) Weak scaling ensemble throughput321285122 0488 192166425610244096GPUsSamples (/s) (d) Strong scaling data throughput↑ higher is better12481632641280.250.250.50.511224488161632326464DP Group SizeForecasts (/s) Fig. 3: Scaling results of BEAST on JUPITER and LUMI for various data‑ and uncertainty-parallel setups, expressed in computational, data, and ensemble throughput as well as peak performance under strong scaling, i.e., increasing DP with increasing node counts, and weak scaling with fixed per-node workload, i.e., increasing UP with increasing full GPU counts. V-B Scaling experiments To evaluate the efficiency of uncertainty parallelism, we scale out the number u of UP groups in a weak scaling manner, i.e., we increase the number of global weight samples evaluated per data batch by increasing the number of UP instances. For JUPITER-2.4B-64K and LUMI-1.1B-64K we conduct this experiment in two regimes: for one, we scale u from 22 to 1616, with a single DP instance, i.e., using 1616 through 128128 GPUs, or GCDs in LUMI. Such a configuration is not very realistic for production training runs, where multiple DP instances would be used to iterate over the dataset faster. Hence, we conduct a second round of this scaling experiment, where we keep the number of DP instances, i.e., the global batch size, fixed at 1616 and increase the number of UP instances from 22 to 6464, i.e., using 256256 to 16 38416\,384 GPUs. For JUPITER-0.7B-1.8K we evaluate UP scaling in both training and inference mode, with a fixed global batch size of 44, increasing the number of UP instances from 88 to a total of 20482048 nodes on JUPITER. The scaling experiments shed light on two interesting aspects of our multidimensional parallelization scheme. For one, DP requires different data to be read from storage by each rank, while for UP, each rank reads the same data. Having both UP and DP instances participate in the training illustrates how our parallelization scheme handles these different I/O streams. Second, both UP and DP utilize PyTorch’s DistributedDataParallel wrapper [16] to handle communication. Thus, the experiment allows us to investigate if these two communication axes potentially compete against one another, thereby hindering performance. The results for JUPITER-2.4B-64K and LUMI-1.1B-64K are shown in Figure 3 (a). The values in the gray section indicate the scaling experiments where the DP size is 11, and the values after the gray section showcase the results of the configuration with the increased DP size of 1616. On LUMI, we observe almost constant throughput for lower node counts up to 30723072 GPUs, whereas on JUPITER, we observe a small degradation of per-GPU performance. On both systems, performance starts to deteriorate notably at larger node counts. This scalability limit of UP can be attributed to the increased stress on the file system and higher communication cost during the loss calculation. Ensemble throughput scaling of JUPITER-0.7B-1.8K is presented in Figure 3 (c). On 20482048 nodes, the training configuration reaches a total of 1293.31 sample/s1293.31\,sample/s, with a scaling efficiency of 52.35 %52.35\,\%, similar to the results shown in Figure 3 (a). Again, we observe a degradation in performance, which can be attributed to the additional communication required to calculate the loss for a larger number of samples generated per forward pass. This effect disappears in inference mode, where the loss is not calculated, and the uncertainty parallel groups become independent of each other. On 20482048, the predictive model reaches 5716.61 sample/s5716.61\,sample/s, with an ensemble scaling efficiency of 82.01 %82.01\,\%. Next, we evaluate scaling out data parallelism in a strong scaling manner: We increase the global batch size by increasing the number of DP instances from 11 to 1616 at a fixed local batch size of 11, thereby reducing the overall number of iterations in an epoch. Similarly to UP weak scaling, we conduct this strong scaling experiment for JUPITER-2.4B-64K and LUMI-1.1B-64K in two setups: one with just 22 UP instances, thus scaling from 22 to 3232 nodes, and one with 1616 UP instances, thus scaling from 6464 to 51205120 nodes. For JUPITER-0.7B-1.8K we keep the number of random weight samples (ensemble members) in the training configuration at 9696, using 4848 UP instances, and scale out DP from 11 to 1616, i.e. 9696 to 15361536 nodes. Results are presented in Figure 3 (b) and (d). Again, the values in the gray sections indicate the scaling experiments where the UP size is 22, i.e., training with a 22 UP instance, and the values after the gray section showcase the results of the configuration with the increased UP size of 1616. We observe near-perfect DP-strong scaling for all three model configurations on both JUPITER and LUMI. On 51205120 nodes on JUPITER, i.e., 20 48020\,480 GH200 GPUs, JUPITER-2.4B-64K runs at a total peak performance of 3.96 EFLOP/s3.96\,EFLOP/s during forward–backward-update step, and a scaling efficiency of 73.99 %73.99\,\%. LUMI-1.1B-64K reaches 463.23 PFLOP/s463.23\,PFLOP/s at 10241024 nodes, corresponding to 92.09 %92.09\,\% scaling efficiency. These results indicate that the parallelism strategy and model configuration space allow for good system portability. The seamless transition between the different UP and DP instance sizes highlights the efficiency of both parallelism axes. JUPITER-0.7B-1.8K yields a higher forecast throughput per data parallel group than the large configurations. At a single DP instance, the predictive configuration achieves a 1.13 forecast/s1.13\,forecast/s, almost double the throughput of JUPITER-2.4B-64K. With 16 DP instances, or 15361536 nodes, it reaches 14.08 forecast/s14.08\,forecast/s, each with 9696 ensemble members, at a sustained performance of 239.88 TFLOP239.88\,TFLOP. 4 ↓ lower is betterU10T2MU850T850000.20.20.40.40.60.60.80.8111.21.21.41.41.61.6VariableRMSE ↓ lower is betterU10T2MU850T850000.20.20.40.40.60.60.80.8VariableCRPS ↑⁣↓ over-/under-dispersiveU10T2MU850T850000.50.5111.51.5VariableSpread/skill ratio Fig. 4: Forecasting skill of BEAST across the four key variables, u component of wind velocity at 10m (U10), 2 meter temperature (T2M), u component of wind velocity at 850 hPa850\,hPa (U850), and temperature at 850 hPa850\,hPa (T850), compared to IFS-ENS and FGN. All three models were evaluated against ERA5, on one year of data (2022). V-C Full model training and predictive skill We train BEAST in the JUPITER-0.7B-1.8K configuration on 384384 nodes on JUPITER, using 44 DP and 4848 UP instances, for 983 350983\,350 optimizer update steps. We use 22 samples per UP instance and keep the local batch size to 11, thus using a global batch size of 44 and drawing 9696 random weight samples per optimizer update step. One epoch on 40 years of data (1979–2020) requires 89 83589\,835 gradient update steps and takes 25.45 h25.45\,h. The sustained training performance is 70.46 PFLOP/s70.46\,PFLOP/s. Predictive metrics of the trained model for four different variables (U10, T2M, U850, T850) for the two forecast time steps 12 h12\,h and 24 h24\,h are shown in Figure 4, in comparison to the state-of-the-art probabilistic model FGN and the IFS ensemble forecast. We observe competitive to superior predictive scores with much better calibrated uncertainty estimates, and stable roll-out for up to ten days. Fig. 5: Probability for cyclone passing within 150 km150\,km for TC Damien, predicted by BEAST (initialized with ERA5 data from 02-01-2020, 96 ensemble members). The reference track is obtained from the International Best Track Archive for Climate Stewardship (IBTrACS). European Heat Wave (August 2022)3.8263.8723.9774.2394.5723.6203.6743.8213.8863.9693.6143.6593.8103.8663.9563.6153.6613.8023.8573.959IFS HRESBEASTn=32n=96n=256135710Lead Time [days]-15015% Diff vs IFS HRES European Freeze (December 2022)5.7305.8556.0376.1966.9215.8756.0706.0455.9695.9555.8656.0456.0035.9875.8945.8686.0215.9895.9985.874IFS HRESBEASTn=32n=96n=256135710Lead Time [days]-20020% Diff vs IFS HRES Fig. 6: RMSE values of BEAST predictions for a heat wave and a freeze, both occurring over Europe in 2022, compared to the IFS HRES. Predictions were generated with 3232,9696 and 256256 ensemble members, rolled out for up to 10-days, with 6 h6\,h- initialization times. Lower values indicate better predictive skill. Furthermore, BEAST can predict extreme events with exceptional skill. For one, it excels at the task of TC track prediction, as shown in Figure 6 for TC Damien. The mean ensemble prediction of BEAST follows the observed track, including the transition from land to ocean and the subsequent landfall. The ensemble members form a spread that encloses the observed track, with uncertainty increasing with forecast lead time as expected. Second, it can forecast both the heat wave and the freeze several days in advance, with skill scores competitive to or even superior to the IFS, as shown in Figure 6. These results show that BEAST’s probabilistic forecasts remain informative not only in aggregate verification, but also for rare events where accurately representing both the expected evolution and its uncertainty is particularly consequential. Increasing the number of random weight samples not only impacts computational performance but also predictive skill. Table I illustrates how larger ensemble sizes lead to better predictive skill in terms of lower RMSE. This result highlights the benefit of our uncertainty parallelism approach, which allows us to efficiently generate large ensembles by scaling out uncertainty parallelism. Generating a 15-day forecast trajectory with BEAST takes about 16 s16\,s to 17 s17\,s on 44 GH200 GPUs, independent of the ensemble size. For comparison, FGN reports inference times of just under 1 minute on a single TPU v5p [9]. In summary, our results demonstrate the strengths of our innovations: They enable us to efficiently utilize the available resources (up to 20 48020\,480 GH200 GPUs) to generate large, accurate ensembles, while allowing for high flexibility regarding the modeling objectives. The combination of domain-, tensor-, uncertainty, and data-parallelism in an orthogonal 4D parallelization scheme can be leveraged to: • prioritize model size beyond the 1B parameter mark towards building bigger models, e.g., foundation models • prioritize weight sample number towards maximizing ensemble generation throughput for better uncertainty quantification • prioritize batch size towards maximizing forecast throughput for faster training and inference Together, these aspects allow us to leverage the full potential of the Bayesian paradigm for smarter AI models in atmospheric forecasting. TABLE I: Inference time for a 15-day forecast, and RMSE values of 2-meter temperature (T2m) for 12 h12\,h and 24 h24\,h predictions generated with BEAST-JUPITER-0.7B-1.8K, evaluated on one year of validation data (2022) with increasingly larger ensembles. Evaluation was conducted by distributing one DTP instance on 44 GH200 GPUs, and a global batch size of 16. # GPUs Ensemble member Inference time [s] T2m RMSE 12h 24h 128 8 15.58 0.7507 0.8661 512 32 16.85 0.7289 0.8390 1536 96 16.98 0.7237 0.8328 3072 768 15.74 0.7226 0.8316 VI Implications Until now, training AI models for global atmospheric forecasting in a fully Bayesian framework has not been possible. While implementing the complex mathematical concepts behind BNNs is already a challenge in itself, the associated computational training cost has ultimately been the prohibitive factor. We present BEAST, the first large-scale Bayesian Swin Transformer for atmospheric forecasting, able to quantify both, epistemic (model) and aleatoric (data) uncertainty. This breakthrough is enabled by our innovative 4D parallelization scheme, consisting of domain-, tensor-, uncertainty-, and data-parallelism. We achieve state-of-the-art performance in both training speed and efficiency, peaking at 3.96 EFLOP/s3.96\,EFLOP/s on 20 48020\,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. After training for nearly one million gradient updates, a 700700-million-parameter model achieves predictive skill scores competitive with the state-of-the-art probabilistic AI model FGN and the IFS ensemble forecast. Yet, through our highly efficient domain-tensor-parallelism approach, BEAST is able to generate large ensembles significantly faster: a single 15-day forecast can be generated on 44 GH200 GPUs in less than 20 s20\,s. On 15361536 nodes, we achieve a forecast generation throughput of 14.08 forecast/s14.08\,forecast/s with 96 ensemble members. This innovation opens a new era for AI-based atmospheric sciences. Under a changing climate, fast and accurate uncertainty quantification in weather forecasting is becoming increasingly important. Existing approaches to uncertainty quantification in AI models follow methods from generative AI, using input randomness to produce ensemble forecasts; however, they neglect the uncertainty of the model itself, making the predictions overconfident [58]. In contrast, our Bayesian model formulation allows for a more principled estimation of the prediction certainty, leading to better calibrated uncertainty estimates. This calibration is crucial, especially for predicting extreme weather events such as heatwaves, floods, and severe storms. Decision makers depend on accurate estimates of the likelihood and severity of extreme events to prepare for possible outcomes and take preventative measures to reduce potential damages to people, infrastructures, and the economy. As demonstrated by our results, BEAST can deliver upon this promise by predicting extreme events, such as heat waves or tropical cyclones, with high accuracy. Our contribution thus directly addresses the challenges we are facing under a changing climate. Next to these direct implications, our approach to enable fast and efficient training of large-scale BNNs also holds the potential to revolutionize Earth system sciences beyond atmospheric forecasting. While the lack of sufficiently large, well-curated datasets has thus far hindered true breakthroughs in AI-based models of Earth system components other than the atmosphere—such as the ocean—our Bayesian approach can overcome this barrier. With its capabilities to learn the underlying data distribution and to capture the associated uncertainty in its distributional model weights, a BNN is particularly well-suited to leverage small, high-quality datasets, e.g., from km-scale climate simulations [59] to build the next generation of climate emulators. Our approach to flexibly allocate GPU memory through highly efficient domain–tensor-parallelism puts us in a prime position to increase the size of input data items beyond the current resolutions towards multi-timestep forecasting and km-scale modeling. Our breakthroughs were enabled by leveraging top-tier supercomputers with thousands of GPUs. Going forward, the results presented in this work mark only the beginning of exploring high-performance computing for BNNs. By introducing uncertainty parallelism as an orthogonal scaling axis, we have successfully overcome the current barriers in scaling out neural network training through plain data-parallelism, which is inevitably limited due to large-batch effects. Scaling out uncertainty parallelism goes beyond merely accelerating neural network training—it creates scientific value. For future work, we aim to leverage the power of current and future exascale systems by improving stability for lower precision computations and seeking out additional parallelization axes for scaling. Acknowledgment This project received access to the JUPITER supercomputer through the JUPITER Research and Early Access Program (JUREAP). JUPITER is funded by the EuroHPC Joint Undertaking, the German Federal Ministry of Research, Technology and Space, and the Ministry of Culture and Science of the German state of North Rhine-Westphalia. Computations on the LUMI supercomputer were enabled by resources provided by CSC - IT Center for Science, Finland. LUMI is co-funded by the EuroHPC Joint Undertaking and the LUMI Consortium. We would like to express our gratitude to Christelle Piechurski from NVIDIA for the continuous support and encouragement to pursue this project from the beginning. We would like to thank Benedikt von St. Vieth and Damian Alvarez from JSC for their assistance and support, and Christian Schiffer for the exchanges on extreme-scale AI. We are also thankful to Jesse Harrison and Jarmo Mäkelä from CSC for supporting this submission, and to Gaurav Naithani and Billy Braithwaite from CSC for their continuous input and insightful discussions on the topics of Bayesian neural networks and FlashAttention. This research is supported by the German Federal Ministry of Research, Technology and Space (BMFTR) under the 01LK2313A SMARTWEATHER21-SCC-2 grant, by the Carl Zeiss Foundation under the Breakthroughs 2025 “WOW – a World model of Our World” project, by the European High Performance Computing Joint Undertaking (JU) and the the German Federal Ministry of Research, Technology and Space (BMFTR), the Ministry of Culture and Science of North Rhine-Westphalia (MKW NRW) and the Hessian Ministry of Science and Research, Arts and Culture (HMWK) under grant agreement No 101250682 (JUPITER AI Factory, JAIF). References [1] Y. Yu, L. Huang, A. Calotoiu, and T. Hoefler (2026) Scaling laws of global weather models. External Links: 2602.22962, Link, Document Cited by: §I, §I. [2] R. Lam, A. Sanchez-Gonzalez, M. Willson, P. Wirnsberger, M. Fortunato, F. Alet, S. Ravuri, T. Ewalds, Z. Eaton-Rosen, W. Hu, A. Merose, S. Hoyer, G. Holland, O. Vinyals, J. Stott, A. Pritzel, S. Mohamed, and P. Battaglia (2023) Learning skillful medium-range global weather forecasting. Science 382 (6677), p. 1416–1421. External Links: Document, Link, https://w.science.org/doi/pdf/10.1126/science.adi2336 Cited by: §I. [3] K. Bi, L. Xie, H. Zhang, X. Chen, X. Gu, and Q. Tian (2023) Accurate medium-range global weather forecasting with 3d neural networks. Nature 619 (7970), p. 533–538. External Links: Document Cited by: §I. [4] T. Kurth, S. Subramanian, P. Harrington, J. Pathak, M. Mardani, D. Hall, A. Miele, K. Kashinath, and A. Anandkumar (2023) Fourcastnet: accelerating global high-resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the platform for advanced scientific computing conference, p. 1–11. External Links: Document Cited by: §I, §I. [5] B. Bonev, T. Kurth, C. Hundt, J. Pathak, M. Baust, K. Kashinath, and A. Anandkumar (2023) Spherical fourier neural operators: learning stable dynamics on the sphere. In Proceedings of the 40th International Conference on Machine Learning, ICML’23. External Links: Document Cited by: §I. [6] A. Mahesh, W. D. Collins, B. Bonev, N. Brenowitz, Y. Cohen, J. Elms, P. Harrington, K. Kashinath, T. Kurth, J. North, T. O’Brien, M. Pritchard, D. Pruitt, M. Risser, S. Subramanian, and J. Willard (2025) Huge ensembles – part 1: design of ensemble weather forecasts using spherical fourier neural operators. Geoscientific Model Development 18 (17), p. 5575–5603. External Links: Link, Document Cited by: §I. [7] A. Mahesh, W. D. Collins, B. Bonev, N. Brenowitz, Y. Cohen, P. Harrington, K. Kashinath, T. Kurth, J. North, T. A. O’Brien, M. Pritchard, D. Pruitt, M. Risser, S. Subramanian, and J. Willard (2025) Huge ensembles – part 2: properties of a huge ensemble of hindcasts generated with spherical fourier neural operators. Geoscientific Model Development 18 (17), p. 5605–5633. External Links: Link, Document Cited by: §I. [8] I. Price, A. Sanchez-Gonzalez, F. Alet, T. R. Andersson, A. El-Kadi, D. Masters, T. Ewalds, J. Stott, S. Mohamed, P. Battaglia, R. Lam, and M. Willson (2024) Probabilistic weather forecasting with machine learning. Nature 637 (8044), p. 84–90. External Links: Document Cited by: §I. [9] F. Alet, I. Price, A. El-Kadi, D. Masters, S. Markou, T. R. Andersson, J. Stott, R. Lam, M. Willson, A. Sanchez-Gonzalez, and P. Battaglia (2025) Skillful joint probabilistic weather forecasting from marginals. External Links: 2506.10772, Link, Document Cited by: §I, §I-B, §IV-E, §V-C. [10] B. Bonev, T. Kurth, A. Mahesh, M. Bisson, J. Kossaifi, K. Kashinath, A. Anandkumar, W. D. Collins, M. S. Pritchard, and A. Keller (2025) FourCastNet 3: a geometric approach to probabilistic machine-learning weather forecasting at scale. External Links: 2507.12144, Link, Document Cited by: §I. [11] S. Lang, M. Alexe, M. C. Clare, C. Roberts, R. Adewoyin, Z. Ben Bouallègue, M. Chantry, J. Dramsch, P. D. Dueben, S. Hahner, P. Maciel, A. Prieto-Nemesio, C. O’Brien, F. Pinault, J. Polster, B. Raoult, S. Tietsche, and M. Leutbecher (2026) Aifs-crps: ensemble forecasting using a model trained with a loss function based on the continuous ranked probability score. npj Artificial Intelligence 2 (1), p. 18. External Links: Document Cited by: §I, §I-B. [12] S. Lee, S. Purushwalkam, M. Cogswell, D. Crandall, and D. Batra (2015) Why m heads are better than one: training a diverse ensemble of deep networks. External Links: 1511.06314, Link, Document Cited by: §I. [13] B. Lakshminarayanan, A. Pritzel, and C. Blundell (2017) Simple and scalable predictive uncertainty estimation using deep ensembles. In Proceedings of the 31st International Conference on Neural Information Processing Systems, Vol. 30, p. 6405–6416. External Links: Link Cited by: §I. [14] J. Gawlikowski, C. R. N. Tassi, M. Ali, J. Lee, M. Humt, J. Feng, A. Kruspe, R. Triebel, P. Jung, R. Roscher, M. Shahzad, W. Yang, R. Bamler, and X. X. Zhu (2023) A survey of uncertainty in deep neural networks. Artificial intelligence review 56 (Suppl 1), p. 1513–1589. Cited by: §I. [15] L. V. Jospin, H. Laga, F. Boussaid, W. Buntine, and M. Bennamoun (2022) Hands-on bayesian neural networks—a tutorial for deep learning users. IEEE Computational Intelligence Magazine 17 (2), p. 29–48. Cited by: §I. [16] S. Li, Y. Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, J. Smith, B. Vaughan, P. Damania, and S. Chintala (2020) PyTorch distributed: experiences on accelerating data parallel training. Proceedings of the VLDB Endowment 13 (12), p. 3005–3018. External Links: Link, Document Cited by: §I, §I-D, §V-B. [17] T. Ben-Nun and T. Hoefler (2019) Demystifying parallel and distributed deep learning: an in-depth concurrency analysis. ACM Computing Surveys (CSUR) 52 (4), p. 1–43. External Links: Document Cited by: §I. [18] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang (2017) On large-batch training for deep learning: generalization gap and sharp minima. In International Conference on Learning Representations, p. 1–16. External Links: Link Cited by: §I, §I-C. [19] T. Nguyen, J. Brandstetter, A. Kapoor, J. K. Gupta, and A. Grover (2023) Climax: a foundation model for weather and climate. External Links: 2301.10343, Link, Document Cited by: §I. [20] C. Lessig, I. Luise, B. Gong, M. Langguth, S. Stadtler, and M. Schultz (2023) AtmoRep: a stochastic model of atmosphere dynamics using large scale representation learning. External Links: 2308.13280, Link, Document Cited by: §I. [21] C. Bodnar, W. P. Bruinsma, A. Lucic, M. Stanley, A. Allen, J. Brandstetter, P. Garvan, M. Riechert, J. A. Weyn, H. Dong, J. K. Gupta, K. Thambiratnam, A. T. Archibald, C. Wu, E. Heider, M. Welling, R. E. Turner, and P. Perdikaris (2025) A foundation model for the earth system. Nature 641 (8065), p. 1180–1187. External Links: Document Cited by: §I. [22] Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. X. Chen, D. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and Z. Chen (2019) GPipe: efficient training of giant neural networks using pipeline parallelism. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, Red Hook, NY, USA, p. 103–112. External Links: Link Cited by: §I. [23] A. Harlap, D. Narayanan, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, and P. B. Gibbons (2018) PipeDream: fast and efficient pipeline parallel DNN training. CoRR abs/1806.03377. External Links: Link, 1806.03377, Document Cited by: §I. [24] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro (2020) Megatron-lm: training multi-billion parameter language models using model parallelism. External Links: 1909.08053, Document Cited by: §I. [25] S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, R. Y. Aminadabi, S. L. Song, S. Rajbhandari, and Y. He (2024) System optimizations for enabling training of extreme long sequence transformer models. In Proceedings of the 43rd ACM Symposium on Principles of Distributed Computing, p. 121–130. External Links: Document Cited by: §I, §I. [26] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He (2020) Zero: memory optimizations toward training trillion parameter models. In SC20: international conference for high performance computing, networking, storage and analysis, p. 1–16. External Links: Document Cited by: §I, §I. [27] Y. Zhao, A. Gu, R. Varma, L. Luo, C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y. Hao, A. Mathews, and S. Li (2023) PyTorch fsdp: experiences on scaling fully sharded data parallel. External Links: 2304.11277, Link, Document Cited by: §I. [28] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He (2024) \zero-Offload\: democratizing \billion-scale\ model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), p. 551–564. External Links: Link Cited by: §I. [29] T. Chen, B. Xu, C. Zhang, and C. Guestrin (2016) Training deep nets with sublinear memory cost. External Links: 1604.06174, Link, Document Cited by: §I. [30] X. Wang, S. Liu, A. Tsaris, J. Choi, A. Aji, M. Fan, W. Zhang, J. Yin, M. Ashfaq, D. Lu, and P. Balaprakash (2024) ORBIT: oak ridge base foundation model for earth system predictability. External Links: 2404.14712, Link, Document Cited by: §I, §I-A. [31] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby (2021) An image is worth 16x16 words: transformers for image recognition at scale. External Links: 2010.11929, Link, Document Cited by: §I. [32] X. Wang, J. Choi, T. Kurihaya, I. Lyngaas, H. Yoon, X. Xiao, D. Pugmire, M. Fan, N. M. Nafi, A. Tsaris, A. M. Aji, M. Hossain, M. Wahib, D. Wang, P. Thornton, P. Balaprakash, M. Ashfaq, and D. Lu (2025) ORBIT-2: scaling exascale vision foundation models for weather and climate downscaling. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, New York, NY, USA, p. 86–98. External Links: Link, Document Cited by: §I. [33] V. Hatanpä, E. Ku, J. Stock, M. Emani, S. Foreman, C. Jung, S. Madireddy, T. Nguyen, V. Sastry, R. A. O. Sinurat, S. Wheeler, H. Zheng, T. Arcomano, V. Vishwanath, and R. Kotamarthi (2025) AERIS: argonne earth systems model for reliable and skillful predictions. External Links: 2509.13523, Link, Document Cited by: §I. [34] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo (2021) Swin transformer: hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), p. 10012–10022. External Links: Link Cited by: §I, §I, §IV-A. [35] A. Graves (2011) Practical variational inference for neural networks. In Advances in Neural Information Processing Systems, J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K.Q. Weinberger (Eds.), Vol. 24, p. 1–9. External Links: Link Cited by: §I. [36] M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley (2013) Stochastic variational inference. Journal of Machine Learning Research 14 (1), p. 1303–1347. External Links: ISSN 1532-4435, Link Cited by: §I. [37] S. Kullback and R. A. Leibler (1951) On information and sufficiency. The Annals of Mathematical Statistics 22 (1), p. 79–86. External Links: Link Cited by: §I. [38] F. Wenzel, K. Roth, B. Veeling, J. Swiatkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, and S. Nowozin (2020) How good is the bayes posterior in deep neural networks really?. In Proceedings of the 37th International Conference on Machine Learning, p. 10248–10259. External Links: Link Cited by: §I. [39] D. Kieckhefen, M. Götz, L. H. Heyen, A. Streit, and C. Debus (2026) Combining domain and tensor parallelism to train multi-billion-parameter ai weather models. In Platform for Advanced Scientific Computing Conference, p. 1–13. Cited by: §I-A, §IV-D. [40] T. Dao (2024) FlashAttention-2: faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, p. 1–10. External Links: Link Cited by: §I-A. [41] Transformer engine Note: Version 1.11.0+e7a7f6da, accessed 2026-08-12 External Links: Link Cited by: §I-A. [42] A. K. Özdemir, L. H. Heyen, A. Weyrauch, A. Streit, M. Götz, and C. Debus (2026) Sampling parallelism for fast and efficient bayesian learning. In Platform for Advanced Scientific Computing Conference, p. 1–13. Cited by: §I-B. [43] F. Brakel, U. Odyurt, and A. Varbanescu (2024) Model parallelism on distributed infrastructure: a literature review from theory to llm case-studies. External Links: 2403.03699, Link, Document Cited by: §I-B. [44] M. Chen, X. Tan, B. Li, Y. Liu, T. Qin, S. Zhao, and T. Liu (2021) Adaspeech: adaptive text to speech for custom voice. External Links: 2103.00993, Link, Document Cited by: §I-B. [45] P. Huber, D. Li, J. P. G. H. Muriedas, D. Kieckhefen, M. Götz, A. Streit, and C. Debus (2026) Energy consumption in parallel neural network training. In Proceedings of the 2026 SIAM Conference on Parallel Processing for Scientific Computing (P), p. 46–59. External Links: Document, Link Cited by: §I-C. [46] H. Hersbach, B. Bell, P. Berrisford, S. Hirahara, A. Horányi, J. Muñoz-Sabater, J. Nicolas, C. Peubey, R. Radu, D. Schepers, et al. (2020) The era5 global reanalysis. Quarterly journal of the royal meteorological society 146 (730), p. 1999–2049. External Links: Document Cited by: §IV-B. [47] S. Rasp, S. Hoyer, A. Merose, I. Langmore, P. Battaglia, T. Russell, A. Sanchez-Gonzalez, V. Yang, R. Carver, S. Agrawal, M. Chantry, Z. Ben Bouallegue, P. Dueben, C. Bromberg, J. Sisk, L. Barrington, A. Bell, and F. Sha (2024) WeatherBench 2: a benchmark for the next generation of data-driven global weather models. Journal of Advances in Modeling Earth Systems 16 (6), p. e2023MS004019. External Links: Document, Link Cited by: §IV-B, §IV-E. [48] A. Herten, S. Achilles, D. Alvarez, J. Badwaik, E. Behle, M. Bode, T. Breuer, D. Caviedes-Voullième, M. Cherti, A. Dabah, S. E. Sayed, W. Frings, A. Gonzalez-Nicolas, E. B. Gregory, K. H. Mood, T. Hater, J. Jitsev, C. M. John, J. H. Meinke, C. I. Meyer, P. Mezentsev, J. Mirus, S. Nassyr, C. Penke, M. Römmer, U. Sinha, B. v. St. Vieth, O. Stein, E. Suarez, D. Willsch, and I. Zhukov (2024) Application-driven exascale: the jupiter benchmark suite. In SC24: International Conference for High Performance Computing, Networking, Storage and Analysis, p. 1–45. External Links: Document Cited by: §IV-C. [49] G. M. Kurtzer, V. Sochat, and M. W. Bauer (2017) Singularity: Scientific containers for mobility of compute. PLOS ONE 12 (5), p. e0177459. External Links: Link, Document Cited by: §IV-D. [50] NVIDIA (2026) NVIDIA NGC PyTorch. External Links: Link Cited by: §IV-D. [51] sfantao (2026) LUMI-containers/RecipesDocker/pytorch/build-rocm-6.2.4-python-3.12-pytorch-v2.7.1.docker at main · sfantao/LUMI-containers. GitHub. External Links: Link Cited by: §IV-D. [52] A. Weyrauch, L. H. Heyen, J. P. G. H. Muriedas, P. Hsia, A. K. Özdemir, A. Streit, M. Götz, and C. Debus (2026) Torch_blue: a flexible python package for bayesian neural networks in PyTorch. Journal of Open Source Software 11 (117), p. 9415. External Links: ISSN 2475-9066, Link, Document Cited by: §IV-D. [53] J. P. Gutiérrez Hermosillo Muriedas, K. Flügel, C. Debus, H. Obermaier, A. Streit, and M. Götz (2023) Perun: benchmarking energy consumption of high-performance computing applications. In European Conference on Parallel Processing, p. 17–31. External Links: Document Cited by: §IV-D. [54] M. Zamo and P. Naveau (2018) Estimation of the continuous ranked probability score with limited information and applications to ensemble weather forecasts. Mathematical Geosciences 50 (2), p. 209–234. External Links: Document Cited by: §IV-E. [55] European Centre for Medium-Range Weather Forecasts (ECMWF) (2024) IFS documentation cy49r1-part v: ensemble prediction system. Technical report European Centre for Medium-Range Weather Forecasts (ECMWF). External Links: Document Cited by: §IV-E. [56] A. McGovern, T. Mandelbaum, D. Rothenberg, N. Loveday, C. Potvin, M. Flora, L. Magnusson, E. Gilleland, and J. Allen (2026) Extreme weather bench: a framework and benchmark for evaluation of high-impact weather. arXiv preprint arXiv:2605.01126. Cited by: §IV-E. [57] C. Debus, M. Piraud, A. Streit, F. Theis, and M. Götz (2023) Reporting electricity consumption is essential for sustainable ai. Nature Machine Intelligence 5 (11), p. 1176–1178. External Links: Document Cited by: §V-A. [58] S. Nath and T. Palmer (2026) Can ai models reliably forecast extreme weather events?. Nature 651 (8106), p. 583–584. Cited by: §VI. [59] D. Klocke, C. Frauen, J. F. Engels, D. Alexeev, R. Redler, R. Schnur, H. Haak, L. Kornblueh, N. Brüggemann, F. Chegini, et al. (2025) Computing the full earth system at 1km resolution. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, p. 125–136. External Links: Document Cited by: §VI.

Instant research discovery