SLANG: Fast Structured Covariance Approximations for Bayesian Deep Learning with Natural Gradient

TL;DR

SLANG introduces a fast, low-rank natural-gradient approximation for Bayesian deep learning, improving uncertainty estimates with less computational cost than traditional methods.

Contribution

It proposes a novel stochastic low-rank covariance approximation method that enhances uncertainty estimation in large deep models using only gradient information.

Findings

Faster and more accurate uncertainty estimation than mean-field methods

Comparable performance to state-of-the-art Bayesian inference techniques

Requires fewer gradient computations than existing approaches

Abstract

Uncertainty estimation in large deep-learning models is a computationally challenging task, where it is difficult to form even a Gaussian approximation to the posterior distribution. In such situations, existing methods usually resort to a diagonal approximation of the covariance matrix despite, the fact that these matrices are known to result in poor uncertainty estimates. To address this issue, we propose a new stochastic, low-rank, approximate natural-gradient (SLANG) method for variational inference in large, deep models. Our method estimates a "diagonal plus low-rank" structure based solely on back-propagated gradients of the network log-likelihood. This requires strictly less gradient computations than methods that compute the gradient of the whole variational objective. Empirical evaluations on standard benchmarks confirm that SLANG enables faster and more accurate estimation of uncertainty than mean-field methods, and performs comparably to state-of-the-art methods.