Accelerated Execution of Bayesian Neural Networks using a Single Probabilistic Forward Pass and Code Generation

AI summary

Authors: Bernhard Klein, Falk Selker, Hendrik Borras (Heidelberg University); Sophie Steger, Franz Pernkopf (Graz University of Technology); Holger Fröning (Heidelberg University) arXiv: 2511.23440v1 [cs.LG], 28 Nov 2025

Overview

Research area: Efficient inference for Bayesian neural networks, deep learning compilers, and uncertainty quantification on embedded hardware.

Technical level: Advanced. The paper assumes familiarity with Bayesian inference, variational methods, deep learning compilers, and operator-level optimization.

Scope: This paper presents an end-to-end pipeline for training, compiling, tuning, and deploying Probabilistic Forward Pass (PFP) Bayesian neural networks on embedded ARM CPUs by adding custom probabilistic operators to the TVM deep learning compiler.

What This Paper Is About

Bayesian neural networks can express uncertainty about their predictions, which matters for safety-critical uses, but they normally require sampling weights and running many forward passes per prediction, making them expensive. PFP sidesteps this by assuming both weights and activations are Gaussian, allowing a single closed-form forward pass that propagates probability distributions instead of running many samples. The problem is that PFP relies on non-standard operators that no standard deep learning library provides, so the authors build those operators into TVM and optimize them for embedded ARM processors.

Key Contributions

1. A training pipeline based on Stochastic Variational Inference (SVI) that shows PFP reaches comparable performance to SVI-BNNs in uncertainty estimation and out-of-domain detection on the "Dirty-MNIST" dataset.

2. An extension of the TVM deep learning compiler with essential probabilistic operators for both multi-layer perceptron (MLP) and convolutional neural network (CNN) architectures, enabling efficient execution of PFP-based models on ARM processors.

3. Optimized execution schedules plus manual and automatic tuning strategies for the computationally intensive PFP operators, targeting high performance on resource-constrained ARM CPUs.

4. An empirical evaluation on embedded ARM processors demonstrating substantial performance improvements that support deploying BNNs for uncertainty-aware edge and IoT applications. The operator library is released publicly at https://github.com/UniHD-CEG/PFP-Operator-Library.

Main Findings

• PFP matches SVI on quality metrics: On Dirty-MNIST, PFP-based BNNs achieve performance comparable to SVI-BNNs in classification accuracy, uncertainty quantification, and out-of-domain detection, per the abstract.

• Large speedups at small mini-batch sizes: Ablation studies show PFP consistently outperforms SVI in computational efficiency. For small mini-batch sizes, which the authors identify as critical for low-latency applications, the approach achieves speedups of up to 4200x.

• Sample count drives SVI uncertainty metrics: Figure 1(b) shows that Softmax Entropy (aleatoric uncertainty) remains stable as the number of predictive samples grows, while Total Predictive Uncertainty and Mutual Information (epistemic uncertainty) require more samples for reliable OOD detection, particularly for out-of-domain data such as Fashion-MNIST.

• A quantified limitation of the Gaussian assumption: In an artificially constructed high epistemic uncertainty scenario built by assigning random one-hot encoded class predictions, total uncertainty is identical between the Gaussian approximation and the sample-based estimate of the true distribution, but the Gaussian approximation underestimates Mutual Information by 44%. The authors describe this as an extreme case that highlights the limitation of the Gaussian assumption when predictive distributions deviate substantially from a Gaussian.

• Custom operator support in TVM is feasible: Although the PFP concept relies on highly non-standard operators, the authors report that casting the problem within a mapping framework such as TVM enables deployment without specialized libraries while also allowing performance optimizations with reasonable effort.

• Hyperparameters used for the reported SVI training: The SVI-BNNs were trained for 1000 epochs using Adam with a constant learning rate of 0.001, variational posterior weights initialized at mean 0.08 and standard deviation 0.0001, and mini-batch size 100. KL annealing scales a factor from 0 to a maximum of 0.25 over training epochs, which the authors report mitigates sensitivity to initialization and removes the need for non-probabilistic pretraining.

• Not reported in the available content: Absolute wall-clock latencies, energy figures, and the specific ARM processor model or models used for benchmarking are not given in the text provided. Numerical classification accuracy values on Dirty-MNIST are described qualitatively as "comparable" rather than stated as figures. The paper's Sections 5 through 7, which cover operator implementation, tuning, profiling, and benchmarking details, are referenced but their contents are not included in the truncated content.

Methodology in Plain English

The researchers start from an existing baseline: Stochastic Variational Inference (SVI), the common approach where each network weight is described by a Gaussian distribution. To make a prediction, SVI draws many sets of weights and runs the network once per set, then aggregates the resulting logits — accurate, but slow because of all those forward passes.

PFP is described as an extreme simplification of SVI. Where SVI assumes Gaussian-distributed weights only, PFP additionally assumes Gaussian-distributed activations. That lets the network push means and variances through each layer analytically, so prediction and uncertainty come out of one forward pass instead of many. The catch is that the operators need to work on distributions, not scalars: for example, the ReLU nonlinearity has to be replaced by a moment-matching operator that uses the error function to compute the mean and second moment of a truncated Gaussian.

The authors train BNNs with Pyro's SVI — two architectures, a multi-layer perceptron with one hidden layer of 100 neurons and a LeNet-5 — treating all weights probabilistically with Gaussian priors and a mean-field assumption. They then convert the trained weight means and variances into the form PFP needs, applying a calibration factor that reweights variances globally.

To make PFP run on hardware, the authors implement the custom probabilistic operators inside the TVM compiler for MLP and CNN architectures. Because TVM can generate code for different backends and supports auto-tuning, they write optimized schedules and combine manual tuning with automatic tuning to find implementations suited to ARM CPUs. Evaluation uses the Dirty-MNIST construction: MNIST as in-domain data, Ambiguous-MNIST for aleatoric uncertainty, and held-out Fashion-MNIST for epistemic uncertainty / OOD detection. Since PFP has no natural sample dimension, the authors introduce an artificial sampling dimension from PFP-predicted means to make the comparison with SVI fair.

Why This Matters

Impact on research: The work bridges two communities that rarely meet — probabilistic machine learning and compilers for embedded hardware. It shows that a mathematically non-standard Bayesian approximation can be expressed as a set of compiler-level operators, and that doing so yields practical speedups without sacrificing uncertainty quality. It also provides a concrete, quantified example of where the Gaussian assumption behind PFP breaks down (the 44% Mutual Information underestimate), which is useful for anyone deciding when such approximations are safe.

Real-world applications:

• Medical diagnostics and monitoring on embedded or wearable hardware, where a model must flag unfamiliar inputs rather than silently misclassify.
• Autonomous driving and robotics, where perception systems benefit from knowing when a scene falls outside their training distribution.
• Industrial IoT and sensor networks, where devices must run under strict latency and power budgets but still report confidence.
• Weather forecasting and other safety-relevant prediction tasks that need calibrated uncertainty rather than raw confidence scores.

Industry relevance: The result is directly about deployment. By targeting ARM CPUs and building on TVM, the approach fits commodity embedded hardware rather than requiring specialized accelerators, ASICs, or FPGAs of the kind proposed in prior BNN hardware work (ShiftBNN, B²N², VIBNN, photonic and resistive-memory designs). The reported speedups at small mini-batch sizes address the low-latency regime that matters for interactive and real-time edge applications, and the released operator library lowers the barrier for others to reuse the implementation.

Future Directions

1. Characterizing when the Gaussian activation assumption is safe. The authors show a 44% Mutual Information underestimate in a deliberately extreme case; a practical criterion for detecting such failures in real models, and mitigations for them, remains open.

2. Extending the operator library beyond the evaluated architectures. The current library targets multilayer perceptrons and convolutional networks; generalizing to architectures such as transformers and recurrent networks is not addressed.

3. Comparing against other efficient BNN approximations on hardware. The paper compares PFP against SVI, but Monte Carlo Dropout, Deep Ensembles, and last-layer methods are discussed as alternatives without a head-to-head embedded-hardware comparison.

4. Broadening the deployment targets and the evaluation. The work focuses on embedded ARM CPUs; the same code-generation approach could be examined for other backends, and open questions remain about accuracy, latency, and energy trade-offs on larger-scale datasets and models, which the truncated content does not report.

Target Audience

This paper suits researchers and engineers working on Bayesian deep learning, uncertainty quantification, or probabilistic inference who care about deployment cost rather than theory alone. It is also relevant to compiler and systems researchers interested in how non-standard operators can be mapped to embedded hardware through frameworks like TVM, and to practitioners in edge, IoT, and safety-critical domains who need uncertainty-aware models under tight latency and resource budgets. Readers without background in variational inference or compiler internals will find the methodology sections demanding.

Authors’ abstract

Machine learning models perform well across domains such as diagnostics, weather forecasting, NLP, and autonomous driving, but their limited uncertainty handling restricts use in safety-critical settings. Traditional neural networks often fail to detect out-of-domain (OOD) data and may output confident yet incorrect predictions. Bayesian neural networks (BNNs) address this by providing probabilistic estimates, but incur high computational cost because predictions require sampling weight distributions and multiple forward passes. The Probabilistic Forward Pass (PFP) offers a highly efficient approximation to Stochastic Variational Inference (SVI) by assuming Gaussian-distributed weights and activations, enabling fully analytic uncertainty propagation and replacing sampling with a single deterministic forward pass. We present an end-to-end pipeline for training, compiling, optimizing, and deploying PFP-based BNNs on embedded ARM CPUs. Using the TVM deep learning compiler, we implement a dedicated library of Gaussian-propagating operators for multilayer perceptrons and convolutional neural networks, combined with manual and automated tuning strategies. Ablation studies show that PFP consistently outperforms SVI in computational efficiency, achieving speedups of up to 4200x for small mini-batches. PFP-BNNs match SVI-BNNs on Dirty-MNIST in accuracy, uncertainty estimation, and OOD detection while greatly reducing compute cost. These results highlight the potential of combining Bayesian approximations with code generation to enable efficient BNN deployment on resource-constrained systems.