EKFAC (Eigenvalue-Corrected Kronecker Factorization)

For each neural network layer l with input activation covariance A = E[h h^T] = U_A S_A U_A^T and backpropagated pre-activation gradient covariance B = E[delta delta^T] = U_B S_B U_B^T, EKFAC forms the orthogonal Kronecker-Factored Eigenbasis (KFE) U_A (x) U_B and defines the preconditioning matrix G_EKFAC = (U_A (x) U_B) S* (U_A (x) U_B)^T, where the diagonal entries S*_{ii} = s*_i = E[((U_A (x) U_B)^T grad_theta)_i^2] track the exact second moments of the gradient coordinates in the KFE basis, computing basis projections efficiently via (U_A (x) U_B)^T vec(C) = vec(U_B^T C U_A).

While recomputing the eigendecompositions of A and B is amortized every n mini-batches, EKFAC (and its exponential moving average variant EKFAC-ra) updates the diagonal scaling vector s* in the fixed KFE basis at every mini-batch, preventing the eigenspectrum of the curvature estimate from drifting away from the true empirical Fisher matrix between basis recomputations.

Lemma 1 and Theorems 2-3 prove that for any positive semi-definite matrix G = E[grad_theta grad_theta^T] and orthogonal KFE basis Q = U_A (x) U_B, the diagonal matrix D with entries D_{ii} = (Q^T G Q)_{ii} = E[(Q^T grad_theta)_i^2] uniquely minimizes ||G - Q D Q^T||_F over all diagonal matrices D, guaranteeing ||G - G_EKFAC||_F <= ||G - G_KFAC||_F everywhere.

Applied to training an 8-layer deep MNIST auto-encoder and VGG11 and ResNet-34 convolutional neural networks on CIFAR-10 (using Spatially Uncorrelated Activations for convolutional layers), EKFAC and EKFAC-ra achieve lower training loss per epoch and faster wall-clock convergence than amortized KFAC and batch-normalized SGD with momentum.

EKFAC optimization performance is sensitive to the choice and initialization of the damping hyperparameter epsilon and learning rate eta, requires careful variance-preserving weight initialization when training deep residual networks without batch normalization, and shows diminished wall-clock speedup over KFAC at small mini-batch sizes on smaller networks.