用单次前向传播实现高效贝叶斯神经网络,适合嵌入式设备部署。
Accelerated Execution of Bayesian Neural Networks using a Single Probabilistic Forward Pass and Code Generation
- 通过高斯假设将采样变为单次确定性前向传播,大幅降低计算开销。
- 在小批量下速度提升最高达4200倍,准确率与不确定性估计媲美传统方法。
- 专为嵌入式ARM CPU设计,结合代码生成与自动调优,适合边缘计算场景。
机器学习模型在诊断、天气预测、自然语言处理和自动驾驶等领域表现优异,但其不确定性建模能力有限,限制了在安全关键场景的应用。传统神经网络难以识别域外数据(OOD),可能输出高置信度的错误预测。贝叶斯神经网络(BNN)通过概率估计缓解此问题,但因需多次采样和前向传播导致计算成本高昂。本文提出的概率前向传播(PFP)方法假设权重与激活服从高斯分布,实现解析式的不确定性传播,以单次确定性前向传播替代采样,显著提升效率。我们构建了端到端流程,涵盖训练、编译、优化与部署,基于TVM深度学习编译器,实现适用于多层感知机与卷积神经网络的高斯传播算子库,并结合手动与自动化调优策略。消融实验表明,PFP在计算效率上持续优于变分推断(SVI),小批量下最快提升4200倍;在Dirty-MNIST数据集上,PFP-BNN的准确率、不确定性估计与域外检测性能均与SVI-BNN相当,同时大幅降低计算开销。结果表明,结合贝叶斯近似与代码生成,可实现资源受限系统上的高效BNN部署。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。