Reading the order one layer at a time

Reading the order off the Fisher decay worked for one direction. A real network has billions of parameters, and its full Fisher information matrix is that number squared: far too large to form, store, or take eigenvalues of. So the slope reading from earlier cannot be run on the whole network at once. The way through is that the Fisher has structure, and that structure is what makes the order readable at scale.

The Fisher of a layer factorises

Take a single weight matrix $W$ in the network, mapping an input activation $a$ to a pre-activation output. Its block of the Fisher information has a remarkable approximate form, a result from optimization theory we take as the starting point here and the basis of K-FAC:1

$$ F_W \;\approx\; A \otimes G. $$

Here $A = \mathbb{E}[a\,a^\top]$ is the covariance of the layer's inputs (a forward quantity), and $G = \mathbb{E}[\delta\,\delta^\top]$ is the covariance of the gradient arriving at the layer's output (a backward quantity). Both are small: their size is the layer's width, not the parameter count. The cost of looking at curvature drops from one giant matrix to two tiny ones.

The Kronecker product has a clean spectrum. Every eigenvalue of $F_W$ is a product of one eigenvalue from each factor,

$$ \lambda_{ij}(F_W) \;=\; \lambda_i(A)\,\mu_j(G), $$

with eigenvector $v_i \otimes u_j$. The curvature in a layer is laid out as a grid: forward modes along one axis, backward modes along the other, and each cell their product.

The demo below makes that grid literal. Two small factors, $A$ on the forward side and $G$ on the backward side, build the whole block between them: the Kronecker product drops a scaled copy of $G$ into every entry of $A$. Drag any cell of a factor and watch the copies of it inside the product move with it.

Shrink a diagonal entry of either factor toward zero and a whole copy of the other factor fades with it. A small mode on one side drags a block of the layer's curvature down, and that is a dead direction read off the small factor rather than the full block.

Forward death and backward death

A dead direction is a cell whose product is zero, and a product is zero when either factor is. This gives two ways for a layer to go flat, and they are mirror images.

A forward mode dies when some direction in the input activations carries no variance, for example a feature that has saturated to a constant. Its eigenvalue $\lambda_i(A)$ goes to zero, and the entire row of products it sits in collapses: paired with every backward mode, it is dead.

A backward mode dies when some direction in the output gradient vanishes, for example an output the loss has stopped constraining. Its eigenvalue $\mu_j(G)$ goes to zero, and the entire column collapses.

This is the A–G duality: the forward, activation side and the backward, gradient side are two readings of the same curvature, and a dead direction can come from either.

Because a dead direction can sit on either side, finding one means measuring both. You take the two cheap spectra, the forward $A$ and the backward $G$, and read off which one has a collapsed mode. That side is the correct one to measure the dead direction in, and a layer that looks perfectly healthy from one side can still be degenerate when you check the other.

Starve a mode on the forward side or the backward side and the verdict names where the dead direction lives. Reading only one spectrum would miss a degeneracy hiding on the other, which is why the method always measures both, and each spectrum is small enough to be cheap to take.

The same reading, layer by layer

A real network stacks many such layers, and each one carries its own $A$ and its own $G$, so its own pair of spectra and its own order. To read where a deep network has gone flat you never assemble its enormous Fisher. You walk the layers, and at each one factor the local curvature into its small forward and backward parts and read which modes have gone slack.

Pick a layer and the two spectra change; the side that is flatter, where this layer's dead direction lives, shifts with depth. Train the network and the spectra spread as the model commits to some directions and lets others go slack. The slope reading that needed one clean direction now climbs the network one rung at a time, combining the per-layer readings into a picture of where the whole network has gone flat.

Still open / underspecified. $F_W \approx A \otimes G$ is an approximation, exact for linear networks, approximate for nonlinear ones. The two spectra you read are exact objects either way: the rates of $A$ and $G$ are statements about those factors themselves, with no dependence on the approximation. What the approximation governs is the lift, how faithfully the per-layer block picture reflects the full network's Fisher, and when cross-layer coupling is strong that lift can miss or invent whole-network degeneracy. Characterising when the lift is trustworthy is active work. Two pieces that once sat in this caveat now have explainers of their own: the two factors are enough to construct a candidate dead direction outright (reading at a standstill#21), and the per-direction readings assemble into a single global coefficient (adding up the orders#23).


  1. The Kronecker factorisation of a layer's Fisher is from James Martens and Roger Grosse, Optimizing Neural Networks with Kronecker-factored Approximate Curvature, ICML 2015. ↩︎