Optimizing Neural Networks with Kronecker-Factored Approximate Curvature

Standard gradient descent computes the steepest-descent direction in Euclidean space — it asks "which small change to the weights reduces the loss the most?" But this question depends on how you measure "small." In Euclidean space, a step of 0.01 in any direction looks the same. In the loss landscape, it doesn't: moving 0.01 along a steep valley wall might launch you out of the valley entirely, while moving 0.01 along the valley floor barely changes the loss.

The curvature of the loss — how fast the gradient itself changes — varies enormously across different parameter directions. First-order methods are blind to this curvature, which forces practitioners to use small learning rates that accommodate the steepest direction, wasting potential progress in all other directions.

Amari (1998) proposed a better question: instead of "which parameter change of Euclidean length ε reduces the loss the most?", ask "which parameter change that moves the model's output distribution by at most ε reduces the loss most?" This is the natural gradient.

The natural gradient replaces the identity matrix in the gradient update with the Fisher information matrix $F$, which measures how sensitive the model's output distribution is to each parameter direction. The update becomes:

$\theta_{t+1} = \theta_t - \eta F^{-1} \nabla L(\theta_t)$

By pre-multiplying the gradient by $F^{-1}$, the natural gradient stretches the step along flat directions (where the Fisher eigenvalues are small) and compresses it along steep directions (where they are large). The result is reparameterization invariant — the optimization behaves the same regardless of how you parameterize the network.

Think of the Fisher as a ruler that measures distances in "distribution space" rather than "parameter space." Two parameter vectors that are far apart in Euclidean terms but produce nearly identical output distributions should be considered close; two that are nearby in parameters but produce very different distributions should be considered far apart. The Fisher encodes exactly this geometry.

K-FAC makes two structural approximations to the Fisher:

Approximation 1 — Block diagonal. Assume parameters in different layers are statistically independent. This turns the giant $N \times N$ Fisher into a block-diagonal matrix where each block covers one layer. Cross-layer interactions are ignored — a reasonable trade-off because most curvature structure is within layers.

Approximation 2 — Kronecker factorization. For a fully-connected layer with weight matrix $W \in \mathbb{R}^{m \times n}$, the Fisher block is an $mn \times mn$ matrix. K-FAC approximates it as the Kronecker product of two much smaller matrices: $\hat{F}_\ell \approx A_{\ell-1} \otimes G_\ell$ where $A_{\ell-1} = \mathbb{E}[a_{\ell-1} a_{\ell-1}^\top]$ is the $n \times n$ covariance of layer inputs (activations), and $G_\ell = \mathbb{E}[g_\ell g_\ell^\top]$ is the $m \times m$ covariance of backpropagated gradients.

The Kronecker factorization is equivalent to assuming that the activations $a$ and the backpropagated gradients $g$ are statistically independent. While not exactly true, this assumption captures the dominant structure: $A$ encodes which input features co-activate and $G$ encodes which output gradients co-vary. Their Kronecker product naturally models how combinations of input correlations and gradient correlations shape the curvature.

The Kronecker product $A \otimes G$ of an $n \times n$ matrix and an $m \times m$ matrix produces an $mn \times mn$ matrix — but you never need to form it explicitly. The critical identity is:

$(A \otimes G)^{-1} = A^{-1} \otimes G^{-1}$

Instead of inverting one $mn \times mn$ matrix (cost $O(m^3 n^3)$), you invert an $n \times n$ matrix and an $m \times m$ matrix separately (cost $O(n^3 + m^3)$). For a layer with 512 inputs and 256 outputs, this reduces the inversion from a $131{,}072 \times 131{,}072$ problem to a $512 \times 512$ plus a $256 \times 256$ problem — a factor of roughly one billion savings.

For each layer $\ell$ with weight matrix $W_\ell$, the K-FAC update has three stages:

Stage 1 — Accumulate statistics. During training, maintain running averages of the activation covariance $A_{\ell-1} = \mathbb{E}[a_{\ell-1} a_{\ell-1}^\top]$ and the gradient covariance $G_\ell = \mathbb{E}[g_\ell g_\ell^\top]$. These are computed from the same forward and backward passes you already do for SGD — almost no extra work.

Stage 2 — Invert the factors. Periodically (not every step), invert $A_{\ell-1}$ and $G_\ell$ separately. Each is a small matrix, so this is cheap and can run asynchronously on CPU while the GPU does the next forward pass.

Stage 3 — Apply the preconditioned gradient. The natural gradient update for the vectorized weight is $F^{-1} \text{vec}(\nabla L)$. Using the Kronecker structure, this simplifies to a sandwich: $\Delta W_\ell = G_\ell^{-1} \, (\nabla_{W_\ell} L) \, A_{\ell-1}^{-1}$ This is just two matrix multiplications — the same computational pattern as a forward pass.

The Fisher approximation isn't perfect, and the Kronecker factors can have very small eigenvalues, leading to dangerously large update steps. K-FAC uses Tikhonov damping — adding a small constant $\lambda$ to the diagonal before inverting:

$(\hat{F}_\ell + \lambda I)^{-1}$

But naively adding $\lambda I$ to $A \otimes G$ destroys the Kronecker structure and makes inversion expensive again. K-FAC's clever solution is factored Tikhonov damping: distribute the damping across the two factors using a balancing scalar $\pi$:

$(A + \pi\sqrt{\lambda}\,I)^{-1} \otimes (G + \frac{1}{\pi}\sqrt{\lambda}\,I)^{-1}$

The scalar $\pi$ is chosen to minimize the approximation error. This preserves the Kronecker structure, so the cheap inversion identity still applies.

The paper's experiments on deep autoencoders and convolutional networks showed K-FAC converging in significantly fewer iterations than SGD with momentum — often reaching the same loss in 3–10× fewer updates. While each K-FAC step costs several times more than an SGD step (due to the covariance estimation and periodic inversions), the net effect is faster wall-clock convergence.

Unlike Hessian-free optimization — the previous best second-order method — K-FAC works well in highly stochastic settings with small mini-batches. The running-average covariance estimates naturally smooth out noise, and the factored damping prevents the approximation from becoming unstable.