How does the Shampoo optimizer change a neural network’s learning updates? Instead of dividing each gradient coordinate by its own running scale, it uses matrix statistics along the dimensions of a weight tensor to construct a structured, preconditioned update. For a two-dimensional weight matrix, a simplified Shampoo step multiplies its gradient from the left and right by inverse fourth roots of accumulated positive-definite matrix statistics. The goal is to account for relationships among coordinates without storing one enormous dense matrix for every parameter in the model.
Shampoo is a family of practical implementations building on Gupta, Koren and Singer’s 2018 ICML research. Its value is not that every model learns faster or that the optimizer knows the right answer. It changes the geometry of the training step; useful gains depend on numerical stability, matrix size, precision, loss scaling, hardware, implementation details and a fair experimental comparison. Here the examples show exactly how the matrices, dimensions, memory and step scales are calculated.
Start with the short mechanics if you are learning optimisation. Engineers can follow the stability, batching and distributed-training sections; evaluators can go directly to the worked comparisons, failure diagnosis and independent transfer questions. The numerical examples are deliberately specified teaching cases, not benchmarks measured on eduKate systems or promises of accuracy.
Contents and reader routes
For first principles, follow The 60-second route: what changes from gradient descent?; First principles: parameters, gradients and preconditioners; Why diagonal statistics cannot see the same correlations; Build the row and column statistics with a small example; Calculate an inverse fourth root without guessing; Why the exponent is one quarter for a matrix; Why damping and positive definiteness are mandatory. Work the two-by-two example before attempting a real model.
For practitioners, continue through A checked memory example: 1,024 rows and 4,096 columns; The price of inverse matrix roots; Blockwise Shampoo and the geometry it discards; Go from matrices to genuine higher-rank tensors; Cumulative statistics versus exponential moving averages; Grafting: borrow step magnitude without losing direction; Momentum and weight decay are separate controls; Precision and conditioning: the numerical failure layer; When should the preconditioner be refreshed?; Distributed Shampoo: ownership, communication and checkpoints; A fair comparison with AdamW, SGD and Adafactor. These are practical algorithm choices and resource boundaries.
For evaluation, read A failure-to-repair map for actual training runs; Original miniature experiment: two steps, two observed directions; Transfer problem: what if the matrix shape changes?; What Shampoo can and cannot establish. Finish with answered questions, original research and implementation sources and verified related eduKate reading.
The 60-second route: what changes from gradient descent?
Ordinary gradient descent takes a negative step along the estimated loss gradient. If two model coordinates respond differently to the same raw step size, their updates may be poorly balanced. Optimizers such as Adam scale coordinates with running first- and second-moment estimates. This is powerful, but a diagonal scale does not directly represent a relationship between different coordinates in one weight matrix. Shampoo uses a matrix-aware preconditioner to represent some of those relationships.
Imagine a sheet of paper covered by long, slanted contour lines. A raw gradient tells you a local direction of steepest rise in the current coordinate system, while a preconditioner changes the notion of distance used to choose a step. It can rotate and rescale the update. A good transformation might reduce zigzags on a difficult landscape; a poorly estimated one might magnify noise or make training unstable. Matrix preconditioning is a choice to be tested, not a shortcut to a guaranteed minimum.
The method keeps separate statistics for the rows and columns of a matrix, and more generally for each mode of a tensor. The statistics are computed from past gradients, then regularised and transformed with inverse matrix powers. The resulting update is multiplied by a learning rate, often with additional techniques such as momentum and grafting in modern implementations. A useful slogan is ‘respect tensor shape, but verify the resulting step’.
First principles: parameters, gradients and preconditioners
Let W have m rows and n columns. Its gradient G has exactly the same shape and units: a row entry describes how the loss changes when that parameter changes, holding the others locally fixed. A scalar learning rate controls the size of an ordinary update W_new = W_old minus eta times G. The first useful diagnostic is to confirm that the gradient’s dimensions match W before discussing any optimizer mathematics.
Shampoo builds one row-side matrix L of shape m by m and one column-side matrix R of shape n by n. These are not learned weight matrices and not necessarily Hessians. They are accumulated outer-product-style statistics derived from the gradients. In a simple cumulative matrix case L becomes the sum over steps of G times G-transpose, and R becomes the sum of G-transpose times G, with regularisation added to make inverse powers well behaved.
With appropriate symmetric positive-definite statistics, a simplified update direction is P = L raised to minus one-quarter times G times R raised to minus one-quarter. Its shape is m by n because the matrix multiplications are (m by m) times (m by n) times (n by n). The fractional exponents are matrix functions: they act on eigenvalues in an eigenbasis, not by taking the fourth root of each displayed entry.
Why diagonal statistics cannot see the same correlations
Suppose two parameters have gradients that often rise and fall together. A diagonal method records the magnitude information separately, but not the cross term telling you those coordinates tend to move together. An outer product retains entries relating pairs of coordinates. For example, the row statistic G times G-transpose has entry (i,j) equal to the dot product between gradient row i and gradient row j. Off-diagonal values can therefore represent relationships that coordinatewise squares miss.
This should not be confused with knowing the true loss curvature. Shampoo uses gradient correlations as preconditioning information; it does not magically observe every second derivative, and the statistics may reflect minibatch noise. A row correlation from one dataset, label mix or phase of training can become less relevant later. Learning benefits come from how useful the structured estimate is for the current optimization problem.
One practical reason shape matters is that neural network parameters already form matrices and tensors. Flattening a 1,024 by 4,096 weight matrix into a four-million-element vector throws away the convenient row-column factorisation when building a generic dense preconditioner. Shampoo exploits that two-dimensional organisation. It is still an approximation of richer parameter interactions: row and column factors cannot reproduce every possible full-vector covariance relationship.
Build the row and column statistics with a small example
Take an illustrative two-by-two gradient G whose first row is (2,0) and second row is (0,1). Its transpose is the same matrix. Multiplying G by G-transpose gives a row-side statistic L equal to diag(4,1); multiplying G-transpose by G gives R equal to diag(4,1). The first row and column directions have accumulated stronger squared-gradient evidence than the second. All numbers here are dimensionless model gradients for a made-up exercise.
At a first observation, these statistics are rank-complete and positive definite; later observations would add additional outer products if using cumulative statistics. If the gradient changed to a non-diagonal matrix, the products could include off-diagonal entries. Whether to accumulate the new matrices, exponentially average them or partition them into blocks is a distinct implementation decision. Do not silently interchange these variants when comparing logs.
An inspection test is to recompute one entry by hand: L(1,1) equals 2 times 2 plus 0 times 0, which is 4, while L(1,2) equals 2 times 0 plus 0 times 1, which is 0. A common error is to calculate only G elementwise-squared and call it the row statistic. Elementwise squares may be useful for a diagonal optimizer, but they are not the full matrix multiplication defining L here.
Calculate an inverse fourth root without guessing
For a diagonal positive matrix, matrix inverse powers are simple. In the preceding toy example, L raised to minus one-quarter is diag(4 to the minus one-quarter, 1), which is diag(1 divided by square-root-of-two, 1). R has the same result. Multiply these on both sides of G, and the new direction P is diag(1,1), because the first gradient entry 2 receives two factors each equal to 1 divided by square-root-of-two.
The raw gradient would have had diagonal entries (2,1), while this one-step illustrative preconditioned direction has (1,1). That is a precise, hand-checkable example of coordinate rescaling, not evidence of better learning. If eta equals 0.01 parameter-units per gradient-unit, the corresponding update magnitudes are 0.01 and 0.01 parameter units for this simplified Shampoo direction. The raw update would have been 0.02 and 0.01.
For a non-diagonal matrix, do not take powers of the four displayed entries. The mathematical object must be decomposed or otherwise transformed using a numerically appropriate matrix-function algorithm. If L = Q diag(lambda_1,lambda_2) Q-transpose for an orthogonal eigenbasis Q, then L to the negative one-quarter is Q diag(lambda_1 to the negative one-quarter,lambda_2 to the negative one-quarter) Q-transpose. Near-zero eigenvalues demand regularisation and careful floating-point treatment.
Why the exponent is one quarter for a matrix
An AdaGrad-style inverse square root corrects an accumulated squared-gradient statistic. A matrix gradient has two separate mode factors, one on the left and one on the right. In the original Shampoo construction, each factor uses a fourth-root inverse so that their combined action balances the desired structured scaling. Multiplying two fourth-root factors in a compatible diagonal example yields a square-root-scale effect rather than applying a full inverse square root twice.
Do not generalise the one-quarter exponent blindly to every tensor rank. For a rank-k tensor, the paper’s structured construction applies a separate inverse power with exponent related to 1 divided by 2k to each mode. For rank two, 1 divided by 4 appears. For a rank-four tensor, the corresponding basic exponent is 1 divided by 8. Real implementations can override these exponents, reshape dimensions or use variants, so the exact configuration matters.
Nor should one mistake ‘fourth-root’ for a fourth-order derivative method. The power describes a matrix transformation, not a derivative of the loss. Shampoo is often discussed alongside second-order or full-matrix preconditioning because of its use of richer gradient statistics, but the simplified update does not evaluate the entire Hessian. The explanation of what information was measured must remain separate from the exponent used to turn it into a step.
Why damping and positive definiteness are mandatory
A preconditioner built from outer products is positive semidefinite, but it can have zero eigenvalues when a direction has not been excited by the gradients. An inverse negative power of zero is undefined or numerically explosive. Implementations therefore add a positive multiple of the identity or use a related damping strategy before taking inverse powers. The size and scaling of this adjustment can have a strong effect on updates.
Consider an artificial row statistic diag(0,16). Without damping, the zero direction cannot be inverted. Add an identity matrix multiplied by 1, obtaining diag(1,17); both eigenvalues are now positive. Its inverse fourth root has diagonal entries 1 and 17 to the negative one-quarter, approximately 1 and 0.4925. The damping parameter is one unit of the statistic’s squared-gradient scale in this example, not an arbitrary dimensionless token.
Changing the damping to 0.000001 would yield a much larger inverse factor in the previously zero direction; applying that factor to a noisy new gradient can be hazardous. Conversely too much damping makes the preconditioner close to a uniform rescaling, losing the geometry one hoped to exploit. Diagnose failures by logging smallest eigenvalues, condition indicators, chosen epsilon, data type and resulting step norms, not by changing the learning rate in isolation.
A checked memory example: 1,024 rows and 4,096 columns
Take a single weight matrix with 1,024 rows and 4,096 columns. It contains 1,024 multiplied by 4,096, or 4,194,304 parameters. A hypothetical dense matrix over every flattened parameter would have 4,194,304 rows and the same number of columns, or 17,592,186,044,416 entries. At four bytes per float32 entry, that is 70,368,744,177,664 bytes: exactly 64 binary tebibytes. That is not a plausible per-layer optimizer state in an ordinary accelerator training setup.
Shampoo’s unblocked row and column statistics for the same weight matrix have 1,024 squared plus 4,096 squared entries. Those are 1,048,576 plus 16,777,216, equalling 17,825,792 stored numbers. At four bytes each, the two statistics alone require 71,303,168 bytes, exactly 68 MiB. This excludes roots, momentum, weights, gradients, temporary workspaces and optimizer metadata, so it must not be reported as the total training-memory footprint.
A diagonal second-moment accumulator needs roughly one number per parameter, or 4,194,304 float32 values and 16 MiB for this one statistic. The comparison reveals the trade-off: two dense mode matrices are astronomically smaller than a full-vector dense matrix but heavier than a diagonal accumulator. The right denominator changes with the question. If asking ‘mode-state overhead versus a diagonal state’, use 68 divided by 16, about 4.25 times; if asking total GPU memory, include all other states and activations.
The price of inverse matrix roots
Storage is not the only cost. Eigen decompositions and general matrix root operations can scale roughly cubically with the matrix dimension when performed using dense methods. For a 4,096 by 4,096 matrix, 4,096 cubed equals 68,719,476,736 in a simple cubic-size scaling expression. That is not a measured flop count: actual solvers have different constant factors, precision requirements and algorithms. It explains why a large column-side factor can dominate implementation planning.
Computing new inverse roots at every minibatch might spend more time in optimizer algebra than the model saves in training steps. Computing them less frequently can reduce immediate cost but allow the geometry to lag behind recent gradients. Different workloads need different refresh strategies. A useful report separates time to update statistics, time to compute inverse roots, time to apply them and time spent in the model’s forward and backward passes.
The original paper’s headline experimental comparisons are evidence for its particular evaluated workloads, not a general theorem that Shampoo always beats Adam, SGD or every modern variant. Hardware parallelism, block sizes, memory bandwidth and fusion alter wall-clock outcomes. Always inspect a currently maintained implementation’s precision choices and default parameters before claiming its cost matches one simplified equation written in an article.
Blockwise Shampoo and the geometry it discards
A practical implementation may split a large tensor into blocks before building mode statistics. For example, partitioning a 1,024 by 4,096 parameter matrix into blocks of 128 by 128 yields 8 row-blocks times 32 column-blocks, or 256 blocks. Each block can maintain two 128 by 128 statistics. One block contains 16,384 entries in each statistic, so both together contain 32,768; multiplying by 256 yields 8,388,608 entries.
At four bytes per entry, this illustrative blockwise state uses 33,554,432 bytes, or 32 MiB, for those statistics. That is lower than the unblocked example’s 68 MiB, but it discards correlations between different blocks. The precise memory of an implementation may differ because it also stores inverse roots, grafting state, momentum, padding, layouts or buffers. A partition is a modelling choice as well as a resource choice.
The cubic scaling argument is similarly instructive. A 128 cubed dimension measure is 2,097,152; two roots per block across 256 blocks gives an aggregate scale of 1,073,741,824 such units, compared with the large unblocked factors. Do not call this an exact observed operation count or neglect distributed scheduling overhead. The trade-off is smaller local matrix problems at the price of a coarser view of global structure.
Go from matrices to genuine higher-rank tensors
A convolutional parameter array may have four axes for spatial dimensions and channel dimensions. Treating it as a matrix requires some reshape decision, while mode-wise Shampoo can preserve a higher-rank structure. Each mode statistic contracts the gradient over all of the other modes and leaves a square matrix for that one axis. This is a way to exploit the natural coordinates rather than take an outer product of the entire flattened array.
A miniature tensor with shape 2 by 3 by 4 has twenty-four parameters. Its three basic mode statistics have shapes 2 by 2, 3 by 3 and 4 by 4, requiring 4 plus 9 plus 16, or 29 entries. A full dense parameter-space statistic would have 24 squared, or 576 entries. This memory comparison ignores workspace and duplicated buffers; its purpose is to show why mode dimensions, rather than just total parameter count, govern the structured state.
Some distributed implementations combine small dimensions before building preconditioners to improve statistical information and hardware efficiency. For example, a convolution with dimensions 3 by 3 by 128 by 128 may be treated as 9 by 128 by 128 for certain calculations. That changes the mode structure and exponent in the corresponding approximation. State the exact reshaping and exponent whenever presenting a result as Shampoo rather than suggesting there is one immutable code path.
Cumulative statistics versus exponential moving averages
In a cumulative formulation the row-side statistic includes every observed gradient product so far, perhaps with a numerical stabiliser. Older examples therefore never disappear from the sum. In an exponential moving average, newer gradients receive more weight and old contributions decay. These approaches estimate different histories. They should not be silently mixed in a worked derivation or fair optimizer comparison.
Consider a scalarised teaching analogue where squared gradient evidence is 4 on one step and 16 on the next. The cumulative sum is 20 after two steps. An exponential moving average starting at zero with coefficient 0.9 has first value 0.4 and second value 0.9 times 0.4 plus 0.1 times 16, or 1.96. Their magnitudes are not directly comparable without accounting for update definition and bias correction; this example is not a claim about a specific Shampoo default.
In real training, changing from cumulative statistics to exponential weights alters responsiveness to distribution drift and can affect numerical scale. A long-memory estimate may be stable yet slow to reflect changed geometry, while a short-memory estimate can track new data but be noisy. Diagnose the statistical horizon together with the preconditioning interval and learning-rate schedule. Otherwise two experiments labelled ‘Shampoo’ might be solving different optimisation problems.
Grafting: borrow step magnitude without losing direction
A well-tuned first-order optimizer may already have an appropriate update magnitude for a layer. Some Shampoo implementations use ‘grafting’: construct the Shampoo direction but rescale it to match a direction norm produced by SGD or AdaGrad. This allows practitioners to separate the learned geometry or direction from part of the step-length calibration. It is not identical to simply averaging two optimizers’ updates.
Suppose the Shampoo direction for a layer is vector (3,4), with Euclidean norm 5, while the trusted reference direction from SGD has norm 2. To graft the norm, multiply the Shampoo direction by 2 divided by 5, producing (1.2,1.6) with norm 2. The example checks by squaring and adding: 1.2 squared plus 1.6 squared equals 4. The direction ratio 3 to 4 remains, while its length changes.
Grafting can be useful when the preconditioned step’s raw scale would otherwise require difficult retuning, but it also has limitations. A reference optimizer’s scale can itself be poorly calibrated; matching its norm does not ensure safe updates. Inspect the chosen norm, possible clipping, sparse gradients, decoupled weight decay and the learning-rate schedule. Do not report a grafted run as the unmodified algorithm from a paper without explaining the difference.
Momentum and weight decay are separate controls
Momentum smooths or combines gradient-derived update directions across training steps; it does not create the Shampoo mode matrices. A system may apply momentum before or after some preconditioning operations depending on its exact variant. Those ordering choices change the mathematics. Reusing familiar optimizer names without recording the actual update order makes an experiment hard to reproduce.
Weight decay acts as regularisation and can be coupled to the gradient or decoupled from the gradient-derived step. A decoupled decay applies a parameter shrinkage operation separately rather than letting that term distort the statistic in exactly the same way as task gradients. Its behavior needs to be specified, especially for biases, embedding tables and normalisation parameters that may receive different treatment.
Do not repair divergence by indiscriminately turning on every stabiliser. If a run becomes unstable, examine effective learning-rate scale, damping, the preconditioned direction norm, update clipping, mixed precision and momentum one at a time. A careful ablation can show which component actually fixed the issue. The broader lesson is that Shampoo is a geometry-bearing component inside an optimizer stack, not a substitute for disciplined training configuration.
Precision and conditioning: the numerical failure layer
Matrix inverse roots can be sensitive when eigenvalues span many orders of magnitude. A positive-definite statistic with one eigenvalue near zero and another very large may be mathematically invertible yet difficult to process accurately in low precision. The reciprocal power amplifies uncertainty along the small-eigenvalue directions. Damping, normalisation, solver tolerance and precision therefore belong in the core implementation, not in a footnote.
The Google Research scalable Shampoo implementation documents a matrix epsilon, data-type choices, matrix-inverse-root failure thresholds, block sizes and controls for how often statistics and roots are computed. Those are concrete signs that a theoretically elegant formula still needs numerical engineering. One project may prefer float32 roots with suitable stabilisation; another may use higher precision when hardware and resources make that feasible.
For a debugging exercise, imagine an inverse-root routine returns NaNs after a sudden spike in gradient norm. Do not infer that the loss landscape has become fundamentally unlearnable. Inspect the input statistic for non-finite values, the damping applied, whether the root algorithm converged, and whether automatic mixed-precision scaling altered gradient magnitudes. An immediate rollback to a previously stable checkpoint may be safer than allowing invalid optimizer state to propagate.
When should the preconditioner be refreshed?
A Shampoo iteration can update gradient statistics more frequently than it recomputes expensive inverse roots. This creates a two-clock system: the matrix remembers recent gradients, while the transformation actually applied to today’s step may correspond to an earlier snapshot. It is incorrect to describe a run as computing fresh matrix roots every step unless its configuration really does that.
Suppose a training job takes 1,000 steps and computes a new inverse root at steps 1, 101, 201 and so forth. The nominal interval is 100 steps, yielding about ten refresh events under this illustrative schedule. If each refresh takes 200 milliseconds, its isolated cost is about two seconds, excluding overlap, distributed communication and other optimizer costs. This arithmetic says nothing about whether the staler roots harm convergence.
The right refresh frequency depends on how quickly the gradient geometry changes and on hardware throughput. During an early unstable stage, slow refresh may miss meaningful shifts; later, an infrequently refreshed stable preconditioner may be sufficient. Benchmark both convergence per step and elapsed time to a fixed validation metric. Avoid comparing only the optimizer kernel’s microbenchmark while ignoring the time needed to reach the desired result.
Distributed Shampoo: ownership, communication and checkpoints
In distributed training, gradients, optimizer states and model parameters may be sharded across accelerator devices. Shampoo adds matrix states and root computations that must be assigned to devices and combined consistently with the chosen data-parallel or model-parallel layout. A method that fits on one accelerator can still fail in a distributed job because of communication patterns or uneven distribution of a few large matrices.
Google Research’s scalable Shampoo materials describe splitting large tensors, computing inverse pth roots on accelerators and distributing work. Those implementation decisions should be read as a specific engineering approach, not proof that any network topology will produce identical speed. If a 4,096-square column factor is assigned to one device while others process many small factors, load imbalance may erase some of the expected scaling gains.
A checkpoint must preserve enough optimizer state to resume the same algorithm. Restoring weights alone and reconstructing empty preconditioners changes the subsequent trajectory. Store or deliberately reset accumulated statistics, inverse-root refresh metadata, momentum and grafting state according to the documented recovery contract. Verify a restore on a short reproducible run before assuming resuming at step 50,001 is equivalent to uninterrupted training.
A fair comparison with AdamW, SGD and Adafactor
AdamW uses coordinatewise running moments and decoupled weight decay, often with mature, highly optimised implementations. SGD with momentum can be cheap in state and compute and may respond well to a carefully tuned learning-rate schedule. Adafactor reduces some optimizer-state memory by factoring statistics, particularly for large matrices. Shampoo is distinguished by dense mode-wise preconditioning and its associated matrix-function computation, not by the mere fact that it uses two arrays.
Compare training algorithms under a shared budget and objective. Fix the dataset splits, evaluation metric, model architecture, batch semantics and permitted tuning resources. Tune learning rates and stabilisers fairly instead of giving one optimizer a carefully selected configuration and the others defaults. If compute cost differs significantly, provide both steps-to-quality and total time-to-quality, along with peak memory and the relevant accelerator model.
A five-percent reduction in training steps does not imply a five-percent reduction in wall-clock training time. Suppose optimizer A needs 1,000 steps at 0.5 seconds each, or 500 seconds, while B needs 900 steps at 0.6 seconds each, or 540 seconds. B reduces steps by ten percent yet takes eight percent more time. The denominator for an efficiency claim must be explicit: steps, seconds, energy, hardware cost or achieved metric.
A fairness check also asks whether the model has enough training time to show adaptation benefits. Some algorithms look attractive at one learning-rate schedule or early checkpoint and worse at another. Generalisation, sensitivity to random seeds and final calibration may differ. An optimizer that makes the training loss fall faster is not necessarily the winner on untouched validation data or a downstream transfer task.
A failure-to-repair map for actual training runs
If the optimizer diverges immediately, first verify tensor dimensions and inverse-power implementation, then inspect finite gradients, damping and learning-rate magnitude. A row-side matrix intended to be m by m must never be silently broadcast as n by n. In a matrix example, check the shapes of both intermediate products and confirm that both statistics are symmetric before calling the inverse-root routine.
If a run is numerically stable but slower than AdamW, measure where time goes. Break out statistical updates, eigen or root computation, application of preconditioners, communication and model steps. Try smaller blocks or less frequent root refresh under a controlled experiment, recognising that either choice may change convergence. The goal is not to force a win, but to explain which cost or approximation is limiting the method.
If validation accuracy improves only on the source domain, examine data distribution, labels and evaluation conditions. Shampoo is an optimizer; it does not automatically fix dataset shift, leakage, a mislabeled benchmark or inappropriate objective. A better-conditioned training trajectory cannot replace a correct experiment design. Keep data curation and statistical reliability checks separate from the mathematics of a parameter update.
If memory spikes or a rank-deficient factor causes failure, identify the exact layer and matrix size. Some tensors are better handled by diagonal fallbacks, blocked factors or excluded parameter groups, especially when dimensions are enormous or extremely narrow. Document exclusions as part of the optimizer definition. It is misleading to say the whole network uses Shampoo if critical large layers were silently switched to a different algorithm.
Original miniature experiment: two steps, two observed directions
Begin with W a two-by-two weight matrix, and ignore regularisation and momentum solely to expose the algebra. At the first step the gradient is diagonal (2,1), producing L and R equal to diag(4,1). As already calculated, the simplified inverse-fourth-root preconditioned direction becomes diag(1,1). Apply eta = 0.01 to reduce W’s two diagonal coordinates by 0.01 each. The update amounts are parameter units, not training accuracy percentages.
Now let the second gradient be diag(1,2). Cumulative row and column statistics become diag(4+1,1+4), or diag(5,5). Their inverse fourth roots are both diagonal with entries 5 to the negative one-quarter. Applying one on each side of G multiplies its entries by 5 to the negative one-half, approximately 0.447214. Thus P is approximately diag(0.447214,0.894427). At eta = 0.01, the proposed update magnitudes are about 0.00447214 and 0.00894427 parameter units.
The arithmetic reveals why the new step did not simply reverse the previous step-size allocation. The statistics now contain both gradient observations, and each direction has accumulated the same squared-gradient total of 5. The second raw gradient has twice the first coordinate’s magnitude, so the second preconditioned entry is twice as large in this toy case. Neither step claims a loss decrease because no loss function or curvature is specified.
Independent check: compute the second-step L as G1 G1-transpose plus G2 G2-transpose. The diagonal entries are 4+1 and 1+4, and the off-diagonal entries remain zero. Confirm that both left and right inverse factors are 5 to the negative one-quarter. Multiplying these two factors yields 1 divided by square root of 5; applying that to (1,2) gives the stated values. The two updates are consequently compatible with the simplified formula.
Transfer problem: what if the matrix shape changes?
A second layer has shape 16 by 64, rather than 1,024 by 4,096. Calculate the number of parameters, the sizes of the mode statistics and the memory required for two float32 statistics. Then suppose you block this layer into eight-by-eight submatrices. Compare the number of blocks and state entries without claiming anything about empirical speed.
The first answer is 16 multiplied by 64, or 1,024 parameters. The unblocked row factor has 16 squared, or 256 entries; the column factor has 64 squared, or 4,096 entries. Together they have 4,352 entries, which occupy 17,408 bytes in float32, or 17 KiB. A single diagonal float32 accumulator over all parameters occupies 4,096 bytes, or 4 KiB. These counts exclude roots and other optimizer states.
Blocking into eight-by-eight submatrices creates 2 blocks along rows and 8 along columns, or 16 blocks. Each block has two eight-by-eight factors, together 128 entries. Sixteen blocks therefore require 2,048 entries, or 8,192 float32 bytes, exactly 8 KiB for the two statistics. Blocking saves state in this example but removes correlations spanning different blocks. The right decision remains workload-dependent and must be tested for accuracy and throughput.
Changed-context answer: if the original matrix were tall rather than wide, interchange the left and right shapes but do not alter the total number of parameters. If the tensor had three modes, replace the two-factor structure with three mode matrices and examine the rank-dependent exponent. The generalisable skill is to derive state size from the actual dimensions, not assume all optimizers cost the same number of bytes per parameter.
What Shampoo can and cannot establish
Shampoo is a structured gradient preconditioning method with theoretical results and empirical evaluations in specific settings. Those results do not prove global convergence for every deep neural network, every data distribution or every hyperparameter choice. A nonconvex loss can have saddle points, noisy gradients and multiple useful solutions. A superior training trajectory on one architecture is an experiment result rather than a law of nature.
Matrix preconditioning does not tell you whether a training dataset is legal to use, whether predictions are fair, whether output is truthful or whether an AI agent should be trusted with tools. Those require separate governance, evaluation and application controls. An optimizer can make model fitting more efficient and still amplify existing data problems if the task and observations are flawed.
Be explicit about versions. The 2018 Shampoo paper, later scalable and distributed implementations and more recent related optimizers such as SOAP do not share every detail. SOAP, for example, explores running Adam-style statistics in the preconditioner’s eigenbasis; that is not identical to the original cumulative two-sided update. Report exactly which algorithm, code version and features were used before attributing a result to ‘Shampoo’.
Frequently asked questions
Is Shampoo a shampoo-related computer analogy?
Only its name is playful: the 2018 authors describe it as pre-conditioning. In machine learning, conditioning here concerns scaling parameter updates using matrix statistics. It has no relationship to products that clean hair, and the term is not evidence that the optimizer physically washes data or removes bias.
Does Shampoo always use an inverse fourth root?
For the simplest rank-two matrix construction each of the left and right factors uses the inverse fourth root. Higher-rank tensors use rank-dependent exponents in the basic derivation, and implementations can override powers or reshape dimensions. Always read the version-specific configuration before reproducing the formula.
Is the preconditioner the Hessian?
No. The worked construction accumulates gradient-derived row and column outer-product statistics. That supplies richer directional information than a diagonal accumulator but is not an exact Hessian or its inverse. Some optimizer literature uses second-order language for matrix preconditioning; such labels do not change what the algorithm actually computes.
Why can matrix roots fail with finite input gradients?
Positive semidefinite statistics may have extremely small eigenvalues, and finite-precision matrix-function solvers can struggle with their inverse powers. Damping, normalisation, precision and convergence criteria matter. Failure requires investigating eigenvalues and numerical residuals rather than concluding that the neural-network model itself is impossible to train.
Does Shampoo reduce optimizer memory compared with Adam?
Not in general. Dense row and column statistics can use more memory than diagonal Adam moments, especially for very large dimensions. They remain far smaller than an impossible full dense parameter-space matrix. Blockwise and distributed variants change the cost, and total memory must include the rest of the optimizer and training system.
Why is training-step count an incomplete benchmark?
Each step may have a different cost. If one optimizer reduces the number of iterations but spends more time computing or communicating preconditioners, its total time can be worse. Compare wall-clock time to a fixed validation target, accelerator utilisation, peak memory and repeated-run variability with explicit denominators.
Can I use Shampoo without changing my model architecture?
It generally acts through optimizer state and parameter updates, so many model architectures can be trained with an appropriate implementation. However, layer types, tensor dimensions, sparse gradients, distributed layouts and mixed precision may require exclusions, partitioning or special handling. Read the exact library contract before enabling it across a production system.
How do I know whether a published comparison applies to my workload?
Check its model, dataset, batch size, precision, hardware, training budget, optimizer version and tuning protocol. Reproduce a small controlled experiment on your own workload using common validation criteria. Published examples motivate testing; they do not transfer automatically to every new architecture or data regime.
Is SOAP the same as original Shampoo?
No. SOAP is a later related optimisation method that explores an Adam-style update in an eigenbasis informed by Shampoo preconditioning. Its additional moving moments and update choices distinguish it from the simplified original two-sided matrix transformation. Use the method’s own paper and software documentation before claiming equivalence.
What should a student calculate before running code?
Derive the row and column statistics for one tiny matrix, check their shapes and symmetry, add damping if needed and transform the gradient using the relevant inverse powers. Then calculate bytes for actual tensor dimensions. Those checks catch misunderstandings that an apparently working optimizer library can otherwise conceal.
Primary research and implementation documents
Gupta, Koren and Singer (2018), Shampoo: Preconditioned Stochastic Tensor Optimization, Proceedings of Machine Learning Research, Volume 80, is the original peer-reviewed account. It specifies the tensor-structured preconditioning idea and analyses the algorithm in a stochastic convex setting. It is the authoritative starting point for the matrix derivation, not a claim that every modern variant uses identical defaults.
The Google Research publication record for Shampoo identifies the original researchers and describes the method’s motivation. For operational implementations, read the Google Research scalable Shampoo documentation and its published PyTorch implementation. These primary developer files discuss block size, inverse-root precision, grafting and refresh decisions.
For an explicitly different later design, see Vyas and colleagues (2024), SOAP: Improving and Stabilizing Shampoo using Adam. It is a research preprint describing an Adam-in-eigenbasis variant and should not be silently substituted for the original Shampoo update. Original numerical calculations in this article are teaching exercises rather than statistics or experimentally measured speed-ups from the papers.
Verified nearby eduKate reading routes
For the broader learning-mechanism map, return to How AI Works. For a quite different optimisation problem involving opposing task losses and gradient projection, see PCGrad Gradient Surgery. Shampoo changes gradient geometry; PCGrad changes the interactions among task gradients. Their mechanisms and validation questions must not be conflated.
For a related but distinct kind of model-side representation learning, the Barlow Twins explanation concerns cross-correlation in a training objective rather than an optimizer’s accumulated row and column preconditioners. Recognising the difference between a loss design and an update rule is an important transfer skill.