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.

Note: the definition above is for a generative model $p(x; \theta)$. For supervised learning, the relevant object is the conditional-predictive Fisher $F(\theta) = \mathbb{E}_{x \sim \mathcal{D},\, y \sim p(\cdot \mid x; \theta)}[\nabla_\theta \log p(y \mid x; \theta) \nabla_\theta \log p(y \mid x; \theta)^\top]$, where the outer expectation is over the data distribution and the inner over the model's predictive distribution. K-FAC is explicitly designed as a Kronecker-factored approximation to this Fisher (or the closely related generalized Gauss-Newton matrix). Shampoo is not derived from the Fisher; it is built from accumulated outer products of the observed gradients, more in the spirit of full-matrix AdaGrad. The two are both block-Kronecker preconditioners but capture different second-order objects.

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$ exponent comes from a Löwner-order (matrix AM-GM) bound, not from any vec-vs-matrix conversion. If $H_t = \sum_{s \le t} g_s g_s^\top$ is the full-matrix AdaGrad preconditioner on the vectorized gradient, and $L_t, R_t$ are obtained from partial traces of the $g_s g_s^\top$ terms, then Gupta, Koren, and Singer (2018) show $L_t \otimes R_t \succeq H_t$ in the Löwner order. By operator monotonicity of $X \mapsto X^{-1/2}$, we get $(L_t \otimes R_t)^{-1/2} \preceq H_t^{-1/2}$. Using $(R \otimes L)^{-1/2} = R^{-1/2} \otimes L^{-1/2}$ and the identity $(R^{a} \otimes L^{a}) \text{vec}(G) = \text{vec}(L^{a} G R^{a})$, applying $(L_t \otimes R_t)^{-1/2}$ as a matrix-shaped update and symmetrizing between the two factors yields $\Delta W_t = L_t^{-1/4} G_t R_t^{-1/4}$. The Kronecker vec identity preserves exponents; there is no square-root introduction from switching between matrix and vectorized form.

Scope caveat: the $1/4$ power is a design choice, not uniquely determined. The Löwner bound only establishes that $(L \otimes R)^{-1/2}$ upper-bounds the AdaGrad preconditioner; splitting this symmetrically across the two factors gives $1/4$, but other exponents are defensible. Recent work (Morwani, Shah, Cohen, Mhaskar, 2024; Anil, Gupta, Koren, Regan, Singer, 2020 Distributed Shampoo; SOAP) often prefers $1/2$, which has shown comparable or better results in large-scale deployments. Empirically, intermediate exponents $p \in (0, 1/2]$ trade off preconditioning strength.

Shampoo is a full-matrix AdaGrad approximation, not a Fisher approximation. It directly accumulates $GG^\top$ (left) and $G^\top G$ (right) from the observed gradients themselves. The left preconditioner captures correlations between output rows of the weight matrix; the right captures correlations between input columns. The relevant object being approximated is the AdaGrad full-matrix preconditioner $\sum_s g_s g_s^\top$ on the vectorized gradient, upper-bounded in the Löwner order by $L \otimes R$. K-FAC, by contrast, targets the Fisher / Gauss-Newton matrix using a different factorization (input-activation covariance Kronecker pre-activation-gradient covariance). Both are block-Kronecker preconditioners but they approximate different second-order quantities.

The matrix fourth root $L^{-1/4}$ arises from the operator-inequality chain described in the statement: the symmetric split of $(L \otimes R)^{-1/2}$ across the two factors. Empirically, intermediate exponents $p \in (0, 1/2]$ trade off preconditioning strength against stability, with $1/4$ and $1/2$ both in active use.

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's Distributed Shampoo (Anil et al., 2020) is the production variant: preconditioner statistics $L$ and $R$ are sharded across data-parallel workers, the matrix powers $L^{-p}$ and $R^{-p}$ are recomputed every $k$ steps (typically $k = 100$ to $1000$) on dedicated CPU or GPU workers in parallel with training, and only the small matrix powers (not the full statistics) are communicated back. This amortizes the $O(m^3)$ eigendecomposition cost across many training steps and hides it behind the forward-backward pass. 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.