PaperlayerView exact revision on arXiv

Source-linked AI summary

Eigenvalue Corrected Noisy Natural Gradient

Juhan Bae, Guodong Zhang, Roger Grosse

arXiv:1811.12565v1cs.LGstat.ML

TL;DR

Expressive variational posteriors are difficult to fit, and a matrix-variate Gaussian cannot capture accurate diagonal variance. The paper introduces noisy EK-FAC, which tracks diagonal variance in the K-FAC eigenbasis; across regression and classification evaluations, it achieves higher ELBO and test accuracy.

  • Problem

    Expressive variational posteriors are difficult to train, and matrix-variate Gaussian posteriors cannot capture accurate diagonal variance.

    Source: §1 · p. 1

  • Method

    Noisy EK-FAC uses EK-FAC to track diagonal variance in the K-FAC eigenbasis, producing a more flexible variational posterior.

    Source: §4 · pp. 4–5

  • Results

    Across regression and classification evaluations, noisy EK-FAC achieves higher ELBO and test accuracy.

    Source: §7 · p. 7

  • Takeaways & Limitations

    The evaluations demonstrate noisy EK-FAC’s effectiveness in representing weight correlations and accurate diagonal variance under the K-FAC eigenbasis.

    Source: §7 · p. 7

    The regularization parameter λ is 1 for exact Bayesian inference but can be tuned in practical settings.

    Source: §3 · p. 3

Abstract

from arXiv · show

Variational Bayesian neural networks combine the flexibility of deep learning with Bayesian uncertainty estimation. However, inference procedures for flexible variational posteriors are computationally expensive. A recently proposed method, noisy natural gradient, is a surprisingly simple method to fit expressive posteriors by adding weight noise to regular natural gradient updates. Noisy K-FAC is an instance of noisy natural gradient that fits a matrix-variate Gaussian posterior with minor changes to ordinary K-FAC. Nevertheless, a matrix-variate Gaussian posterior does not capture an accurate diagonal variance. In this work, we extend on noisy K-FAC to obtain a more flexible posterior distribution called eigenvalue corrected matrix-variate Gaussian. The proposed method computes the full diagonal re-scaling factor in Kronecker-factored eigenbasis. Empirically, our approach consistently outperforms existing algorithms (e.g., noisy K-FAC) on regression and classification tasks.

1 Introduction

Variational Bayesian neural networks offer deep-learning flexibility with Bayesian uncertainty estimates, but fitting expressive posteriors is difficult. The paper introduces noisy EK-FAC to improve the posterior flexibility of noisy K-FAC.

1 Introduction

  • Expressive variational posteriors can be difficult to train because they involve strong and complicated posterior dependencies.Fully factorized approximations are among the restricted assumptions used for posterior distributions.Source: §1 · p. 1
  • Noisy K-FAC fits a matrix-variate Gaussian posterior with minimal changes to ordinary K-FAC, but that posterior cannot capture accurate diagonal variance.Noisy natural gradient adds adaptive weight noise to regular natural-gradient updates.Source: §1 · p. 1
  • Noisy EK-FAC computes diagonal variance in the K-FAC eigenbasis, yielding a more expressive posterior distribution.The method builds on noisy K-FAC and EK-FAC.Source: §1 · p. 1

2 Natural Gradient

The section develops natural-gradient updates and their K-FAC approximation, then interprets curvature scaling in an eigenbasis to motivate EK-FAC’s more flexible diagonal rescaling.

2 Natural Gradient

  • Natural gradient descent implements steepest descent in distribution space, using predictive-distribution KL divergence as its distance measure.Its preconditioned gradients are invariant to smooth, invertible model reparameterizations.Source: §2 · p. 2

2.1 Kronecker-Factored Approximate Curvature

  • K-FAC approximates the inverse Fisher with Kronecker products, avoiding full-Fisher storage while enabling efficient inverse-Fisher products.Its layerwise update uses matrix multiplications comparable in size to the weight matrix.Source: §2.1 · p. 2

2.2 An Alternative Interpretation of Natural Gradient

  • The natural-gradient update can be viewed as projecting into the Fisher eigenbasis, rescaling coordinates, and projecting back; K-FAC uses its factored eigenbasis.The K-FAC eigenbasis is a more accurate approximation to the full Fisher eigenbasis, but its rescaling remains Kronecker-constrained and omits estimated variance along that basis.Source: §2.2 · p. 3
  • A diagonal Fisher approximation uses the identity basis and gradient second moments for rescaling, but its crude basis introduces significant approximation error.Obtaining an accurate eigenbasis is difficult.Source: §2.2 · p. 3

2.3 Eigenvalue Corrected Kronecker-Factored Approximate Curvature

  • EK-FAC computes the gradient’s second moment in the K-FAC eigenbasis, giving the diagonal rescaling factor n × p degrees of freedom instead of K-FAC’s n + p.Its rescaling matrix minimizes the approximation error in Frobenius norm.Source: §2.3 · p. 3

3 Variational Bayesian Neural Networks

Variational Bayesian neural networks fit an approximate posterior by maximizing the ELBO, while richer posterior families bring computational challenges. Noisy natural gradient and its Noisy K-FAC variant offer structured Gaussian approaches to fitting these posteriors.

3 Variational Bayesian Neural Networks

  • Variational Bayesian methods fit an approximate posterior q(w) by maximizing the evidence lower bound (ELBO).The regularization parameter λ is 1 for exact Bayesian inference but can be tuned in practical settings.Source: §3 · p. 3
  • Bayes By Backprop uses a fully-factorized Gaussian posterior approximation, with variational parameters updated by stochastic gradients using the reparameterization trick.Its variational parameters are the mean and diagonal variance.Source: §3 · pp. 3–4
  • Matrix-variate Gaussian posteriors capture correlations between weights, but gradient computation and positive semi-definite constraints make inference challenging.Prior methods use additional structures such as diagonal covariance or products of Householder transformations for efficient updates.Source: §3 · p. 4

3.1 Noisy Natural Gradient

  • Noisy natural gradient fits multivariate Gaussian posteriors by adding adaptive weight noise to ordinary natural gradient updates.At each iteration it samples weights from the variational posterior; computational intractability necessitates a structured covariance restriction, equivalently a restriction on the Fisher matrix.Source: §3.1 · p. 4

3.2 Fitting Matrix-Variate Gaussian Posteriors with Noisy K-FAC

  • Noisy K-FAC efficiently fits a matrix-variate Gaussian posterior without the additional approximations used by existing matrix-variate methods.It is a tractable noisy natural gradient method using a Kronecker-factored Fisher approximation; π keeps the Kronecker product representation compact.Source: §3.2 · p. 4

4 Methods

Noisy EK-FAC uses EK-FAC’s improved Fisher approximation to address the inaccurate diagonal variance of matrix-variate Gaussian posteriors. Its EMVG posterior retains weight correlations while representing full diagonal variance in the K-FAC eigenbasis.

4 Methods

  • Noisy EK-FAC replaces the K-FAC Fisher approximation with EK-FAC to improve the posterior’s diagonal variance representation.The matrix-variate Gaussian covariance does not capture an accurate diagonal variance; EK-FAC computes a more accurate diagonal re-scaling factor in the K-FAC eigenbasis.Source: §4 · pp. 4–5
  • Noisy EK-FAC estimates EK-FAC’s covariance factors and re-scaling matrix online using exponential moving averages.These updates resemble those used for noisy natural gradient; the method tracks the diagonal variance in the K-FAC eigenbasis.Source: §4 · pp. 4–5
  • EMVG combines compact weight-covariance representation with full diagonal variance in the K-FAC eigenbasis.It is obtained by applying the EK-FAC approximation to the matrix-variate Gaussian covariance factorization.Source: §4 · p. 5
  • The covariance factorization uses intrinsic damping, and extrinsic damping may be added to the re-scaling matrix for stable training.The paper identifies γ as intrinsic damping and states that additional extrinsic damping may be advantageous.Source: §4 · p. 6
  • Unlike standard EK-FAC, noisy EK-FAC samples weights from the variational posterior q during training.The paper interprets the posterior mean as a weight point estimate and its covariance as correlated Gaussian noise for each training example.Source: §4 · p. 6

5 Related Work

Related work spans fully factorized and non-Gaussian variational posteriors, dropout interpretations of Bayesian neural networks, and structured covariance approximations. Structured methods trade richer covariance representations against computational cost or restrictions on network depth.

5 Related Work

  • Graves [2011] used fully factorized Gaussian posteriors with a biased gradient estimator, while Blundell et al. [2015] used an unbiased reparameterization-trick estimator.The reparameterization trick is attributed to Kingma and Welling [2013].Source: §5 · p. 6
  • Other variational BNN approaches include Multiplicative Normalizing Flows [Louizos and Welling, 2017], implicit distributions [Shi et al., 2018], and dropout interpretations [Gal and Ghahramani, 2016, Gal et al., 2017].Source: §5 · p. 6
  • Natural-gradient VI with diagonal Fisher approximations yields fully factorized Gaussians, while K-FAC yields matrix-variate Gaussians [Zhang et al., 2017, Khan et al., 2018].Mishkin et al. [2018] instead use a diagonal-plus-low-rank approximation.Source: §5 · p. 6
  • Mishkin et al.’s [2018] diagonal-plus-low-rank method shares this work’s spirit, but its computational cost confines its application to shallow, two-layer networks.Source: §5 · p. 6

6 Experiments

The experiments assess noisy EK-FAC’s predictive performance and scalability in regression and classification. It reports higher test log-likelihood on UCI regression benchmarks and the highest CIFAR10 test accuracy across settings, including a 1.55% improvement over noisy K-FAC without extra regularization tricks.

6 Experiments

  • The study asks whether noisy EK-FAC improves prediction over existing methods and scales to large datasets and modern convolutional architectures.Classification evaluates a modified VGG16 on CIFAR10 against SGD with momentum, K-FAC, and noisy K-FAC.Source: §6 · p. 6, §6.2 · p. 7

6.1 Regression

  • Regression tests use UCI benchmarks, a Gamma prior and posterior for Gaussian-likelihood precision, and mostly single-hidden-layer networks.Training and test splits are 90% and 10%; the two largest datasets use fewer repeated splits, and network width differs for Protein and Year.Source: §6.1 · p. 6
  • Figure 3 presents training curves for BBB, noisy K-FAC, and noisy EK-FAC, paired respectively with fully factorized, matrix-variate, and eigenvalue-corrected matrix-variate posteriors.Source: Figure 3 · p. 6
  • Noisy EK-FAC yields higher test log-likelihood than the compared methods and higher ELBO than noisy K-FAC on UCI regression benchmarks.The reported test metrics are RMSE and log-likelihood; Table 1 summarizes average values for the UCI benchmarks.Source: §6.1 · p. 7, Table 1 · p. 7

6.2 Classification

  • 1.55% improvement over noisy K-FAC without extra regularization tricks; noisy EK-FAC achieves the highest test accuracy across CIFAR10 settings.The evaluation uses a modified VGG16, with batch normalization and data augmentation settings specified for the experiments.Source: §6.2 · p. 7, Table 2 · p. 7

7 Conclusion

Noisy EK-FAC extends the matrix-variate Gaussian posterior to capture correlations between weights and accurate diagonal variance in the K-FAC eigenbasis. It achieves higher ELBO and test accuracy in regression and classification evaluations.

7 Conclusion

  • Noisy EK-FAC achieves higher ELBO and test accuracy in both regression and classification evaluations.The conclusion presents these results as evidence of the method's effectiveness.Source: §7 · p. 7
  • The eigenvalue corrected matrix-variate Gaussian extends the matrix-variate Gaussian to represent the posterior with more flexibility.It captures correlations between weights while computing accurate diagonal variance under the K-FAC eigenbasis.Source: §7 · p. 7
  • The paper introduces a modified training method for variational Bayesian neural networks.Source: §7 · p. 7

A Matrix-Variate Gaussian

A matrix-variate Gaussian models a matrix-valued random variable with a mean and row and column covariances. After vectorization, its covariance is a Kronecker product of the column and row covariances.

A Matrix-Variate Gaussian

B Eigenvalue Corrected Matrix-Variate Gaussian

The eigenvalue corrected matrix-variate Gaussian extends the matrix-variate Gaussian by considering full diagonal variance in the Kronecker-factored eigenbasis. Its rescaling matrix has positive diagonal entries, and sampling is a special case of multivariate Gaussian sampling.

B Eigenvalue Corrected Matrix-Variate Gaussian

C Derivation of EK-FAC Update

The EK-FAC update derivation uses the weight gradient and activation and pre-activation covariance matrices, with diagonal rescaling in the K-FAC eigenbasis. It also specifies element-wise division and inverse vectorization in its notation.

C Derivation of EK-FAC Update

D Pseudo-Code for Noisy K-FAC

Algorithm 2 specifies Noisy K-FAC’s settings and layer-wise update loop, including its damping terms and scheduled operations.

D Pseudo-Code for Noisy K-FAC

  • The procedure assumes zero momentum, defines each layer’s weights and means as vectorized W_l and M_l, and lists its required tuning parameters.These include the stepsize, covariance-factor averaging parameter, KL weighting, prior variance, extrinsic damping, and statistics and inverse update intervals.Source: D Pseudo-Code for Noisy K-FAC · p. 11
  • Intrinsic damping is γ_in = λNη, and total damping is γ = γ_in + γ_ex.Source: D Pseudo-Code for Noisy K-FAC · p. 11
  • The loop updates covariance factors and includes scheduled operations using equations 11 and 12, followed by a displayed gradient and mean update.The displayed gradient term subtracts γ_in times W_l from the likelihood gradient; the mean update is shown using A and γ.Source: D Pseudo-Code for Noisy K-FAC · p. 11

Related papers

Loading 1811.12565v1…