Preconditioned Optimizers: Shampoo, K-FAC, and Natural Gradient

The natural gradient of a loss $\mathcal{L}(\theta)$ is:

$\tilde{\nabla} \mathcal{L}(\theta) = F(\theta)^{-1} \nabla \mathcal{L}(\theta)$

where $F(\theta) = \mathbb{E}_{x \sim p(\cdot; \theta)}[\nabla \log p(x; \theta) \nabla \log p(x; \theta)^\top]$ is the Fisher information matrix.

The natural gradient update is:

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

This is steepest descent in the KL divergence geometry: it finds the direction that decreases the loss most per unit of distributional change $D_{\text{KL}}(p_{\theta + \delta} \| p_\theta)$.

Euclidean gradient descent treats all parameter directions equally: a step of size $\epsilon$ in weight 1 is the same as a step of size $\epsilon$ in weight 1000. But in a neural network, some parameters have much larger effects on the output distribution than others. The Fisher information matrix captures these sensitivities.

Multiplying by $F^{-1}$ rescales the gradient so that the update produces equal change in the output distribution in all directions. This is the same idea as Newton's method (which uses the Hessian instead of the Fisher), but adapted for probabilistic models.

The natural gradient is parameterization-invariant: it gives the same update regardless of how you parameterize the model. Reparameterizing a neural network (e.g., scaling a weight matrix by a constant and dividing the next layer by the same constant) changes the Euclidean gradient but not the natural gradient. This is why natural gradient methods can be more robust to architecture choices.

Practically, natural gradient methods converge in fewer steps than Adam on problems with ill-conditioned Fisher matrices (which is most neural networks). The challenge is computing $F^{-1}$, which costs $O(d^3)$ for $d$ parameters.

The Fisher matrix has $d^2$ entries for $d$ parameters. For a model with 100M parameters, storing $F$ requires $10^{16}$ entries, which is impossible. Every practical preconditioned optimizer is an approximation to the natural gradient: diagonal (Adam), block-diagonal (K-FAC), Kronecker-factored (Shampoo), or low-rank.

Shampoo (Gupta, Koren, Singer, 2018) maintains left and right preconditioners for each weight matrix:

$L_t = (1 - \beta) L_{t-1} + \beta \, G_t G_t^\top \in \mathbb{R}^{m \times m}$ $R_t = (1 - \beta) R_{t-1} + \beta \, G_t^\top G_t \in \mathbb{R}^{n \times n}$

The preconditioned update is:

$\Delta W_t = L_t^{-1/4} \, G_t \, R_t^{-1/4}$

The $-1/4$ power (rather than $-1/2$) arises because the Kronecker approximation $F \approx L \otimes R$ means $F^{-1/2} \approx L^{-1/2} \otimes R^{-1/2}$, but applied to a matrix (not a vectorized matrix), this becomes $L^{-1/4} G R^{-1/4}$.

Shampoo is K-FAC for general weight matrices, simplified. Instead of separating the Fisher into input covariance and gradient covariance, it directly accumulates $GG^\top$ (left) and $G^\top G$ (right) from the gradient itself. The left preconditioner captures correlations between output neurons; the right captures correlations between input features. Together, they approximate the full Fisher.

The matrix fourth root $L^{-1/4}$ is the geometric mean of the identity and the inverse square root. It provides a balance between no preconditioning ($I$) and full natural gradient ($F^{-1/2}$).

Shampoo has shown strong empirical results on transformer training, sometimes matching or exceeding Adam with fewer steps (though each step is more expensive due to the matrix power computation). Google has used Shampoo-derived optimizers (Distributed Shampoo) for production training of large models. The connection to Riemannian optimization is direct: Shampoo's update can be interpreted as Riemannian gradient descent on the space of weight matrices with a specific metric.

Computing $L^{-1/4}$ requires eigendecomposition of $L$, costing $O(m^3)$ per step. For large layers ($m = 4096$), this is expensive. Practical implementations amortize this cost by recomputing the preconditioner every $k$ steps (e.g., $k = 100$). The matrix power can also be approximated using Newton-Schulz iterations (same technique as Muon). Memory cost is $O(m^2 + n^2)$ per weight matrix for the preconditioners.