Shampoo: Preconditioned Stochastic Tensor Optimization

By Vineet Gupta, Tomer Koren, Yoram Singer. 2018

492 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computational Mathematics, Computational Mechanics, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mathematics

In Computational Mechanics

In Optimization & Training

Citations in the corpus

Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges (2019, background): “…including meth-ods like gradient checkpointing and low precision training (Courbariaux et al., 2014; Ott et al., 2018), memory efﬁcient optimizers (Shazeer and Stern, 2018; Gupta et al., 2018) and frameworks supporting model parallelism (Shazeer et al., 2018; Harlap et al.; Huang et al., 2018).”
Memory-Efficient Adaptive Optimization for Large-Scale Learning (2019, methodology): “Also related to our work is the Shampoo algorithm for optimization over tensor structures [12].”
Pathways: Asynchronous Distributed Dataflow for ML (2022, background): “Concurrent (time-multiplexed or overlapping) ML task execution (Gupta et al., 2018; Xiao et al., 2020; Bai et al., 2020; Yu and Chowd-hury, 2020; Wang et al., 2021; Lim et al., 2021) helps harvest idle resources within accelerators.”
Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (2023, unknown): “In general, previous second-order optimizers do not achieve a speed-up on large language models in wall-clock time or total compute (Gupta et al., 2018; Yao et al., 2021) (see more discussions in Section C).”
A Generalizable Approach to Learning Optimizers (2021, background): “…that LHOPTs could greatly assist other subﬁelds of research: there have been many promising optimization primitives emerging that so far have been unsuccessful but perhaps just need the right hyperparameters to get working such as SVRG [5], hypergradients [27], and second-order methods [17, 4, 28].”
The Case for Full-Matrix Adaptive Regularization (2018, background): “[GKS18] propose a way to do AdaGrad with Kronecker products of full-matrix preconditioners, a more limited setting which requires knowledge of the model’s structure.”
Large-Scale Differentially Private BERT (2021, background): “An interesting direction is to improve the efficiency of DP training further by leveraging recent advances such as higher-order methods (Anil et al., 2020; Gupta et al., 2018), local loss optimization (Amid et al.”
On the Convex Behavior of Deep Neural Networks in Relation to the Layers' Width (2019, methodology): “tell us quite a bit about the behavior of the optimizer, and how we could modify our algorithms to make them converge faster and to better solutions (Gupta et al., 2018; Martens & Grosse, 2015).”
Practical tradeoffs between memory, compute, and performance in learned optimizers (2022, methodology): “…Newtons method has led to a number of approximations – e.g. diagonal approximations (Duchi et al., 2011), block diagonal (Martens and Grosse, 2015; Gupta et al., 2018), as well as a large family of quasi newton methods (Dennis and Moré, 1977) (e.g. BFGS (Broyden, 1970) and its low memory…”
Symbolic Discovery of Optimization Algorithms (2023, background): “Other related works include numerous handcrafted optimizers (Anil et al., 2020; Bernstein et al., 2018; Dozat, 2016; Duchi et al., 2011; Gupta et al., 2018; Kingma and Ba, 2014; Liu et al., 2020; Reddi et al., 2018; Shazeer and Stern, 2018; Zhuang et al., 2020), which we discuss in Section 3.”
Does Knowledge Distillation Really Work? (2021, unknown): “While we do not pursue Shampoo [16, 1] specifically, Beyer et al.”
Extreme Tensoring for Low-Memory Preconditioning (2019, methodology): “[GKS18] Let g 1 , g 2 , · · · , g T be a sequence of vectors, and let M t = diag( =1 g s g s ) .”
ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning (2020, methodology): “A naive use of Fisher information has computational and memory overhead, but it is possible to also approximate the Fisher information matrix using low rank decomposition [20].”