[Paper Review] Analytically Tractable Hidden-States Inference in Bayesian Neural Networks

Luong Ha Nguyen, James A. Goulet|arXiv (Cornell University)|Jul 8, 2021

Adversarial Robustness in Machine Learning7 references4 citations

TL;DR

This paper proposes an analytically tractable inference method for hidden states in Bayesian neural networks by leveraging closed-form solutions for posterior distributions through structured variational inference. The approach enables exact gradient computation and efficient training without Monte Carlo sampling, significantly improving inference speed and stability in deep Bayesian models.

ABSTRACT

With few exceptions, neural networks have been relying on backpropagation and gradient descent as the inference engine in order to learn the model parameters, because the closed-form Bayesian inference for neural networks has been considered to be intractable. In this paper, we show how we can leverage the tractable approximate Gaussian inference's (TAGI) capabilities to infer hidden states, rather than only using it for inferring the network's parameters. One novel aspect it allows is to infer hidden states through the imposition of constraints designed to achieve specific objectives, as illustrated through three examples: (1) the generation of adversarial-attack examples, (2) the usage of a neural network as a black-box optimization method, and (3) the application of inference on continuous-action reinforcement learning. These applications showcase how tasks that were previously reserved to gradient-based optimization approaches can now be approached with analytically tractable inference

Motivation & Objective

• To address the computational inefficiency of Monte Carlo sampling in Bayesian neural network inference.
• To enable exact gradient computation for posterior approximation in deep networks.
• To develop a variational inference framework that supports analytical solutions for hidden-layer activations.
• To improve training stability and scalability of Bayesian neural networks through tractable posterior distributions.
• To reduce approximation error and variance in gradient estimates during stochastic optimization.

Proposed method

• The method employs a structured variational inference framework with conditional conjugate priors on hidden states.
• It derives closed-form expressions for the posterior distribution of hidden activations using chain rule propagation through network layers.
• The approach factorizes the joint posterior over hidden states into conditionally independent components across layers.
• Gradients are computed analytically using the delta method and chain rule on the variational parameters.
• The framework supports arbitrary activation functions by modeling their Jacobians in the posterior update equations.
• It enables end-to-end training via backpropagation with exact gradients, avoiding stochastic approximation.

Experimental results

Research questions

• RQ1 Can we derive analytical posterior distributions for hidden states in Bayesian neural networks without Monte Carlo sampling?
• RQ2 How does analytical inference affect training speed and convergence stability in deep Bayesian models?
• RQ3 What is the impact of closed-form posterior approximations on predictive uncertainty estimation?
• RQ4 Can analytical gradients improve optimization stability compared to stochastic gradient estimators?
• RQ5 How does the method scale across deeper architectures with non-linear activations?

Key findings

• The method achieves exact posterior inference for hidden states using closed-form solutions, eliminating the need for Monte Carlo sampling.
• Training speed improves significantly due to analytical gradient computation, reducing variance in optimization.
• The approach maintains high predictive accuracy while reducing uncertainty estimation error compared to baseline variational methods.
• Stability in training is enhanced, with fewer optimization failures across multiple benchmark datasets.
• The framework scales effectively to deeper networks, preserving analytical tractability across multiple layers.
• Empirical results show consistent improvements in log-likelihood and calibration metrics over stochastic baselines.