Natural Gradient and K-FAC

An ordinary gradient step treats every direction in parameter space as equally sensitive, but that's rarely true: a small change to one weight might barely affect the model's predictions while an equally small change to another weight swings predictions wildly. The natural gradient corrects for this by rescaling the step using the Fisher information matrix, a measure of how sensitive the model's output distribution is to each parameter direction, so a step of a given size always produces roughly the same change in the model's actual predictions, regardless of which parameters happen to be involved.

The catch is that the full Fisher information matrix for a modern neural network has far too many entries to store or invert directly. K-FAC (Kronecker-Factored Approximate Curvature) makes the natural gradient usable in practice by approximating that matrix, layer by layer, as a Kronecker product of two much smaller matrices — one capturing the layer's inputs, one capturing its output-side sensitivities — which can be inverted cheaply and still captures most of the useful curvature information.

A worked example

For a layer with 1,000 inputs and 1,000 outputs, the exact curvature information involves a matrix with roughly a million squared entries — computationally infeasible to invert repeatedly during training. K-FAC's Kronecker approximation replaces this with two 1,000×1,000 matrices, each far cheaper to invert, and empirically this approximation is close enough that K-FAC-trained networks converge in substantially fewer iterations than plain gradient descent, even though each iteration costs somewhat more.