[Paper Review] Wide Bayesian neural networks have a simple weight posterior: theory and accelerated sampling

Jiri Hron, Roman Novak+2 more|arXiv (Cornell University)|Jun 15, 2022

Gaussian Processes and Bayesian Inference4 citations

TL;DR

This paper introduces repriorisation, a data-dependent reparameterisation that transforms the posterior of wide Bayesian neural networks (BNNs) so its Kullback–Leibler (KL) divergence to the standard normal prior vanishes as width increases. By enabling faster mixing in Markov chain Monte Carlo (MCMC) sampling—achieving up to 50x higher effective sample size—the method significantly accelerates posterior inference in both fully-connected and residual networks, especially at wider widths.

ABSTRACT

We introduce repriorisation, a data-dependent reparameterisation which transforms a Bayesian neural network (BNN) posterior to a distribution whose KL divergence to the BNN prior vanishes as layer widths grow. The repriorisation map acts directly on parameters, and its analytic simplicity complements the known neural network Gaussian process (NNGP) behaviour of wide BNNs in function space. Exploiting the repriorisation, we develop a Markov chain Monte Carlo (MCMC) posterior sampling algorithm which mixes faster the wider the BNN. This contrasts with the typically poor performance of MCMC in high dimensions. We observe up to 50x higher effective sample size relative to no reparametrisation for both fully-connected and residual networks. Improvements are achieved at all widths, with the margin between reparametrised and standard BNNs growing with layer width.

Motivation & Objective

• To address the limited theoretical understanding and high computational cost of Bayesian neural networks (BNNs), especially in parameter space.
• To close the gap in understanding wide BNNs by characterizing their posterior behavior in weight space, complementing existing function-space results like the NNGP limit.
• To develop a practical, efficient MCMC sampling method for BNNs that scales favorably with width, overcoming typical high-dimensional sampling challenges.
• To enable faster and more effective posterior inference in BNNs by transforming the posterior to be increasingly close to a standard normal distribution as width grows.

Proposed method

• Introduces repriorisation, a reparameterisation map $ T( heta) $ that transforms BNN weights $ heta $ into a new parameterization $ heta = T( heta) $, such that the KL divergence between the reparametrized posterior and the standard normal prior vanishes as layer width increases.
• Proves theoretically (Theorem 2.1) that the KL divergence $ \mathrm{KL}(\mathcal{N}(0,I_d) \| p(\phi|\mathcal{D})) \to 0 $ as the minimum layer width $ d_{\min} \to \infty $, establishing the posterior's convergence to the prior in parameter space.
• Proposes a computationally efficient implementation using Cholesky decomposition to simultaneously compute the reparametrization and its Jacobian determinant, enabling tractable MCMC sampling.
• Applies Langevin Monte Carlo (LMC) to the reparametrized posterior, leveraging the fact that the gradient of the log-posterior concentrates around that of the standard normal, improving mixing speed.
• Uses a second-order Taylor approximation and Hutchinson's estimator to efficiently approximate the log-determinant of the Jacobian, reducing memory usage and enabling scalable gradient computation.
• Employs JAX’s `vjp` and `jvp` functions to compute gradients efficiently, with a custom forward-pass strategy to avoid out-of-memory errors during backpropagation.

Experimental results

Research questions

• RQ1 How does the posterior distribution of wide Bayesian neural networks behave in parameter space as network width increases?
• RQ2 Can a reparameterisation be designed such that the KL divergence between the reparametrized posterior and a standard normal prior vanishes in the infinite-width limit?
• RQ3 Does such a reparameterisation lead to faster mixing in MCMC sampling for BNNs, especially in high-dimensional weight spaces?
• RQ4 Can the method be efficiently implemented and scaled across different architectures, including residual networks and varying dataset sizes?
• RQ5 To what extent do the benefits of repriorisation persist outside the strict NNGP regime, particularly in practical, finite-width settings?

Key findings

• Repriorisation reduces the KL divergence between the reparametrized BNN posterior and the standard normal prior to zero in the infinite-width limit, providing a theoretical foundation for simplifying the posterior in parameter space.
• The method enables up to 50x higher effective sample size (ESS) in MCMC sampling compared to standard BNNs, with improvements observed across all widths and growing with increasing layer width.
• For a 3-hidden-layer fully-connected network with 1024 units per layer on CIFAR-10, a 10x improvement in ESS was observed, even far from the NNGP regime.
• On ResNet-20, a 10x improvement in ESS was achieved only when the top-layer width was comparable to or larger than the number of training examples, indicating regime dependence.
• The computational efficiency is enhanced via Cholesky-based Jacobian computation and memory-optimized gradient estimation using Hutchinson’s estimator and JAX’s `vjp`.
• The method remains effective with non-Gaussian likelihoods, as demonstrated with a categorical likelihood on a smaller experiment, where the reparametrized version achieved ~99% acceptance rate with a higher step size than the standard parametrization.