Wide Bayesian neural networks have a simple weight posterior: theory and accelerated sampling
Jiri Hron, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 2da87782-6aca-4264-8aa9-5f627bbf53bbCited by top-tier papers1
Ask how each one uses itBuilds on11
- What Are Bayesian Neural Network Posteriors Really Like?Pavel Izmailov, Sharad Vikram, Matthew D. Hoffman, Andrew Gordon WilsonICML 2021 · 458 citations
- Neural Tangents: Fast and Easy Infinite Neural Networks in PythonRoman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee et al.ICLR 2020 · 254 citations
- Efficient and Scalable Bayesian Neural Nets with Rank-1 FactorsMichael Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-An Ma et al.ICML 2020 · 239 citations
- Infinite attention: NNGP and NTK for deep attention networksJiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, Roman NovakICML 2020 · 147 citations
- Bayesian Deep Learning via Subnetwork InferenceErik A. Daxberger, Eric T. Nalisnick, James Urquhart Allingham, Javier Antorán et al.ICML 2021 · 108 citations
Related papers
- Parameter Expanded Stochastic Gradient Markov Chain Monte CarloHyunsu Kim, Giung Nam, Chulhee Yun, Hongseok Yang et al.ICLR 2025
- Posterior Refinement Improves Sample Efficiency in Bayesian Neural NetworksAgustinus Kristiadi, Runa Eschenhagen, Philipp HennigNeurIPS 2022 · 17 citations
- Learning to Explore for Stochastic Gradient MCMCSeunghyun Kim, Seohyeon Jung, Seonghyeon Kim, Juho LeeICML 2024 · 2 citations
- Efficient Low Rank Gaussian Variational Inference for Neural NetworksMarcin Tomczak, Siddharth Swaroop, Richard E. TurnerNeurIPS 2020 · 37 citations
- Reparameterization invariance in approximate Bayesian inferenceHrittik Roy, Marco Miani, Carl Henrik Ek, Philipp Hennig et al.NeurIPS 2024 · 20 citations
