Machine learning models, and deep neural networks in particular, are increasingly deployed in risk-sensitive domains such as healthcare, environmental forecasting, and finance, where reliable quantification of predictive uncertainty is essential. However, many uncertainty quantification (UQ) methods remain difficult to apply due to their substantial computational cost. Sampling-based Bayesian learning approaches, such as Bayesian neural networks (BNNs), are particularly expensive since drawing and evaluating multiple parameter samples rapidly exhausts memory and compute resources. These constraints have limited the accessibility and exploration of Bayesian techniques thus far. To address these challenges, we introduce sampling parallelism, a simple yet powerful parallelization strategy that targets the primary bottleneck of sampling-based Bayesian learning: the samples themselves. By distributing sample evaluations across multiple GPUs, our method reduces memory pressure and training time without requiring architectural changes or extensive hyperparameter tuning. We detail the methodology and evaluate its performance on a few example tasks and architectures, comparing against distributed data parallelism (DDP) as a baseline. We further demonstrate that sampling parallelism is complementary to existing strategies by implementing a hybrid approach that combines sample and data parallelism. Our experiments show near-perfect scaling when the sample number is scaled proportionally to the computational resources, confirming that sample evaluations parallelize cleanly. Although DDP achieves better raw speedups under scaling with constant workload, sampling parallelism has a notable advantage: by applying independent stochastic augmentations to the same batch on each GPU, it increases augmentation diversity and thus reduces the number of epochs required for convergence.
Figures & tables
Figure 1. BNN Training with Parallel Sampling
Figure 2. Distributing Training with Hybrid Parallelization
Figure 3. Simplified illustration of the SWIN transformer architecture used for the weather forecasting task; patch embedding and recovery is performed via a 2d convolution or transpose convolution layer with kernel size = stride = (2,2) over the latitude and longitude dimensions
Figure 4. Speed-up (fixed-sample scaling, top) and efficiency (proportional-sample scaling, bottom) of training the Bayesian ViT on CIFAR10. For fixed-sample scaling, we compare different sampling parallelism, data parallelism, and a hybrid of both, using 16 random samples and either a global batch size of 256 or an increasing global batch size of 256 times number of GPUs. For proportional-sample scaling we use a fixed global batch size of 1024 and scale the number of random samples from 1 to 16.
Figure 5. Accuracy of the Bayesian ViT on CIFAR10 with two different global batch sizes and 16 random samples, using sampling parallelism, DDP and hybrid parallelism, over the course of training, for different numbers of GPUs.
Figure 6. Accuracy of sampling parallelism over the course of training, using 16 random samples on 16 GPUS and two data augmentation settings: one in which all GPUs apply identical augmentations (blue), and one in which each GPU applies its own stochastic augmentations (orange).
Figure 7. Accuracy of the Bayesian ViT on CIFAR10 with two different global batch sizes and 16 random samples, using sampling parallelism, DDP, and hybrid parallelism, with respect to wall-clock time, for different numbers of GPUs.
Figure 8. Negative log likelihood during training of the ViT on CIFAR-10 trained on 16 GPUs with 16 samples labeled with parallelization strategy and global batch size.
Method/GBS
256
512
1024
2048
DDP
0.1407
0.1399
0.1753
0.1525
SP
0.1211
0.1121
0.0657
0.0534
Table 1. Mean Absolute Calibration Error (MACE ↓ ) of the ViT on CIFAR-10 trained on 16 GPUs with 16 samples given parallelization strategy and global batch size.
Figure 9. Speed up (fixed-sample scaling, top) and efficiency (proportional-sample scaling, bottom) of MCD on MLP with the ENTSO-E dataset. For fixed-sample scaling, we compare different sampling parallelism, data parallelism, and a hybrid of both, using 16 random samples and either a global batch size of 64 or an increasing global batch size of 64 times number of GPUs. For proportional-sample scaling we use a fixed global batch size of 1024 and scale the number of random samples from 1 to 16.
Figure 10. Scaling efficiency of the Bayesian SWIN Transformer with the ERA5 dataset, where the number of random samples equals 2 times the number of GPUs, and the local batch size per GPU is 1.
The practical adoption of sampling-based inference (SAI) in Bayesian neural networks (BNNs) remains limited, partly due to persistent misconceptions about the feasibility and efficiency of sampling. This position paper argues that SAI has achieved computational parity with optimization-based methods and is at the verge of superseding such methods for effective and efficient inference in BNNs. This development should be in the interest of the whole community, promoting BNNs as a principled paradigm with its long-standing yet unfulfilled promise of providing principled uncertainty quantification for neural networks. SAI can even do more -- yielding superior prediction performance through model averaging, serving as the foundation for a plethora of possible downstream tasks, and providing crucial insights into the landscape of BNNs. In order to make such a change happen and unfold the potential of sampling, overcoming current misconceptions is a necessary first step. The next step is to realign research efforts toward addressing remaining challenges in SAI. In particular, the community must focus on two core problems: sufficient exploration of the posterior landscape and high-fidelity distillation of posterior samples for efficient downstream inference. By addressing conceptual and practical obstacles, we can unlock the full potential of SAI and establish it as a central tool in Bayesian deep learning.
Emanuel Sommer, David Rügamer
Department of Statistics, LMU Munich, Munich, Germany · Munich Center for Machine Learning, Munich, Germany.
Sampling-based methods offer a principled approach to uncertainty quantification in Bayesian neural networks. Their practical use, however, is often challenged by the computational cost of exploring high-dimensional and multimodal posterior distributions. To overcome these difficulties, Bayesian Deep Ensembles, i.e., warmstarting the sampling from several optimized solutions, have proven to be an effective strategy. In this paper, we demonstrate that curvature estimates computed during the warmstart as a byproduct in adaptive optimizers such as AdamW can inform the sampling phase at negligible additional cost. Specifically, our proposed preconditioned sampling strategy based on optimizer-derived geometries can substantially reduce or even eliminate the need for a lengthy sampling burn-in phase and leads to greater numerical stability. This approach consistently maintains or improves predictive performance and uncertainty quantification without any additional computational costs. We confirm the consistency of our findings across various datasets and network architectures.
Moritz Schlager, Emanuel Sommer, Thomas Möllenhoff +1
Modern deep learning models remain notoriously prone to overconfidence, limiting their reliability in high-stakes applications. Bayesian methods aim to counter this by learning a distribution over model parameters, and recent advances now make this feasible for large-scale architectures at costs comparable to AdamW. However, a challenge remains at test time: predictions must be averaged across many forward passes with weights sampled from the posterior, which is prohibitively expensive. Variance propagation offers an efficient alternative, computing layer-wise analytical approximations of uncertainty in a single forward pass. While such techniques are effective for MLPs, their extension to modern architectures remains challenging, due to increased depth and diversity of layer types. To fill this gap, we propose Calibrated Variance Propagation (CVP), which introduces a new propagation method for normalization layers, combines it with recent techniques for handling activation functions, and absorbs residual error through a light calibration step. CVP yields comparably accurate uncertainty estimates to MC sampling across transformers and CNNs, at a fraction of the cost. Against prior variance propagation work, CVP improves coverage at 0.5% risk from 8.2% to 14.6% with BEiT-3 on Visual Reasoning (NLVR2) and from 2.6% to 10.8% with ViLT on VQAv2, with gains extending to convolutional architectures.
Tobias Jan Wieczorek, Leon de Andrade, Thomas Möllenhoff +1
TU Darmstadt & hessian.AI, Darmstadt, Germany · RIKEN Center for Advanced Intelligence Project, Tokyo, Japan