PaperlayerView exact revision on arXiv

Source-linked AI summary

Shampoo: Preconditioned Stochastic Tensor Optimization

Vineet Gupta, Tomer Koren, Yoram Singer

arXiv:1802.09568v2cs.LGmath.OCstat.ML

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.

    Source: §1 · p. 1

  • Method

    Shampoo retains tensor structure and maintains separate full preconditioner matrices for each tensor dimension, updated online from accumulated-gradient second-order statistics.

    Source: §1 · p. 2

  • Results

    Experiments report that Shampoo converges considerably faster than commonly used optimizers, while its per-step runtime is comparable to SGD, AdaGrad, and Adam.

    Source: §6 · p. 13

  • Takeaways & Limitations

    Shampoo provides a structure-aware preconditioning approach whose moderate-sized matrices make preconditioned optimization practical for large-scale learning problems.

    Source: §1 · p. 2, §1.1 · p. 3

    The analysis may require projecting iterates onto a Frobenius-norm-bounded set, but this projection is computationally expensive and is rarely performed in practice.

    Source: §3 · p. 7

Abstract

from arXiv · show

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.Source: §1 · p. 1

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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §1.1 · p. 3
  • 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.Source: §1.1 · p. 3

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.Source: §1.1 · p. 3

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.Source: §2.1 · p. 4

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.Source: §2.2 · pp. 4–5
  • 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.Source: §2.2 · p. 5

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.Source: §2.3 · p. 5

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.Source: §2.4 · p. 6

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.Source: §3 · pp. 6, 9
  • 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.Source: §3 · p. 7

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.Source: §3 · p. 7
  • 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.Source: §3 · p. 7

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.Source: §4.2 · p. 11, §4.3 · pp. 11–12

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.Source: §4.2 · p. 11
  • 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.Source: §4.1 · p. 10, §4.2 · p. 11

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.Source: §4.2 · p. 11

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*.Source: §4.2 · p. 11, §4.3 · p. 12
  • 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.Source: §4.2 · p. 11

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.Source: §5 · p. 12
  • 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.Source: §5 · p. 12
  • 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.Source: §5 · p. 12

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.Source: §6 · pp. 12–13
  • 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.Source: §6 · p. 13

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.Source: §6 · p. 13

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.Source: §6 · p. 13, Table 1 · p. 13

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.Source: Figure 2 · p. 13, §6.1 · p. 14, §6.2 · p. 14

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.Source: A Diagonal Shampoo · p. 16

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.Source: A Diagonal Shampoo · pp. 16–17

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.Source: A Diagonal Shampoo · p. 16

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}.Source: A Diagonal Shampoo · p. 17

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.Source: A Diagonal Shampoo · p. 17

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.Source: B Tensor case: Technical proofs · pp. 17–18
  • 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.Source: B Tensor case: Technical proofs · pp. 17–18
  • 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.Source: B Tensor case: Technical proofs · p. 18

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.Source: §B.1 · pp. 18–19

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).Source: §B.1 · p. 19
  • 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.Source: §B.1 · p. 20
  • 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.Source: §B.1 · p. 20

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.Source: §B.2 · p. 20

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.Source: §B.2 · p. 20

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.Source: §B.2 · p. 20

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.Source: §B.2 · p. 20

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.Source: §C.1 · p. 21

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.Source: §C.1 · p. 21

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.Source: §C.2 · p. 21

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.Source: §C.2 · p. 21

Related papers

Loading 1802.09568v2…