[Paper Review] Variational Dropout and the Local Reparameterization Trick

Diederik P. Kingma, Tim Salimans+1 more|UvA-DARE (University of Amsterdam)|Jun 8, 2015

Gaussian Processes and Bayesian Inference23 references461 citations

TL;DR

This paper introduces a local reparameterization trick that reduces variance in stochastic gradient variational inference for Bayesian neural networks, enabling efficient, parallelizable training with significantly faster convergence. The method generalizes Gaussian dropout by learning adaptive dropout rates, achieving state-of-the-art performance with up to 200x speedup and lower variance than standard baselines.

ABSTRACT

We investigate a local reparameterizaton technique for greatly reducing the variance of stochastic gradients for variational Bayesian inference (SGVB) of a posterior over model parameters, while retaining parallelizability. This local reparameterization translates uncertainty about global parameters into local noise that is independent across datapoints in the minibatch. Such parameterizations can be trivially parallelized and have variance that is inversely proportional to the minibatch size, generally leading to much faster convergence. Additionally, we explore a connection with dropout: Gaussian dropout objectives correspond to SGVB with local reparameterization, a scale-invariant prior and proportionally fixed posterior variance. Our method allows inference of more flexibly parameterized posteriors; specifically, we propose variational dropout, a generalization of Gaussian dropout where the dropout rates are learned, often leading to better models. The method is demonstrated through several experiments.

Motivation & Objective

• Address the high variance in stochastic gradient variational inference (SGVB) for Bayesian neural networks, which hinders efficient optimization.
• Improve the efficiency and scalability of variational inference by reparameterizing global parameter uncertainty into local noise per data point.
• Generalize Gaussian dropout by learning optimal dropout rates instead of fixing them, enabling more flexible and adaptive regularization.
• Demonstrate that the proposed method achieves faster convergence and better generalization than standard dropout and baseline variational inference.

Proposed method

• Introduce a local reparameterization trick that transforms uncertainty in global model weights into independent, local noise per data point in a minibatch.
• Use the reparameterization trick to derive a stochastic gradient estimator with variance inversely proportional to minibatch size, enabling efficient and parallelizable optimization.
• Formulate variational dropout as a generalization of Gaussian dropout, where dropout rates are learned parameters of the posterior distribution over weights.
• Apply the local reparameterization to the variational lower bound (ELBO), enabling backpropagation through the noise using local, data-point-specific random variables.
• Use a scale-invariant prior and fixed posterior variance to recover Gaussian dropout as a special case of the proposed framework.
• Implement the method using a single random sample per weight per data point, drastically reducing the number of random variables compared to standard reparameterization.

Experimental results

Research questions

• RQ1 Can local reparameterization reduce the variance of stochastic gradients in variational Bayesian inference for neural networks, leading to faster convergence?
• RQ2 How does the proposed method compare to standard reparameterization and baseline dropout in terms of training speed and test accuracy?
• RQ3 Can variational dropout, with learned dropout rates, outperform fixed-rate dropout and Gaussian dropout in generalization?
• RQ4 What is the impact of the local reparameterization trick on computational efficiency and parallelization in minibatch training?
• RQ5 Does the method maintain or improve performance across different network architectures and dataset sizes?

Key findings

• The local reparameterization estimator achieves variance scaling as $1/M$, where $M$ is the minibatch size, leading to significantly lower gradient variance than standard estimators.
• The method enables a 200x speedup in training: 1635 seconds per epoch with standard SGVB vs. 7.4 seconds per epoch with the local reparameterization trick on a modern GPU.
• Variational dropout with adaptive dropout rates outperforms standard dropout and non-adaptive Gaussian dropout, especially in smaller networks where dropout often causes underfitting.
• The adaptive variational dropout method achieves lower test error than all baselines on both MNIST and CIFAR-10, with the best performance observed when the KL divergence term is down-scaled by a factor of 3.
• The variance of the stochastic gradient estimator is lowest with the local reparameterization method, and it remains stable across training, outperforming estimators that sample weights per data point or per minibatch.
• The method generalizes Gaussian dropout and provides a principled Bayesian interpretation of dropout as variational inference with local noise.