Optimizing Neural Networks with Kronecker-factored Approximate Curvature

• SGD with momentum remains preferred for large-scale neural-network training, despite curvature-aware methods making more progress per iteration.Hessian-free optimization is costly because conjugate-gradient iterations require repeated curvature-matrix products and must use a fixed curvature estimate, limiting data processing in stochastic settings.Source: §1 · pp. 1–2
• Highly non-diagonal neural-network curvature motivates direct inversion methods, because diagonal approximations provide limited practical improvement and conjugate gradients can lose their advantage on difficult spectra.An efficient high-quality non-diagonal approximation could combine HF-like update power with stochastic-gradient-like computational cost.Source: §1 · p. 2
• K-FAC is proposed as an optimizer that can be much faster than highly tuned SGD with momentum on standard neural-network optimization benchmarks.Its central approximation is neither diagonal nor low-rank, yet can be inverted efficiently and estimated online from arbitrarily large data subsets without increasing inversion cost.Source: §1 · p. 3
• The method supports large-window online curvature estimates and practical optimization through theoretically grounded damping and an effective momentum variant.The paper also develops efficient cost reductions, complete pseudocode, and analyses of computational requirements and invariance to network transformations.Source: §1 · p. 3
• K-FAC first partitions the Fisher by layer and approximates each large layer block as a Kronecker product of two smaller matrices.This construction corresponds to assumptions about gradient statistics while using substantially larger blocks than earlier block-based methods.Source: §1 · p. 3
• K-FAC then approximates the inverse Fisher as block-diagonal or block-tridiagonal, justified through inverse covariances, tree-structured graphical models, and linear regression.The justification applies to the inverse rather than the Fisher itself, and experiments indicate that this structure is approximate for the inverse Fisher but not for the Fisher.Source: §1 · p. 3

2.1 Neural Networks

• A feed-forward network maps input x through ℓ layers using weight matrices, nonlinear activations, homogeneous coordinates for implicit biases, and concatenated parameters θ.At layer i, weighted sums are transformed element-wise by φ_i; the final column of W_i represents the bias vector.Source: §2.1 · p. 4
• Training minimizes the expected loss L(y, f(x, θ)) under the training distribution, with negative-log-likelihood models covering least-squares and cross-entropy objectives.The learned conditional model P_y|x(θ) supports a maximum-likelihood interpretation, while the training objective approximates expectation under the true data distribution.Source: §2.1 · pp. 4–5

2.2 The Natural Gradient

• The natural gradient F^-1∇h(θ) chooses the objective-improving direction per unit KL-divergence change in the model, unlike the Euclidean-metric standard gradient.The Fisher is defined for the network’s conditional model and is estimated using the training input distribution because the true distribution and exact expectation are generally unavailable.Source: §2.2 · p. 6
• In important settings, the Fisher equals the generalized Gauss–Newton matrix, linking natural-gradient methods to positive-semidefinite Hessian approximations and motivating the paper’s optimizer.This equivalence holds when the network is linearized up to the loss and the loss is a negative log probability from an exponential-family model with natural-parameter outputs.Source: §2.2 · pp. 6–7

3 A block-wise Kronecker-factored Fisher approximation

• K-FAC approximates each Fisher block with a Kronecker product, exploiting (A ⊗ B)^−1 = A^−1 ⊗ B^−1 for efficient inversion.This approximation yields significant storage and inversion savings for computing an approximate natural gradient.Source: §3 · p. 8
• The Kronecker-factorization is a major approximation because expectations of Kronecker products generally do not factorize, although it captures the Fisher’s coarse structure in practice.The supplied example compares the exact Fisher with the approximation for a partially trained deep network in Figure 2.Source: §3 · p. 8, Figure 2 · p. 9

3.1 Interpretations of this approximation

• The approximation assumes statistical independence between products of unit activities and products of unit input derivatives.For individual weight pairs, the Fisher approximation replaces their joint expectation with the corresponding factorized form.Source: §3.1 · pp. 8–9
• The approximation-error bound is small when higher-order cumulants are small; in the example, total error was 2894.4 versus an upper bound of 4134.6.Higher-order cumulants vanish for jointly multivariate-Gaussian variables, though the authors give no practical argument that this distributional condition generally holds.Source: §3.1 · p. 10

4 Additional approximations to ˜F and inverse computations

• Because no efficient general inversion method exists for Khatri-Rao products, the section introduces two structured inverse approximations and efficient matrix-vector products for natural-gradient computation.The block-tridiagonal structure is less restrictive and more accurate, but requires additional computational complexity.Source: §4 · p. 10

4.1 Structured inverses and the connection to linear regression

• Linear-regression analysis suggests that ˜F⁻¹ is approximately block-diagonal or block-tridiagonal because entries within the same or adjacent gradient layers are most useful for prediction.The block-tridiagonal alternative reflects forward and backward information flow between adjacent layers and is exact for corresponding Gaussian graphical models.Source: §4.1 · p. 11

4.2 Approximating ˜F⁻¹ as block-diagonal

• The block-diagonal approximation preserves layer blocks, whose Kronecker structure reduces inversion to 2ℓ smaller matrix inversions and enables efficient matrix-vector products.Unlike earlier unit-level block approximations in TONGA (Le Roux et al., 2008), these blocks contain all parameters in a layer.Source: §4.2 · p. 13

4.3 Approximating ˜F⁻¹ as block-tridiagonal

• The block-tridiagonal approximation constructs ˆF to match ˜F’s tridiagonal blocks while enforcing a block-tridiagonal inverse, using a Gaussian graphical-model representation and efficient factored products.Multiplication requires conditional covariance inversions; these are differences of Kronecker products and are handled with techniques described in Appendix B.Source: §4.3 · pp. 13–15

4.4 Examining the approximation quality

• ˆF⁻¹ is significantly more accurate than ˘F⁻¹, while ˆF also approximates ˜F’s off-tridiagonal blocks well.The inverse-Fisher comparison is especially relevant because natural-gradient approximation quality is roughly proportional to inverse-Fisher approximation quality.Source: §4.4 · p. 16, Figure 5 · p. 17, Figure 6 · p. 18

• The approximations require selected Ā_i,j and G_i,j statistics: block-diagonal blocks use i = j, whereas block-tridiagonal blocks use j ∈ {i, i + 1}.Source: §5 · p. 18
• Ā statistics average over the training input distribution, while G statistics require averaging over inputs and the network’s predictive distribution rather than training targets.Using training labels would produce the empirical Fisher, which lacks the relevant Generalized Gauss–Newton equivalence and natural-gradient interpretation and may perform worse in practice.Source: §5 · pp. 18–19
• Because forming complete G_i,j matrices is prohibitively expensive for mini-batches, the method estimates predictive expectations by sampling targets and rerunning backward propagation.Matrix-vector products can be computed efficiently for a given input, but constructing the full matrix would require work comparable to matrix-matrix multiplication per example.Source: §5 · p. 19
• One modified backward pass usually suffices, since the required outer-product averages resemble those already computed during ordinary gradient backpropagation.The same input data used for ∇h can be reused, with only additional backward passes and outer-product averaging required.Source: §5 · p. 19
• For online optimization, exponentially decaying averages track Ā_i,j and G_i,j, using ϵ = min{1 − 1/k, 0.95} at iteration k because older estimates become stale as parameters change.Equal-weight averaging is inappropriate because the curvature statistics depend on the evolving network parameters.Source: §5 · p. 19
• Unlike exact-Fisher methods such as HF, these compact running statistics can incorporate data beyond a single manageable mini-batch, avoiding curvature estimates restricted by full-matrix storage.Exact Fisher representations independent of dataset size would require impractically large explicit matrices, limiting HF-like methods in stochastic optimization.Source: §5 · p. 20

6.1 Background and motivation

• K-FAC’s approximate quadratic model can be inaccurate beyond short distances, so damping is essential during exploration while adapting away near local convergence.Natural-gradient directions are optimal directions rather than discrete updates; adaptive damping addresses local-model breakdowns and can turn off when the quadratic approximation becomes accurate.Source: §6.1 · pp. 20–21

6.2 A highly effective damping scheme for K-FAC

• The effective K-FAC damping scheme first applies adaptive Tikhonov regularization to its approximate Fisher, then rescales the proposal using one exact-Fisher matrix-vector product.This compensates for the approximate model’s intrinsic second-order inaccuracy while retaining an efficient computation.Source: §6.2 · p. 22, §6.4 · pp. 24–25

6.3 A factored Tikhonov regularization technique

• Factored Tikhonov regularization modifies the Kronecker factors separately, preserving single-Kronecker-product structure and efficient inversion; empirically, it often outperforms exact regularization.The method uses a factor-scaling parameter π_i, with a trace-norm choice based on relative average eigenvalues, while retaining the computations used for the original approximation.Source: §6.3 · pp. 23–24

6.4 Re-scaling according to the exact F

• The exact-Fisher rescaling minimizes a one-dimensional quadratic model for δ = α∆, using the same mini-batch gradient and efficient scalar curvature estimate.Only the scalar ∆^T F ∆ is required, rather than the extended quadratic-model optimization used by HF.Source: §6.4 · pp. 24–25
• Re-scaling the factored-Tikhonov proposal using the exact Fisher produces much larger and more effective updates than using the raw proposal.Figure 7 shows that raw updates may fail to improve the objective unless damping is very large, whereas re-scaling permits smaller damping and stronger updates.Source: §6.4 · p. 25, Figure 7 · p. 26

6.5 Adapting λ

• Levenberg–Marquardt adaptation adjusts λ toward the smallest value that preserves model reliability, allowing the implicit trust region to expand and λ to vanish near convergence.In the reported experiments, λ started at 150 and was updated every five K-FAC iterations, though these settings may be application-dependent.Source: §6.5 · pp. 25–26

6.6 Maintaining a separate damping strength for the approximate Fisher

• K-FAC maintains γ separately from λ because γ controls the quality of the approximate-Fisher proposal, whereas λ controls trustworthiness of the exact quadratic model.Every T2 iterations, three candidate γ values are compared using the quadratic-model quality metric; this metric performs similarly to alternatives at essentially no extra cost.Source: §6.6 · p. 27

Method

• K-FAC uses δ = α∆ + µδ0, choosing α and µ to minimize M(δ) and build a better local quadratic solution across iterations.Here δ0 is the previous iteration’s final update, and M uses the exact Fisher F; the approach is analogous to Matrix Momentum and HF.Source: §7 · p. 28

Implementation

• Computing the momentum coefficients requires only two forward passes, equivalent to the cost of one matrix-vector product with F.The calculation evaluates four required scalars involving F∆ and Fδ0.Source: §7 · p. 28

Results

• Momentum provides substantial acceleration when the gradient signal has a low noise-to-signal ratio, typically early to mid optimization or later with sufficiently large mini-batches.These findings agree with convex optimization theory and earlier neural-network optimization work.Source: §7 · p. 28

Method

• The momentum decay coefficient µ is computed on the fly, avoiding schedules or heuristics commonly needed for SGD.Source: §7 · p. 28
• For deterministic quadratic objectives, this momentum makes K-FAC equivalent to preconditioned linear CG on M(δ), using the approximate Fisher as preconditioner.The equivalence follows because linear CG jointly chooses α and µ to minimize M(δ) at each iteration.Source: §7 · p. 28

Per-update tasks

• The main stochastic-gradient tasks cost C2ℓd^2m, the randomized additional backward pass costs C1ℓd^2m, and updating curvature estimates costs 2C2ℓd^2m.These tasks are measured as sequential arithmetic operations, although the implementation parallelized over units and training cases.Source: §8 · pp. 28–29

Cost drivers

• In a naive implementation, matrix inversion or SVD can dominate per-iteration cost because these operations are several times costlier than matrix multiplication and harder to parallelize.Tasks 5 and 6 can run across layers, but task 5 is especially costly when m is comparable to or smaller than d.Source: §8 · pp. 29–30

Mini-batch subsampling

• Using subsets of sizes τ1m and τ2m reduces covariance-estimation and exact-F matrix-vector costs; τ1 = 1/8 and τ2 = 1/4 had negligible update-quality effects in experiments.When smaller batches or severe overfitting created deterioration, τ2 had to be restored to 1; T1 = 5 made the λ-adjustment pass relatively insignificant.Source: §8 · pp. 29–30

Inverse-cost reductions

• Recomputing the approximate Fisher inverse every T3 iterations amortizes inversion, with T3 = 20 used experimentally and only a modest observed decrease in update quality.Iterative Stein solvers, iterative matrix methods, hot-starting, and parallel layerwise computation provide additional alternatives; low-rank-plus-diagonal approximations may degrade updates.Source: §8 · p. 30

Low-rank gradient computation

• Exploiting the rank-m stochastic gradient reduces block-diagonal task-6 cost from C5ℓd^3 to C5ℓd^2m, but standard ℓ2 weight decay is incompatible with this trick.The weight-decay contribution is generally not low-rank and may require separate or occasional computation, or an alternative regularizer.Source: §8 · pp. 30–31

Overall cost

• The efficiency-adjusted average cost combines ℓd^2m terms with inverse-related ℓd^3 terms and, for block-diagonal inversion, an ℓd^2 min{d, m} term controlled by momentum, inverse type, and update intervals.The cost expression uses flags χmom and χtri and intervals T1, T2, and T3.Source: §8 · p. 31

Algorithm 2

• Algorithm 2 provides high-level pseudocode while leaving the computational details of each major step to their respective sections.Source: §9 · p. 32
• K-FAC initializes parameters and damping, then iterates until satisfactory by sampling mini-batches, estimating gradients and curvature factors, and generating candidate damping values.The procedure uses an additional backward pass with random targets to update the required activation and gradient-factor estimates.Source: §9 · p. 32
• The algorithm periodically recomputes a factored-Tikhonov-damped approximate Fisher inverse, multiplies it by the gradient estimate, and forms final updates using curvature-vector products.Inverse recomputation occurs initially and at scheduled iterations; an efficiency trick is suggested for layers with size d < m.Source: §9 · p. 32
• K-FAC evaluates candidate updates and damping values using M(δ), optionally updates λ with a Levenberg-Marquardt rule, and applies the selected update to θ.Momentum may be used when computing the final update.Source: §9 · p. 32

• K-FAC updates are equivalent across the original and transformed networks under Theorem 1, for either Fisher approximation, via an invertible linear reparameterization.The transformed networks represent the same function, and corresponding additive updates satisfy ζ(θ† + δ†) = θ + δ.Source: §10 · p. 35
• The K-FAC optimization path through predictive-distribution space is identical under fixed invertible layer transformations, given equivalent initialization, negligible damping, no momentum, and parameterization-independent learning rates.Source: §10 · p. 35
• This invariance includes switching between logistic sigmoid and tanh activations and applying arbitrary affine transformations to inputs, under the same initialization and damping conditions.The input case covers many common training-data preprocessing transformations.Source: §10 · p. 35
• K-FAC is approximately invariant to smoothly varying centering and normalization transformations, and to stronger transformations that center units and whiten their covariance to I.Invariance becomes exact as the learning rate approaches zero; approximation depends on transformation smoothness and update size.Source: §10 · p. 35
• Without damping, block-diagonal K-FAC updates equal standard gradient-descent updates in a transformed network whose unit activities and unit-gradients are both centered and whitened.Source: §10 · pp. 35–36

Hessian-free optimization

• Compared with Hessian-free optimization, K-FAC directly inverts its curvature approximation and estimates curvature from large online data averages, avoiding CG products and small fixed mini-batches.This trades exact curvature for an inexact but more efficiently computed approximation.Source: §11 · p. 36

TONGA

• K-FAC uses two layer-level Kronecker factors instead of TONGA’s per-unit blocks, making it less computationally intensive while arguably providing a more accurate Fisher approximation.TONGA further uses low-rank-plus-diagonal blocks to reduce inversion costs, introducing additional approximation error.Source: §11 · p. 36

Centering methods

• Centering methods reduce correlations by centering unit-wise quantities, whereas K-FAC additionally whitens activities and back-propagated gradients to account for within-layer correlations without requiring skip connections.The centering argument assumes strong independence-like structure, represented by diagonal Gi,i and rank-1-plus-diagonal Āi,i approximations.Source: §11 · pp. 36–37

Unitwise Fisher approximations

• Unlike Ollivier’s unitwise Fisher approximation, which uses an expensive biased deterministic estimate, K-FAC uses an unbiased stochastic approximation that performs nearly as well as exact values in practice.Ollivier’s blocks require k backpropagations, where k is the number of output units.Source: §11 · p. 37

Heskes’ method

• K-FAC extends Heskes’ Kronecker-factored Fisher method with dynamically adapted damping, exact-Fisher-based re-scaling, stochastic factor estimation, and additional inverse and optimization mechanisms.Basic fixed-parameter Tikhonov damping alone produced poor updates, even when its parameter was optimized at each iteration.Source: §11 · pp. 37–38

Povey et al. (2015)

• Compared with Povey et al. (2015), K-FAC uses the standard Fisher, adaptive re-scaling, and momentum, while Povey et al. use the empirical Fisher, basic damping, and online low-rank-plus-diagonal factor estimates.The alternative online factor-estimation strategy may still suit very wide layers, although similar experiments performed poorly relative to occasional inverse refreshing.Source: §11 · p. 38

Interpretation

• Heskes (2000) identifies the block-diagonal Fisher approximation as the curvature matrix of a quadratic function measuring changes from θ to θ′.The auxiliary quantities determined by θ are independent of θ′.Source: §12 · p. 38
• The quadratic measure is a reweighted sum of squared changes in each s_i, with G_i,i approximately translating those changes into changes in the predictive distribution over y.The weighting uses an expected Fisher rather than the more specific Fisher conditioned on s_i.Source: §12 · pp. 38–39

Intrinsic measures

• The resulting measures are intrinsic because their terms quantify changes in output distributions, subject to dividing θ into independent groups parameterizing separately conditioned predictive distributions.Using the specific Fisher instead would yield an exact layerwise block-diagonal approximation without the Kronecker-factorizing approximation.Source: §12 · p. 39

Limitations

• It remains unclear whether the Kronecker-factored approximation is the Hessian of an equally intrinsic measure, although both approximations produce updates with strong invariance properties.The proposed quadratic D(θ′, θ) is not obviously intrinsic, so this interpretation does not simplify the proof of Theorem 1.Source: §12 · p. 39

Experimental setup

• The experiments used standard MNIST, CURVES, and FACES deep-autoencoder benchmarks, a calibrated Nesterov-momentum SGD baseline, shared initialization, and training reconstruction error as the evaluation metric.The study targeted optimization speed rather than generalization and included ℓ2 regularization with η = 10−5.Source: §13 · pp. 39–40

First experiment

• As mini-batch size increased on MNIST, momentum K-FAC’s per-iteration progress grew superlinearly, unlike the baseline and momentum-less K-FAC, whose growth was smaller or near-linear.This indicates gradient noise was a major late-stage limitation for momentum K-FAC and suggests benefits from massively parallel implementations.Source: §13 · p. 40, Figure 9 · p. 41

Mini-batch schedule

• The exponentially increasing schedule m_k = min(m_1 exp((k −1)/b), |S|), with m_1 = 1000 and m_500 = |S|, was designed to exploit K-FAC’s scaling with mini-batch size.The authors note that slower or capped schedules may suit larger datasets, and adaptive mini-batch selection is another option.Source: §13 · p. 42

Second experiment

• On CURVES, MNIST, and FACES, momentum K-FAC achieved orders-of-magnitude higher per-iteration progress and much higher per-second progress than the baseline despite more expensive iterations.The comparison used training reconstruction error, with K-FAC employing an exponentially increasing mini-batch schedule and the baseline a fixed mini-batch size.Source: §13 · p. 42, Figure 10 · p. 43, Figure 11 · p. 44

Momentum

• Without momentum, K-FAC was not significantly faster than the strongly momentum-based baseline, confirming that momentum was essential on these autoencoder problems.The authors attribute K-FAC’s momentum benefit to compensating for approximate-Fisher inexactness, while conventional momentum performed better for SGD.Source: §13 · pp. 42, 45

K-FAC variants

• Block-tridiagonal K-FAC improved per-iteration progress by 25% to 40% over block-diagonal K-FAC, but its higher iteration cost made the overall per-second advantage only moderate.The results support a more accurate inverse-Fisher approximation for the block-tridiagonal variant, while favoring block-diagonal K-FAC overall because of its simplicity and comparable speed.Source: §13 · p. 45

14 Conclusions and future directions

• K-FAC’s complete optimization algorithm uses quadratic-model damping or regularization and is described as robust and virtually free from hyper-parameter tuning.Its Fisher approximation is efficiently invertible and was motivated by theoretical and empirical analysis of neural-network gradient statistics.Source: §14 · pp. 45–46
• K-FAC, combined with momentum and increasing mini-batch sizes, far surpasses well-tuned SGD with momentum on difficult deep auto-encoder benchmarks while requiring orders of magnitude fewer updates.The method preserves invariance to certain network reparameterizations and is suited to distributed settings where synchronization is the main bottleneck.Source: §14 · p. 46
• Future work includes a more principled treatment of gradient stochasticity than a predetermined increasing mini-batch schedule and extensions to recurrent or convolutional architectures.Such architectural extensions may require specialized approximations to their associated Fisher matrices.Source: §14 · p. 46

A Derivation of the expression for the approximation from Section 3.1

• The Kronecker-factor approximation follows from cumulant identities and a lemma asserting conditional uncorrelatedness between forward intermediates and backward derivatives under the predictive distribution.The required independence holds when targets are sampled from P_y|x, whereas derivative-dependent choices involving g^(k) are invalid.Source: A Derivation of the expression for the… · pp. 50–51

B Efficient techniques for inverting A ⊗B ± C ⊗D

• For symmetric positive semi-definite factors, SVD- and matrix-square-root-based decompositions make repeated products with (A ⊗B + C ⊗D)^−1 efficient after fixed preprocessing.The method reduces the remaining inverse to a diagonal element-wise operation, avoiding additional SVD computations for later vectors.Source: B Efficient techniques for inverting A… · pp. 52–53

C Computing v⊤Fv and u⊤Fv more efficiently

• Efficient Fisher-vector products use the sequence Jv, F_R(Jv), and J⊤ applied through linearized forward, cheap Fisher, and backward passes over each mini-batch.Because F_R is usually diagonal or diagonal plus rank-1, its matrix-vector multiplication is cheap.Source: C Computing v⊤Fv and u⊤Fv more… · pp. 53–54
• Computing v⊤Fv, u⊤Fv, and u⊤Fu from Jv and Ju instead of separate Fisher-vector products reduces the associated scalar-computation cost by roughly half.The identities are v⊤Fv=(Jv)⊤F_R(Jv), u⊤Fv=(Ju)⊤F_R(Jv), and u⊤Fu=(Ju)⊤F_R(Ju).Source: C Computing v⊤Fv and u⊤Fv more… · p. 54

Theorem 1

• The transformed network is an invertible linear reparameterization that preserves its output and objective, with f†(x, θ†) = f(x, ζ(θ†)) and h†(θ†) = h(ζ(θ†)).The equivalence follows by induction over the network layers and equality of the original and transformed outputs.Source: D Proofs for Section 10 · p. 54
• Under an invertible affine reparameterization, appropriately related curvature matrices produce equivalent additive updates in the original and transformed parameters.The proof invokes the stated equivalence ζ(θ† + δ†) = θ + δ for updates based on B and B†.Source: D Proofs for Section 10 · p. 55

Corollary 3

• Choosing transformations that make the transformed factors identity yields an identity Kronecker-factored curvature approximation.The construction uses whitening and centering properties of transformed activations and gradients, together with suitable choices of Ωi and Φi.Source: D Proofs for Section 10 · pp. 57–58
• With the transformed approximation equal to the identity, the curvature update −α F̌^−1∇h† becomes the standard gradient-descent update −α∇h†.The identity follows from the Kronecker products of identity factors.Source: D Proofs for Section 10 · p. 58