Source-linked AI summary
Bayesian Dark Knowledge
Anoop Korattikara, Vivek Rathod, Kevin Murphy, Max Welling
TL;DR
The paper addresses Bayesian parameter estimation for neural networks when accurate predictive uncertainty matters but standard approaches are costly or limited. It distills an online MCMC posterior predictive distribution into a single neural network and reports better performance than SGD, EP, and VB, with reduced test-time cost.
Problem
Neural-network point estimates can ignore parameter uncertainty, while online MCMC requires storing and evaluating many parameter samples.
Method
The method trains a student neural network online to approximate the posterior predictive distribution produced by an SGLD teacher by minimizing KL(q(y|x)||S(y|x, w)).
Results
The method improves test-set log likelihood scores over SGD and recently proposed EP and VB approaches.
Takeaways & Limitations
Distillation retains Bayesian predictive behavior in a single network, reducing SGLD’s storage and test-time costs without sacrificing much accuracy in the MNIST experiment.
Takeaways & Limitations
The methods were not evaluated uniformly across all datasets because EP code supported only regression, BBB results were quoted, and HMC was too slow for large problems.
Abstract
from arXiv · showhide
We consider the problem of Bayesian parameter estimation for deep neural networks, which is important in problem settings where we may have little data, and/ or where we need accurate posterior predictive densities, e.g., for applications involving bandits or active learning. One simple approach to this is to use online Monte Carlo methods, such as SGLD (stochastic gradient Langevin dynamics). Unfortunately, such a method needs to store many copies of the parameters (which wastes memory), and needs to make predictions using many versions of the model (which wastes time). We describe a method for "distilling" a Monte Carlo approximation to the posterior predictive density into a more compact form, namely a single deep neural network. We compare to two very recent approaches to Bayesian neural networks, namely an approach based on expectation propagation [Hernandez-Lobato and Adams, 2015] and an approach based on variational Bayes [Blundell et al., 2015]. Our method performs better than both of these, is much simpler to implement, and uses less computation at test time.
1 Introduction
Bayesian inference addresses overconfident neural-network predictions, but existing approximations trade accuracy, speed, memory, or implementation simplicity. The paper proposes distilling online MCMC predictions into one neural network and reports improved test log likelihoods over SGD, EP, and VB.
- DNNs often produce overconfident predictions, undermining uncertainty estimates needed for active learning, reinforcement learning, and classifier fusion.
- Point estimates such as MAP and plug-in predictions ignore parameter uncertainty, losing much of Bayesian inference’s benefit.
- EP and VB can struggle with nonfactorized, multimodal, or skewed posteriors, while also increasing predictive computation, memory use, and implementation complexity.
- Online MCMC methods provide potentially more accurate posterior approximations and simpler implementation, but require S times more test computation and parameter storage.SGLD adds Gaussian noise to SGD updates, but predictions average across sampled models.
- The proposed method trains a student network online to minimize KL(q(y|x)||S(y|x, w)) against an SGLD-generated teacher.This distills the Monte Carlo posterior predictive distribution into a single parametric model.
- The approach combines online MCMC with model distillation and reports improved test-set log likelihoods compared with SGD, EP, and VB.The student uses deep neural networks and online training, supporting larger datasets than earlier batch or mixture-model distillation approaches.
2 Methods
The method trains a single student neural network to approximate a Monte Carlo teacher’s Bayesian predictive distribution, using online training and teacher-provided probability labels. It applies this framework to classification and regression while avoiding the need to retain all parameter and data samples.
- The student neural network approximates the Bayesian predictive distribution produced by a Monte Carlo ensemble of teacher neural networks.
- The teacher uses SGLD to estimate the posterior predictive distribution while the student is trained online to minimize KL divergence from it.
- Student training uses an unlabeled dataset whose probability-distribution labels come from the teacher, with sampling choices defining the domain of accurate predictions.For higher-dimensional problems, inputs can be sampled near the training data by perturbing them slightly.
- The stochastic algorithm uses one posterior parameter sample and a minibatch of student inputs at each step instead of precomputing and storing all samples.
- 2.1 Classification: For classification, both teacher and student use softmax outputs, making the estimated distillation loss the standard cross-entropy loss.Student gradients are computed with respect to log softmax outputs and back-propagated through the network.
- 2.2 Regression: For regression, the student predicts the mean and log variance of a Gaussian predictive distribution, doubling the teacher’s output parameters to capture data-dependent variance.Predicting log variance via e^α avoids positivity constraints on the variance during training.
3 Experimental results
Experiments compare SGLD and distilled SGLD with plugin SGD, EP/VB-based methods, and HMC across classification and regression tasks. Distillation closely preserves predictive performance while reducing storage and test-time computation.
- Experimental setup: The experiments compare SGD, SGLD, distilled SGLD, PBP, BBB, and HMC across classification and regression problems.The configurations include toy classification, MNIST, and regression experiments.
- Toy 2d classification: SGLD closely approximates the HMC posterior predictive distribution on the toy 2d classification task.HMC uses 20k samples, while SGLD uses about 1,000 retained samples after burn-in and thinning.
- Toy 2d classification: The 2-10-10-2 distilled student performs best among the tested student architectures on the toy 2d task.The 2-10-2 student is too simple, the 2-100-2 student works better, and the two-hidden-layer 2-10-10-2 model performs best.
- MNIST classification: 1.271% test error is obtained by SGLD on MNIST, while distilled SGLD obtains 1.307%, close to the teacher's performance.The SGLD result is better than the reported SGD, dropout, and BBB results; distillation uses a single student network.
- Computational costs: 10,000-model prediction storage and test-time costs are reduced by a factor of 10,000 through distillation, without sacrificing much accuracy.Per-iteration training times are 1.3 ms for SGD, 1.6 ms for SGLD, and 3.2 ms for distilled SGLD.
- Predictive quality: Distilled SGLD incurs little predictive loss relative to SGLD and produces better-calibrated MNIST predictions than SGD.On the toy 1d regression problem, distillation incurs little accuracy loss while saving computational cost; MNIST log-likelihood results favor SGLD and distilled SGLD over SGD.
4 Conclusions and future work
The paper presents a simple method that combines online MCMC with model distillation for scalable Bayesian inference, and reports better performance than EP- and VB-based alternatives.
- The method combines online MCMC with model distillation to provide a simple, scalable approach to Bayesian inference for neural networks and other models.
- The authors report that their approach works better than recently proposed alternatives based on expectation propagation and variational Bayes.
- Future work includes testing predictive uncertainty in contextual bandits or active learning.
- Future directions also include reducing algorithmic variance, improving student-data generation, and testing confident false predictions on adversarial inputs.