Gibbs Sampling the Posterior of Neural Networks

TL;DR

This work tackles sampling from the neural-network weight posterior $P(W|X,y)$ by introducing an intermediate-noise generative model that augments activations with Gaussian noise and a Gibbs sampler tailored to the augmented posterior. It introduces a teacher–student thermalization criterion to certify sampling equilibrium and derives conditional Gaussian updates that enable efficient Gibbs steps, including handling non-Gaussian activation variables via truncation. Empirically, the Gibbs approach on the intermediate-noise posterior is competitive with HMC and MALA for small architectures and offers a hyperparameter-free, parallelizable alternative, though memory costs limit scalability. Overall, the paper provides a principled diagnostic and practical Bayesian inference pathway for uncertainty estimation in compact neural networks.

Abstract

In this paper, we study sampling from a posterior derived from a neural network. We propose a new probabilistic model consisting of adding noise at every pre- and post-activation in the network, arguing that the resulting posterior can be sampled using an efficient Gibbs sampler. For small models, the Gibbs sampler attains similar performances as the state-of-the-art Markov chain Monte Carlo (MCMC) methods, such as the Hamiltonian Monte Carlo (HMC) or the Metropolis adjusted Langevin algorithm (MALA), both on real and synthetic data. By framing our analysis in the teacher-student setting, we introduce a thermalization criterion that allows us to detect when an algorithm, when run on data with synthetic labels, fails to sample from the posterior. The criterion is based on the fact that in the teacher-student setting we can initialize an algorithm directly at equilibrium.

Figures (5)

• Figure 1: Comparison of different thermalization measures. In the legend, next to each method we write between parentheses the initialization (or pair of initializations) the method is applied to. The circles on the $x$ axis represent the thermalization times estimated by each method. Left: We compare the predictions for the thermalization time of the zero-initialized MCMC. The red $y$ scale on the right refers uniquely to the lines in red. All the other quantities should be read on the black $y$ scale. Right: We compare the predictions for the thermalization time of two chains initialized independently at random. The pink $y$ scale refers uniquely to the pink line. All other quantities should be read on the black logarithmic scale. The randomly initialized runs fail to thermalize and their test MSEs get stuck on a plateau. However, $\hat{R}$, whose time series on the plateau is stationary and close to 1, fails to detect this lack of thermalization.
• Figure 2: Thermalization experiments on synthetic data. Left: Proportion of the 72 runs that thermalize plotted against the equilibrium test MSE. Right: Example of the dynamics of the test MSE in a particular run where all four algorithms thermalize. In order to get a similar equilibrium test MSE in the classical and intermediate noise posteriors, we pick respectively $\Delta=10^{-3}$ and $\Delta=4.64\times10^{-4}$. The transparent lines represent the informed initializations.
• Figure 3: Gibbs on the intermediate noise posterior and HMC, MALA both on the classical posterior, compared on MNIST. Left: MLP with one hidden layer with $12$ hidden units. Right: CNN network.
• Figure 4: Percentiles of $\hat{R}$ as a function of time. A line marked with the number $k$ in the legend represents how the $k$-th percentile of $\hat{R}$ changes throughout the simulation. The data comes from the same simulation that was used for computing the average $\hat{R}$ in figure \ref{['fig:convergence_measures']}. The red dashed horizontal line is placed at a height of 1, the value that $\hat{R}$ should approach when the chains are close to each other. Left: percentiles of $\hat{R}$, computed on two chains with respectively zero and informed initialization. Right: percentiles when computing $\hat{R}$ from two chains independently initialized at random.
• Figure 5: CNN architecture used in the experiments of section \ref{['sec:real_data_experiments']}. The convolutional layer is composed of the filter $W^{(1)}$ with shape $2\times1\times4\times4$ and a output channel bias $b^{(1)}\in\mathbb R^2$. The final layer instead has weights $W^{(2)}\in\mathbb R^{72\times 10}$ and bias $b^{(2)}\in \mathbb R^{10}$.