[Paper Review] Weight Uncertainty in Neural Networks
This paper introduces Bayes by Backprop, a differentiable, backpropagation-compatible algorithm that learns a probability distribution over neural network weights using variational inference. It achieves performance comparable to dropout on MNIST, improves generalization in regression via uncertainty-aware predictions, and enables natural exploration in reinforcement learning through Thompson sampling, all while scaling efficiently with standard deep learning hardware.
We introduce a new, efficient, principled and backpropagation-compatible algorithm for learning a probability distribution on the weights of a neural network, called Bayes by Backprop. It regularises the weights by minimising a compression cost, known as the variational free energy or the expected lower bound on the marginal likelihood. We show that this principled kind of regularisation yields comparable performance to dropout on MNIST classification. We then demonstrate how the learnt uncertainty in the weights can be used to improve generalisation in non-linear regression problems, and how this weight uncertainty can be used to drive the exploration-exploitation trade-off in reinforcement learning.
Motivation & Objective
- To address overfitting and poor uncertainty estimation in standard neural networks by introducing principled weight uncertainty.
- To develop a scalable, backpropagation-compatible method for Bayesian neural networks that enables uncertainty-aware predictions.
- To demonstrate how weight uncertainty improves generalization in regression and drives effective exploration in contextual bandit tasks.
- To provide a unified framework for regularisation, model averaging, and exploration using a single probabilistic weight distribution.
Proposed method
- Applies variational inference to approximate the true posterior over network weights using a parameterized distribution, minimizing the variational free energy (evidence lower bound).
- Uses reparameterization gradients to enable unbiased, differentiable optimization of the variational objective, compatible with standard backpropagation.
- Represents each weight as a distribution (e.g., Gaussian), allowing Monte Carlo sampling for stochastic forward passes and gradient estimation.
- Employs a compression cost (variational free energy) that regularizes weights by favoring simpler, more compressible distributions that explain the data well.
- Enables model averaging via Monte Carlo sampling from the weight posterior, without explicitly training multiple networks.
- Supports non-Gaussian priors and is compatible with standard deep learning frameworks, enabling efficient training on GPUs.
Experimental results
Research questions
- RQ1Can a differentiable, backpropagation-compatible method learn meaningful weight uncertainty in deep neural networks?
- RQ2Does weight uncertainty via variational inference improve generalization in regression tasks compared to standard networks?
- RQ3Can weight uncertainty naturally drive exploration in contextual bandit problems, outperforming epsilon-greedy strategies?
- RQ4How does the performance of Bayes by Backprop compare to established regularisation techniques like dropout on standard benchmarks?
Key findings
- Bayes by Backprop achieves test accuracy on MNIST comparable to dropout, demonstrating effective regularisation through weight uncertainty.
- In non-linear regression, the Bayesian network produces wider prediction intervals in data-sparse regions, reflecting uncertainty, unlike standard networks that are overconfident.
- The Bayesian agent in the mushroom contextual bandit task explores effectively from the start, achieving near-optimal regret much faster than epsilon-greedy agents.
- The epsilon-greedy agent with 5% exploration over-explores, while the greedy agent initially fails to explore due to poor initial policy, but eventually improves.
- The Bayes by Backprop agent maintains low regret throughout, showing that uncertainty-driven exploration is more systematic and efficient than epsilon-greedy.
- The method scales efficiently using standard deep learning hardware and optimisation schemes like asynchronous SGD.
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.