Source-linked AI summary
Shampoo: Preconditioned Stochastic Tensor Optimization
Vineet Gupta, Tomer Koren, Yoram Singer
TL;DR
Full-matrix preconditioning is powerful but requires prohibitively large matrices for large-scale machine learning. Shampoo exploits tensor structure with separate per-dimension preconditioners, and the paper analyzes its convergence in stochastic convex optimization. Experiments report faster convergence than commonly used optimizers, while per-step runtime remains comparable to SGD, AdaGrad, and Adam.
Problem
Full-matrix preconditioning can improve convergence but is impractical for large machine-learning problems because its memory and computation requirements are prohibitive.
Method
Shampoo retains tensor structure and maintains separate full preconditioner matrices for each tensor dimension, updated online from accumulated-gradient second-order statistics.
Results
Experiments report that Shampoo converges considerably faster than commonly used optimizers, while its per-step runtime is comparable to SGD, AdaGrad, and Adam.
Takeaways & Limitations
Shampoo provides a structure-aware preconditioning approach whose moderate-sized matrices make preconditioned optimization practical for large-scale learning problems.
Takeaways & Limitations
The analysis may require projecting iterates onto a Frobenius-norm-bounded set, but this projection is computationally expensive and is rarely performed in practice.
Abstract
from arXiv · showhide
Preconditioned gradient methods are among the most general and powerful tools in optimization. However, preconditioning requires storing and manipulating prohibitively large matrices. We describe and analyze a new structure-aware preconditioning algorithm, called Shampoo, for stochastic optimization over tensor spaces. Shampoo maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions. We establish convergence guarantees in the stochastic convex setting, the proof of which builds upon matrix trace inequalities. Our experiments with state-of-the-art deep learning models show that Shampoo is capable of converging considerably faster than commonly used optimizers. Although it involves a more complex update rule, Shampoo's runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.
1 Introduction
Shampoo addresses the scaling limits of full-matrix preconditioning by exploiting tensor structure, maintaining separate moderate-sized preconditioners for each tensor dimension. The paper develops the method, its structured matrix interpretation, and its practical implementation and evaluation context.
- Motivation: Preconditioned methods can improve convergence, but full-matrix preconditioning is impractical for large machine-learning problems because of memory and computation costs.Diagonal approximations scale favorably, while sketched and estimated alternatives often require super-linear resources.
- Shampoo: Shampoo retains tensor structure and maintains a separate preconditioner matrix for each gradient dimension.Each preconditioner is full but moderately sized and updated online using second-order statistics of accumulated gradients, similarly to AdaGrad.
- Analysis: The paper analyzes Shampoo in stochastic and online convex optimization, using matrix trace and geometric-mean inequalities to establish convergence guarantees.The tensor extension requires additional matrix-analysis arguments.
- Matrix case: For matrix parameters, Shampoo maintains left and right matrices containing accumulated-gradient second-moment information instead of one mn × mn full preconditioner.The left and right matrices multiply the gradient from the corresponding sides.
- Matrix case: In the matrix case, Shampoo requires m2 + n2 storage and O(m3 + n3) preconditioner construction, versus m2n2 storage and O(m3n3) for full-matrix methods.These savings arise because the algorithm manipulates separate matrices for the two parameter dimensions.
- Structured interpretation: Shampoo implicitly applies a structured preconditioner equivalent to a Kronecker-product preconditioner after flattening the gradient, without forming the full matrix.This structure is connected to the full unstructured preconditioner used by methods such as AdaGrad.
2 Background and technical tools
The paper introduces online convex optimization, adaptive mirror descent, vectorization, Kronecker-product identities, and matrix inequalities as tools for analyzing Shampoo. These tools connect adaptive preconditioned updates to structured tensor operations and support the convergence analysis.
- Online convex optimization: Online convex optimization generalizes stochastic convex optimization by comparing cumulative loss against the best fixed point in a convex domain.Regret-minimizing algorithms can be converted to stochastic optimization algorithms through an online-to-batch conversion.
- Adaptive regularization: Adaptive online mirror descent receives a loss, computes its gradient, and updates using a time-dependent positive definite matrix H_t.When the domain is R^d, the update is equivalent to a preconditioned gradient step.
- Adaptive regularization: For constrained domains, the adaptive mirror-descent update can be expressed as a projection onto the convex set using the norm induced by H.The accompanying regret lemma bounds performance for arbitrary positive definite matrix sequences.
- Kronecker products: Vectorization flattens an m × n matrix into an mn-dimensional column vector, enabling matrix operations to be related to vector operations.The paper uses this representation throughout its Kronecker-product analysis.
- Matrix inequalities: The matrix-inequality analysis uses operator monotonicity of geometric means and the map x ↦ x^α for α ∈ [0,1].The geometric-mean result is stated first for commuting positive semidefinite matrices, while a stronger noncommuting result is noted but not required.
3 Analysis of Shampoo for matrices
The matrix analysis shows that Shampoo implicitly approximates full-matrix preconditioning while retaining tractable structure, and obtains sublinear regret under rank and convexity conditions.
- Regret guarantee: The matrix regret theorem assumes gradients of rank at most r and bounds regret against any comparator W* in R^(m×n).The bound also depends on the distance parameter D and trace terms involving the preconditioners.
- Regret guarantee: The unprojected algorithm has a looser bound because D may grow with T, although projection can address this issue.Projection in the preconditioner-induced norm is computationally expensive in large-scale problems and is rarely used in practice.
- Structured preconditioning: The Kronecker product of Shampoo’s preconditioners is lower bounded by a full mn × mn preconditioning matrix.This comparison supports Shampoo’s implicit approximation of full-matrix preconditioning.
- Structured preconditioning: Small eigenvalues, which are important for effective preconditioning, do not vanish under Shampoo’s implicit approximation.The result follows from the matrix inequality used in the analysis.
4 Shampoo for tensors
For tensors, Shampoo preserves the gradient’s multidimensional structure by maintaining one moderate-sized full preconditioner per mode and applying them through tensor contractions.
- Tensor analysis: The general tensor algorithm is analyzed over convex functions on the vector space of order-k tensors.The analysis is presented in the online convex optimization framework.
- Tensor algorithm: Shampoo maintains a separate ni × ni preconditioning matrix for each tensor dimension.Each matrix is updated online from accumulated gradient second-order statistics, similarly to AdaGrad.
- Tensor algorithm: On each step, Shampoo multiplies the gradient along every mode by the corresponding preconditioner raised to the power −1/(2k), then takes a gradient step.The tensor-matrix products may be applied in any order because products along distinct modes commute.
- Tensor implementation: Tensor contractions and tensor-matrix products can be implemented with standard scientific-computing operations supported by modern machine-learning frameworks.The implementation uses tensor contraction for the operators Api and M ×i A.
- Tensor analysis: Theorem 10 assumes rank bounds ri for every gradient matricization and defines a combined rank parameter r from their product.Under these assumptions, the theorem gives a regret bound against any tensor comparator W*.
- Tensor analysis: Under standard assumptions, each trace term in the tensor regret bound scales as O(T^1/(2k)), and their product determines the overall bound.The supplied passage states the trace-term scaling and its role in the overall regret guarantee.
5 Implementation details
The TensorFlow implementation applies Shampoo independently to each model tensor and selectively falls back to diagonal preconditioning when a mode’s matrix is too costly to store or compute.
- Implementation: The implementation uses TensorFlow tensordot for tensor contractions and tensor-matrix products, and SVDs to compute matrix powers.These operations are supported directly in TensorFlow.
- Implementation: Each tensor in the model is optimized independently, making the optimizer architecture-oblivious and equivalent to a block-diagonal preconditioner.The blocks correspond to the model’s different tensors.
- Implementation: Shampoo automatically uses a diagonal variant for any tensor dimension whose preconditioner is too large to store or whose SVD is too expensive.Other dimensions of the same tensor can continue using non-diagonal preconditioning when feasible.
6 Experimental results
Shampoo was evaluated on image-classification and language-modeling benchmarks using existing training code, with optimizer-only substitutions and mini-batches of 128. Despite greater computation per step, its practical runtime was usually only slightly slower than SGD, AdaGrad, and Adam.
- Experimental setup: Experiments covered image classification on CIFAR-10/100 and statistical language modeling on LM1B using standard deep neural-network models.The experiments replaced only the TensorFlow optimizer in existing training code.
- Experimental setup: All experiments used mini-batches of 128, updating Shampoo's preconditioners once per batch from the averaged gradient.The averaged gradient, rather than individual-example gradients, was used for each preconditioner update.
- Implementation adjustments: Shampoo used delayed preconditioner updates every 20–100 steps and momentum with α = 0.9 to improve amortized runtime and convergence.The delayed updates had almost no impact on accuracy, while momentum slightly improved convergence.
- Runtime: Shampoo's runtime per step was typically only slightly slower than SGD, AdaGrad, and Adam, and was faster in some ResNet-55 cases.Runtime was measured as average steps per second on a Tesla K40 GPU with batch size 128.
- Benchmarks: The benchmarks included a 32-layer residual network and a 20-layer inception network on CIFAR-10, a 55-layer residual network on CIFAR-100, and an Attention model on LM1B.The CIFAR-10 models had 2.4M and 1.65M trainable parameters; the CIFAR-100 and LM1B models had 13.5M and 9.8M, respectively.
A Diagonal Shampoo
The diagonal Shampoo variant replaces each full preconditioning matrix with a diagonal one when dimensions are too large for full storage or matrix-power computation. It stores only diagonal elements and admits a rank-dependent regret bound.
- Diagonal variant: The diagonal variant uses diagonal preconditioners for dimensions that are too large to store or manipulate as full matrices.Other dimensions can independently retain non-diagonal preconditioning.
- Updates: Its preconditioners update by accumulating diagonal second-moment terms from GtGT and GTGt.The matrix-case updates replace full preconditioner updates with diagonal updates.
- Complexity: The diagonal implementation stores O(m + n) numbers and each matrix-case update can run in O(mn) time.This is linear in the number of matrix parameters.
- Guarantee: Theorem 13 bounds the regret of diagonal Shampoo for gradients whose matrices have rank at most r.The bound is stated relative to any comparator W* in R^{m×n}.
- Proof: The proof uses diagonal positive-semidefinite inequalities and order-preserving properties of diag(·), together with the preceding theorem's proof strategy.The relevant matrices are positive semidefinite and commute, enabling the matrix trace inequality step.
B Tensor case: Technical proofs
The tensor-case proofs reduce tensor operations to matrix and vectorized representations, establishing identities for rank-one tensors and extending them to arbitrary tensors by linearity.
- Tensor identities: Lemma 14 relates tensor vectorization to Kronecker products of its factor vectors.The identity is first stated for rank-one tensors and then extended through linearity.
- Tensor identities: The matricization of a rank-one tensor along dimension i isolates the corresponding factor while combining the remaining factors through a Kronecker product.This provides the matrix representation used in the tensor analysis.
- Tensor identities: Mode multiplication transforms the corresponding factor of a rank-one tensor by M, and its matricization is multiplied by M.The property extends to arbitrary tensors because the operators are linear.
B.1 Proof of Lemma 11
The proof of Lemma 11 compares tensor and matricized quadratic forms, then aggregates dimension-wise inequalities and applies a matrix inequality to obtain the result.
- Technical lemma: Lemma 15 supplies the technical inequality connecting a tensor G, an order-k tensor X, and an ni × ni matrix B.The proof begins by expressing X in the standard tensor-product basis.
- Proof of Lemma 11: The proof shows that the two compared left-hand sides are equal because both equal the squared dot product of G and X.The dot product is represented through vec(X).
- Proof of Lemma 11: Applying Lemma 9 to mat_i(Gt) yields a dimension-wise matrix inequality involving G^(i) and an identity factor.The proof then sums these inequalities over t and adds εI_n.
- Proof of Lemma 11: The resulting right-hand-side matrices are positive semidefinite and mutually commuting, allowing Lemma 5 to complete the argument.This commutativity is the condition used for the final matrix inequality step.
B.2 Proof of Lemma 12
The proof proceeds by induction on the tensor order k, using tensor slices, vectorization, and mode-wise matrix products to establish the identity.
- Induction: The argument uses induction on k, with the base case k = 2 established previously in Lemma 4.The induction step introduces H and applies the induction hypothesis to mode-wise transformed tensors.
- Conclusion: The induction step combines the transformed slices and their stacked representations to obtain the desired matrix identity.The relation vec(G) = (H ⊗ Mk) connects the tensor and matrix formulations.
- Slice decomposition: The k-th-dimensional slices of G are treated as order-(k − 1) tensors whose transformed forms are analyzed row by row.Each slice contributes a row of mat_k(G), and the resulting vectors are stacked into a matrix.
- Tensor identities: Vectorization identities relate mode-wise tensor contractions to Kronecker products and matrix operations.The proof substitutes G1 = G ×1 M1 ×2 M2 ··· ×k−1 Mk−1 and uses the stated vectorization relation.
C.1 Proof of Lemma 2
The proof invokes the Follow-the-Leader / Be-the-Leader lemma and proceeds by induction on T. It first verifies the relevant identity for rank-one matrices, then extends it by linearity.
- FTL-BTL argument: The proof is framed as an instance of the Follow-the-Leader / Be-the-Leader lemma of Kalai and Vempala.The target inequality is rewritten before the induction argument begins.
- Induction: Induction on T starts from the trivial base case T = 0 and uses that H_T−1 minimizes the relevant objective.The inductive step derives the second inequality from the minimizer property and completes the result by adding the same squared term to both sides.
- Rank-one case: For rank-one G = uv^T, vectorization and Kronecker-product identities reduce both sides to (Lu) ⊗ (R^Tv).The proof applies identities from Lemma 3 to establish equality for rank-one matrices.
- General case: The identity extends from rank-one matrices to arbitrary matrices because every matrix is a sum of rank-one matrices and the operators are linear.This linearity argument completes the general case.