Second-Order Optimization and K-FAC

Plain gradient descent takes a step in the steepest downhill direction, but it has no idea how the loss surface curves — it can crawl painfully slowly through flat valleys and overshoot in steep ones. Second-order methods fix this by also using curvature (the Hessian, or in practice a related matrix called the Fisher information), but for a modern neural network with millions of parameters, that curvature matrix is far too large to even store, let alone invert.

K-FAC (Kronecker-Factored Approximate Curvature) makes second-order updates tractable by approximating each layer's curvature matrix as the Kronecker product of two much smaller matrices — one built from the layer's inputs, one from the gradients flowing back through its outputs. Both of those smaller matrices are cheap to estimate and invert, and their Kronecker product can be "divided into" the gradient without ever forming the full curvature matrix explicitly.

To see why the approximation is worth making, imagine a layer with 1,000 inputs and 1,000 outputs. Its full curvature matrix over the million weights would itself be a million-by-million object — computationally hopeless to store or invert every step. K-FAC instead estimates two 1,000-by-1,000 matrices (one from inputs, one from output gradients), each of which is small enough to invert in a fraction of a second, and combines them mathematically to stand in for the impossible million-by-million version.

In practice K-FAC shows up in settings where training runs are expensive enough that shaving off training steps is worth the extra bookkeeping per step — large-scale image classifiers and reinforcement learning policies being the classic examples — while smaller, cheaper models usually stick with first-order optimizers like Adam, for which the per-step cost is lower even if it needs more steps overall to converge.