Kronecker-Factored Curvatures in Neural Networks
- Kronecker-Factored Curvatures are structured approximations that factorize large Fisher blocks into Kronecker products of smaller matrices encoding activation and gradient statistics.
- They enable efficient inversion and evaluation of quadratic forms by reducing computational complexity while preserving key parameter correlations compared to diagonal methods.
- This approach underpins scalable natural-gradient optimization, model compression, influence estimation, and efficient training across diverse neural network architectures.
Kronecker-Factored Curvatures are structured approximations to second-order geometry in neural networks in which a layerwise curvature block—typically a Fisher Information Matrix block, a generalized Gauss–Newton block, or a closely related positive semidefinite surrogate—is represented as a Kronecker product of two much smaller matrices, one encoding input-activation statistics and one encoding output-gradient or sensitivity statistics. This construction makes storage, inversion, and quadratic-form evaluation tractable for deep models while preserving correlations that are absent in diagonal approximations, and it has become a central mechanism for scalable natural-gradient methods and related curvature-aware procedures (Martens et al., 2015, Firouzi, 2018).
1. Geometric basis and historical development
The geometric motivation is the natural gradient. For parameters and loss , the expected Fisher may be written as
Under regular conditions in the maximum-likelihood setting, the Fisher equals the Hessian of the negative log-likelihood at the optimum, $F(\theta^\*)=H(\theta^\*)$. This makes Fisher a curvature proxy and explains why preconditioning by yields the steepest descent direction in the statistical manifold rather than in Euclidean parameter space (Chekalina et al., 23 May 2025).
Historically, Kronecker-Factored Approximate Curvature (K-FAC) was introduced by Martens and Grosse in 2015 to make natural-gradient and second-order optimization practical for deep networks. The central observation was that large Fisher blocks associated with entire layers are neither diagonal nor low-rank in general, yet can often be approximated by Kronecker products of small factors that are cheap to estimate and invert. The original formulation emphasized fully connected feedforward networks, damping, step-size selection, and momentum, and showed that the storage and inversion cost depends on layer widths rather than on the amount of data used to estimate curvature (Martens et al., 2015).
The same geometric logic was subsequently extended to other settings. Convolutional layers motivated Kronecker Factors for Convolution (KFC), which adapted the factorization to weight sharing and spatial structure (Grosse et al., 2016). In reinforcement learning, ACKTR used K-FAC within a trust-region actor–critic method, and later work applied K-FAC to the control-variate network in RELAX/LAX estimators to improve sample efficiency and reduce gradient variance (Wu et al., 2017, Firouzi, 2018). This progression established “Kronecker-factored curvature” not as a single algorithm, but as a family of structured approximations to curvature that trade exactness for tractability.
2. Canonical layerwise factorization
For a linear layer with weight matrix , input activation vector , and pre-activation gradient vector , the per-sample gradient with respect to is
and vectorization gives
0
This bilinear structure is the algebraic source of Kronecker-factored curvature. Under the standard independence approximation between activations and backpropagated derivatives, the Fisher block for the layer is approximated as
1
with 2 and 3 (Firouzi, 2018, Chekalina et al., 23 May 2025).
The preconditioned update then exploits the Kronecker inverse identity
4
so that for the vectorized gradient one may write
5
Because 6 and 7 are only 8 and 9, their inversion is far cheaper than inverting the full block of size 0 (Firouzi, 2018).
A second important identity is the conversion from parameter-space quadratic forms to weighted Frobenius norms. With 1 and 2,
3
This shows that the Kronecker factors define a curvature-weighted metric on perturbations of the weight matrix, not merely a preconditioner for optimization. Recent compression methods make direct use of this equivalence by minimizing curvature-weighted reconstruction error rather than ordinary Frobenius error (Chekalina et al., 23 May 2025).
The computational advantage is substantial. For a layer 4, the full Fisher block has size 5 and is prohibitive to form or store, whereas the Kronecker structure stores only the two factors, with memory 6. This sits between diagonal curvature, which has lower storage but discards cross-parameter correlations, and exact second-order methods, which are generally intractable at modern scales (Chekalina et al., 23 May 2025).
3. Weight sharing and architecture-specific generalizations
Convolutional networks require additional assumptions because gradients aggregate over spatial locations and weights are shared. KFC treats each spatial location as contributing a local linear mapping, unfolds input patches with an im2col view, and derives a Kronecker factorization from three approximations: independent activations and derivatives, spatial homogeneity, and spatially uncorrelated derivatives. In this setting, the activation factor is formed from unfolded patches, the derivative factor from pre-activation derivatives across spatial positions, and the resulting updates preserve invariance to commonly used reparameterizations such as centering of activations (Grosse et al., 2016).
A broader formulation was later developed for modern architectures whose linear components exhibit generic weight sharing. This framework distinguishes two regimes. In the “expand” setting, the loss expands across the shared dimension and the weight-sharing axis is treated analogously to an enlarged batch; in the “reduce” setting, the shared dimension is aggregated before the loss and the Kronecker factors are built from reduced statistics. These two constructions, K-FAC-expand and K-FAC-reduce, are exact for deep linear networks with weight sharing in their respective settings. Empirically, both variants were able to reach a fixed validation metric target in 7–8 of the number of steps of a first-order reference run for a graph neural network and a vision transformer, with a comparable wall-clock improvement (Eschenhagen et al., 2023).
Physics-informed neural networks introduce a different form of sharing: the same weights act not only on activations but also on derivative channels arising from Taylor-mode automatic differentiation. In that setting, the preferred curvature is the Gauss–Newton or Fisher built from residual linearization, and the interior loss block can be approximated by Kronecker factors that average over both batch and shared derivative channels. The resulting KFAC-based optimizers were reported to be competitive with expensive second-order methods on small problems, to scale more favorably to higher-dimensional neural networks and PDEs, and to consistently outperform first-order methods and LBFGS (Dangel et al., 2024).
Normalization layers expose a limitation of the standard assumptions. Batch normalization creates inter-example dependence because each example depends on mini-batch mean and variance. Extended K-FAC addresses this by introducing additional activation and derivative blocks that capture cross-example relations, yielding a Fisher-based approximation that remains positive semidefinite and is suitable for continual learning with batch-normalized architectures (Lee et al., 2020).
4. Algorithmic variants and spectral refinements
Plain K-FAC is only one point in a larger design space. Several later methods keep the Kronecker skeleton but modify either the spectrum, the block approximation, or the way the inverse curvature is represented.
| Variant | Core modification | Reported consequence |
|---|---|---|
| EKFAC | Keeps the Kronecker-factored eigenbasis and estimates diagonal second moments in that basis | Provably closer than K-FAC in Frobenius norm |
| TKFAC | Uses a scaled Kronecker block 9 with trace restriction | Generally tighter Frobenius error bound than K-FAC |
| K-BFGS / K-BFGS(L) | Uses Kronecker-factored Hessian blocks with quasi-Newton inverse updates | Memory comparable to first-order methods |
| KrADagrad$F(\theta^\*)=H(\theta^\*)$0 | Approximates the inverse empirical Fisher directly in Kronecker form using positive roots | Avoids inverse matrix roots and 64-bit precision |
EKFAC is centered on the Kronecker-factored eigenbasis $F(\theta^\*)=H(\theta^\*)$1. Instead of using the K-FAC eigenvalues $F(\theta^\*)=H(\theta^\*)$2, it tracks the diagonal variance of the gradient in that basis,
$F(\theta^\*)=H(\theta^\*)$3
and uses this diagonal as the curvature approximation. The method is provably better than K-FAC in Frobenius norm because, for a fixed orthogonal basis, the optimal diagonal approximation is obtained by matching the diagonal second moments in that basis. The same idea underlies later EK-FAC parameterizations for influence estimation in transformer LLMs (George et al., 2018, Bao et al., 8 May 2025).
TKFAC modifies the block itself. Each block is approximated as $F(\theta^\*)=H(\theta^\*)$4, where $F(\theta^\*)=H(\theta^\*)$5 is a trace coefficient chosen so that the approximate block preserves the trace of the exact block. The paper provides an upper bound on the approximation error and states that, in general, the TKFAC bound is smaller than the corresponding KFAC bound except in the degenerate case where all per-example traces are equal. It also introduces a CNN-specific damping scheme intended to prevent normal damping from overwhelming second-order structure in convolutional blocks (Gao et al., 2020).
Kronecker-factored quasi-Newton methods replace Fisher- or Gauss–Newton-based factor inversion with BFGS or L-BFGS updates on Kronecker factors that approximate Hessian blocks directly. For multilayer perceptrons and convolutional neural networks, these methods were designed to have memory requirements comparable to first-order methods and substantially lower per-iteration time complexity than earlier Hessian-based quasi-Newton methods, while performing comparably to second-order state-of-the-art baselines in the reported experiments (Ren et al., 2021).
KrADagrad$F(\theta^\*)=H(\theta^\*)$6 sits further away from classical natural gradient. It uses a Kronecker Approximation-Domination construction to update matrices that directly approximate the inverse empirical Fisher, thereby avoiding inverse matrix fractional powers. The reported motivation is numerical: inverse matrix roots of ill-conditioned matrices usually require 64-bit precision, whereas KrAD uses positive fractional powers and remains stable in 32-bit settings with computational costs similar to Shampoo (Mei et al., 2023).
5. Uses across optimization, compression, attribution, and model editing
In reinforcement learning, Kronecker-factored curvature became a practical trust-region mechanism. ACKTR applies K-FAC to both actor and critic, using the policy Fisher for the actor and a Gauss–Newton/Fisher surrogate for the critic, and reports higher rewards and a $F(\theta^\*)=H(\theta^\*)$7- to $F(\theta^\*)=H(\theta^\*)$8-fold improvement in sample efficiency on average compared to previous state-of-the-art on-policy actor–critic methods. KF-LAX and KF-RELAX use the same layerwise Kronecker structure not for the policy itself but for the control-variate network $F(\theta^\*)=H(\theta^\*)$9, preconditioning gradients of the variance-reduction objective to accelerate optimization of the surrogate (Wu et al., 2017, Firouzi, 2018).
In scientific computing and quantitative finance, the same structure has been used outside conventional supervised classification. For PINNs, KFAC-based optimizers were reported to be competitive with expensive second-order methods on small problems and to scale to networks with hundreds of thousands to over a million parameters while consistently outperforming first-order methods and LBFGS (Dangel et al., 2024). In deep hedging, K-FAC was applied to the fully connected output layer of an LSTM-based hedging model and yielded a 78.3% reduction in transaction costs, a 34.4% decrease in profit-and-loss variance, and a Sharpe ratio of 0.0401 versus 0 for Adam, with comparable reported training duration (Enkhbayar, 2024).
Large-language-model compression has introduced a distinct use of Kronecker-factored curvature: weighting approximation error by task curvature rather than by Euclidean norm. “Generalized Fisher-Weighted SVD” constructs a Kronecker-factored approximation to the observed Fisher information of a single weight matrix and solves
1
by forming 2 and taking its truncated SVD. On LLaMA 2–7B at the most aggressive setting in the reported table, where 80% of parameters were retained, average MMLU was 0.32 for GFWSVD versus 0.27 for FWSVD, 0.29 for SVD-LLM, and 0.26 for ASVD; on PTB, perplexity was 50.50 for GFWSVD versus 1523.00, 98.91, and 241.57, respectively (Chekalina et al., 23 May 2025).
Influence estimation for billion-parameter transformers uses EK-FAC for efficient inverse-Hessian-vector-like actions. In the reported GPT-NeoX-like setup, EK-FAC achieved Spearman 3 against conjugate-gradient ground truth with overhead 4 h and pair-wise time 5 s, while an MLP-only EK-FAC retained 6 with reduced pair-wise time 7 s. The same paper argues that MLP parameters contributed disproportionately to influence, accounting for approximately 85.5% of influence while comprising 75% of analyzed parameters (Bao et al., 8 May 2025).
Model editing and compression also use Kronecker-factored curvature as a regularizer or metric rather than as a preconditioner. In task arithmetic, a dataless regularizer based on K-FAC approximates the generalized Gauss–Newton penalty that measures representation drift, aggregates task-specific factors into a single surrogate per layer, and achieves constant complexity in the number of tasks while reporting state-of-the-art results in task addition and negation (Porrello et al., 19 Feb 2026). In neuron-based pruning, KFC approximates the Hessian block-diagonally as 8 and is used to penalize spectral radius during simultaneous training and pruning, with the reported result that the method improves the state of the art on neuron compression and obtains smaller networks with small accuracy degradation (Ebrahimi et al., 2021).
6. Assumptions, limitations, and open directions
The defining approximations remain the main source of both tractability and error. Standard K-FAC ignores inter-layer curvature and factorizes within-layer second-order statistics into independent activation and gradient terms. For architectures with strong cross-position interactions, this can be imperfect. In transformers, attention introduces cross-token couplings, so neither the expand nor the reduce formulation is exact even though both can work well empirically (Eschenhagen et al., 2023).
Estimation quality is another recurring issue. With limited or unrepresentative calibration or training data, Kronecker factors may be noisy, ill-conditioned, or singular. Recent LLM compression work explicitly reports singular estimates in some layers and recommends damping, diagonal loading, shrinkage toward the identity, or low-rank truncation; PINN work similarly reports that slower resampling can stabilize factor estimates (Chekalina et al., 23 May 2025, Dangel et al., 2024). Batch normalization adds a more structural problem: standard K-FAC assumes example independence within a minibatch, an assumption violated by BN statistics, which is why extended K-FAC introduces additional blocks to capture inter-example terms (Lee et al., 2020).
Memory and compute remain modest relative to exact second-order methods, but not negligible. Storage scales as 9 per block, which can be substantial for very wide layers. This has motivated module selection, block-diagonalization, quantization, truncated SVD, and the empirical practice of focusing first on MLP blocks in large transformers (Bao et al., 8 May 2025, Porrello et al., 19 Feb 2026).
The contemporary literature suggests a broadening of purpose. Kronecker-factored curvatures now serve as preconditioners for natural-gradient optimization, curvature metrics for low-rank compression, surrogates for influence functions, penalties for representation disentanglement, and proxies for flatness in pruning. A plausible implication is that the concept has evolved from a single optimizer design into a general structured language for second-order reasoning in large neural systems.