Towards Practical Second Order Optimization for Deep Learning
Authors: Rohan Anil
Vineet Gupta
Tomer Koren
Kevin Regan
Yoram Singer
Venue:  (2021)
Keywords: large scale distributed deep learning
second order optimization
bert
resnet
criteo
transformer
machine translation

Abstract: 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.

Under review as a conference paper at ICLR 2021
TOWARDS PRACTICAL SECOND ORDER OPTIMIZATION FOR DEEP LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
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 attempt to bridge this gap between theoretical and practical optimization, we present scalable implementation of second-order preconditioned method (coneretely, 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 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-ard on very large learning tasks such machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.
INTRODUCTION
Second order methods are among the most powerful algorithms in mathematical optimization.
Algorithms in this family often use preconditioning matrix to transform the gradient before applying each step. Classically, the preconditioner the matrix of second-order derivatives (i.e., the Hessian) in the context of exact deterministic optimization (e.g., Fletcher, 2013; Lewis & Overton, 2013; Nocedal, 1980). While second-order methods often have significantly better convergence properties than first-order methods, the size of typical problems prohibits their use in practice, as they require quadratic storage and cubic computation time for each gradient update. Approximate algorithms such as quasi-Newton methods are aimed at significantly reducing these requirements; nonetheless, they still impose non-trivial memory costs equivalent to storing several copies of the model (and often quadratic computation, as in the popular two-loop recursion (Nocedal, 1980)), which severely limits their use at the immense scale of present-day deep learning.
Arguably, one of the greatest challenges of modern optimization is bridge this gap between theoretical and practical optimization towards making second-order methods feasible to implement and deploy at immense scale. Besides the compelling scientific and mathematical developments may stimulate, this challenge has also a clear real-world significance: recent practice of training deep learning models suggests that the utility of common first-order methods quickly reaching plateau, in large part because their time-per-stey is already negligible (compared to other parts of the computation) and cannot be optimized further; thus, the only way obtain faster training performance is by drastically reducing the number of update steps. To this end, utilizing second-order methods seem very natural and promising approach.
In this paper we attempt to narrow the gap between theory and practice of second-order methods, focusing on second-order adaptive methods for stochastic optimization. These methods can be thought of as full-matrix analogues of common adaptive algorithms such as AdaGrad (Duchi et al., 201 I; McMahan & Streeter, 2010) and Adam (Kingma & Ba, 2014): they precondition each gradient with second moment matrix, akin to covariance matrix, that accumulates the outer products of the stochastic gradients. Full-matrix versions are potentially more powerful than first-order methods as they can exploit statistical correlations between (gradients of) different parameters; geometrically,

---

Under review as a conference paper at ICLR 2021
they can scale and rotate gradients whereas first order methods only scale gradients. However they suffer from similar prohibitive runtime and memory costs as Hessian-based methods.
Recent developments in the space of second-order methods, on which focus on in this paper, include the K-FAC (Heskes, 2000; Martens & Grosse, 2015) and Shampoo (Gupta et al., 2018) algorithms that exploit the structure of deep networks (and more generally, models described by collection of tensors) for mitigating the space and runtime costs of full-matrix second-orde algorithms. These methods approximate each preconditioning matrix using factored representation that stems from the network structure. However, in very large applications, such algorithms are still impractical due to a number of numerical and infrastructural pitfalls and are difficult to parallelize.
Contributions. We provide solutions to practical concerns and challenges that arise in implementing and using second-order methods at large scale. Our focus will be on the Shampoo algorithm, but most of the challenges we address are relevant to the implementation of many other second-order methods. These include:
We design and implement an pipelined version of the optimization algorithm, critically exploiting the heterogeneity and computing power of CPU-Accelerator coupled architectures; We extend Shampoo in a number of ways so to make it applicable to larger range of deep architectures; in particular, the extensions allow Shampoo to be used for training very large layers such as embedding layers ubiquitous in language and translation models;
We replace expensive spectral decompositions (c.g.. SVD) used for manipulating preconditioners with efficient and numerically-stable iterative method for computing roots of PSD matrices; We describe practical challenges and limitations we faced in our design, which we argue could be useful for the design considerations of next-generation accelerator hardware architectures.
Our distributed implementation demonstrates significant improvements in performance, both in terms of number of steps, and often in actual wall-clock time, on some extremely large deep learning tasks:
Machine translation: we train Transformer models (Vaswani et al., 2017) on the WMT' 14 English to French translation task (Bojar et al., 2014) in half as many steps compared to state-of-the-art (well tuned Adam), resulting with to 45% reduction in wall-time.
Language modeling: we trained BERT (Devlin et al., 2018) in 16% fewer steps and achieve higher masked- LM accuracy compared to state-of-the fit-the-ar art optimizer (You et al., 2019) at 32K batch size; overall wall-time decreased by 4% from 3.8 3.65 hours. (For this task, our system has not yet been tuned for performance; we discuss several possible optimizations below.)
Click-Through Rate (CTR) prediction: we trained the DLRM model (Naumov et al., 2019) on the terabyte Criteo dataset (Criteo Labs, 2015) at 64K batch size in half many steps as the current wtat-offthheat optimizer, with a wall-time reduction of 37.5%. We achieve a new tifee-offthhee art performance of 80.56% AUC (s 0.3% improvement) improvement of 0.1 % is considered significant; see Rong et al., 2020; Wang et al., 2017.) Image classification:
this task. (An
achieve MLPerf target accuracy of 75.9% (Mattson et al., 2019) at 32K batch size on the standard ResNet-50 ImageNet benchmark in 10% fower sleps than previous state-of-the-arts Here do not see wall-time gains, mainly because the problem is too small (only few thousand steps for convergence which does not allow for amortization of costs). However, we expect that one would be able better exploit parallelism via improved software and hardware support.
We note that one of our main points in this work was to demonstrate wall-time speedups with second- order methods implemented on real-world distributed setup being used to train oftt--of-the-an deep models. In our view, this is important for influencing future hardware accelerator design and runtime software. Indeed, first-order methods have received huge investments in tuning, implementation, platform support and tailored accelerator hardware over the last decade; we believe there are numerous opportunities to improve the per-step time performance of preconditioned methods as well. For example, our results provide a concrete justification for incorporating 64bit accumulation units in hardware for distributed training, adding larger on-chip memory, better model parallelism and tighter coupling between accelerators and CPUs, which would make second order methods feasible across more domains and models.
Related work. Classic techniques for addressing the high storage and computation costs of second- order methods mostly belong the quasi-Newton or the trust-region families of algorithms (Conn et al., 2000; Nocedal & Wright, 2006). Traditionally, these methods need accura- gradients in
2

---

Under review as a conference paper at ICLR 2021
order to construct useful quadratic approximations and implement reliable line searches, rendering them as suitable for training with very large batch sizes, and resulting in expensive iterations that make the overall algorithm slow compared with stochastic first-order methods (see, e.g,, Bollapragada et al., 2018 for a recent account). Hence, our focus in this paper is on adaptive second-order methods which are directly applicable in stochastic setting. That said, our effort could be relevant to quasi-Newton and trust-region methods as well: e.g., each iteration of typical trust-region methods amounts to solving a certain generalized eigenvalue problem, which presents numerical difficulties of similar nature to those encountered in matrix root/inverse computations, being addressed here.
Various approximations to the preconditioning matrix have been proposed the recent literature (e.g., Gonen & Shalev-Shwartz, 2015; Erdogdu & Montanari, 2015; Agarwal et al., 2016; Xu et al., 2016; Pilanci & Wainwright, 2017). However, so far the only prevalent and pragmatic approximation is the diagonal approximation. Some recent approaches for approximating full-matrix preconditioner are K-FAC (Martens & Grosse, 2015), Shampoo (Gupta et al., 2018) and GGT (Agarwal et al., 2018). K-FAC uses factored approximation of the Fisher-information matrix as a preconditioner. While our focus in this paper is on Shampoo, we believe that many of the techniques presented here could also be applied to make K-FAC practical in large scale (see Appendix C). GGT uses clever trick to compute a low-rank approximation to the AdaGrad preconditioner. However, GGT maintains several hundred copies of the gradient in memory, which too expensive even for mid-sized models.
Ba et al. (2017) took a first important step at experimenting with distributed K-FAC for training deep models, using a single machine with 8 GPUs to simulate distributed environment for training. In contrast, a main thrust of our work is to demonstrate wall-time speedups with second-order methods on real-world distributed setup used for training state-of-the-art deep models, that call for design considerations crucially different than in (Ba et al., 2017). More recently, Osawa et al. (2019) scaled up K-FAC for training convolutional networks, but fell short of reaching the accuracy of first order methods, despite making changes to data augmentation and model architecture.
2 PRELIMINARIES
Adaptive preconditioning methods. First order methods iteratively update the parameters solely based on gradient information: Wr+1 w, - 778r where w, and & are (column) vectors in Rd. Here g denotes a linear combination of the current and past gradients .r. where different algorithms use different combinations. Preconditioned methods take the form Wr+1 w,, p,g, where P; is an d x d matrix. Whereas in Newton-type methods this matrix is related to the Hessian matrix of second-order derivatives, adaptive preconditioning based on
correlations.
The parameters of deep network are structured as a set of tensors of order two (i.c., matrix), three, or four. For simplicity of presentation we focus on the matrix case-however our design, analysis, and implementation hold for tensors of arbitrary order. We denote the space of parameters by the matrix W Rxx and an estimate of its gradient by G. Full matrix Adagrad flattens W, W.G. to vectors of dimension mn, it thus requires m²n" space to store the preconditioner and m²n²² time to perform the update. m and are in the 1 1000's in state-of-the-art models, thus rendering full-matrix preconditioning impractical. For this reason, both AdaGrad and Adam constrain the preconditioning matrices to be diagonal. Shampoo bridges the gap between full matrix preconditioning and the diagonal version by approximating the matrices.
The Shampoo algorithm. We describe Shampoo in the context of the Online Convex Optimization (OCO) framework, which generalizes stochastic optimization (see, e.g., Shaaev-Shwartz. 2012; Hazan, 2016). In OCO, learning progresses in rounds where on round / the learner receives an input X, and then uses the parameters W, to form prediction denoted fr. After making the prediction, the true outcome y, is revealed. The discrepancy between the true and predicted outcomes is assessed by a loss function é which takes values in R., The learner then uses the discrepancy to update the matrix to Wr+1 and prepare for the next round. For instance, the input on round can be an example x, E R. for which the learner predicts & f(W,, where f R - R. and the loss a function L: R x R - R, such as E(f,y) (y - ?) or E(f.y) log(1 + exp(-y9))
Stochastic gradient methods use the gradient C, = (wW(f((,x))) thus G, E m if the parameters are shaped as matrix W E RR For matrix-shaped parameters, Shampoo tracks two statistics over the course of its run, l,, and R,, which are defined as follows:
L cl + E'= G.G.
R₂ cla + E'=1 G.G,-
3

---

Under review as a conference paper at ICLR 2021
Note that L, E Rmxm and E R These are used to precondition the gradient and update W::
W+1 = W, L, -,,,,,,,
The primary complexity of Shampoo arises from the computation of 1, and R. which was naively implemented using spectral decompositions (i.c., SVD).
3 FULL-MATRIX PRECONDITIONING: CHALLENGES
We discuss the main challenges and design choices in the development of the distributed implementation of Shampoo. These largely arose from the fact that modern accelerators are highly optimized for training using first-order optimizers, which have low computational and memory requirements. The Shampoo algorithm is computationally expensive. The extra computation Shampoo compared to standard first-order methods is in the following steps:
Preconditioner statistics computation: 1 = 11-1 G,G.P and R R₁-1 + G.G.: Inverse tth root computation: 4, and R R-1/4,
Preconditioned gradient computation: L, 1/4 G, 1/4
Preconditioner statistics and gradient computations are expensive for large fully connected as well as embedding layers, we address these below. For other layers we show in Section 5 that they do not add significantly to the runtime of each step. Computing the inverse ''th roots is very slow-as much as 100 times the step time some cases- and performing these without slowing down training was a key challenge in our system.
3.1 ALGORITHMIC CHALLENGES
Modern ML architectures often use very large embedding layers, where the longer dimension can be in the millions. For example, DLRM (Naumov et al., 2019) on Criteo- uses vocabulary with ~186 million hash buckets, while Transformer models (Shazeer et al., 2018) 8) the largest layer can have up to 65536 units per dimension. This makes preconditioning impossible due to O(d²) memory and O(d²) computational complexity. We show how to extend Shampoo to overcome these problems; we provide proofs and convergence results in Appendix B.
Large layers. For embedding layers specifically, we extend the Shampoo algorithm to allow us use only one of the preconditioners, in case both preconditioners are too expensive compute. Our choice is empirically supported by the experiments shown in Figs. 2b, 3a and 5a which suggest that there is a benefit from preconditioning one dimension of the large softmax and embedding layers with minimal increase in time. The following result allows us to choose a subset of preconditioners:
LEMMA 1. Let GI...., G, E mx be matrices of rank at most r. Let & = vec(G,) and define H, = el + +Z'= 8xg. Let L,, R2 be defined as above: L, = el E'= G,G, R₁ = EI,+E,=1 G.G.s. Then for any p.g > 0 such that 1/p 1/q 1, we have 11 L @ R,14.
A consequence is that for any p.q > 0 such that 1/p 1/q = 1, the full AdaGrad preconditioned gradient H1- gr is approximated by (L, @ 1//4)-1/22, gr. giving us G, = L,-1/2p R, 1/24, Now, by choosing (p.4) (1, oo) and (p.4) = (oo, 1) we obtain the simple preconditioned gradients: G,R.1/2 and G,, Theorem shows that Lemma can be used to prove a regret bound for this extended Shampoo in the online convex optimization setting this provides intuitive justification for the usefulness of this approximation. We further optimize the computation of these preconditioned gradients for embedding layers by taking advantage of the sparse inputs, see details in Appendix D.
Preconditioning blocks from large tensors. addition to embedding layers, large models occasionally have large fully connected layers. To reduce the computational cost of computing statisties and preconditioned gradient: we divide the tensor into blocks and treating individual block as separate tensor. Concretely this would entail dividing tensor W R kmokkn into Wi, Wm,n such that Wi. R j.. Shampoo still converges in this case in the convex setting (Theorem 4), showing that the extension is justified.
LEMMA 2. Assume that 81. E mk are vectors, and let gi = [g, -8. where gij E E
Define A, clma + Σ g, and let B, E Rimxomk be the block diagonal matrix with k m x m blocks, where the j-th block is B(4) elm Ei=1 8.,85.j Then A, k B,,
4

---

Under review as a conference paper at ICLR 2021
We performed experiments study the effect of partitioning intermediate layers into blocks, in which we observed that the latter had minimal impact on quality of the solution while providing faster step time well as reduced memory overheads; see Fig. 3b.
Delayed preconditioners. As remarked above, computing the preconditioners is the most expensive computation in every Shampoo step. In Fig. 3c we show that we can compute the preconditioners once every few hundred steps without a significant effect on the accuracy which indicates that the loss function landscape does not change significantly with each step. We observe that there is a performance//uality tradeoff' here - in our experiments we set the frequency of computing preconditioners the smallest value that does not degrade performance, i.e. the number of training steps that can be completed in the amount of time needed to compute the largest preconditioner. The only way to increase the frequency of computing preconditioners with better hardware/software support.
3.2 NUMERICAL CHALLENGES
Inverse p roots (where typically p = 2,4,8) can be computed using SVD, but there are efficient iterative algorithms such the coupled Newton iteration algorithm (Guo & Higham, 2006; lannazzo, 2006) that can compute the inverse p'th root via a sequence of matrix-vector and matrix-matrix products, which are highly optimized on modern accelerators. However, our experiments suggest that on real workloads the condition numbers of the L, R, matrices are very large (see Fig. 6 in Appendix E) so both SVD and the coupled iteration must be run in double-precision. but this is very expensive on accelerators. We applied several further optimizations to speedup the coupled Newton iteration in implementation; these are described in Appendix E.
3.3 INFRASTRUCTURAL CHALLENGES
Heterogeneous training hardware. Neural network accelerators are custom designed to run machine learning workloads faster and at lower cost. Accelerator design is trending towards preferring lower-precision (8-bit/16-biij arithmetic that satisfy both of these goals on existing benchmarks. Our method demands double-precision arithmetic as described above, which makes running computation on accelerators non-starter, and therefore we had to design the system to leverage the existing underutilized CPUs attached to the accelerators (Section 4).
API inflexibility. Deep learning libraries such as TensorFlow (Abadi et al., 2016) offer APIs for optimizer implementation that are well suited for first-order optimizers and for mini-batch training. Our design requires that we interact with the training loop in non-standard ways, which requires framework level changes. Our Transformer experiments were carried out in the Lingvo (Shen et al., 2019) TensorFlow framework, while BERT- DRLM, well ResNet-50 used the MLPerf v0.7 Tensorflow baselines (Mattson et al., 2019). Experimentation required changes to the training loop such gathering statistics at regular intervals, distributing computation across all the CPUs available in the cluster without blocking the TPU training, as well as updating the preconditioners. We anticipate that this proof-of-concept for full-matrix preconditioning will encourage the development of more flexible API's to fully utilize heterogeneous hardware.
DISTRIBUTED SYSTEM DESIGN
We present our distributed system design of the modified Shampoo algorithm. Our method designed to run effectively on modern neural network accelerators such as TPUs (Jouppi et al., 2017) or GPUs. We first describe the standard paradigm of data parallelism used in training models on these accelerators (Dean et al., 2012). Parameters are replicated on each core of the accelerator, and each core computes forward propagation and back propagation on a sub-batch (a subset which itself is small randomly selected subset of the training set) of input examples. These gradients are averaged across all cores via all-reduction to get the average gradient for the mini-bate batch. Each core uses the average mini-batch gradient to update its copy of the parameters.
mini-batch,
All-reduction adds a barrier as all the cores need to synchronize to compute the mini-batch gradient. In Fig. 2b measure the overhead of each of the steps on a Transformer model (Vaswani et al., 2017) described in the experiment section. We observe that the overheads from all-reduction and weight updates are minor part (< 5%) of the overall step time.
The overall design of our implementation is illustrated by the timeline in Fig. 1. As discussed in the previous section the preconditioner computation (inverse pth root) is expensive and requires double
5

---

Under review as conference paper at ICLR 2021
Devices
Preconditioner computation Layer
Preconditioerr computation Layer
Preconditioners
Statistics
onc computation Layer
Preconditione computation Layer
Step
Transters
Step N+1
Stop
Transters
Accelerator
Time
Figure I: Timeline illustrating the design optimization algorithm. Preconditioner statistics (I,, and R₁) are computed at each step by the accelerators. Preconditioners (1,7 and R!/A) are only computed every N steps and this computation is distributed to all available CPU cores.
precision, also we need do this computation once every few hundred steps. These observations naturally suggested using the often underutilized CPUs on the machines to which the accelerators such GPUs or Cloud TPUs are attached. CPUs offer double precision arithmetic but are slower than GPUs or Cloud TPUs, which makes them perfect choice to run the preconditioner computation without adding any extra cost to the training run, as the computation is pipelined and runs asynchronously without blocking the training loop.
Preconditioners need be computed for every layer of the network so we distribute the computation across all the CPUs that are part of the training system. As result, the most expensive step in Shampoo adds almost nothing to the overall training time. Moreover, preconditioned gradient is independent of the batch size. Thus, increasing the batch size allows to linearly decrease the overhcad making Shampoo practical for very large scale training setups. On smaller problems such CIFAR-10, 10, we find that our design still results in training time improvements (Appendix G.3) as preconditioner computations take very little time.
computational overhead of
5 EXPERIMENTS
We compare our method against various widespread optimization algorithms for training large state-of-the-arr deep models for machine translation, language modeling, recommendation systems as well as image classification. Details of the experiments are given in Appendix G and we will opensource our code before publication.
5.1 MACHINE TRANSLATION WITH TRANSFORMER
We demonstrate the effectiveness of our implementation on the standard machine translation dataset from WMT 14 English to French (en-ffr) with 36.3M sentence pairs. We used the state-of-the-art Transformer architecture (Vaswani et al., 2017). This architecture contains 93.3M parameters and consists of 6 layers for its encoder and decoder. Each layer is composed of 512 model dimensions, 2048 hidden dimensions, and attention heads. The model makes use of sub-word vocabulary that contains 32K word pieces (Schuster Nakajima, 2012). The experiment was run on 32 cores of Cloud TPU v3 Pod, and the implementation of the optimizer was carried out in the Lingvo (Shen al., 2019) sequence to sequence modeling based on TensorFlow. Our results are shown in Fig. 2a: our algorithm achieves the same accuracy as AdaGrad or Adam in about half as many steps.
Preconditioning of embedding and softmax layers. Following the first extension in Section 3.1 the algorithm preconditions the large layers with only one of the preconditioners (G,R, 1/2 L, 1-1/2,, to make it tractable. Fig. 2b shows the increase in step time only 6% while Fig. 3a shows that we can reduce the number of steps to convergence x 20%.
Reducing overhead in fully- connected layers. Following the second extension Section 3.1 we ran two experiments where partitioned fully connected layer of size [512, 2048] into two blocks of size [512. 1024] and four blocks of size [512, 512]. Our experiments show no drop in quality under this approximation with small reduction in runtime (<%...

---

Under review as conference paper at ICLR 2021
Forward Gradient
All reduction of gradients
Precondiitoner statistics
Precondilloned gradient
Parameter updates
3.60
Diagenal Adagrad
Our algorithm for layers
Our algorthm on all layers
2.30
0 ms
20 ms
40 ms
60 ms
80 ms
100 ms
120 ms
140 ms
ms
(a)
(b)
Figure 2: Results Transformer model WMT" en-ffr, trained with batch of 1536. (a) Test pog-perpyyxity vs, number steps: algorithm converges 1.95x faster steps, while being only 16% slower per step. This allows the method to attain a particular log- perplexity in 40% less wall-time. (b) Detailed breakdown of latency of single step (Appendix G.6). Diagonal AdaGrad optimizer: 134ms, Shampoo: 145ms (all layers except embedding and softmax layers) and 155ms (all layers). Preconditioner computation is pipelined and distributed over CPUs, thus adding any overhead, and transfer latency (< (=100ms) amortized over hundreds of steps.
7,00
1.35
2.35
L200
1,57
105
1524
180%
200
80C 100
155%
1200
100
1902
(a)
(b)
(c)
Figure 3: Impact of Shampoo extensions WMT" 14 en fr training: (a) preconditioning applied to all layers except embedding and softmax layers, vs. applied to all layers; (b) preconditioning with fully-connected layers partitioned into sub-blocks: (c) varying interval between preconditioner updates.
5.2 TRANSFORMER- BIG MODEL
We also ran experiments with larger Transformer model with 375.4M parameters, consisting of layers for its encoder and decoder. Each layer is composed 1024 model dimensions, 8192 hidden dimensions, and 16 attention heads. Results are presented in Fig. 4a where again an improvement the end-to-end wall-clock time. For the softmax, embedding and the projection fully-connected layer (with 8192 hidden dimensions) only make of the left preconditioner. We note that step time is dominated by the preconditioned gradient computation which can reduced by sub-blocking the layers.
On the overhead of the optimizer. We show the computational and memory complexity of the Shampoo extensions described in Section 3.1 in Table in the appendix. We note that the overhead from computing the statistics, as well as from computing the preconditioned update for single step of training, can be further reduced by increasing the batch sizes (indeed, these overheads are independent of the batch size) as shown in Fig. 4b where the overhead dramatically reduces from 40% to 19%.
5.3 ADS CLICK-THROUGG RATE (CTR) PREDICTION
We trained the Deep Learning Recommendations Model (DLRM) of Naumov al. (2019) on the terabyte Criteo click logs dataset for online advertisement click-through-rate prediction task (Criteo Labs, 2015). We compared Shampoo against the highly tuned SOTA baseline from MLPerf v0.7 training benchmarks (Wu et al., 2020). We trained the model with a batch size of 65536 for 64000 steps (1 epoch). We trained a version of the model where Shampoo is applied only to the hidden layers well exact same setup as the baseline. We found that Shampoo achieves the target accuracy of 80.25% in only 30.97K steps compared to 64K steps for the baseline. Moreover, Shampoo achieves
where we apply it for all layers. We only tune the learning rate, and keep the
7

---

Under review as a conference paper at ICLR 2021
210
2.50
Adepud
Shampoo Same AddOted
270
2.4
2.50
250
2.40
2.40
2:20
2.40
3K6
100C
33000
2004
250C
137
2014
6000
1016
1106
3200
3.00
3870
3.90
(a) Batch size: 384
(b) Batch size: 1536
Figure 4: Test log- log-perplexity of a Transformer-Bis model on WMT 14 en-frr. (a) Shampoo converges faster than AdaGrad (= 2x faster in steps), and allows larger learning rates; duc to the large overhcad in step time, this results in only 30% improvement in wall-time. (b) Larger batch sizes reduce the optimizer overhead from 40% to 19%, resulting in an end-to-end improvement of % wall-time convergence.
state-of-the performance of 80.56% AUC (an = 0.3% improvement) on this dataset, note that an improvement of 0.1 % is considered significant in this task; see Rong et al., 2020; Wang et al., 2017. Here preconditioning embedding layers further reduced the number of steps needed to reach the target accuracy from 39.96K to 30.97K.
LMPE
66.00
19
02.00
18.50
18.00
Hurrea (oely baynn
10C
S/c
(a) Test AUC on the Criteo- Sitteo-1Tb Ibb dataset.
(b) Masked Language accuracy BERT-Large.
Figure 5: (a) Shampoo reaches a target AUC of 80.25% in half as many steps with preconditioning embedding layers improving the results, and achieves new state-of-the-art AUC of 80.56%; (b) Shampoo converges in 16% fewer steps, and achieves 1% higher MLM accuracy than the baseline BERT-Large.
5.4 LANGUAGE MODELING
We trained BERT-Large (the Bidirectional Encoder Representation architecture Devlin et al., 2018) for the language modeling task on the concatenation of Wikipedia and BooksCorpus, with 2.5B and 800M words respectively. BERT-Large is a large bidirectional transformer model containing 24 transformer blocks with 1024 hidden dimensions and 16 self attention heads. It has 340M parameters (a) masked language model (Masked-L.M) loss where the task is to predict masked tokens based surrounding context, and (b) next sentence prediction (NSP) loss where the task is to predict whether two given sentences are consecutive the text. In Fig. 5b compare et al., 2019). Models were trained with batch size 16K; in these experiments we replaced the Adam update rule in Lamb that produces the preconditioned gradient with Shampoo. Both experiments used existing well-tuned hyperparameters of the baseline.
and set up to jointly optimize two objectives:
results against the current state of the art in training BERT (You
5.5
IMAGE CLASSIFICATION
We trained a ResNet-50 model (He et al., 2016) the ImageNet-2012 (Russakovsky et al., 2015) dataset and compared against the state-of- ffite-of-theee -tteh-art baseline using SGD+Momentum. We base our experiments off the Tensorflow baseline available from Mattson et al. (2019) where the target criteria reaching 75.9% accuracy. See results in Table 1; in particular, we find that Shampoo reaches the target accuracy in fewer steps than the current state of the art. Tuning details in Appendix G.4,
8

---

Under review as a conference paper at ICLR 2021
OPTIMIZER
BATCH SizE
EPOCHS
StepS
SGD+Momentum
4096
85
26586
LARS
4096
45
14040
LARS
32768
64
2512
Shampoo 4096
45
14040
Shampoo
16384
48
3744
Shampoo 32768
58
2262
Table 1: Epochs and steps MI.Perf target accuracy of 75.9% with ResNect-00.
6 CONCLUDING REMARKS
We have presented an implementation of second order optimizer, and demonstrated step time as well as wall time improvements on multiple large tasks in different domains in each case our implementation performed as well or better than state-of-the-ar optimizers specialized for each domain. The main point our work is to demonstrate that second order methods implemented on real-world distributed setup can be used train state-of-the-art deep models. We hope that this work will influence future hardware accelerator design and runtime software first order methods have received large investments in tuning, implementation, platform support and hardware tailored for them, and we believe there are several opportunities to improve the per-step time performance of second order methods well:
Most second order methods use symmetric matrices, but we haven't found support for typing operands symmetric, which can reduce compute flops and storage by upto 50%.
Several optimizations that are currently tuned towards first order methods could be extended second order methods. For example, weight update sharding pattern matches first order methods (Xu et al., 2020) and dramatically reduces the time spent in the update step as well as memory used. This change can also be applied to Shampoo with blocked preconditioners - but we do not have support for it yet as it requires compiler level support, and is not expressible at the program layer. Currently every core must update all layers which is quite inefficient.
Mixed precision algorithms may work for inverse pth roots and can help increase frequency of preconditioner computation.
Increased memory per chip can allow larger preconditioners.
Hardware support for high-precision arithmetic in accelerators can allow more frequent preconditioner computation. The benefits high precision arithmetic optimization run counter the prevailing wisdom in ML' which has led to the focus on low-precision formats such as bíloat16 (Wang & Kanwar, 2019).
Hardware support for storing/packing and using upper/lower triangular matrices efficiently, as available in LAPACK.
Our hope is that these suggestions could result in innovations that would make second-order methods practical across more domains and models, especially in data limited regimes where we may not able amortize the latency added in the data transfer between the accelerator and the CPU.
Forr example, (Gupta et al., 2015) say "It is well appreciated that the presence of statistical approximation and estimation errors, nigh-precision computation in the context of learning is rather unnecessary (Bottou & Bousquet, 2007)" and (Higham & Pranesh, 2019) say "... machine learning provides much of the impetus for the development of half precision arithmetic hardware...

---

Under review as conference paper at ICLR 2021
REFERENCES
Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: A system for large-scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16), pp. 265-283, 2016.
Naman Agarwal, Brian Bullins, and Elad Hazan. Second order stochastic optimization in linear time. arXiv preprint arXiv: /ttp/ 2016.
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. The case for full-matrix adaptive regularization. CoRR, abs/1806.022) 02060, 2018.
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang, Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002. 11003 2020.
Tsuyoshi Ando, Chi-Kwong Li, and Roy Mathias. Geometric means. Linear algebra and its applications, 385:305-533333 -3,434,,,,,,,, 2004.
Jimmy Ba, James Martens, and Roger Grosse. Distributed second-order optimization using kronecker- factored approximations. International conference on machine learning, pp. 2408-2417, 2017.
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ale Tamchyna. Findings the 2014 workshop on statistical machine translation. In Proceedings of the Niuth Workshop on Statistical Machine Translation, pp. 12-58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics. URL http:/ httt::/wwwww /wwww//wwwww/////o////.// aclweb org/ anthology/ W/W14/W14 3302.
Raghu Bollapragada, Jorge Nocedal, Dheevatsa Mudigere, Hao-Jun Shi, and Ping Tak Peter Tang. A progressive batching 1-bfgs method for machine learning. In International Conference on Machine Learning, pp. 220-629, 2018.
Andrew Conn, Nicholas IM Gould, and Philippe Toint. Trust region methods. SIAM, 2000.
Criteo Labs. Criteo releases industry's largest-ever dataset for machine learning academic community, July 2015. URL https: /www. criteo. htt//new//preee reieessrrleeass/221155 industrys-
-arelaasee-mtutttts------------- wawprsss--rreeeeeeeeeeeee ever dataset/.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthicu Devin, Mark Mao, Mare' Aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, Quoc V. Lc, and Andrew Y. Ng. Large scale distributed deep networks. Advances in Neural Information Processing Systems 25, 2012.
Jacob Devlin, Ming Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint arXiv: 810.00805.0 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Reseurch, 12(lu))2121-2155) 2011.
Murat A Erdogdu and Andrea Montanari. Convergence rates of sub-sampled newton methods. In Proceedings of the 28th International Conference on Neural Information Processing Lextems-Volume 2, pp. 3052-3060. MIT Press, 2015.
Roger Fletcher. Practical methods of optimization. John Wiley Sons, 2013.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in Kronecker factored eigenbasis. In Advances in Neural Information Processing Systems, pp. 9550-9560, 2018.
Alon Gonen and Shai Shalev-Shwartz. Faster sgd using sketched conditioning. arXiv preprint arXiv: 1506.0200602449 2015.
10

---

Under review as a conference paper at ICLR 2021
Chun-Hua Guo and Nicholas J Higham. A Schur-Newton method for the matrix p'th root and its inverse. SIAM Journal On Matrix Analysis and Applications, 28(3):788-80)) 2006.
Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In International Conference on Machine Learning, pp. 1737-1746, 2015.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 1842-1850, 2018.
Elad Hazan. Introduction to online convex optimization. Foundations and Trends Optimization, (3-4):157-325, 2016.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Tom Heskes. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12(4):881-901, 2000.
Nicholas J Higham and Srikara Pranesh. Simulating low precision floating-poin arithmetic. SIAM Journal on Scientific Computing, 41(5):C55--662, 2019.
Bruno lannazzo. On the Newton method for the matrix p-th root. SIAM journal on matrix analysis and applications, 28(2):503-523 2006.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Computer Architecture (ISCA), 2017 ACM/IEEE 44th Annual International Symposium on, pp. 1-12. IEEE, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. In Advances in Neural Information Processing Systems, pp. 4156-4167, 2019.
Adrian S Lewis and Michael L Overton. Nonsmooth optimization via quasi-newton methods. Mathematical Programming, 141(1-2)135--66, 2013.
James Martens and Roger Grosse. Optimizing neural networks with Krotecker-factored approximate 2015.
curvature. In International conference machine learning, pp.
Peter Mattson, Christine Cheng, Cody Coleman, Greg Diamos, Paulius Micikevicius, David Patterson, Hanlin Tang, Gu-Yeon Wei, Peter Bailis, Victor Bittorf, et al. MIperf training benchmark. arXiv preprint 10Xiv:990.001500, 2019.
H Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. COLT 2010, pp. 244, 2010.
Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang. Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint 2019.
Jorge Nocedal. Updating quasi-newton matrices with limited storage. Mathematics of computation, 35(151):773-72, 1980.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Science & Business Media, 2006.
11

---

Under review as a conference paper at ICLR 2021
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-seale distributed second-order optimization using sonecker-factorce approximate curvature for deep convolutional neural networks. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 12359-12367, 2019.
Michael L Overton. Numerical compuling with IEEE floating point arithmetic. SIAM, 2001.
Mert Pilanci and Martin J. Wainwright. Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence. SIAM Journal on Optimization, 27(1):205 -455, 2017.
Haidong Rong. Yangzihao Wang, Feihu Zhou, Junjie Zhai, Haiyang Wu, Rui Lan, Fan Li, Han Zhang. Yuekui Yang. Zhenyu Guo, et al. Distributed equivalent substitution training for large-scale recommender systems. In Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 911 911-920, 2020.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211225 2015.
Mike Schuster and Kaisuke Nakajima. Japanese and Korean voice search. ICASSP, pp. 5149-5152.
IEEE, 2012.
Shai Sawarr- Online learning and online convex optimization. Foundations and Trends in Machine Learning, 4(2):107-194, 2012.
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et Mesh-tenssrfflow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems, pp. 10414 10423, 2018
Jonathan Shen, Patrick Nguyen, Yonghui Wu, Zhifeng Chen, et Lingvo: a modular and scalable framework for sequence-to- -o-suuennce modeling, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Ruoxi Wang. Bin Fu, Gang Fu, and Mingliang Wang. Deep cross network for ad click predictions. In Proceedings of the ADKDD'I7 17, pp. 2017.
Shibo Wang and Pankaj Kanwar Bfloat16: The secret to high performance -elaniin-
on cloud tpus. https://cloud. gwwgle. bfloat16 the secret-o- hioo-higg hwwh-higghieeeeeeeeeeeeeeeeeeeee performance cloud- tpus, 2019.
Carole-Jean Wu, Robin Burke, Ed Chi, Joseph Konstan, Julian McAuley, Yves Raimond, and Hao Zhang. Developing recommendation benchmark for miperf training and inference. arXiv preprint arXiv:2003.07336. 2020.
Peng Xu, Jiyan Yang, Farbod Koosta-Khorasen Christopher Ré, and Michael W Mahoney, Sub- sampled newton methods with non-uniform sampling. In Advances in Neural Information Processing Systems, pp. 3000-3008, 2016.
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Hongjun Choi, Blake Hechtman, and Shibo Wang. Automatic cross-replica sharding of weight update in data-parallel training. arXiv preprint arXiv:2004. 13336, 2020.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.0096, 2019.
12

---

Under review as a conference paper at ICLR 2021
A NOTATION
We use lowercase letters to denote scalars and vectors, and uppercase letters to denote matrices. ||||| denotes the Frobenius norm A, i.e. ||A||; Li, At.. A denotes the Hadamard or element-wise product of A and B which have the same shape, so C = A o B = Cij = Aij D is the element-wise power, (Do = Di
We use N to denote the Loewner order: given square symmetric matrices A, B, we write A = B iff A is positive semidefinite (PSD).
Given symmetric PSD matrix A, and E R. A" is defined as follows: let A UDU" be the singular A, where a unitary matrix and is diagonal matrix (with Dj 2 0 as A is PSD), then A = UD*U* where (D)) Dj. If < 0, this is defined for positive definite matrices only, where Du > O.
value decomposition
matrix A: if A has rows a)....,mm then We use vec(A) to denote the flattening of the am)" A @ B denotes the Kronecker product vec(A) is the mn x column vector vec(A) = (d).....mm) of two matrices A and B, and we will use the identities (A @ B) A & Ba for R, and (A @ B) vec(C) = vec(ACB").
B DEFERRED PROOFS
PROoF (or LEMMA 1). Lemma 8 in Gupta et al. (2018) shows that H, x rL, @ I, and H, rm @ R By using Ando's incquality (Ando et al., 2004), we get
H, r(L, @ In) (Im (P R2)
r(L,"P @ 1m @ R!
rL, 1/ @ R 19
which concludes the proof.
This lemma immediately allows to prove regret bound for Shampoo with extended exponents: THEOREM 3. Assume that the gradients CI....,yy Gr are matrices rank at most r. Then the regret of Shampoo with extended exponents compared to any W* mmm is bounded as follows,
Σ ((W)) 1=1
Σ f(W*) VErDTT(L )TT((( E=1
where
G,G., Rr eln Σ G.G., D max |[, W* re[T]
L
elm
1=0
and 1/p+ 1/9 = 1.p.9
PRooF. The proof follows the proof of Theorem in Gupta et (2018). Let H, = 1,, @ R," Then the update rule of the extended Shampoo algorithm is equivalent to W,+1 w, HH, 8: Since 0 N L1 .... < Lr and 0 x Ri I = Ry, standard properties of the Kronecker product and the operator monotonicity of the function x + xa for a (an immediate consequence of Ando's inequality) ensure that 0 N H1
1 llr.
Following the aforementioned proof, we have the regret bound
Σ ((W.) Σ f(W*) 1 Ti(Hr) [=1
D2
T Σ [=]
l1g. 177
2m
where D max, ||W, ll. Define & vec(G,) and A, (cl,, + ES 8xg,) then Lemma 1 shows that f, = VrH, using operator monotonicity. Using this equation twice, along with Equation (6) from the proof of Theorem 7, we have
11rr lee 5 VF Σ
ll 11 $2VFTr(fhr) s 2r Tw(IT)).
7-1l
13

---

Under review as conference paper at ICLR 2021
This gives
Σ f((,)) Σ f((*)=
D²
Tr(Hr) mリ Tr(Hy).
2n
1=1
Setting 1) D/V2r and observing that Tr(H,) = Tr(L, 112) Tr( 1/0g gives us the required bound.
PROOF (oF LEMMA 2). Let x E R and x = 12. xxl, where xj E R Then
x = + Σ x &ggggx e||x|12
Σ (g]x)² e|1x112
Σ 8 B
kel|x112 + E
(oT
xx)) Σ
eelxzy133 Σ
x 2788,,88...))
k (elm
Σ 88.881
Σ j=1
B xj
kx -x..,...
xj
Here used the inequality (Z)=1 a²)² = k -/= a) which follows from the convexity of r x2 (or from the fact that variance of random variable is non-negative).
This lemma once again allows us to prove regret bound, exactly following the proof of the regret bound above:
THEOREM 4. Assume that the gradients .7 Rmn and let gi Igi. gi.A where 8i,j R Then the regret of Shampoo with blocking compared to any w* R is bounded as follows:
Σ f( Σn(w> V2kD
Σ j=|
Tr
mm(eemm clm
Σ 81.88 S1.8BTJ)
The two regret bounds can be combined to show that Shampoo with both extensions also converges.
C COMPARISON WITH K-FAC
K-FAC is natural gradient algorithm, and approximates the curvature of the loss using the Fisher Information Matrix:
F=
E p(x(0)
Vlog Vlog p(x)) p(xx=000)))))))))
E p(xi0)
8p(x10) 8 8p(x10)
For fully connected layer with W Ruxx where Wx s, the gradient for the layer G, Rmxn can be written via the chain rule as G, V,l(s,, x))x and in vectorized form as: v,l(s, yı) @ We can then write the Fisher information matrix as:
== p(x(0)
E (VV.lss,.,) p(x(0)
(xxx))
Assuming independence between v,l(s,,)) and x, K-FAC rewrites the Fisher in tractable form as:
F E
E E
we let / = (V. and X E x the update rule then becomes:
WA11 W, D G
We note some of the differences and similarities between the two updates here. KFAC preconditioners use exponent of -1 (as original Fisher is inverted) whereas Shampoo uses - 112pp where is the rank of the tensor. KFAC computes statisties based gradients with labels sampled from the model's
14

---

Under review as a conference paper at ICLR 2021
predictive distribution (hence requiring strictly more computation) where Shampoo relies on the gradient of the mini-batch.
Now we can compute each term in the Shampoo preconditioners as:
G,G. V,
G.G. = x, ( V,, )))x; = 3x) Dividing by the scale, and taking expectations on both sides:
G,G.
E [V D:
||
G,
E
E x X.
|V, D112
This shows that K-FAC preconditioners closely related Shampoo preconditioners, especially when one uses the empirical Fisher (Kunstner et al., 2019).
The main difficulty in implementing K-FAC 1x. ((),))) to the optimizer, so K-FAC difficult to send additional information such as ||x implementations have register the structure of each layer. Moreover, due to the dependence of K-FAC on the structure of the network, it is difficult to implement standard operators like batch norm, weight norm, layer norm, etc., which are prevalent in the tasks and models we considered. For example, if we write fully connected layer with weight norm as = Wx/|[W]], then the gradient
model is that current optimizer APIs make
G = ||i|| Val(s,, V,l(s,, yd) Wx W, ||W]
so rewriting E[vec(G,) ve as Kronecker product is not an easy task.
The similarity between K-FAC and Shampoo preconditioners also allows us to use techniques explored by the K-FAC community for Shampoo. One of the extensions for KFAC is the E-KFAC algorithm (George et al., 2018) which constructs a better approximation of the Fisher matrix by using the eigenbasis computed from the Kronecker approximation, but rescaling the eigenvalues to match the diagonal of the Fisher matrix in this eigenbasis. This method produces provably better approximation, and can immediately be applied to Shampoo too with a simple modification:
Let ft, = 1,,/2 @ R,/22 Let the singular value decompositions of the factors be 1,1/2 UDUT and R, = VD'V.. The L @ R!/2 = (U @ V)D a D')(U @ v)T. Now the EKFAC correction replaces D D' by the optimal diagonal
diag((U @ V) A(UU
cl +
diag((U @ v)" vec(G.) (U
diag(vec(U" G,V) ve) G,v)))
cl +
s=l
cl
vee(U" 02
s=1
Thus we can approximately compute A++1 = + GU)22, and the new update becomes: W,+1 W, ( This technique does have the disadvantage that it requires computing the singular value decompositions (which we already observed are much slower than coupled Newton iterations), and doubles the number of matrix multiplications in the preconditioned gradient computation. At this time our experiments did not show significant improvements over the standard Shampoo implementation, but we plan to explore this further.
D SHAMPOO FOR EMBEDDING LAYERS
In modern networks, embedding layers are usually very large, and even computing the left precondi- tioner described in Section 3.1 can be prohibitively expensive. However can take advantage
15

---

Under review as a conference paper at ICLR 2021
of the fact that the inputs to the network are very sparse, and use this to reduce the computation significantly.
Let our input example such network consist of set of categorical features: each feature such as user language, user country etc consists of one out of a set of options. Then the output of the embedding layer is the concatenation of the embeddings for each such feature. If the embeddings are of width d and there are N such embeddings, then the embedding layer is W E RdxN The input can be represented as e RN- where is the number of categorical features, and each column is one-hot: if the k-th feature is x(k), then xjk = =x)x())) The output of the layer is = Wx.
Now G = vwvl= = v,lx'. so GG = Vyex'x can compute the preconditioner for W by computing on the output of the embedding layer, and this is a much smaller computation since y is of dimension b x m. this computation is O(dmm) rather than N). Note that sparse multiplication would also be m)) but accelerators usually implement sparse operations by densifying the tensors.
But x x Im. so GG VveV,e". Thus we
If each column of x is multi-hot, as is the case when the features are words and their embeddings are averaged, x'x x is diagonal matrix, where each diagonal entry is a function of the number of ones in each column of x. Computing GG = Vv,((* x)V, e is still m) & N).
E A COUPLED NEWTON ITERATION FOR COMPUTATION OF INVERSE p-TH ROOTS
The Newton method for solving the matrix equation X A 0 produces the iteration X2+1 H I)Xx xft1|A A], where we take Xo = LT. This iteration satisfies X - A as k co, but
it is not numerically stable. Introducing the matrix M₂ X A. we get
(p + 1)/- Me
1
X = X1
Xo I,
and
(p 1)/ - Me
(p + 1)/
M
1
MR+1
XP 2+1 A
XPA=
MR.
Mo
A.
p
P
cP
since Xz, My and commute with each other. This is the coupled Newton iteration for computing inverse p-th roots, and was shown to be numerically stable in (Guo & Higham, 2006; lannazzo, 2006).
We implemented the following optimizations to the coupled Newton iteration method:
Warm Start: The coupled Newton iteration to compute l/p starts with X = L, G and maintains the invariant M X while driving M - 1, resulting in X G-lip_ We need to find the p-th root of a sequence G, so we instead set X = G, 1/P M = X G since the difference between G, and G,+1 is small, this ensures that M is already close to 1. In our experiments warmstart improves convergence (by upto 4x fewer steps).
Projecting top singular values: In practice our G, matrices have large condition numbers, which sometimes leads to inaccurate results. As a rule of thumb, computing G-1/p leads to a loss of log2( { x()) bits of precision (Overton, 2001), where x(G) is the condition number of the G. However we also find that usually there is a sharp falloff within the first few singular values, in order to reduce the condition number, we project away the largest singular values, since these are the least important after taking inverses. We find the top-k singular values A, and their associated singular vectors using a standard iterative method, and replace each with AR+1- This produces a better approximation than adding e/ to each G,: the latter can wash out the smallest (and most crucial) singular values, see Fig. where the smallest singular value for layer can be as small as 10-10 to 10-6 during the course of optimization.
Dynamic uning of projection: We dynamically tune the number of singular values we need to project in the previous step, by computing the condition number x(G,) and using it to estimate the smallest singular value of G as Amax (G (G,). We then keep projecting out singular values of G until get an acceptable condition number.
Correcting for projection: If G, L; A. then G. G:liP- =; V... Projection above means replacing M As by A but since we have already computed the corresponding
Vi.... Vk. we correct the approximate p-th root by adding (2, 1 -1/p > This is a small effect, but adding it is a straightforward modification (details deferred to Appendix E).
16

---

Under review as conference paper at ICLR 2021
Algorithm I A coupled Newton iteration procedure for computing inverse p-th roots of PSD matrix, with warm start and singular value projection
1: procedure MAXSV(G)
2: Parameters: > 0, Matep
3: E R", where G E R
4: i= O, error = 00,
5: while < Nstep and error > e do
6:
= v/llv|l
7:
Gf
8:
Aold 1:A VTv
9:
error 11 Aotal;i =
return v/v/lv/l
12:
procedure PROJECT(G, (optional), kJ (optional), *pooj (optional))
14:
Á =
v, MAXSV(G)
16:
Amax = A
17:
while À Kd Ama or < Alrroj do
18:
G = G - Avv
vy
20:
A, = MaxSV(G)
i= -ii+1
22:
return G + AA
23:
24:
procedure COUPLEDITERATTIONG.. p N, X (optional), (optional)
25:
Parameters: O, Kdi Пproj
26:
Outputs: G-1/p
27:
G (ROJECT(G,X, K, K, (proj)
28:
a x=
29:
if X is provided then
30:
M XºG
31:
else
1tp
32:
z 2||G
X 20
1
34:
M = ZG
35:
while |M I||s> > do
36:
M (1 a)I + aM
37:
X = XM1
38:
M M. M
39:
return X
IMPLEMENTATION DETAILS OF SHAMPOO
Our implementation of the Shampoo algorithm for fully-connected layers is described in Algorithm II. The algorithm can use heavy-ball momentum for its updates, as well an exponential moving average over the preconditioners, like Adam. The configuration parameter T denotes the number of steps between subsequent fetches of the latest available preconditioner by the accelerator. must be set sufficiently high so that there enough time for the CPU to complete the computation of preconditioner asynchronously and pipeline it efficiently, but otherwise its setting does not have significant effect on convergence. The configuration parameter T2 (default value 1) determines the frequency of gathering gradient statistics update L,, R, every 72 steps only for efficiency.
E.1 COMPUTATION COST SHAMPOO
We capture the computational and memory complexity under various schemes described in Section 3.1 of handling large layers in Table 2.
17

---

Under review as conference paper at ICLR 2021
SVD
Exclatoo conditio numbers
530.0
430.0
20
230.0
120.0
10
0.0
1200
30.00
4002
3032
0030
23K 4010
Emmendann ofmarda c1
$1655
Figure Benchmarks on computing inverse-pti -ehrse-phhheee root statistics varying dimensions (left), and condition numbers for L, of layer in the transformer model over time (right). We find that the coupled Newton iteration method can effectively utilize the CPUs and give large walltime improvements compared to SVD (that relies on bidiagonal srriid--and-oonquer These upto by reducing the number of iterations the solution. These were measured Intel Skylake CPUs. Note that since log2( xx(,)) bits of precision are lost in computing p-th roots, 64-bit arithmetic becomes necessary.
measured without warmstart which provides additional speedup of
Algorithm II Sketch of the Shampoo algorithm parameters: learning rate Tr,, momentum: BI, P2
2: for I.... do
3: Receive stochastic gradients G, for each layer
4: if r % = 0 then
5:
if B2 then
6:
L. + B2 14-12 + +1 (1 (f-)))) G.G.
7:
R, + f2 R1-r, (1 B2) G.G.
8:
else
9:
L, +- L- + G
10:
R - R + C.G.,
D, D,-1 G, G,
M, < 91 M,-1 +(1 (1-1,)) B1) 1/-1/2 G,
13:
if % = then
14:
Gather preconditioners
1/4 R- 1-14 ((-r))) (T-T) from CPUs
Send L, R, to CPU host compute L,"
11/ R 1/4
16:
if > T then
P - P1P.-1 + (1- (1) L, 1/4 G, 1/4
17:
17 - no0| P
18:
W, W,- P.
20:
else
π7 < 770
22:
W, = W,- Mr
G FURTHER DETAILS ON EXPERIMENTS
Layer wise learning rates. As seen in Fig. 7 the step size scale for each layer is dependent on the operator norm of the preconditioners (inverse-pll -erre root of the smallest singular value of the statistics
TYPE
COMPUTATION
MEMORY
All preconditioner W,: [n, m]
O + m²n)
O(n + m²)
Left only preconditioner for Wy: [n.m] Preconditioner: block size
O(n'm)
O(n")
O(mbb)
O(mn)
Table 2: Computational and memory complexity of variants of Shampoo.
18

---

Under review as a conference paper at ICLR 2021
matrix) has large spread in its range which results in optimization instabilities in practice. Moreover, as statistics as well as preconditioner computation are amortized across many steps the norm does not grow at every step. Hence, rely learning rate schedule based the update directions of well tuned first order optimizer (in our experiments we use diagonal AdaGrad for Transformers in machine translation, well as Criteo, layer-wise scaling heuristic proposed in LARS/LAMB optimizer, where each layer's s learning rate is set to be | | for BERT and ResNet training. For example, when used with diagonal AdaGrad: Shampoo is used to determine the direction of the update, and AdaGrad to determine its magnitude.
This procedure termed Grafting (Agarwal al., 2020) allows us to bootstrap a reasonable learning rate schedule for a specific problem that is well tuned, and study the effect of preconditioned gradient directions in isolation. The weight matrix W, is updated as W, W A,S, where:
D,
Σ s=1
G. C5; =
= no ||, 1-1 GillF
(Adagrad magnitude)
L. R, 1/4
S,
(Shampoo direction).
|L, 1/4 G, 14
10
-3
10-9
Figure 7: Minimum (dashed) and maximum (solid) singular values for statistics matrices of the embedding, softmax and intermediate attention query projection layers.
G.1 TRANSFORMER MODEL WMT 14 EN-FRR
For all optimizers, we make use a warmup schedule where the learning rate increased from 0.0 over 40k steps. For the smaller transformer experiments, larger transformer experiments we use a linear warmup. We found that quadratic warmup improves all optimizers equally and provides better log-perplexity. For the Adam optimizer experiments, we use a learning rate decay schedule of the form 7) /d/rr, following the suggestion of Vaswani et al. (2017). For the smaller Transformer experiments, we tuned the hyperparameters for each algorithm over 100 trials. We took the best settings for the momentum and recond-moment parameters, and tuned the learning rates until either the model became unstable, did not increase performance. For Shampoo, we used a per layer learning rate derived from AdaGrad (see Appendix G for details), and found that for the exact same hyperparameter settings AdaGrad, Shampoo provides modest improvement in performance. Moreover, Shampoo allows for larger learning rates than AdaGrad does, shown in Fig. 4a.
quadratic warmup, and for the
G.2 STEP TIME BERT-LARGE
Our current implementation showed a 14% increase in step time for BERT-Large, nearly wiping out all the gains from reduced number of steps (16%). We note that due amount of resources it would require to tune BERT, we used Shampoo with exact same to understand the effect of preconditioner. Moreover, step time be optimized considerably the current implementation is not heavily optimized. For example, larger batch sizes help amortize the preconditioning overhead, and reduce overall wall time to reach the same accuracy. Furthermore,
as LAMB with grafting
19

---

Under review as a conference paper at ICLR 2021
EXPERIMENT (TPU CORES)
OPTIMIZER
BATCH
OPTIMIZER PARAMETERS
WARMUP
Adam
1536
1 0.000225, f1 0.9, P 0.98
40k steps
Transformer (32)
Adagrad
1536
7 7= 0.125,, B1 0.95
40k steps
Shampoo
1536
? = 0.225. f1 = 0.95, K = 500
40k steps
= 1000, T2
Adam
384
7 0.000154, A1 0.9, P22 = 0.999
40k steps
Transformer- Big (32)
Adagrad
384
7 0.03, 61 = 0.9
40k steps
Shampoo
384
1 0.06, P1 0.9, =500
40k steps
F 1000, T2
Adagrad Shampoo
1536
1 0.06, Pj 0.9
40k steps 40k steps
Transformer-Bip (32)
1536
π 0.08, P1 0.9, 500
T = 1000, T2 =
LAMB Shampoo
16384
r 0.0060 PI 0.9, & 0.999
6.4k steps 6.4k steps
16384
7 0.0060 P1 = 0.9, f2 0.999,
Bert-Large (256)
A 10- T1 400, T2 10
Block size: 1024
SGD Shampoo
65536
7 0.1, poly dccay(p=2) at 38k steps
2k steps 2k steps
65536
1 0.1 poly decay(p=2) ß1 = 0.9, T = 999, T2 10
steps
DLRM (32)
(w/ embd)
65536
7emmdd = 0.31
Table Hyperparameter setup used in our experiments.
in current implementation, all TPU cores compute all the preconditioning statistics and the preconditioned gradients, which involves over hundred 1024 x 1024 matrix multiplications. This repeated work can be avoided by cross-replica sharding of weight update (Xu et al., 2020), which distributes this computation across cores, and should save at least half the step time overhead.
G.3 CIFAR-10
We train a ResNet-50 model on CIFAR-10 10 (Krizhevsky et al., 2009) with 2 cores of CloudTPU-v2 at batch size 2048. Our baseline achieves 93.45% accuracy 300 epochs, where Shampoo reaches the same accuracy in 143 epochs. We see an overall training time reduction of 42% (1428 seconds to 827 seconds). As it smaller problem, the time taken for preconditioner inverse computation for the largest preconditioning matrix is less than Ims on the CPU. We use total of 8 CPU cores to run these inverses.
G.4 IMAGENET
For SGD with Momentum, the learning rate warmed up over the first 5 epochs from 1.6, followed by 10x drops of the learning rate at 30, 60 and 80 epochs. For LARS, we use warmup learning rate over 20 epochs for 4K and 16K batch sizes, 25 epochs for 32K batch size with polynomial decay (p=2) until end of training. For Shampoo we use the same layer-wise heuristics and hyperparameters as LARS with Grafting such that the direction changed to the one computed by Shampoo. We make use weight decay with value: A = 21 and label smoothing of 10-!
G.5 DETAILED RESULTS FOR EXPERIMENTS
Approximate wall clock times for the various tasks are follows:
Task
Model
Baseline
Shampoo
Recommendations: Criteo-1TT
DLRM
13 min
8.2 min
Translation: WMT-14 En-Fr
Transformer
2 12 hrs
6.5 hrs
Translation: WMT-14 En-Fr
Transfomer-Bi
47 hrs
29.5 hrs
Language Modeling: Wikipedia+Books
BERT-Large
228 mins
219 mins
G.6
BREAKDOWN OF STEP-TIME IN FIG. 2B
Each step of training consists of the following phases, whose times are shown in Fig. 2b.
20

---

Under review as conference paper at ICLR 2021
Forward Pass: Each core independently computes the predictions for each training example in its sub-batch.
Gradient:
The gradient for sub-batch is computed using the back-propagation algorithm.
All reduction: The gradients for the sub-batches from all are averaged compute the gradient for the minibatch. This is then sent back to each core.
The preconditioner statistics for adaptive algorithms are updated, e.g. for AdaGrad, set H; == H; + 8 for all parameters, while for Shampoo, we set L GG etc.
Preconditioner statistics:
Preconditioned gradient: The preconditioned gradient is computed e.g. for AdaGrad, compute Ri/HH, while for Shampoo, we compute I. 1--//GGR 114
Parameter updates: The parameters are updated using the preconditioned gradients. This step is the same for all algorithms: W := W rτ,, where G is the preconditioned gradient.
Note that the Shampoo computation of the preconditioners L-1/4 R is pipelined on the host CPU, so does not show up the step times.
21


Reviews:
[{"review_id": "f6GMzmS5vG", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "######################################################################\n\n1.  Paper Summary \n\nThis work addresses practical challenges in applying full matrix pre-conditioner methods (such as Shampoo) on problems involving large datasets and architectures trained using a distributed setup.  In particular, this work presents a practical extension for the Shampoo algorithm by (1) using only a left or right preconditioner for large layers (2) computing inverse pth roots via coupled Newton iteration algorithms (3) distributing preconditioner computation across CPU cores in a CPU-GPU/TPU cluster and (4) delaying preconditioner computation to occur only once per several steps.  The proposed modifications lead to an implementation of Shampoo that consistently decreases the number of training steps and in certain cases provides a direct wall time improvement over Adagrad/Adam.    \n\n######################################################################\n\n2. Strengths\n\n2.1. The proposed modifications do seem to make Shampoo practical when training large models on large datasets using a distributed setup, and the authors present evidence that their implementation does converge faster (in number of steps and occasionally wall time) than Adam and Adagrad.  The authors also present an example of CTR prediction where their method produces state of the art results.  \n\n2.2. A few of the modifications are mostly well founded: namely, distributing preconditioner computation across unused CPUs and using the coupled Newton iteration for computing the inverse pth root.  \n\n######################################################################\n\n\n3. Limitations/Questions\n\n3.1. Some of the modifications seem to be heuristic and I found it unclear why these should work well in practice.  In particular, while I found Lemma 1 an interesting theoretical result, it was unclear to me why the proposed preconditioner was actually a good approximation for the full preconditioned matrix from Adagrad.  Does this somehow follow directly from the regret bound provided in the appendix? It appears that these regret bounds only show that Shampoo with the provided extensions converges, but there could be still be a gap between the performance of the modified Shampoo and original Shampoo.  It would be helpful if the authors could comment on this.  I felt similarly about the delayed preconditioner computation: The fact that computing the preconditioner only once every several epochs is interesting, but I didn't quite follow why this should have worked in the first place.  Is there some further intuition the authors could provide for this phenomenon?  \n\n3.2. While the authors demonstrate that their modifications worked well on the settings they described on the bottom of page 2, it appears that there are a lot of elements to tune (described in the appendix) to ensure that their method works well in these settings.  It would be helpful if the authors could clarify a bit more about the basic elements of their method that require significant tuning.  For example, while I understand that there are several works using learning rate scheduling to achieve state of the art results, in most experiments, Adam can be used with a learning rate of 1e-4.  In the proposed method, how often (in number of steps) should one compute the preconditioner matrix? I understand that there is a plot comparing this in Figure 3c, but on smaller datasets, do we actually need to compute the preconditioner often in order to get convergence? \n\n3.3. Point 3.2 leads to my next concern about the work.  Namely, if the preconditioner does need to be computed more frequently, this would lead to a significant increase in wall time on smaller datasets.  In particular, the authors even mention that there is in fact no wall time improvement on full ImageNet due to a lack of benefit of amortization (I'm assuming this is referring to the frequency of preconditioner computation prior to convergence, but please correct me if I'm mistaken).  Thus, the practicality of the proposed method seems limited to extremely large datasets, and may be of practical use only to large organizations with access to resources capable of handling such large datasets.  However, for several of these organizations, it may be easier to just parallelize a simpler method across more GPUs/TPUs.  Thus, I would be interested in hearing more of a discussion about concrete hardware/software improvements that would boost the performance of the current algorithm over currently used algorithms.  For example, is the main suggested improvement to provide better support for double precision computations? I feel that most other hardware improvements would benefit first order methods anyway.  \n\n\n######################################################################\n\n4. Score and Rationale\n\n\nUpdates: I have changed my score in light of the author responses and edits and now am in favor of accepting this work.  \n\nCurrently, I vote for rejecting.  However, my decision is borderline.  My main concerns with the work are (1) the practical relevance of the work as compared to first order methods and (2) the rationale around the heuristics proposed in the current method.  In particular, I am concerned that while the proposed method consistently improves the number of steps, it does not necessarily improve wall time computations even on ImageNet (which would be considered a large dataset by many organizations).  I am also unsure as to how software/hardware benefits could boost the performance of the current method without significantly boosting the performance of first order methods as well.  I am definitely open to changing my review provided that the authors are able to provide clarification around these and other points in section 3. \n\n\n######################################################################\n\n\n5. Minor Comments\n\n5.1. Apart from the points in section 3, it would be helpful to have a discussion section at the end of the work. \n\n5.2. I believe the y-axis label is missing in Figure 3 (though I believe it should just be -log(perplexity)).  \n\n\n######################################################################\n\nPost-Rebuttal Updates:\n\n6. Thank you for running the experiments on small datasets and for clarifying my additional concerns.  In light of your clarifications, I do find this work to be a practical extension of Shampoo and am in favor of accepting the work.  I am still surprised by the fact that delayed pre-conditioner computation still leads to improved convergence over 1st order methods, but the presented empirical evidence demonstrates a consistent benefit in wall time across a number of settings. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "XacLjho8D8Z", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Main idea: The paper developed a practical second-order preconditioned method, which is based on the Shampoo algorithm, to improve the wall-clock time compared with state-of-the-art first-order methods for training deep networks. \n\n\n1.\tThis paper is heavily drawn from the shampoo algorithm, algorithm designing, and even the later theoretical analysis (such as theorem 3). It appears that the iterative preconditioning method is only incremental improvement over the Shampoo algorithm. However, the iterative method is widely used in the similar spirit of the limited memory of BFGS or LU decomposition.\n2.\tThe main text of this paper has only two lemmas. What is the main theorem of this paper?  \n3.\tEven though the paper mentioned wall-time everywhere, all the improvements shown in the figures, compared with Adagrad and Adam, are in terms of steps, such as Figure 2(a), Figure 3 - 5 and Table 1. The authors should directly show wall-clock time.\n4.\tSome minor comments: fonts in Figure 1 are too small; in other figures, the fonts in x-axis and y-axis are also very small; in figure 2 (b), the author should explain more about how to break down the latency of a single step into the parts mentioned in the figure.\n\n##################################\n\nUpdates: I have changed my score after reading the author's responses.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Um8jFmyJB0V", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper tackles the important problem of making second-order methods practical for stochastic optimization. The main contributions of the paper include an improvement to a previously proposed Shampoo algorithm that is, on a high level, an extension of AdaGrad to include a second matrix moment (such as covariance of gradients).\n\nThe paper addresses challenges along three axes. On the algorithmic level, computing the preconditioners is expensive both in terms of memory and compute and the authors suggest improvements theoretically motivated techniques to make this faster. On a numerical level, the authors put forth a case for investing in 64-bit arithmetic in hardware implementations required for computing the p-th root. This is an important point to make considering that the field has been trending towards the opposite direction and risks a local maxima. On the distributed implementation level, the authors implement the preconditioning computation (done every few steps) asynchronously on the host CPU owing to lack of double precision support on the accelerators attached to the host (GPU/TPU).\n\nExperiments are run on a varied set of tasks (translation, CTR, NLP, vision) and they show fairly significant reduction in the number of optimization steps. Overall, the paper tackles the important problem of second order optimization in a more practical way. The motivations, main contribution and supporting experiments are well laid out.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Dear authors,

I like to topic of your paper very much. Indeed, your work is trying to show that 2nd order methods can be efficiently implemented in a distributed environment and can achieve improvement in training times.

However, having worked on distributed computing for many years, I personally think that reporting running time in your work is not very informative (without mentioning what hardware is used during computation), and one cannot understand the connection or reproduce your results. Also, it was not clear how the baselines were implemented, and how the hyper-parameters were tuned. It is also not clear why you haven't picked better benchmarks to compare your work. 

I think that addressing both the concerns from reviewers and the one mentioned above would improve the paper significantly. I would really like to see if accepted in the near future after these issues are fully addressed.

Decision: Reject