PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam

Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, Akash Srivastava

arXiv:1806.04854v3stat.MLcs.AIcs.LGstat.CO

TL;DR

Deep-learning systems need uncertainty estimates, but Bayesian inference and existing VI methods are costly to implement and run at scale. The paper develops natural-gradient Gaussian mean-field VI algorithms that use weight perturbations within Adam, achieving comparable uncertainty estimates with lower resource and implementation demands.

  • Problem

    Deep-learning applications need reliable uncertainty estimates, while Bayesian inference is computationally challenging and existing VI methods require more effort than MLE.

  • Method

    The paper develops natural-gradient algorithms for Gaussian mean-field VI that perturb network weights during gradient evaluations and derive uncertainty from adaptive learning-rate vectors.

  • Results

    The proposed algorithms obtain uncertainty estimates comparable to existing VI methods while requiring less computational and implementation effort.

  • Takeaways & Limitations

    VI can be performed within Adam-like adaptive optimization with minimal code changes, and weight perturbation may support exploration in reinforcement learning and stochastic optimization.

  • Takeaways & Limitations

    The authors identify generalizing the natural-gradient approach to other distribution families and further validating reinforcement-learning and stochastic-optimization applications as future work.

Abstract

from arXiv · show

Uncertainty computation in deep learning is essential to design robust and reliable systems. Variational inference (VI) is a promising approach for such computation, but requires more effort to implement and execute compared to maximum-likelihood methods. In this paper, we propose new natural-gradient algorithms to reduce such efforts for Gaussian mean-field VI. Our algorithms can be implemented within the Adam optimizer by perturbing the network weights during gradient evaluations, and uncertainty estimates can be cheaply obtained by using the vector that adapts the learning rate. This requires lower memory, computation, and implementation effort than existing VI methods, while obtaining uncertainty estimates of comparable quality. Our empirical results confirm this and further suggest that the weight-perturbation in our algorithm could be useful for exploration in reinforcement learning and stochastic optimization.

1. Introduction

Bayesian deep learning can provide uncertainty estimates needed for reliable decisions, but existing VI methods demand more resources and code changes than MLE. The paper introduces natural-gradient algorithms that integrate with adaptive optimizers through weight perturbation.

  • Motivation: Uncertainty estimates are important for applying deep learning to safety-sensitive areas such as robotics and medical diagnostics.Without them, decisions may be unreliable and sometimes have disastrous consequences.
  • Motivation: VI scales better than MCMC for large models, but existing VI methods require more computation, memory, and implementation effort than MLE.VI commonly uses more parameters, while existing codebases require substantial modification.
  • Contribution: The paper proposes natural-gradient algorithms for Gaussian mean-field VI that require minimal changes to adaptive learning-rate codebases.The main algorithmic change is perturbing network weights during gradient computation.
  • Contribution: The proposed algorithms obtain uncertainty estimates from the vector that adapts the learning rate, reducing memory, computation, and implementation effort.The paper reports uncertainty estimates of comparable quality to existing VI methods.
  • Connections: The approach is related to exploration methods for avoiding local minima and may support exploration in reinforcement learning and stochastic optimization.Related approaches include natural evolution strategies, entropy-SGD, and noisy networks.

2. Gaussian Mean-Field Variational Inference

The paper frames Bayesian neural-network inference as a difficult posterior-estimation problem and uses Gaussian mean-field VI to make the approximation tractable. Direct adaptive VI remains costly, motivating algorithms that fit within Adam-like implementations with minimal code changes.

  • Model: The model treats independently sampled data examples as generated by a probabilistic deep neural network with weights θ.The likelihood may belong to an exponential family whose mean is the network output.
  • Maximum-Likelihood Estimation: MLE estimates a single parameter vector efficiently with stochastic-gradient methods such as RMSProp, AdaGrad, and Adam.These methods are popular partly because their implementations are simple and efficient.
  • Bayesian Inference: Bayesian inference instead estimates the posterior over θ, but its normalization constant is difficult to compute in large, nonconjugate deep models.Both model size and nonconjugacy contribute to the computational challenge.
  • Gaussian Mean-Field VI: Gaussian mean-field VI approximates the posterior with a Gaussian distribution having diagonal covariance and parameters µ and σ.The variational parameters are optimized through a variational objective.
  • Computational Cost: Direct stochastic-gradient VI updates both µ and σ, increasing the parameter and optimization state compared with MLE.Adaptive methods additionally store scaling vectors for both variational parameters.
  • Proposed Direction: The proposed algorithms approximate natural-gradient VI and can be implemented within Adam with minimal code changes.This design addresses the extra memory, computation, and implementation effort of direct adaptive VI.

3. Approximate Natural-Gradient VI

The paper develops approximate natural-gradient methods for Gaussian mean-field VI, progressively adapting them toward RMSprop-like updates that are easier to implement. The approximations trade computational simplicity against uncertainty-estimation accuracy.

  • Natural-gradient VI: Natural-gradient VI uses expectation parameters to avoid directly computing the Fisher information matrix.For Gaussian mean-field VI, this yields the NGVI update, whose learning rate is adapted by the variance.
  • Variational Online-Newton (VON): VON rewrites NGVI using stochastic gradients and an online diagonal-Hessian estimate, enabling backpropagation without variational-objective gradients.Its scaling vector estimates the diagonal Hessian, motivating the name variational online-Newton.
  • Variational Online Gauss-Newton (VOGN): The GGN approximation keeps variance estimates positive when initialized positively, producing the VOGN algorithm.VOGN avoids VON’s negative-variance failure mode, although GGN solutions can differ from VON because sampled variational parameters make approximation error nonzero.
  • Variational Online Gauss-Newton (VOGN): VOGN is difficult to implement in standard deep-learning codebases because it requires individual per-example gradients rather than only minibatch gradient sums.The paper introduces another approximation to address this implementation requirement.
  • Variational RMSprop (Vprop): Gradient magnitude replaces the GGN’s sum of squared gradients with the square of the summed gradient, yielding the variational RMSprop algorithm Vprop.Vprop samples perturbed weights, adds a prior term, and uses the learning-rate adaptation vector to obtain variance-based uncertainty estimates.
  • Analysis of the GM approximation: For minibatch size M = 1, gradient magnitude is unbiased for GGN; for M = N, it contains no second-order information.The paper therefore presents VOGN with M = 1 as a more accurate uncertainty option, while Vprop with M > 1 is easier and faster but may yield slightly worse estimates.

4. Variational Adam (Vadam)

The paper extends its natural-gradient framework with momentum to obtain an Adam-like variational optimizer. The resulting Vadam update is derived from natural momentum and uses sampled weights during gradient evaluation.

  • Natural momentum: Natural momentum replaces ordinary gradients with natural gradients in Polyak’s heavy-ball framework.For Gaussian q(θ), the resulting update can be expressed as a VON update with momentum.
  • Adam-like adaptation: The momentum-based update uses adapted learning rates and is designed to resemble Adam.The derivation includes an approximation rather than being an exact update for the earlier natural-momentum formulation.
  • Variational Adam (Vadam): Applying the Vprop-style derivation produces variational Adam, or Vadam, with pseudocode provided in Fig. 1.The implementation uses Adam learning rates rather than selecting them through the intermediate adaptive-rate expressions.

5. Variational AdaGrad (VadaGrad)

The paper extends variational inference toward optimization by introducing an interpolation parameter and derives VadaGrad as a variational AdaGrad method.

  • 5. Variational AdaGrad (VadaGrad): VadaGrad arises by modifying Vprop and setting τ = 0, yielding a variational version of AdaGrad.The τ parameter interpolates between variational inference and variational optimization.
  • 5. Variational AdaGrad (VadaGrad): Variational optimization minimizes E_q[F(θ)] rather than F(θ) directly, with equality recovered as σ2 → 0.The expectation under q(θ) = N(θ|µ, σ2) serves as a surrogate objective.
  • 5. Variational AdaGrad (VadaGrad): The interpolation parameter τ connects VI at τ = 1 with variational optimization at τ = 0.Intermediate values provide a continuum between the two objectives.
  • 5. Variational AdaGrad (VadaGrad): Unlike Vprop and Vadam, VadaGrad’s scaling vector is a weighted sum of past gradient magnitudes and never decreases.Consequently, its variance estimate never expands.

6. Results

Experiments evaluate uncertainty quality, predictive performance, convergence, and reinforcement-learning exploration across logistic regression, neural networks, and deep RL.

  • 6.1. Uncertainty Estimation in Logistic Regression: On USPS-3vs5, ELBO and log-loss are comparable across methods, while symmetric KL reveals differences in posterior-approximation quality.VOGN-1 performs well on KL, whereas Vadam performs worse because minibatch GM approximations introduce bias for M > 1.
  • 6.1. Uncertainty Estimation in Logistic Regression: As minibatch size decreases, Vadam’s Breast-Cancer performance improves and closely matches VOGN-1 at M = 1.The methods still differ because Vadam uses the square-root over st.
  • 6.2. Uncertainty Estimation in Neural Network: In BNN regression, BBVI and Vadam perform comparably on test RMSE and test log-likelihood, while MC-Dropout outperforms both.The comparison uses the benchmark setup and paired testing described in the experiment.
  • 6.2. Uncertainty Estimation in Neural Network: VOGN and Vadam often converge faster than BBVI, with Vadam performing well at M = 1 on Australian-Scale.Figure 3 varies minibatch size and MC samples; VOGN converges fastest in the displayed Australian-Scale experiments.
  • 6.3. Exploration in Deep Reinforcement Learning: In Half-Cheetah, Vadam and VadaGrad learn faster than SGD-Plain and SGD-Explore according to cumulative rewards over training iterations.Against Adam variants, their advantage remains mainly early because Adam-based methods catch up quickly.

7. Discussion

The discussion presents the algorithms as simple Adam-compatible VI methods, while acknowledging a trade-off between implementation ease and uncertainty-approximation quality.

  • 7. Discussion: The proposed VI algorithms use approximations and natural momentum to achieve Adam-compatible updates with minimal implementation changes.The authors report lower computational and implementation effort than existing VI methods.
  • 7. Discussion: The paper concludes that its methods obtain uncertainty estimates comparable to existing VI methods while requiring less computational and implementation effort.This is the central empirical and practical conclusion stated by the authors.
  • 7. Discussion: The revised conclusions identify a trade-off between ease of implementation and the quality of uncertainty approximation.This qualification narrows the broader simplicity claim.
  • 7. Discussion: The Gaussian mean-field derivation assumes diagonal covariance, and stochastic gradients are used in the resulting update.The update also requires positive precision; Gauss-Newton approximation is one proposed way to enforce this.
  • 7. Discussion: The reparameterization-based Hessian approximation can have high variance and produce negative st, requiring backtracking to remain feasible.The backtracking condition enforces sd + βthd > 0 for every coordinate.

E.1. Adam as an Adaptive Heavy-Ball Method

This section connects Adam’s adaptive heavy-ball form to natural-momentum VI, showing how Gaussian mean-field updates become implementable through Adam-like scaling and perturbation.

  • E.1. Adam as an Adaptive Heavy-Ball Method: Adam can be expressed as an adaptive version of Polyak’s heavy-ball method.This representation provides the optimization structure used to derive the natural-momentum analogue.
  • E.1. Adam as an Adaptive Heavy-Ball Method: Using natural gradients in the heavy-ball formulation yields an update similar to Adam, enabling approximated NGVI methods to use Adam.The resulting natural-momentum term is justified through the derivation.
  • E.1. Adam as an Adaptive Heavy-Ball Method: Natural momentum replaces Euclidean distance with KL divergence in the corresponding mirror-descent optimization problem.The KL term explains the name natural-momentum.
  • E.1. Adam as an Adaptive Heavy-Ball Method: For Gaussian approximations, the derivation expresses natural-momentum updates in terms of the mean µ and covariance Σ.The method then applies mean-field and gradient-magnitude approximations to obtain Vprop- and VadaGrad-like updates.
  • E.1. Adam as an Adaptive Heavy-Ball Method: Vadam becomes implementable through Adam’s update after bias correction, with differences from Adam confined to the perturbation-based variational update.The method samples θt from the Gaussian approximation during gradient evaluation.
  • E.1. Adam as an Adaptive Heavy-Ball Method: Setting τ = 0 and applying the gradient-magnitude and square-root approximations produces VadaGrad with σ2_t := 1/st.This links the optimization endpoint of the interpolation to the variational AdaGrad update.

G. Proof of Theorem 1

The proof establishes fixed-point properties for the proposed updates under full-batch, exact-expectation assumptions, and examines how scaling affects convergence. It also relates minibatch gradient variation to GGN approximations.

  • For minibatch size M = 1, G_M is an unbiased GGN estimator, whereas for M = N it reduces to gradient magnitude information without second-order content.
  • Taking the square root of s_t+1 leaves the fixed points unchanged but can affect convergence and the optimization path.
  • The analysis assumes full-batch gradients and exact expectations under the variational distribution.
  • The fixed-point analysis replaces the Hessian with a G_M approximation to derive the stated update behavior.
  • The toy experiment uses a two-Gaussian mixture with N = 60 examples and one Monte Carlo sample to compare the algorithms.

I.2. Real-Data Experiments

The real-data logistic-regression experiments evaluate uncertainty approximations across randomized dataset splits, while the accompanying table and figure define the datasets and UCI evaluation setup.

  • The logistic-regression experiments use 20 random splits for each dataset.
  • Table 2 lists the logistic-regression datasets together with their training-set sizes.
  • The USPS experiment compares MF-Exact, VOGN-1, and Vadam on 20 random 50-50 splits with prior precision λ = 25.
  • The UCI figure reports test RMSE using 100 Monte Carlo samples, with means and standard errors computed over 20 data splits.

J.1. UCI Regression Experiments

The UCI regression experiments compare Bayesian neural-network methods under repeated train-test splits and tuned hyperparameters, while the surrounding section motivates parameter-based exploration for reinforcement learning.

  • J.1. UCI Regression Experiments: The UCI benchmark uses 8 datasets, 20 provided train-test splits per dataset, and Bayesian optimization to tune prior and noise precisions.
  • J.1. UCI Regression Experiments: The neural networks have one hidden layer with 50 ReLU units and are trained for 40 epochs.
  • J.1. UCI Regression Experiments: Vadam and BBVI use different minibatch and Monte Carlo settings for smaller versus larger datasets, while evaluation uses 100 Monte Carlo samples throughout.
  • Parameter-based Exploration: The reinforcement-learning setup defines policy gradients for stochastic and deterministic policies and introduces parameter-based exploration by sampling policy parameters from a Gaussian distribution.
  • Parameter-based Exploration: Sampling policy parameters enables richer exploratory behavior than injecting noise directly into actions.
  • Parameter-based Exploration: The paper applies its methods to deep deterministic policy gradient and reports better performance from induced parameter exploration than existing methods.

K.3. Parameter-based Exploration Deep Deterministic Policy Gradient via VadaGrad and Vadam

The deep-RL experiments apply VadaGrad and Vadam to parameter-based exploration in DDPG, comparing them with SGD and Adam baselines. Vadam learns faster initially and avoids VadaGrad’s premature-convergence behavior, while exploration benefits diminish later.

  • Experimental setup: DDPG uses two-layer policy and Q-networks with 400 and 300 ReLU units, minibatches of M = 64, and one Monte Carlo sample.
  • Parameter-based exploration: VadaGrad and Vadam sample policy parameters from a Gaussian distribution and update the mean and variance during DDPG training.
  • Method behavior: VadaGrad’s increasing precision can make the policy deterministic too quickly, causing premature convergence and vanishing effective gradient steps.
  • Baselines: The baselines include SGD-Plain, SGD-Explore, Adam-Plain, and Adam-Explore, separating ordinary optimization from parameter-exploration variants.
  • Numerical considerations: An added precision constant prevents ill-conditioned initial sampling because λ = 10^-8 produces uninformative sampled weights and actions.
  • Results: Vadam learns faster early in Half-Cheetah, while Adam-based methods eventually catch up and achieve comparable performance at 3 million time-steps.
  • Results: Vadam and Adam-Explore remain slightly better than Adam-Plain, indicating that parameter-based exploration is still beneficial for DDPG.

M. Experiment on Improving “Marginal Value of Adaptive-Gradient Methods”

The experiment tests whether weight-perturbation improves adaptive-gradient methods, finding that VadaGrad matches SGD on War and Peace while showing only minor improvement on CIFAR-10. In reinforcement learning, Vadam learns faster early and Adam-based methods achieve comparable final performance.

  • Marginal Value of Adaptive-Gradient Methods: VadaGrad matches SGD on War and Peace, unlike AdaGrad and Adam, but its CIFAR-10 improvement is minor.The authors attribute the dataset difference partly to regularization techniques such as batch normalization, batch flip, and dropout.
  • Reinforcement Learning: VadaGrad uses a Gaussian variational distribution for parameter-based exploration and updates its mean and variance during training.Its algorithm samples policy parameters, collects replay-buffer transitions, computes policy gradients, and updates target parameters.
  • Reinforcement Learning: Parameter-based exploration samples policy parameters from a Gaussian distribution before collecting transitions and updating replay-buffer minibatches.The Vadam procedure updates first- and second-order moments, the policy mean, and target-network parameters.
  • Reinforcement Learning: Vadam learns faster than Adam-Plain and Adam-Explore during early Half-Cheetah training.Early learning performance is averaged over 5 trials with standard errors.
  • Reinforcement Learning: Vadam and Adam-based methods achieve comparable final performance on Half-Cheetah, while SGD-based methods do not learn well throughout.VadaGrad also learns well but shows signs of premature convergence.
  • Variational Optimization: The two-dimensional illustration shows Vadam annealed toward VadaGrad reaching a flat global minimum, whereas gradient descent gets stuck at various locations.The comparison uses four runs from different initial values.
Loading 1806.04854v3…