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

TL;DR

This paper introduces an efficient method for deploying Bayesian neural networks using a single deterministic forward pass, significantly reducing computational costs while maintaining accuracy and uncertainty estimation capabilities.

Contribution

The paper presents a novel end-to-end pipeline utilizing probabilistic forward pass approximation and code generation to efficiently deploy Bayesian neural networks on embedded systems.

Findings

PFP outperforms SVI in computational efficiency with up to 4200x speedup.

PFP-BNNs match SVI-BNNs in accuracy and uncertainty estimation.

The approach enables deployment of Bayesian models on resource-constrained devices.

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.