Optimizing Neural Networks with Kronecker-factored Approximate Curvature

James Martens; Roger Grosse

Source-linked AI summary

Optimizing Neural Networks with Kronecker-factored Approximate Curvature

James Martens, Roger Grosse

arXiv:1503.05671v7Mar 19, 2015cs.LGcs.NEstat.ML

TL;DR

Neural-network natural-gradient methods need efficient non-diagonal curvature approximations that remain practical in stochastic optimization. K-FAC approximates layer-wise Fisher blocks with Kronecker products and efficiently inverts the result, achieving much faster practical optimization than tuned SGD with momentum in supported benchmarks.

Problem
Efficiently exploiting high-quality non-diagonal curvature remains difficult because exact or advanced curvature methods can be expensive and poorly suited to highly stochastic optimization.
Method
K-FAC approximates layer-wise Fisher blocks with Kronecker products and uses an efficiently invertible approximation of the resulting Fisher matrix.
Results
K-FAC can be much faster in practice than highly tuned SGD with momentum on certain standard neural-network optimization benchmarks.
Takeaways & Limitations
K-FAC combines non-diagonal curvature information with storage and inversion costs that do not increase with the amount of data used for estimation.
Takeaways & Limitations
The Kronecker-factor approximation is not generally exact, and simple adaptive damping was insufficient to match exact-Fisher update quality.

Abstract

from arXiv · show

We propose an efficient method for approximating natural gradient descent in neural networks which we call Kronecker-Factored Approximate Curvature (K-FAC). K-FAC is based on an efficiently invertible approximation of a neural network's Fisher information matrix which is neither diagonal nor low-rank, and in some cases is completely non-sparse. It is derived by approximating various large blocks of the Fisher (corresponding to entire layers) as being the Kronecker product of two much smaller matrices. While only several times more expensive to compute than the plain stochastic gradient, the updates produced by K-FAC make much more progress optimizing the objective, which results in an algorithm that can be much faster than stochastic gradient descent with momentum in practice. And unlike some previously proposed approximate natural-gradient/Newton methods which use high-quality non-diagonal curvature matrices (such as Hessian-free optimization), K-FAC works very well in highly stochastic optimization regimes. This is because the cost of storing and inverting K-FAC's approximation to the curvature matrix does not depend on the amount of data used to estimate it, which is a feature typically associated only with diagonal or low-rank approximations to the curvature matrix.

1 Introduction

K-FAC addresses the limits of stochastic-gradient and Hessian-free optimization by efficiently approximating and inverting non-diagonal Fisher curvature. The method is designed for stochastic training and is reported to outperform tuned SGD with momentum on some benchmarks.

SGD with momentum remains the standard large-scale neural-network training method despite sophisticated alternatives.
Hessian-free updates can make more progress per iteration, but their costly conjugate-gradient computations and fixed curvature estimates limit stochastic-data throughput.
Highly spread-out curvature eigenvalues can eliminate conjugate gradient’s distinct advantage over well-tuned gradient descent with momentum.
Diagonal curvature approximations generally provide limited practical improvement compared with SGD with momentum.
K-FAC directly approximates non-diagonal curvature so its updates can be powerful like Hessian-free updates while remaining nearly as cheap as stochastic gradients.
K-FAC is reported to be much faster in practice than highly tuned SGD with momentum on certain standard neural-network optimization benchmarks.
K-FAC partitions Fisher blocks by layer, approximates them with Kronecker products, and further approximates the inverse as block-diagonal or block-tridiagonal.

2 Background and notation

The paper formalizes feed-forward networks, their likelihood-based training objective, and the Fisher-based natural gradient. This background connects natural-gradient optimization to curvature methods such as the generalized Gauss–Newton approach.

2.1 Neural Networks: A feed-forward network transforms inputs through layers of weighted sums and nonlinear activation functions, with homogeneous coordinates incorporating biases.
2.1 Neural Networks: The parameter vector concatenates vectorized weight matrices from all network layers.
2.2 Objective: The training objective averages prediction losses over the training distribution and serves as a proxy for the corresponding expectation under the true data distribution.
2.2 Objective: The loss is modeled as negative log probability under a predictive distribution, covering standard least-squares and cross-entropy objectives.
2.2 Objective: Minimizing the objective is equivalent to maximum-likelihood learning of the network’s conditional predictive model.
2.3 Natural Gradient: The natural gradient uses the inverse Fisher to select the direction producing the largest objective change per unit KL-divergence change.
2.3 Natural Gradient: In important cases, the Fisher equals the generalized Gauss–Newton matrix, linking natural-gradient methods to second-order optimization.

3 A block-wise Kronecker-factored Fisher approximation

The paper approximates the Fisher in layer-sized blocks using Kronecker factors, yielding substantial storage and inversion savings. The approximation captures coarse Fisher structure but relies on assumptions that are not generally exact.

Directly computing the inverse Fisher is impractical for large networks with potentially millions of parameters.
The Kronecker product structure enables convenient inversion through (A ⊗ B)^-1 = A^-1 ⊗ B^-1.
The initial approximation represents the Fisher as a layer-wise block matrix whose blocks are Kronecker products of smaller matrices.
The expectation-of-product approximation is not generally exact, but empirically captures the Fisher’s coarse structure in the example network.
The approximation substantially reduces storage and inversion costs, enabling an efficient approximate natural-gradient algorithm.
The approximation can be interpreted as assuming statistical independence between products of unit activities and products of unit input derivatives.
The approximation error is bounded by higher-order cumulant terms and is expected to be small when the relevant joint distribution is close to multivariate Gaussian.

4 Additional approximations to ˜F and inverse computations

The paper introduces further approximations to the Kronecker-factored Fisher and its inverse, using block-diagonal or block-tridiagonal inverse structure to enable efficient natural-gradient computations. The block-tridiagonal inverse is more accurate while remaining efficiently computable.

˜F is difficult to invert efficiently, so the paper introduces additional approximations to obtain an efficient approximation of the inverse Fisher.
The block-diagonal and block-tridiagonal inverse structures are motivated by dependencies among gradient entries within the same layer and between adjacent layers.The block-tridiagonal structure is justified by the forward and backward computations linking neighboring layers.
Figure 3 examines the inverse structure, while Figures 5 and 6 compare the Fisher and inverse-Fisher approximations for an example network.
The block-diagonal approximation requires inverses of 2ℓ smaller matrices and supports efficient matrix-vector products through Kronecker identities.
ˆF is constructed to match ˜F on tridiagonal blocks while ensuring that ˆF −1 is block-tridiagonal; this also determines its off-tridiagonal blocks.
ˆF approximates the off-tridiagonal blocks of ˜F well, and ˆF −1 is significantly more accurate than the block-diagonal inverse approximation.The inverse comparison is especially relevant because natural-gradient quality is roughly proportional to inverse-Fisher approximation quality.

5 Estimating the required statistics

The required curvature statistics are estimated from network activations and gradient factors, using predictive-distribution sampling for output-dependent quantities and exponentially decaying averages during stochastic optimization.

The block-diagonal inverse requires statistics for i = j, while the block-tridiagonal inverse additionally requires statistics for adjacent layers.
Output-independent activation factors are averaged over training inputs, whereas gradient factors require expectations over inputs and the network’s predictive distribution.
Monte Carlo sampling from the network’s predictive distribution estimates the output-dependent expectation, followed by rerunning backpropagation with sampled targets.
The required statistics can be estimated using the same input data as the gradient, with additional backward passes and outer-product averages.
Running estimates use exponentially decaying averaging because older statistics become stale as the network parameters change.
Using the training distribution for output targets would produce an empirical Fisher incompatible with the paper’s theoretical analysis and usual natural-gradient interpretation.
Unlike K-FAC’s compact statistics, exact-Fisher methods cannot efficiently summarize curvature independently of data volume, limiting their use of large data windows in stochastic optimization.

6 Update damping

K-FAC uses damping to make its approximate natural-gradient updates reliable beyond infinitesimal steps. The proposed scheme combines factored Tikhonov damping with exact-Fisher rescaling, improving update quality when the approximate curvature is inaccurate.

Motivation: Natural-gradient updates are useful only while their local quadratic model accurately approximates the objective, motivating damping during optimization's exploration phase.Well-designed damping can adapt to local function properties and turn off when the quadratic model becomes accurate.
Motivation: Powerful second-order methods such as HF and K-FAC require sophisticated damping because they can take much larger steps than first-order or diagonal-curvature methods.The paper characterizes damping as essential for preventing these large updates from failing.
A highly effective damping scheme for K-FAC: A simple adaptive Tikhonov method that works with the exact Fisher has no good λ choice for K-FAC updates of comparable quality.The approximate model lacks a guarantee of second-order accuracy, so large λ can wash out important low-curvature directions.
A highly effective damping scheme for K-FAC: K-FAC first computes a Tikhonov-damped proposal using its approximate Fisher, then rescales it using a quadratic model computed with the exact Fisher.The rescaling step is practical because it adjusts the already-computed proposal rather than directly inverting the exact Fisher.
A factored Tikhonov regularization technique: Factored approximate Tikhonov damping often works better in practice than exact Tikhonov damping, although the authors describe the reason as uncertain.They suggest separately regularizing factors may produce more robust inverse estimates.
Re-scaling according to the exact F: Without rescaling, the proposal may fail to improve the objective unless factored damping is very strong; rescaling enables smaller damping and larger, more effective updates.Figure 7 evaluates objective improvement against the factored Tikhonov strength constant γ.

7 Momentum

K-FAC adds momentum by jointly choosing the current proposal's scale and the previous update's coefficient to minimize its exact-Fisher quadratic model. This can accelerate optimization when gradient noise is relatively low.

Momentum: K-FAC uses δ = α∆ + µδ0, choosing α and µ to minimize the quadratic model built with the exact Fisher.This lets successive updates build a better solution to the local quadratic problem.
Momentum: The required momentum scalars can be computed with two forward passes, equivalent to the cost of one matrix-vector product with the Fisher.The main formula requires products F∆ and Fδ0.
Empirical behavior: Substantial acceleration occurs when the gradient signal has a low noise-to-signal ratio, commonly in early-to-mid stochastic optimization or with sufficiently large mini-batches.The paper relates this behavior to convex optimization theory and earlier neural-network optimization results.
Momentum: K-FAC computes its momentum decay coefficient online, avoiding manually specified schedules or heuristic adjustments used commonly with SGD.The coefficient µ is selected as part of the quadratic-model minimization.
Connection to conjugate gradient: For deterministic quadratic objectives, this momentum scheme is equivalent to preconditioned linear conjugate gradient with the approximate Fisher as preconditioner.The equivalence follows because linear conjugate gradient jointly chooses step size and momentum parameters.

8 Computational Costs and Efficiency Improvements

K-FAC adds gradient-like, curvature-estimation, inversion, and Fisher-product computations, with costs depending on layer width, mini-batch size, update frequency, and approximation choice. Several parallelization, low-rank, and refresh strategies reduce these overheads.

Cost components: A K-FAC iteration includes gradient computation, curvature-statistic updates, approximate-Fisher inversion, approximate-inverse products, exact-Fisher products, and damping adjustments.The listed tasks are expressed using layer width d, mini-batch size m, and implementation-dependent constants Ci.
Cost components: The block-diagonal inverse costs C3ℓd3 for inverses, while the block-tridiagonal inverse costs C4ℓd3 for its SVD-based computation.These are rough serial-operation estimates using naive cubic algorithms.
Cost components: Exact-Fisher matrix-vector products cost 4C1ℓd2m with momentum and 2C1ℓd2m without momentum.These products are computed on the current mini-batch.
Efficiency improvements: Curvature computations can be parallelized across layers, approximate-Fisher inversion can be refreshed every T3 iterations, and the paper used T3 = 20 experimentally.The authors report only a modest decrease in update quality because curvature remains relatively stable, especially later in optimization.
Cost limitations: When m is comparable to or smaller than d, inversion and SVD costs can dominate a naive implementation, especially because they are expensive and harder to parallelize on GPUs.If m ≫ d, tasks 5 and 6 become negligible relative to the other costs.
Efficiency improvements: For block-diagonal inversion, exploiting mini-batch-induced low-rank structure reduces one cost from C5ℓd3 to C5ℓd2m.The trick is based on matrix products involving d × m and m × d matrices, but standard ℓ2 weight decay is incompatible with it.
Overall cost: The average per-iteration cost combines gradient-related ℓd2m terms, inverse-related ℓd3 terms, and block-diagonal products scaling with ℓd2 min{d,m}.The formula distinguishes momentum and block-tridiagonal usage through binary flags χmom and χtri.

9 Pseudocode for K-FAC

K-FAC pseudocode alternates mini-batch gradient estimation, curvature-statistic updates, periodic approximate-Fisher inversion, damped proposal construction, exact-Fisher rescaling, and parameter updates.

Initialization: The algorithm initializes parameters and damping values, sets γ from λ and η, and iterates until the parameters are satisfactory.The initial λ is chosen conservatively large.
Per-iteration computation: Each iteration samples a mini-batch and subsets, computes the gradient, and performs an additional randomized backward pass to estimate curvature statistics.The subsets have sizes τ1|S′| and τ2|S′|.
Curvature and damping: The approximate Fisher inverse is recomputed initially and every T3 iterations, using damped factor estimates and each candidate γ when adjustment is scheduled.The inverse is formed with either the block-diagonal or block-tridiagonal formulas.
Update: K-FAC multiplies the approximate inverse by the gradient to form ∆, then applies the damping and rescaling procedures to obtain the final update δ.The parameter vector is updated as θk+1 ← θk + δ.
Damping control: The algorithm updates λ using a Levenberg-Marquardt-style rule before applying δ and advancing the iteration counter.This adjustment occurs within the pseudocode's damping-control loop.

10 Invariance Properties and the Relationship to Whitening and Centering

K-FAC has invariance properties under a class of invertible network reparameterizations, including activation and input transformations, under stated assumptions. Its block-diagonal updates also correspond to gradient descent on a centered and whitened transformed network.

Limitations: For finite steps, invariance errors decrease as damping diminishes, the reparameterization becomes locally linear, or the update region shrinks.The exact invariance result therefore has a practical boundary for large discrete updates.
Reparameterization invariance: Theorem 1 establishes equivalence between updates in the original and transformed networks under an invertible parameter mapping.The equivalence applies to updates using either the exact or approximate Fisher forms stated in the theorem.
Reparameterization invariance: K-FAC follows the same predictive-distribution optimization path across the default and transformed networks when initialization and algorithmic assumptions are matched.The result assumes fixed transformation matrices, equivalent initialization, negligible damping, no momentum, and parameterization-independent learning rates.
Activation, input, and unit transformations: The invariance extends to sigmoid-versus-tanh choices, arbitrary affine input transformations, and broader centering and whitening transformations.For smoothly varying transformations, invariance is approximate and becomes exact as the learning rate approaches zero.
Whitening and centering: With no damping, K-FAC’s block-diagonal update equals ordinary gradient descent in a transformed network whose unit activities and unit-gradients are centered and whitened.This interpretation is stated formally in Corollary 3.

11 Related Work

K-FAC is positioned against Hessian-free, centering, unit-wise, and prior Kronecker-factored methods. Its distinguishing combination is layer-level Kronecker structure with adaptive damping, stochastic estimation, richer inverse approximations, and momentum.

Hessian-free optimization: Unlike Hessian-free optimization, K-FAC directly inverts its curvature approximation and avoids repeated costly curvature-matrix vector products.The comparison emphasizes computational efficiency rather than a change in the underlying optimization objective.
Unit-wise and layer-wise approximations: K-FAC uses larger layer-wise blocks than TONGA, with two inversions per layer instead of one per unit, while retaining a more accurate Fisher approximation.TONGA reduces inversion costs further through low-rank-plus-diagonal blocks, at the cost of additional approximation error.
Centering and whitening: Centering methods target scalar activity and derivative means, whereas K-FAC’s whitening interpretation accounts for correlations among within-layer activities and back-propagated gradients.The centering argument relies on strong zero-correlation assumptions that correspond to diagonal or rank-1-plus-diagonal factors.
Prior Kronecker-factored methods: K-FAC differs from Heskes’ related Fisher approximation through adaptive damping, exact-Fisher-based rescaling, stochastic estimation, and support for higher-dimensional outputs.Heskes uses fixed hand-set factored damping and computes the layer factors exactly.
Additional contributions: Additional K-FAC components include block-tridiagonal inverse approximations, parameter-free momentum, online factor estimation, and computational-efficiency improvements.The authors report that each added element is important across the settings they study.

12 Heskes’ interpretation of the block-diagonal approximation

Heskes’ interpretation views the block-diagonal Fisher approximation as curvature for a quadratic measure of parameter-induced changes in layer preactivations and predictive distributions. The interpretation is only partly intrinsic because it assumes independent parameter groups, while the Kronecker-factored form lacks an established analogous Hessian interpretation.

Quadratic interpretation: The block-diagonal approximation is the curvature matrix of a quadratic function measuring changes between new and current parameter values.The construction provides an alternative interpretation to the paper’s theoretical analysis.
Layer-wise reweighting: The associated quantity reweights squared changes in each layer’s preactivation by the layer’s Fisher matrix.The reweighting is intended to translate preactivation changes into changes in predictive distributions.
Predictive-distribution geometry: The predictive distribution is parameterized by each layer’s preactivation, while the Fisher factor is averaged over the input-induced distribution of that preactivation.The expected Fisher substitutes for a more specific preactivation-dependent Fisher.
Scope of the interpretation: The interpretation is intrinsic except for assuming that parameters divide into independent groups corresponding to separate layer-wise predictive distributions.This independence assumption is the stated qualification on the intrinsic-measure claim.
Kronecker-factorized limitation: The Kronecker-factored approximation lacks a clearly intrinsic Hessian interpretation, although both proposed Fisher approximations still yield strong invariance properties.The paper notes that establishing such an interpretation would simplify the proof of Theorem 1.

13 Experiments

Experiments on three deep autoencoder benchmarks evaluated K-FAC against momentum-based SGD, varying mini-batch size, momentum, and approximate inverse structure. K-FAC achieved substantially faster optimization, while momentum and implementation cost shaped the best practical variant.

Experimental setup: K-FAC was evaluated on the MNIST, CURVES, and FACES deep autoencoder optimization benchmarks against a calibrated SGD-with-momentum baseline.The experiments used training-set reconstruction error to assess optimization speed rather than generalization.
Mini-batch-size experiment: K-FAC’s per-iteration progress increased superlinearly with mini-batch size, unlike the baseline and momentum-free K-FAC.This suggests stochastic-gradient noise limits momentum-based K-FAC later in optimization and motivates larger distributed implementations.
Mini-batch-size experiment: Using an exponentially increasing mini-batch schedule was motivated by superlinear progress gains alongside roughly linear per-iteration cost.The schedule used m_k = min(m_1 exp((k − 1)/b), |S|), with m_1 = 1000 and b chosen so that m_500 = |S|.
Second experiment: Across all three problems, momentum-enabled K-FAC made orders-of-magnitude more progress per iteration than the baseline and achieved much higher progress per second despite costlier iterations.The comparison used increasing mini-batches for K-FAC and fixed mini-batches for the baseline and momentum-free K-FAC.
Second experiment: Without K-FAC’s momentum technique, K-FAC was not significantly faster than the strong momentum baseline on these problems.The paper attributes K-FAC momentum’s role to progressively improving the solution to the exact Fisher-based quadratic model across iterations.
Approximate inverse structure: The block-tridiagonal K-FAC variant improved per-iteration progress by typically 25% to 40% over block-diagonal K-FAC, but its higher cost yielded only moderately better per-second progress.The authors therefore judged block-diagonal K-FAC probably preferable overall for its simplicity and comparable per-second rate.

14 Conclusions and future directions

K-FAC combines an efficiently invertible Fisher approximation with a complete optimization algorithm, and performs strongly on difficult deep auto-encoder benchmarks. The paper also identifies stochasticity handling and extensions to recurrent or convolutional architectures as future directions.

Contributions: K-FAC approximates large layer-wise Fisher blocks with Kronecker products and develops an efficiently computable approximation to the inverse Fisher.The paper also describes efficient matrix-vector products for these inverse approximations, including decomposition-based methods for repeated evaluations.
Contributions: K-FAC uses quadratic model-based damping and regularization to produce a robust method that is virtually free from hyper-parameter tuning.
Results: K-FAC with momentum and an increasing mini-batch schedule far surpasses well-tuned SGD with momentum on difficult deep auto-encoder benchmarks.The reported setting used a single GPU machine, and K-FAC required orders of magnitude fewer total updates or iterations than SGD with momentum.
Future directions: A better, more principled treatment of gradient stochasticity than a predetermined increasing mini-batch schedule remains a future direction.
Future directions: Extending K-FAC to recurrent or convolutional architectures may require specialized approximations of their associated Fisher matrices.

D Proofs for Section 10

The proofs establish that the network transformations are invertible reparameterizations preserving network outputs and characterize how curvature-based updates transform under such reparameterizations.

Reparameterization: The transformed network computes the same outputs as the original, so it represents an invertible linear reparameterization of the objective.
Update equivalence: Under an invertible affine reparameterization, compatible curvature matrices produce equivalent additive parameter updates in the two coordinate systems.
Proof strategy: The proof uses chain-rule relationships between transformed and original gradients, curvature blocks, and parameter mappings.
Invariance: For a transformed network whose relevant gradient and activation statistics are normalized, the approximate Fisher can reduce to the identity.In that case, the corresponding approximate natural-gradient update becomes a standard gradient-descent update.