SLANG: Fast Structured Covariance Approximations for Bayesian Deep Learning with Natural Gradient
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.
Update: I have considered the author response. I maintain my assessment.  ==========================================================  This paper proposed an interesting idea of learning efficient structured multivariate Gaussian variational distributions. Since directly making a mean-field approximation may make the variational distribution too far from the model posterior distribution while making a full Gaussian approximation may be intractable, the authors proposed that, we could approximate the inverse of covariance matrix of the variational distribution as the sum of a diagonal matrix and a low-rank matrix. To make the inference faster, the authors also built a natural gradient update algorithm based on the method proposed in [1].  Strengths:   In my opinion, this paper addresses an interesting question and many people in the machine learning community may be interested in it. Variational inference is useful since it is fast but the most common variational distribution family, the mean-field family, could not yield satisfactory approximations in some cases. Hence, the problem that this paper focuses on is interesting. The design of their structured variational approximation is guaranteed to work at least as good as the regular mean-field approximation, and could outperform it in many cases, even if the added low-rank approximation only have rank L=1. Moreover, the paper is well-written and it is easy to follow the ideas.   Weakness:  It might show the success of the proposed method better if some more experiment results are presented.   First, for both the Bayesian Logistic regression experiment and the Bayesian neural network experiment, the authors mentioned that a bigger value of L improved the performances, but they only show the results for L=1 (on the metrics in Table 1 and 2). It will be more interesting if the experiment results for larger L could be shown and be compared to the baselines since for both experiments, there are baseline methods work better than the proposed methods. Though we understand that, for the Bayesian Logistic regression experiment, it is not possible for the SLANG method to outperform the full Gaussian method, but the difference between them might be smaller if L is larger.   Second, as shown in the right subfigure in Figure 3, the authors tried to compare the efficiency of SLANG with some baseline methods. It will be great if the authors could also compare the efficiency of their methods with some other second-order variational inference algorithms, such as the Hessian-Free Stochastic Gaussian Variational Inference (HFSGVI) algorithm [2].  In addition to that, it will be good if the authors could provide some theoretical analysis on their method. It will be interesting to see the performance improvements of this method compared to the mean-field approximation.  Minor comment:   Some sentences are a little bit confusing. For example, the sentence "Our goal is to approximate it to bring the cost down to O(D)" on line 110. We know that the time complexity is O(DL^2+DLS), not O(D). Though we understand L and S are small, and the precise time complexity is also mentioned later. It is still a little bit confusing to have this sentence. Probably it is better to say "bring the cost down to be linear in D".  References  [1] Mohammad Emtiyaz Khan and Wu Lin. Conjugate-computation variational inference: Converting variational inference in non-conjugate models to inferences in conjugate models. arXiv preprint arXiv:1703.04265, 2017.  [2] Fan, K., Wang, Z., Beck, J., Kwok, J., & Heller, K. A. (2015). Fast second order stochastic backpropagation for variational inference. In Advances in Neural Information Processing Systems (pp. 1387-1395).