Eigenvalue-Corrected K-FAC (EK-FAC)

• EK-FAC is a second-order curvature approximation that improves K-FAC by integrating empirical eigenvalue corrections to reduce spectral errors.
• It leverages eigendecompositions of layer covariances to form a reliable Kronecker eigenbasis, enabling accurate inverse Hessian-vector product preconditioning.
• The method accelerates optimization and influence-function analysis in deep networks while maintaining lower computational overhead compared to iterative solvers.

Eigenvalue-Corrected Kronecker-Factored Approximate Curvature (EK-FAC) is a second-order curvature approximation technique that significantly improves the accuracy and scalability of inverse Hessian-vector product (IHVP) computations in large neural networks and modern deep learning applications. EK-FAC builds upon the Kronecker-Factored Approximate Curvature (K-FAC) method, augmenting its Kronecker-factor eigenbasis with empirical eigenvalue corrections obtained from the true curvature in that basis, resulting in superior spectral faithfulness and more accurate estimates for preconditioned optimization and influence-function analysis (George et al., 2018, Gao et al., 2020, Bao et al., 8 May 2025, Grosse et al., 2023, Hong et al., 27 Sep 2025).

1. Mathematical Foundation and Motivation

The core challenge addressed by EK-FAC arises from the prohibitive computational and memory costs associated with forming, storing, and inverting curvature matrices (e.g., Fisher information matrix or Generalized Gauss–Newton (GGN) matrix) in high-dimensional neural networks. For a model with $D$ parameters, these matrices are typically $D \times D$, making exact second-order methods infeasible for large-scale models. K-FAC circumvents this via two approximations:

• Block-diagonalization: The curvature is decomposed into independent layer-wise blocks.
• Kronecker-factorization: Each block (for layer $l$) is approximated as a Kronecker product $G_l \approx A_{l-1} \otimes S_l$, where $A_{l-1}$ and $S_l$ are covariances of the layer's inputs and pre-activation gradients, respectively.

However, the Kronecker product assumption leads to systematic spectral errors, as the true covariance in the "mixed" basis deviates from the product of factor spectra. EK-FAC was introduced to remedy this by accurately capturing variances along principal directions in the Kronecker-factor eigenbasis, thus correcting a primary source of K-FAC's approximation error (George et al., 2018, Gao et al., 2020, Hong et al., 27 Sep 2025).

2. EK-FAC Approximation: Definition and Construction

Given a neural network layer with parameter matrix $W_l \in \mathbb{R}^{P \times M}$, and letting $w_l = \mathrm{vec}(W_l)$, the K-FAC approximation to the per-layer GGN block or Fisher matrix is

$G_l \approx A_{l-1} \otimes S_l,$

where $$A_{l-1} = \mathbb{E}[\bar{a}_{l-1}\bar{a}_{l-1}^\top]$$ and $D \times D$0, with $D \times D$1 as the bias-augmented activation and $D \times D$2 as pre-activation pseudo-gradient.

EK-FAC improves on this by performing the following sequence:

1. Eigenbasis formation: Compute the eigendecompositions $D \times D$3 and $D \times D$4.
2. Kronecker Eigenbasis: Form the basis $D \times D$5 spanning the parameter space of the layer.
3. Empirical Eigenvalue Correction: Project per-example gradients (or pseudo-gradients) onto $D \times D$6, and set the diagonal matrix $D \times D$7, where $D \times D$8 is the per-example vectorized gradient for the layer. Thus, the corrected curvature approximation is

$D \times D$9

1. Damping: For well-conditioned inversion, add a diagonal regularizer: $l$0, with $l$1.

The EK-FAC inverse-HVP for a vector $l$2 then follows:

$l$3

This construction produces a preconditioner whose diagonal in the Kronecker eigenbasis matches the true moment matrix, minimizing the Frobenius norm distance to the true curvature among all diagonal-corrected matrices in that basis (George et al., 2018, Gao et al., 2020).

3. Algorithmic Implementation and Computational Complexity

Factor Precomputation:

1. Collect forward activations and pre-activation gradients for each layer over a sample batch.
2. Estimate $l$4 and $l$5 by empirical averaging.
3. Eigendecompose $l$6 and $l$7 to obtain $l$8, $l$9.
4. Project per-example gradients into $G_l \approx A_{l-1} \otimes S_l$0 and estimate their squared magnitudes for the EK-FAC eigenvalue diagonal.

Inverse-HVP Application (per vector):

For each layer $G_l \approx A_{l-1} \otimes S_l$1:

1. Reshape $G_l \approx A_{l-1} \otimes S_l$2 to a $G_l \approx A_{l-1} \otimes S_l$3 matrix $G_l \approx A_{l-1} \otimes S_l$4.
2. Apply $G_l \approx A_{l-1} \otimes S_l$5 (basis transform).
3. Divide elementwise by the corresponding $G_l \approx A_{l-1} \otimes S_l$6 diagonal values.
4. Apply $G_l \approx A_{l-1} \otimes S_l$7 and vectorize for the output.

The major computational bottleneck is the eigendecomposition step, $G_l \approx A_{l-1} \otimes S_l$8 per layer, amortized over many IHVPs. Each inverse-HVP costs $G_l \approx A_{l-1} \otimes S_l$9, matching K-FAC. Memory overhead per layer comprises storage for $A_{l-1}$0, $A_{l-1}$1, and $A_{l-1}$2, scaling as $A_{l-1}$3 per layer (Grosse et al., 2023, Bao et al., 8 May 2025).

A tabular contrast of key computational aspects is given below:

Method	Basis	Eigenvalues	IHVP Complexity per Layer	Memory Overhead
K-FAC	$A_{l-1}$4	$A_{l-1}$5	$A_{l-1}$6	$A_{l-1}$7
EK-FAC	$A_{l-1}$8	$A_{l-1}$9	$S_l$0	$S_l$1

4. Application to Influence Functions and Large Models

Influence functions assess the change in model predictions induced by infinitesimal upweighting of training points, requiring IHVPs with the Hessian or GGN. EK-FAC enables scalable approximate influence-function computation by providing a fast, low-error inverse for the curvature. Compared to iterative solvers such as LiSSA or conjugate gradient, EK-FAC achieves similar or higher accuracy (quantified by Pearson and Spearman correlations with "ground truth" influence values) but at orders-of-magnitude lower wall-clock and compute cost per IHVP (Bao et al., 8 May 2025, Grosse et al., 2023).

EK-FAC has been successfully applied to LLMs with up to 52 billion parameters by:

• Focusing on MLP block parameters (which dominate total count).
• Employing block-diagonalization and blockwise approximations for memory management.
• Leveraging batched query processing and additional filtering (e.g., TF-IDF) to further reduce gradient accumulation overheads (Grosse et al., 2023, Bao et al., 8 May 2025).

5. Comparative Analysis: EK-FAC vs. Alternatives

EK-FAC occupies an intermediate accuracy-efficiency regime between K-FAC and exact (or unfactorized block) curvature inversion.

• Spectral Fidelity: EK-FAC yields a higher overlap between its spectrum and that of the true GGN, capturing 30–50% of the eigenvalue error that K-FAC introduces, particularly in deep or under-trained networks (Hong et al., 27 Sep 2025). The dominant approximation error in K-FAC arises from its Kronecker-product eigenvalues, and EK-FAC’s empirical correction greatly reduces this error source.
• Data Attribution Accuracy: Influence scores estimated using EK-FAC show consistently higher quality than those from K-FAC; however, both remain below unfactorized block-diagonal GGN. The residual gap is due to the fixed Kronecker eigenbasis, which cannot capture off-diagonal structure (Hong et al., 27 Sep 2025).
• Optimization Dynamics: EK-FAC accelerates per-epoch convergence in deep autoencoders, VGG, and ResNet architectures, with no degradation in generalization compared to K-FAC or first-order methods (George et al., 2018, Gao et al., 2020).
• Variants: Trace-restricted EK-FAC (TEKFAC) and alternatives such as TKFAC further refine the approximation, but the fundamental eigenvalue correction of EK-FAC remains central for improved spectral alignment (Gao et al., 2020).

6. Practical Considerations and Limitations

EK-FAC's main practical requirements are amortizing the expensive eigenbasis formation and maintaining accurate running averages of empirical coordinate variances in the Kronecker eigenbasis. Damping hyperparameters are important for numerical stability; empirical values are not highly sensitive as long as the curvature is regularized to avoid singularities (Grosse et al., 2023, Gao et al., 2020). For extremely large layers, further block-diagonalization and careful factor management are necessary to keep memory and computational costs feasible.

Limitations include:

• Residual Kronecker Error: EK-FAC cannot capture cross-layer curvature or non-Kronecker structure, which may dominate in very deep or highly-interdependent models (Hong et al., 27 Sep 2025).
• Linearization Assumptions: EK-FAC is derived under a local linearization (GGN or Fisher), thus cannot account for inherently nonlinear training phenomena such as circuit formation or sharp transitions in capacity (Grosse et al., 2023).
• Empirical Sensitivity: The benefit of EK-FAC over K-FAC may diminish near convergence or in shallow models; for very small batch sizes or small models, eigenvector update cost may outweigh diagonal correction benefits (George et al., 2018, Gao et al., 2020).

7. Impact and Empirical Evidence

EK-FAC has enabled previously intractable large-scale influence-function analyses in billion-parameter LLMs and complex deep architectures. Across several benchmarks and case studies—including GPT-NeoX and Dolly-v2-3b—EK-FAC-based IHVPs provide substantially improved accuracy–efficiency trade-offs compared to both naive dot-product baselines and iterative solvers (Bao et al., 8 May 2025, Grosse et al., 2023). Storage overheads are typically tolerable for models up to the billion-parameter scale.

Empirical studies show that:

• EK-FAC matches the influence-estimation accuracy of more expensive iterative Hessian solvers and outperforms K-FAC for most applications requiring layerwise spectral fidelity.
• The improved curvature approximation enables deeper investigation of generalization, attribution, and robustness phenomena in modern neural networks (Grosse et al., 2023, Hong et al., 27 Sep 2025).
• Its running-average and mini-batch update variants ensure accuracy of second-moment tracking with minimal computational impact (George et al., 2018).

In summary, Eigenvalue-Corrected Kronecker-Factored Approximate Curvature constitutes a crucial advancement for scalable, accurate second-order analysis in deep learning, particularly in influence-function-based model interpretability and diagnostics, and establishes a new standard for tractable large-scale Fisher/GGN matrix approximation (George et al., 2018, Gao et al., 2020, Grosse et al., 2023, Bao et al., 8 May 2025, Hong et al., 27 Sep 2025).