Does Shampoo calculate the exact Hessian?

No. The basic method forms its preconditioners from accumulated gradient products.

Why use matrix inverse roots?

They transform the gradient using the accumulated structured statistics according to the algorithm's update rule.

Is it always faster than a diagonal optimizer?

No. Fewer iterations can be offset by matrix operations, memory usage or implementation overhead.

AI Glossary / Probabilistic machine learning / Advanced

Shampoo optimizer

Definition

Shampoo is an adaptive gradient optimizer that uses the matrix or tensor structure of model parameters to build manageable preconditioners. It accumulates gradient statistics along each tensor dimension and uses matrix inverse roots to transform subsequent updates. This captures richer directional information than elementwise scaling without storing one full matrix over every flattened parameter.

What it means

Ordinary gradient descent uses the gradient directly, scaled by a learning rate. Adaptive optimizers modify that scaling based on past information. A full-matrix preconditioner can account for relationships among parameter directions, but its storage and computation become enormous for large models.

Shampoo keeps smaller matrices associated with tensor axes. For a weight matrix, one matrix describes row-related gradient information and another describes column-related information. This structure is a computational compromise, not a claim that every interaction in a full preconditioner is retained exactly. The name belongs to a family of implementations derived from the original optimization method.

How it works

For the basic matrix case, let G be the current gradient of an m-by-n parameter matrix W. Starting with positive diagonal regularization, accumulate L by adding G G-transpose and R by adding G-transpose G. L is m-by-m and R is n-by-n.

The original update subtracts the learning rate times L to the power -1/4, multiplied by G, multiplied by R to the power -1/4. These are matrix powers, not elementwise fourth roots. Tensor versions apply corresponding preconditioning operations along each axis.

The statistics are built from gradients; Shampoo does not require explicitly computing the exact Hessian. Practical variants may block large tensors, use momentum, change accumulation rules or update inverse roots less frequently. Such choices affect behavior and cost, so the precise implementation matters. Numerical regularization keeps the matrix operations well defined and helps control ill-conditioning.

Three practical examples

1. Neural-network training: precondition the gradients of matrix or tensor weights while retaining their structure. Measure final validation performance and total training time rather than assuming fewer steps means a cheaper run.

2. Large multiclass models: use row and column statistics for a parameter matrix instead of treating all entries as one unstructured vector. Check whether the factor sizes fit the available memory.

3. Optimization research: compare structured preconditioning with diagonal adaptive methods on controlled tasks. Keep learning-rate tuning, data order and stopping criteria comparable when assessing the difference.

Think of it this way

Think of adjusting how movement is scaled along rows and columns of a grid rather than assigning a completely separate interaction rule to every pair of cells. The analogy explains the storage reduction. Real preconditioners can mix directions through matrix operations, and they do not literally measure physical row or column distances.

Worked example

Consider the first update for a two-by-two weight matrix. Set both initial accumulators to the identity matrix, using regularization one for transparent arithmetic. Let the current gradient be diagonal with entries 3 and 1. Then G G-transpose and G-transpose G are both diagonal with entries 9 and 1.

After accumulation, L and R are both diag(10,2). Their inverse fourth roots are diag(10^(-1/4),2^(-1/4)). Multiplying L^(-1/4) G R^(-1/4) gives diag(3/sqrt(10),1/sqrt(2)), approximately diag(0.948683,0.707107).

With learning rate 0.1, the corresponding weight changes are approximately -0.0948683 and -0.0707107 on the diagonal. Plain gradient descent with the same learning rate would change them by -0.3 and -0.1. The larger gradient direction is scaled more strongly in this example. Off-diagonal interactions are absent because the chosen gradient and accumulators are diagonal; the calculation illustrates the matrix-root update without demonstrating its richer non-diagonal behavior. The unusually large regularization is a teaching choice, not a tuning recommendation.

Illustrative case study

Fictional case: Spruce Training Lab tests Shampoo on a model with several large weight tensors. It observes fewer training iterations to reach a selected loss but discovers that frequent inverse-root computations increase wall-clock time. The team benchmarks a documented variant with different preconditioner update frequency.

Reviewers also compare memory consumption and validation quality across repeated runs. They report the implementation settings alongside the result instead of presenting every Shampoo variant as interchangeable. This fictional case illustrates a systems tradeoff, not a claim that a particular update schedule always improves training.

Common misconceptions and mistakes

Shampoo is not simply Adam with a different learning-rate constant. Its matrix statistics and tensor structure change the update. It is also not exact Newton optimization: the accumulated gradient matrices should not be labeled as the exact Hessian.

The original convex analysis does not guarantee a global optimum for a nonconvex neural network. Large axis sizes can still make factor matrices expensive. For an m-by-n matrix, storing the two factors costs m squared plus n squared entries, compared with m squared times n squared for a fully flattened matrix, but this saving does not make every model inexpensive. Numerical precision and inverse-root accuracy remain relevant.

Business use

Consider Shampoo when training cost is significant and there is capacity to evaluate a more complex optimizer. Include preconditioner state, matrix-root work and distributed communication in the cost comparison. A 100-by-50 weight matrix needs 12,500 entries for its two dense factors versus 25 million for a full 5,000-by-5,000 matrix, before other optimizer state is counted.

Record blocking, damping, learning-rate schedule, momentum and preconditioner update frequency. Compare at equivalent quality targets and realistic hardware settings. Treat an optimizer change as an empirical training-system decision; faster loss reduction does not establish better generalization or justify changing evaluation standards.

Check your understanding

Why does the example's first transformed gradient equal 3/sqrt(10)?

Answer: The original gradient 3 is multiplied by 10^(-1/4) from both sides. Together the two factors give 10^(-1/2).

Related terms

Preconditioning
Learning rate
adaptive gradient
AdaGrad
matrix inverse root
tensor
gradient second moment
Kronecker structure

Learn next

Reproduce the diagonal matrix example, then inspect a small non-diagonal gradient to see how matrix powers differ from elementwise operations. Next compare storage requirements and benchmark the complete optimizer implementation.

Sources

Gupta, Koren and Singer, Shampoo: Preconditioned Stochastic Tensor Optimization, matrix algorithm and tensor extension: https://proceedings.mlr.press/v80/gupta18a/gupta18a.pdf

The accumulator and inverse-fourth-root update were checked in the original paper. The numerical update, memory comparison and fictional case are original illustrations. Independent review consulted Algorithm 1 for the matrix update and its factor-storage comparison.

Previous: Shallow learning Next: Shannon coding