[Paper Review] Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam

Mohammad Emtiyaz Khan, Didrik Nielsen+4 more|arXiv (Cornell University)|Jun 13, 2018

Gaussian Processes and Bayesian InferenceComputer Science39 references60 citations

TL;DR

This paper presents natural-gradient variational methods that enable Gaussian mean-field Bayesian deep learning to be implemented inside Adam with minimal code changes by perturbing weights during gradient evaluations, yielding uncertainty estimates with comparable quality to existing VI methods and potential benefits for exploration.

ABSTRACT

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.

Motivation & Objective

• Motivate uncertainty estimation in deep learning for robust decision-making.
• Develop VI methods that are easy to integrate with existing adaptive optimizers like Adam.
• Reduce memory, computation, and implementation effort for Gaussian mean-field VI in large neural networks.
• Provide uncertainty estimates via weight perturbations that can also aid exploration in reinforcement learning and stochastic optimization.

Proposed method

• Propose natural-gradient variational inference (NGVI) for Gaussian mean-field VI.
• Show that NGVI can be implemented as minimal modifications to Adam by perturbing weights during gradient evaluation (Vadam).
• Derive VON (Variational Online-Newton) and VOGN (Variational Online Gauss-Newton) updates as approximations to NGVI, enabling online Hessian information.
• Introduce Vprop (Variational RMSprop) as a tractable RMSprop-like replacement using gradient magnitude as a Hessian proxy.
• Develop Variational Adam (Vadam) by adding natural-momentum in natural-parameter space to produce an Adam-like update.
• Extend to Variational AdaGrad (VadaGrad) by framing VI in a variational-optimization (VO) setting and deriving corresponding updates.

Experimental results

Research questions

• RQ1 Can variational inference for Gaussian mean-field models be implemented with similar simplicity and efficiency as MLE using Adam?
• RQ2 Do weight perturbations during gradient evaluation yield reliable uncertainty estimates with lower memory and computation costs than traditional VI?
• RQ3 How do the proposed approximations (VON, VOGN, Vprop, Vadam, VadaGrad) compare in terms of stability and practicality within standard deep-learning codebases?
• RQ4 Can the resulting uncertainty estimates be leveraged to improve exploration in reinforcement learning and stochastic optimization?

Key findings

• Uncertainty estimates obtained via weight perturbations within Adam are of comparable quality to existing VI methods.
• The proposed approximations enable VI with lower memory, computation, and implementation effort than traditional VI approaches.
• Vadam provides an Adam-like update with natural-momentum that aligns with standard optimizer practices, facilitating integration into existing codebases.
• Vprop offers a practical RMSprop-like method with weight perturbations and an uncertainty-bearing perturbation variance.
• The GM-based approximations (VOGN, Vprop) have predictable biases depending on minibatch size, informing trade-offs between accuracy and efficiency.
• Empirical results suggest weight perturbation can aid exploration in reinforcement learning and stochastic optimization.