Scalable Second Order Optimization for Deep Learning

Rohan Anil; Vineet Gupta; Tomer Koren; Kevin Regan; Yoram Singer

Source-linked AI summary

Scalable Second Order Optimization for Deep Learning

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, Yoram Singer

arXiv:2002.09018v2Feb 20, 2020cs.LGmath.OCstat.ML

TL;DR

Second-order methods can converge strongly but are difficult to use in large-scale deep learning because of computation, memory, communication, and infrastructure costs. The paper develops a scalable Shampoo implementation with structured large-layer preconditioning and CPU-accelerator pipelining, achieving faster convergence and often lower wall-clock time across major benchmarks.

Problem
Second-order optimization methods remain uncommon in large-scale deep learning because their computation, memory, and communication costs are prohibitive.
Method
The paper implements a scalable second-order preconditioned method by extending Shampoo for large layers, delaying preconditioner computation, and exploiting heterogeneous CPU-accelerator hardware.
Results
The implementation improves convergence and often wall-clock time across machine translation, BERT, Criteo CTR prediction, and ResNet-50 ImageNet benchmarks.
Takeaways & Limitations
Second-order methods can be competitive with specialized state-of-the-art optimizers on very large deep-learning tasks when their algorithmic and systems costs are addressed.
Takeaways & Limitations
The implementation still lacks some update sharding support and may be less practical in data-limited regimes where accelerator-to-CPU transfer latency cannot be amortized.

Abstract

from arXiv · show

Optimization in machine learning, both theoretical and applied, is presently dominated by first-order gradient methods such as stochastic gradient descent. Second-order optimization methods, that involve second derivatives and/or second order statistics of the data, are far less prevalent despite strong theoretical properties, due to their prohibitive computation, memory and communication costs. In an attempt to bridge this gap between theoretical and practical optimization, we present a scalable implementation of a second-order preconditioned method (concretely, a variant of full-matrix Adagrad), that along with several critical algorithmic and numerical improvements, provides significant convergence and wall-clock time improvements compared to conventional first-order methods on state-of-the-art deep models. Our novel design effectively utilizes the prevalent heterogeneous hardware architecture for training deep models, consisting of a multicore CPU coupled with multiple accelerator units. We demonstrate superior performance compared to state-of-the-art on very large learning tasks such as machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.

1 Introduction

Second-order methods offer strong convergence but remain difficult to deploy at deep-learning scale because of computational, memory, numerical, and infrastructure costs. This paper develops a scalable Shampoo implementation and reports faster convergence and wall-clock training across large models and domains.

1 Introduction: The design combines a pipelined CPU-accelerator implementation with extensions for large layers and an iterative replacement for expensive spectral decompositions.The paper also documents practical implementation challenges and limitations relevant to future accelerator hardware.
1 Introduction: Across machine translation, language modeling, CTR prediction, and image classification, the distributed implementation improves convergence steps and often wall-clock time on extremely large tasks.The reported evaluations span Transformers, BERT, Criteo CTR prediction, and ResNet-50 ImageNet.
1 Introduction: Transformer training used half as many steps as well-tuned Adam, reducing wall time by 45% for Transformer and 37% for Transformer-Big.Training times fell from approximately 12 to 6.7 hours and from 47 to 29.5 hours, respectively.
1 Introduction: BERT training used 16% fewer steps with higher masked-LM accuracy and a 4% wall-time reduction, although the system was not yet performance-tuned for this task.Wall time decreased from 3.8 to 3.65 hours at batch size 32K.
1 Introduction: CTR prediction reached 80.56% AUC, used half as many steps as the current state-of-the-art optimizer, and reduced wall time by 37.5%.The reduction was approximately 13 minutes to 8.2 minutes; a 0.1% improvement is considered significant for this task [Rong et al., 2020, Wang et al., 2017].
1 Introduction: ResNet-50 reached the 75.9% MLPerf target in 1729 steps, 31.7% fewer than the previous state-of-the-art, with a 13% wall-clock reduction.The comparison used 32K batch size on ImageNet and emulated higher precision [Henry et al., 2019].

2 Preliminaries

Full-matrix preconditioning can exploit parameter correlations but is impractical for large tensor-shaped models because its storage and update costs scale cubically or quadratically with flattened dimensions. Shampoo approximates this preconditioner with structured statistics for matrix and higher-order parameters.

2 Preliminaries: The notation defines elementwise products and powers, Loewner order, PSD matrix powers, vectorization, and Kronecker products used to express Shampoo’s structured preconditioning.For example, (A⊗B) vec(C)=vec(ACB^T).
2 Preliminaries: Preconditioned methods update parameters by multiplying an accumulated gradient estimate by a matrix, while adaptive preconditioning derives that matrix from gradient-gradient correlations.Newton-type methods instead relate the preconditioner to the Hessian.
2 Preliminaries: For W∈R^(m×n), full-matrix Adagrad requires m^2n^2 storage and m^3n^3 update time, motivating Shampoo’s Kronecker-product approximation.Large models can have m and n as large as 10^4, whereas AdaGrad and Adam use diagonal preconditioners.
2 Preliminaries: The formulation applies within online convex optimization, where predictions precede outcome revelation and loss gradients G_t are computed for matrix-shaped parameters W.The gradient has the same m×n shape as W.
2 Preliminaries: Shampoo tracks left and right statistics L_t and R_t for matrix-shaped gradients and approximates the full Adagrad preconditioner with (L_t⊗R_t)^1/2.The resulting parameter update uses the corresponding matrix factors to precondition the gradient.

3 Scaling-up Second Order Optimization

The paper scales Shampoo by addressing its computational, numerical, and infrastructure challenges, including large layers, delayed preconditioning, block structure, and inverse-root computation. These changes make second-order optimization more practical for large deep-learning models.

3 Scaling-up Second Order Optimization: Inverse p-th roots are difficult because they can dominate step time by up to 100× and may require expensive double-precision computation for ill-conditioned matrices.The implementation considers iterative methods as accelerator-friendly alternatives to SVD, but real workloads still impose numerical-cost challenges.
3 Scaling-up Second Order Optimization: Large embedding layers make preconditioning infeasible at O(d^2) memory and O(d^3) computation, motivating Shampoo extensions for these architectures.The paper targets settings such as Criteo-1Tb embeddings with approximately 186 million hash buckets and Transformer layers with up to 65,536 units per dimension.
3 Scaling-up Second Order Optimization: Shampoo can precondition only one dimension of very large embedding or softmax layers, reducing cost while retaining an empirically observed benefit.The choice is supported by experiments showing improved behavior with minimal time increases, and by a convergence result for the approximation.
3 Scaling-up Second Order Optimization: Partitioning large tensors into blocks reduces statistics and preconditioned-gradient costs, while Shampoo still converges in the convex setting.Experiments report minimal solution-quality impact alongside faster steps and lower memory overhead.
3 Scaling-up Second Order Optimization: Preconditioner computation can be delayed by hundreds of steps without significant accuracy loss, creating a performance-quality trade-off governed by computation frequency.The experiments choose the smallest frequency that does not degrade performance, while noting that better hardware or software would permit more frequent updates.
3 Scaling-up Second Order Optimization: The implementation must accommodate heterogeneous hardware and inflexible training APIs, requiring CPU use for double precision and framework-level changes for pipelining.The training loop gathers statistics, distributes computation across CPUs, and updates preconditioners without blocking accelerator training.

4 Distributed System Design

The distributed design pipelines Shampoo across accelerator and CPU resources, using accelerators for model computations and CPUs for expensive preconditioner work. This organization makes the dominant preconditioner computation add almost no overall training time.

4 Distributed System Design: Distributed Shampoo assigns accelerator cores the forward, backward, and synchronized mini-batch-gradient computations under standard data parallelism.Parameters are replicated across accelerator cores, each processes a sub-batch, and all-reduction synchronizes the gradients.
4 Distributed System Design: All-reduction and weight-update overheads are minor, together accounting for less than 5% of Transformer step time.All-reduction still introduces a synchronization barrier because accelerator cores must coordinate to compute the mini-batch gradient.
4 Distributed System Design: Preconditioner computation is distributed across attached CPUs, which provide double precision while accelerators continue training.The design computes inverse p-th roots asynchronously and exploits otherwise underutilized CPU resources.
4 Distributed System Design: The most expensive Shampoo step adds almost no overall training time because preconditioners are distributed across the training system’s CPUs.Preconditioned-gradient overhead is independent of batch size, so increasing batch size linearly reduces its relative overhead.

5 Experiments

Experiments evaluate Shampoo against first-order and second-order baselines across translation, recommendation, language modeling, and image classification, with improvements in convergence and often wall-clock time.

5.2 Machine Translation with a Transformer: Embedding-layer preconditioning increased Transformer step time by only 6% while reducing convergence steps by approximately 20%.Partitioning fully connected layers into sub-blocks caused no quality loss and reduced runtime by less than 3%.
5.3 Transformer-Big model: Increasing batch size reduced optimizer overhead from 40% to 19%, while delayed preconditioner updates every few hundred steps had no significant accuracy effect.The overhead reduction reflects statistics and preconditioned-update costs that are independent of batch size.
5.4 Ads Click-Through Rate (CTR) prediction: Shampoo reached 80.25% AUC in 30.97K steps versus 64K for the baseline and achieved 80.56% AUC on the Criteo task.Preconditioning embedding layers reduced the steps needed to reach the target from 39.96K to 30.97K.
5.6 Image classification: Shampoo reached target accuracy in 1729 steps on ResNet-50, compared with 2512 steps for first-order methods.The target was 75.9% accuracy.

6 Concluding Remarks

The paper presents scalable second-order optimization as effective across large learning tasks while identifying hardware and software limitations that constrain broader deployment.

6 Concluding Remarks: The implementation improved steps to convergence and often wall-clock time across large tasks in multiple domains, matching or exceeding specialized state-of-the-art optimizers.The authors frame the implementation as a basis for influencing accelerator and runtime design.
6 Concluding Remarks: Current systems lack support for symmetric operands, leaving potential reductions of up to approximately 50% in flops and storage unavailable.The limitation concerns second-order methods that use symmetric matrices.
6 Concluding Remarks: Weight-update sharding could reduce update time and memory for Shampoo, but it requires compiler-level support and currently forces every core to update all layers.The authors identify this as inefficient and not expressible at the program layer.
6 Concluding Remarks: The benefits of the approach may be harder to amortize in data-limited regimes because accelerator-to-CPU transfer latency is added to the computation.The authors suggest hardware support for higher precision, larger preconditioners, and efficient triangular-matrix handling as possible remedies.

A Deferred proofs

The deferred proofs establish regret bounds for Shampoo with extended exponents and blocking, then combine them to show convergence when both extensions are used.

Shampoo with extended exponents admits a regret bound under rank-at-most-r gradient matrices.
The proof uses monotonicity of the Kronecker-structured matrices H_t to preserve the ordering required by the regret analysis.
Shampoo with blocking also has a regret bound for gradients partitioned into m-dimensional blocks.
The two regret bounds combine to show that Shampoo with both extensions converges.

B Comparison with K-FAC

The comparison explains how K-FAC and Shampoo relate through structured preconditioners while highlighting implementation and numerical trade-offs that favor Shampoo's scalable design.

K-FAC formulation: K-FAC approximates the Fisher information matrix with Kronecker factors, relying on independence between activation inputs and output gradients.
Comparison: K-FAC and Shampoo use related preconditioners, but K-FAC uses exponent −1 and model-sampled gradients whereas Shampoo uses exponent −1/2p and mini-batch gradients.
Implementation limitations: K-FAC's network-structure dependence makes common operators such as batch normalization, weight normalization, and layer normalization difficult to support.
Extensions: Shampoo's eigenbasis correction can improve the approximation, but it requires SVDs and extra matrix multiplications, with no significant experimental improvement over standard Shampoo.
Embedding layers: For sparse embedding inputs, Shampoo reduces preconditioner computation from O(d^2N) to O(d^2m), avoiding accelerator densification of sparse operations.

E Implementation details of Shampoo

The implementation pipelines Shampoo across CPUs and accelerators, using asynchronous preconditioner computation and numerically careful inverse-root evaluation.

Pipelining: The accelerator fetch interval τ1 must leave enough time for the CPU to compute preconditioners asynchronously and pipeline them efficiently.
Inverse roots: Coupled Newton iterations use CPU-friendly matrix products and deliver wall-time improvements over SVD for inverse-pth-root computation.
Complexity: Table 2 summarizes computational and memory complexity across Shampoo variants.

F Experimental comparison with second order optimizers

The experiments compare Shampoo with other second-order optimizers on autoencoder tasks and describe the algorithmic configuration used for these evaluations.

Autoencoder comparison: Shampoo was evaluated against K-BFGS and K-FAC on MNIST, FACES, and CURVES autoencoder problems using tuned algorithm-specific settings.
Update variants: The standard Shampoo update uses a −1/4 exponent, while experiments also treat the exponent as a tunable parameter α ∈ [0, 1].
Algorithm configuration: The algorithm sketch includes learning rate and momentum parameters and accumulates preconditioner statistics with an exponential moving average.

G Further details on experiments

These experiments describe practical choices for stabilizing and applying Shampoo across architectures, including layer-wise learning rates, grafting, warmup, and CPU-based preconditioning. Results include faster CIFAR-10 training, while BERT-Large exposes implementation overhead that limits wall-clock gains.

Layer wise learning rates: Layer-wise learning rates address optimization instability caused by wide variation in preconditioner operator norms across layers.The statistics and preconditioners are amortized across multiple steps, so their norms do not grow at every step.
Grafting: Grafting uses Shampoo for update directions and a well-tuned optimizer such as diagonal AdaGrad for update magnitudes.This isolates the effect of preconditioned directions while bootstrapping a reasonable learning-rate schedule.
Learning-rate schedules: All optimizers use 40k-step warmup, with quadratic schedules for smaller Transformers and linear schedules for larger ones.Adam additionally uses a learning-rate decay schedule, while Shampoo uses per-layer rates derived from AdaGrad.
G.2 BERT-Large: 14% higher BERT-Large step time nearly offsets the 16% reduction in training steps.The implementation computes preconditioning statistics and gradients redundantly across TPU cores, although sharding and larger batches could reduce overhead.

G.4 Detailed results for experiments

This section decomposes each training step into distributed computation, adaptive-statistics updates, preconditioning, and parameter updates. Shampoo’s inverse preconditioner computation is pipelined on host CPUs rather than appearing in reported step times.

Experiment setup: Table 4 records the experiment configurations, including TPU-core counts and optimizer hyperparameters.The supplied table passage identifies the setup scope but does not provide individual parameter values.
Training-step phases: Each training step comprises forward computation, back-propagation, gradient averaging, statistics updates, preconditioning, and parameter updates.The phases separate model computation and communication from optimizer-specific operations.
Parameter updates: Parameter updates use the same rule across algorithms: W := W − η G̃, where G̃ is the preconditioned gradient.The optimizer differences arise in how the preconditioned gradient is formed.
Preconditioner computation: Shampoo computes L^-1/4 and R^-1/4 on the host CPU in a pipeline, so this preconditioner work does not appear in step-time measurements.The reported timing therefore excludes that pipelined computation from the displayed step phases.