Home / Curated / Algorithms / Shampoo (Preconditioned Stochastic Tensor Optimization)

Shampoo (Preconditioned Stochastic Tensor Optimization)

Algorithms AIMech:ShampooPreconditionedStochasticTensorOptimization DRAFT

Aliases: Shampoo, Shampoo Optimizer, Preconditioned Stochastic Tensor Optimization

A structure-aware preconditioned stochastic optimization algorithm over order-k tensor parameter spaces R^{n_1 x ... x n_k} that maintains a separate positive-definite preconditioning matrix H_t^i in R^{n_i x n_i} for each tensor mode i in [k] by accumulating mode-i gradient contractions H_t^i = H_{t-1}^i + mat_i(G_t) mat_i(G_t)^T (with H_0^i = epsilon I_{n_i}) and updates parameters by multiplying the gradient tensor along every mode i by (H_t^i)^{-1/(2k)}, yielding W_{t+1} = W_t - eta L_t^{-1/4} G_t R_t^{-1/4} in the matrix case (k = 2).

Identity Notes & Concept Scope Distinct from diagonal AdaGrad and Adam in that Shampoo maintains full per-mode covariance matrices H_t^i rather than purely elementwise diagonal accumulators, and distinct from full-matrix AdaGrad in that it implicitly preconditions the vectorized gradient via the Kronecker product bigotimes_{i=1}^k (H_t^i)^{1/(2k)} with O(sum_i n_i^2) storage and O(sum_i n_i^3) inversion cost instead of O(prod_i n_i^2) and O(prod_i n_i^3). Distinct from K-FAC in that Shampoo is model-oblivious and relies solely on tensor-contracted outer products of observed gradients rather than architecture-specific activation/backpropagation Fisher factorizations.

Claims & Grounded Evidence

4 Claims

Mechanism #c-r588-kronecker-preconditioning-equivalence

By Lemma 4 and Lemma 12, preconditioning an order-k gradient tensor G_t along each mode i via the mode-i tensor-matrix product G_t x_1 (H_t^1)^{-1/(2k)} ... x_k (H_t^k)^{-1/(2k)} is algebraically equivalent on the vectorized parameter wt = vec(W_t) to a preconditioned gradient step wt+1 = wt - eta H_t^{-1} gt with Kronecker preconditioner H_t = bigotimes_{i=1}^k (H_t^i)^{1/(2k)}, reducing matrix-case memory from m^2 n^2 to m^2 + n^2 and SVD/inversion work from O(m^3 n^3) to O(m^3 + n^3).

supports Source: Shampoo: Preconditioned Stochastic Tensor Optimization · Locator: Section 1.1, p. 3; Lemma 4, p. 6; Section 4.2-4.3 & Lemma 12, pp. 11-12

Mechanism #c-r588-loewner-dominance-and-regret-bound

Under operator monotonicity of the geometric mean of commuting positive semidefinite matrices (Lemma 5, Lemma 8, Lemma 11), the Kronecker product of Shampoo's mode preconditioners upper-bounds the full-matrix AdaGrad covariance in Loewner order, epsilon I_n + sum_{t=1}^T g_t g_t^T <= r bigotimes_{i=1}^k (epsilon I_{n_i} + sum_{t=1}^T G_t^{(i)})^{1/k}, yielding an O(sqrt(T)) online convex optimization regret bound sum_{t=1}^T (f_t(W_t) - f_t(W^*)) <= sqrt(2r) D prod_{i=1}^k Tr((H_T^i)^{1/(2k)}) (Theorem 7 and Theorem 10).

supports Source: Shampoo: Preconditioned Stochastic Tensor Optimization · Locator: Lemma 5-6, p. 6; Theorem 7 & Lemma 8-9, pp. 6-8; Theorem 10 & Lemma 11, pp. 11-12

Application #c-r588-delayed-roots-and-diagonal-fallback

Recomputing the fractional matrix roots (H_t^i)^{-1/(2k)} via SVD once every 20 to 100 steps amortizes per-step runtime to be comparable to SGD, AdaGrad, and Adam on deep CNNs and Transformer attention models, while dimensions exceeding a size threshold (e.g. 1200) can independently fall back to diagonal preconditioning (Algorithm 3, Theorem 13).

supports Source: Shampoo: Preconditioned Stochastic Tensor Optimization · Locator: Section 5-6, pp. 12-14; Appendix A, Theorem 13, pp. 16-17

Limitation #c-r588-block-diagonal-and-large-dimension-limitation

Shampoo treats each parameter tensor independently as a block-diagonal preconditioner that ignores inter-tensor parameter correlations, and storing full n_i x n_i mode preconditioners or computing their SVDs becomes prohibitively expensive when individual tensor dimensions n_i are very large.

supports Source: Shampoo: Preconditioned Stochastic Tensor Optimization · Locator: Section 5, p. 12; Appendix A, p. 16

⚠️ Curated Limitations & Negative Results

Captures only intra-tensor parameter correlations via block-diagonal per-tensor preconditioning and ignores correlations across distinct parameter tensors.
Requires O(sum_i n_i^2) memory and O(sum_i n_i^3) spectral decomposition work per tensor, necessitating diagonal fallback or delayed root updates when mode dimensions are large.

Artificial Intelligence Ontology (AIO) Grounding

Mapped AIO Classes

AIO Target Class	Mapping Relation	Mapping Rationale
`AIO:ArtificialNeuralNetwork` Artificial Neural Network	related_match	Shampoo is a Kronecker-factored tensor preconditioning optimizer for training artificial neural networks.

Evidence Sources & Primary Literature

3 Sources

Source ID	Publication / Document	Type / Access	Reference & Checksum
`r588-primary-pdf`	Shampoo: Preconditioned Stochastic Tensor Optimization Vineet Gupta, Tomer Koren, Yoram Singer	primary paper full_text	ID: https://dblp.org/rec/conf/icml/0001KS18 sha256:dedd4fdb7f59…
`r588-arxiv-abs`	Shampoo: Preconditioned Stochastic Tensor Optimization (arXiv abstract) Vineet Gupta, Tomer Koren, Yoram Singer	primary paper abstract_only	ID: https://arxiv.org/abs/1802.09568 sha256:f6dba65683c6…
`r588-metadata-cr`	Shampoo: Preconditioned Stochastic Tensor Optimization (Crossref metadata) Vineet Gupta, Tomer Koren, Yoram Singer	primary paper abstract_only	ID: https://dblp.org/rec/conf/icml/0001KS18 sha256:3f2f5f53a0cd…

Ontology & Grounding Topology

Open in interactive network graph →

Curation Provenance

Curator: Antigravity

Method: ai_assisted

Model: Gemini

Created: 2026-10-19T11:11:00Z

Updated: 2026-10-19T11:11:00Z

Source File: data/curated/ShampooPreconditionedStochasticTensorOptimization.yaml · Schema: curated_concept.yaml

Explore across X-Mech

Related computational & biological Mechs: Embedding Map Network Graph ProteinTraitsMech (ESM-2 pLM) DUFMech (Domain pLM) TraitMech CultureMech All 14 X-Mech Repositories