Source-linked AI summary
Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, Pascal Vincent
TL;DR
Natural-gradient-style methods require curvature information that is too large for modern deep networks. This paper introduces EKFAC, which tracks diagonal variance in a Kronecker-factored eigenbasis and enables partial updates; experiments report faster optimization than KFAC across several architectures.
Problem
Gradient-covariance and natural-gradient methods can provide better descent directions, but full covariance matrices are too large for models with many parameters.
Method
EKFAC tracks diagonal gradient variance in a Kronecker-factored eigenbasis and periodically updates that basis while cheaply updating the diagonal scaling.
Results
EKFAC is provably more accurate than KFAC in Frobenius norm and achieves competitive optimization by iterations and wall-clock time on deep-network tasks.
Takeaways & Limitations
Partial curvature updates make curvature-based optimization competitive while retaining accurate eigenvalue estimates between eigenbasis updates.
Takeaways & Limitations
The authors leave robustness to the damping hyperparameter, alternative eigenbasis strategies, and more reliable translation of optimization gains into generalization for future work.
Abstract
from arXiv · showhide
Optimization algorithms that leverage gradient covariance information, such as variants of natural gradient descent (Amari, 1998), offer the prospect of yielding more effective descent directions. For models with many parameters, the covariance matrix they are based on becomes gigantic, making them inapplicable in their original form. This has motivated research into both simple diagonal approximations and more sophisticated factored approximations such as KFAC (Heskes, 2000; Martens & Grosse, 2015; Grosse & Martens, 2016). In the present work we draw inspiration from both to propose a novel approximation that is provably better than KFAC and amendable to cheap partial updates. It consists in tracking a diagonal variance, not in parameter coordinates, but in a Kronecker-factored eigenbasis, in which the diagonal approximation is likely to be more effective. Experiments show improvements over KFAC in optimization speed for several deep network architectures.
1 Introduction
Deep networks achieve strong performance but are costly to train, while imbalanced curvature limits SGD. The section motivates tractable curvature approximations and introduces EKFAC as a more accurate, partially updateable alternative to KFAC.
- Deep networks achieve state-of-the-art performance, but top-performing systems can require days of training and substantial computational power.
- Imbalanced curvature in highly non-convex neural-network loss surfaces limits the efficiency of first-order methods such as SGD.
- Second-order methods can correct curvature, but their nθ × nθ preconditioning matrices are too large to compute and invert for modern networks.
- Diagonal methods ignore covariance structure, whereas KFAC captures parameter correlations using efficiently invertible Kronecker-factored blocks.
- EKFAC tracks diagonal variance in a Kronecker-factored eigenbasis, is provably more accurate than KFAC, and supports cheap partial curvature updates.
- Experiments on deep auto-encoders and CIFAR-10 convolutional networks report optimization improvements over KFAC.
2 Background and notations
The paper formulates neural-network training through stochastic gradients and curvature-based preconditioning. It then describes block-diagonal and Kronecker-factored approximations that make large curvature matrices tractable.
- Training minimizes empirical risk over neural-network parameters using losses averaged across dataset examples.
- SGD updates parameters with stochastic gradients, while second-order methods first transform those gradients using the inverse preconditioning matrix G.
- Natural Gradient, Generalized Gauss-Newton, and TONGA use covariance or second-moment forms of stochastic gradients to define preconditioning matrices.
- Because G is too large to compute and invert, practical methods commonly approximate it block-diagonally by treating network layers separately.
- KFAC approximates each layer block as A ⊗ B, replacing a d_in d_out × d_in d_out matrix with two smaller factors of sizes d_in × d_in and d_out × d_out.
- For convolutional layers, constructing the Kronecker factors requires summing over spatial locations and corresponding receptive fields.
3 Proposed method
EKFAC applies diagonal natural-gradient rescaling in a Kronecker-factored eigenbasis, using gradient second moments that better match those directions than KFAC’s constrained scales. Its basis is amortized while the diagonal variance can be cheaply tracked between recomputations.
- Motivation: diagonal rescaling in different bases: Full natural-gradient preconditioning rescales gradient coordinates in the eigenbasis of G, whereas diagonal preconditioning uses the initial parameter basis.The full update projects into the eigenbasis, applies S^-1, and projects back; diagonal rescaling divides parameter coordinates by estimated second moments.
- Motivation: diagonal rescaling in different bases: KFAC uses a Kronecker product A ⊗ B whose eigenvectors form the KFE basis, but its rescaling factors are products of factor eigenvalues rather than directly estimated moments along those axes.The Kronecker eigendecomposition supplies the basis and associated product eigenvalues, enabling a basis-rescaling interpretation of KFAC.
- Eigenvalue-corrected Kronecker Factorization: EKFAC replaces KFAC’s constrained diagonal with S*, the directly tracked second moments of gradient coordinates in the KFE.The vector s* can be efficiently estimated and stored, while the KFE itself remains computationally expensive and is recomputed less frequently.
- Eigenvalue-corrected Kronecker Factorization: S* is the Frobenius-optimal diagonal rescaling in the KFE, guaranteeing ∥G−GEKFAC∥F ≤ ∥G−GKFAC∥F.The guarantee concerns approximation error to G measured by the Frobenius norm, not necessarily the quality of the inverse preconditioner.
- Eigenvalue-corrected Kronecker Factorization: EKFAC’s update projects the gradient into the KFE, applies diagonal natural-gradient rescaling there, and projects the result back to parameter space.This provides an alternative interpretation of EKFAC as a diagonal method in a basis where the diagonal approximation is more accurate.
- Eigenvalue-corrected Kronecker Factorization: Figure 3 reports less gradient correlation in the KFE and slower drift of EKFAC’s approximation error between basis or inverse recomputations.The figure uses a small MNIST classifier and recomputes the KFE basis and KFAC inverse every 100 iterations.
4 Experiments
Experiments on MNIST and CIFAR-10 evaluate EKFAC and its running-average variant against KFAC and standard baselines. EKFAC variants generally improve optimization speed, while preserving good validation performance under amortized curvature updates.
- MNIST Deep Auto-Encoder: EKFAC and EKFAC-ra minimize MNIST auto-encoder training loss faster per epoch than KFAC and other baselines.EKFAC-ra also transfers this optimization benefit to faster wall-clock training.
- MNIST Deep Auto-Encoder: EKFAC preserves optimization performance better than KFAC when inverse or eigendecomposition recomputation is less frequent.Frequent diagonal reestimation helps EKFAC variants track the empirical Fisher spectrum more closely between basis recomputations.
- CIFAR-10: EKFAC-ra shows optimization and wall-clock improvements over baselines on VGG11 when curvature recomputation is amortized over 500 iterations.The reported validation results indicate that this optimization benefit is maintained without losing generalization capability.
- CIFAR-10: The CIFAR-10 experiments evaluate EKFAC methods on VGG11 and ResNet34 using the SUA approximation without batch normalization for KFAC/EKFAC models.Hyperparameters were selected through grid and random search, with constant learning rates in the reported experiments.
- CIFAR-10: On ResNet34, EKFAC-ra outperforms amortized KFAC and SGD with momentum and batch normalization in optimization per epoch and compute time.The gain appears robust across different batch sizes while maintaining good generalization capability.
5 Conclusion and future work
The paper introduces EKFAC as a computationally manageable, accurate approximation to the empirical Fisher Information Matrix. Its partial curvature updates are competitive for deep-network optimization, while future work targets broader adaptive methods, eigenbasis alternatives, damping robustness, and generalization.
- Conclusion: EKFAC is provably more accurate than KFAC in Frobenius-norm approximation of the empirical Fisher Information Matrix.The method keeps an up-to-date eigenvalue estimate while holding the eigenbasis fixed.
- Conclusion: Cheap partial curvature updates make EKFAC competitive in deep-network optimization by both iteration count and wall-clock time.The approach updates eigenvalues frequently without recomputing the more expensive eigenbasis at every parameter update.
- Future Work: Future work includes applying other diagonal adaptive algorithms in the KFE, finding alternative eigenbasis strategies, improving damping robustness, and developing regularization for generalization.The paper specifically identifies Adagrad, RMSProp, and Adam as possible KFE-based alternatives.
A.1 Proof that EKFAC does optimal diagonal rescaling in the KFE
EKFAC chooses an optimal diagonal rescaling in the Kronecker-factored eigenbasis, yielding an approximation to G that is at least as accurate as KFAC in Frobenius norm.
- Optimal diagonal rescaling: For a given orthogonal basis Q, the optimal diagonal approximation sets each diagonal entry to the corresponding diagonal entry of Q⊤GQ.This minimizes the Frobenius-norm approximation error because diagonal rescaling cannot change off-diagonal terms.
- Kronecker-factored eigenbasis: The Kronecker product UA ⊗ UB is an orthogonal eigenbasis for KFAC’s Kronecker factorization A ⊗ B.The associated diagonal eigenvalue matrix is SA ⊗ SB.
- EKFAC construction: EKFAC applies the optimal diagonal rescaling in the Kronecker-factored eigenbasis, rather than using KFAC’s factored eigenvalue products.Its diagonal entries are selected to minimize the Frobenius norm of the approximation error.
- Comparison with KFAC: EKFAC’s approximation error is no larger than KFAC’s: ∥G − GEKFAC∥F ≤ ∥G − GKFAC∥F.The result follows from the optimality of EKFAC’s diagonal rescaling in the shared basis.
B Residual network initialization
Residual networks without batch normalization require careful initialization, using variance-scaled weights adapted to whether layers are preceded by nonlinearities or skip connections.
- Initialization procedure: Layers directly preceded by ReLU use He initialization: W ∼ N(0, 2/n), with zero biases.Here n denotes the layer’s fan-in.
- Initialization procedure: Layers not directly preceded by an activation function use W ∼ N(0, 1/n), with zero biases.This includes convolutions in skip connections and follows from treating the identity as the activation.
- Initialization procedure: The last convolution in each residual block is initialized with W ∼ N(0, 0.2/n), reducing its scale by a factor of 10.The procedure is intended to preserve variance and ease optimization at the beginning of training.
- Damping: Both KFAC and EKFAC are highly sensitive to damping, and EKFAC performs better when damping is applied to the activation and gradient covariance eigenvalues.The initialization sets diag(ϵ′) using ϵI together with terms involving SA, SB, and their square-root-scaled identity matrices.
D.1 Impact of batch size
Batch size affects EKFAC’s optimization and computational trade-offs: its advantage over KFAC can diminish for VGG11 at smaller batches, while remaining consistent for Resnet34.
- VGG11: For VGG11, EKFAC’s optimization gain over KFAC diminishes as batch size decreases.The comparison uses training loss selected by the best value at each epoch.
- VGG11: For VGG11, larger batches increase EKFAC’s optimization benefit over KFAC by reducing inverse/eigendecomposition operations per epoch.The gain does not translate into faster computation-time training in this relatively small network.
- Resnet34: For Resnet34, EKFAC’s optimization gain over KFAC remains consistent across batch sizes.The experiment is conducted on CIFAR-10 with model selection based on the best training loss at each epoch.
D.2 Learning rate schedule
With a learning-rate schedule that decays the learning rate by 2 every 20 epochs, EKFAC continues to show optimization benefits relative to the baseline on CIFAR-10 experiments.
- Experimental setting: The learning-rate schedule decays the learning rate by 2 every 20 epochs.The experiments use a setting similar to the CIFAR-10 experiments without the schedule.
- Results: EKFAC still shows optimization benefits relative to the baseline when combined with a learning-rate schedule.This result is reported for the VGG11 and Resnet34 experiments shown in Figures D.4 and D.5.
- Evaluation: Figures D.4 and D.5 evaluate VGG11 and Resnet34 on CIFAR-10 under the scheduled-learning-rate setting.ED_freq and Inv_freq denote eigendecomposition and inverse frequencies, respectively.