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.
Figures & tables
Figure 1. ( 1(a) ) Exemplary predictions from a SVI-BNN (blue), its Gaussian representation, and PFP ( Roth and Pernkopf, 2016 ) , showcasing MNIST ( Lecun et al., 1998 ) and Ambiguous-MNIST ( Mukhoti et al., 2022 ) as in-domain examples, and Fashion-MNIST ( Xiao et al., 2017 ) as an out-of-domain sample. Variability in class predictions demonstrates aleatoric uncertainty (higher Softmax Entropy, SME), whereas variability across predictive samples indicates epistemic uncertainty (higher Mutual Information, MI). SVI and PFP effectively quantify these uncertainties. Note: Only three SVI samples shown, insufficient for robust estimation. ( 1(b) ) Influence of predictive sample count on uncertainty metrics. Softmax Entropy (aleatoric uncertainty) remains stable, while Total Predictive Uncertainty and Mutual Information (epistemic uncertainty), especially for out-of-domain data (Fashion-MNIST), require more samples for reliable OOD detection.
Figure 2. Illustration of Gaussian moment matching for ReLU activations. The true distribution (solid line) is approximated as a Gaussian (dashed line). Reproduced with permission from ( Roth and Pernkopf, 2016 ) .
Architecture
Dataset
Method
Calib.
Acc. (%)
AUROC (%)
MLP
Dirty-MNIST
SVI
–
96.30
81.2
PFP-only training (100 epochs)
-
97.71
85.4
SVI → PFP, calibrated
0.3
96.30
85.8
SVI → PFP, calibrated + fine-tuned (6 epochs)
0.3
97.08
90.4
LeNet-5
Dirty-MNIST
SVI
–
98.70
98.6
PFP-only training (200 epochs)
–
97.64
80.5
Table 1. Classification accuracy and OOD-detection AUROC for SVI and PFP-based training variants.
Figure 3. Comparison of SVI and PFP uncertainty predictions. For MNIST, both uncertainties are expected to be low; Ambiguous-MNIST exhibits higher aleatoric uncertainty (Softmax Entropy), and Fashion-MNIST, as OOD data, shows higher epistemic uncertainty (Mutual Information). Both methods effectively assign the majority of images to their respective domains.
Figure 4. Disentangled Epistemic (MI) and Aleatoric (SME) Uncertainty. Comparison of SVI and PFP uncertainty predictions. While PFP performs comparably to SVI in estimating total uncertainty, its ability to disentangle aleatoric and epistemic uncertainties is somewhat limited, as anticipated. Nevertheless, the practical distinction remains sufficiently robust for most cases.
Figure 5. Performance comparison of operator implementations, evaluating the reformulation from Equation 5 to 11 and the use of separate vs. joint operators for mean and variance paths on a ARM Cortex-A72.
Figure 6. Execution time distribution across operator types for three PFP-BNN architectures, measured on a Cortex-A72 with a mini-batch size of 10. While dense layers dominate latency in the MLP and convolutions in ResNet-18, LeNet-5 shows a more balanced distribution. Notably, otherwise trivial operators such as ReLU and Max Pool exhibit increased computational cost. The tooling and profiling overhead is typically negligible but becomes noticeable in the low-latency MLP.
Optimizations
Latency
Speedup
Name
Other Opt.
without Opt.
with Opt.
Baseline (no tuning)
OFF
3.760 ms
-
-
Baseline (min. tuning)
OFF
3.681 ms
-
-
Tiling 1
OFF
3.672 ms
0.747 ms
4.91 ×
Loop Reordering
OFF
3.681 ms
1.940 ms
1.90 ×
Vectorization
OFF
3.681 ms
8.837 ms
0.42 ×
Table 2. Comparison of manual optimization techniques for the PFP dense operator. Measurements on a Cortex-A72, with a 3-layer MLP and a mini-batch size of 10.
Architecture
Implementation
Auto-tuning
Latency
Max Pools
Entire Network
LeNet-5
Generic Max Pool
No
12.09 ms
29.13 ms
LeNet-5
Generic Max Pool
All operators
5.04 ms
10.74 ms
LeNet-5
Generic Max Pool
All except Max Pool
11.92 ms
17.82 ms
LeNet-5
Vect. Max Pool k=2
No
3.54 ms
21.10 ms
LeNet-5
Vect. Max Pool k=2
All operators
27.28 ms
33.42 ms
Table 3. Comparison of Max Pool implementations. Measurements on a Cortex-A72 with a mini-batch size of 10.
Architecture
Layer
Baseline
Tuned Implementation
Speedup
Latency
Fraction
Latency
Fraction
MLP
Dense 1
2.931 ms
62.8 %
0.642 ms
33.7 %
4.6 ×
Dense 2
0.570 ms
12.2 %
0.125 ms
6.6 %
4.6 ×
ReLU 1
0.195 ms
4.2 %
0.185 ms
9.7 %
1.1 ×
Dense 3
0.063 ms
1.3 %
0.036 ms
1.9 %
1.8 ×
Sum
3.846 ms
82.4 %
1.078 ms
56.5 %
3.6 ×
Table 4. Profiling of PFP neural architectures to determine most costly layers on a Cortex-A72 with a mini-batch size of 10.
Figure 7. Speedup and latency in relation to mini-batch size. The SVI-BNN, evaluated using 30 samples, exhibits high latency and poor scalability for small mini-batch sizes, resulting in significantly higher per-image latency. In contrast, the PFP implementation maintains consistent performance across all mini-batch sizes due to targeted tuning. The observed speedup highlights the advantages of PFP, particularly for small mini-batch sizes, which are critical for low-latency processing in embedded systems.
Architecture
batch-size
Processor
Deterministic NN
SVI
PFP
Speedup
not tuned
tuned
not tuned
tuned
MLP
1
Cortex-A53
1.49 ms
0.396 ms
-
1.53 ms
0.548 ms
Cortex-A72
0.49 ms
0.094 ms
598.25 ms
0.42 ms
0.152 ms
3935.9 ×
Cortex-A76
0.18 ms
0.044 ms
226.76 ms
0.20 ms
0.062 ms
3657.4 ×
MLP
10
Cortex-A53
14.02 ms
0.933 ms
-
15.26 ms
4.989 ms
Cortex-A72
4.59 ms
0.186 ms
734.74 ms
3.75 ms
0.742 ms
990.2 ×
Table 5. Algorithm Performance Comparison on Embedded ARM Processors, using vectorized Max Pool for Dirty-MNIST.
Dataset
Prob.
SVI
PFP Calibrated
PFP Fine-tuned
PFP Latency
Layers
Acc.
AUROC
Calib.
Acc.
AUROC
Acc.
AUROC
Not tuned
Tuned
CIFAR-10
All
91.5%
89.2%
0.40
91.5%
94.1%
91.2%
95.9%
2366.3 ms
1024.7 ms
Last-only
92.5%
89.0%
0.05
89.3%
86.9%
92.0%
95.6%
1207.1 ms
36.1 ms
CINIC-10
All
81.5%
86.2%
0.45
62.7%
75.7%
63.9%
78.3%
2366.3 ms
1024.7 ms
Last-only
82.4%
81.0%
0.05
56.0%
50.3%
37.7%
60.2%
1207.1 ms
36.1 ms
Table 6. Quality and PFP inference latency of ResNet-18 variants with probabilistic partiality (Cortex-A72, mini-batch size of 1)
Bayesian neural networks place probability distributions over neural-network parameters or functions to quantify predictive uncertainty and support robust decisions.