[Paper Review] Building Bayesian Neural Networks with Blocks: On Structure, Interpretability and Uncertainty
This paper introduces a block-based framework for constructing Bayesian Neural Networks (BNNs) that enhances interpretability and uncertainty estimation by structuring models as additive or hierarchical components. By leveraging computation skeletons and doubly stochastic variational inference, the approach enables efficient posterior approximation, with a specialized Bayesian Additive Neural Network (AddNN) variant demonstrating superior performance and sparsity in capturing interactions on benchmark datasets.
We provide simple schemes to build Bayesian Neural Networks (BNNs), block by block, inspired by a recent idea of computation skeletons. We show how by adjusting the types of blocks that are used within the computation skeleton, we can identify interesting relationships with Deep Gaussian Processes (DGPs), deep kernel learning (DKL), random features type approximation and other topics. We give strategies to approximate the posterior via doubly stochastic variational inference for such models which yield uncertainty estimates. We give a detailed theoretical analysis and point out extensions that may be of independent interest. As a special case, we instantiate our procedure to define a Bayesian {\em additive} Neural network -- a promising strategy to identify statistical interactions and has direct benefits for obtaining interpretable models.
Motivation & Objective
- To develop a modular, interpretable framework for building Bayesian Neural Networks (BNNs) by composing structured blocks.
- To enable efficient posterior approximation in BNNs using doubly stochastic variational inference, supporting mini-batch training and scalability.
- To establish theoretical and practical links between BNNs, Deep Gaussian Processes (DGPs), deep kernel learning, and random features.
- To design a Bayesian Additive Neural Network (AddNN) that explicitly models sparse additive interactions for improved interpretability.
- To evaluate the method on benchmark datasets, demonstrating competitive predictive performance and uncertainty quantification.
Proposed method
- The framework uses a computation skeleton abstraction to compose BNNs from reusable blocks, enabling modular design of model structure.
- Each block corresponds to a probabilistic component (e.g., a sub-network), with priors placed on weights to enable Bayesian inference.
- Doubly stochastic variational inference is employed to approximate the posterior, allowing mini-batch training and backpropagation.
- The AddNN variant constructs the final output as a sum of small, independent sub-networks, promoting additive structure and sparsity.
- Group Lasso regularization is applied to the first layer of each sub-network to encourage sparse, interpretable input feature selection.
- The method supports uncertainty estimation through posterior sampling and provides empirical $β$-norm interaction measures for detecting feature interactions.
Experimental results
Research questions
- RQ1Can a block-based construction of BNNs enable better interpretability while maintaining uncertainty quantification?
- RQ2How do different block types relate to existing models such as Deep Gaussian Processes (DGPs) and deep kernel learning?
- RQ3Can a Bayesian additive neural network structure effectively learn sparse, interpretable interactions in high-dimensional data?
- RQ4Does the proposed doubly stochastic variational inference scheme scale efficiently to large datasets while preserving model accuracy?
- RQ5How do interaction measures derived from the model help in identifying meaningful feature dependencies?
Key findings
- The Bayesian Additive Neural Network (AddNN) achieves state-of-the-art or competitive performance on benchmark regression tasks, with RMSE of 3.03±0.12 on Boston, 0.06±0.00 on Kin8nm, and 3.68±0.03 on Power.
- AddNN outperforms MC dropout and PBP in several datasets, including Energy (0.65±0.03 vs. 1.66±0.04) and Protein (4.07±0.01 vs. 4.36±0.01).
- The model learns a sparse additive structure, as visualized by the input-to-subnetwork weight matrix, which shows selective, non-uniform feature engagement.
- Computationally, AddNN scales as $k^{*}d^{2L-1}$, requiring ${k^{*}}^{2(L-1)}$ times less computation than a standard DNN with equivalent depth and width.
- The use of $\mathbb{E}^n$ as a baseline for interaction measures enables better empirical comparison and detection of feature interactions than $\delta(\mathbf{x}^0)$.
- The framework successfully unifies BNNs with DGPs, deep kernel learning, and random features under a common block-based structure with theoretical grounding.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.