Extreme Tensoring for Low-Memory Preconditioning 
Authors: Xinyi Chen
Naman Agarwal
Elad Hazan
Cyril Zhang
Yi Zhang
Venue:  (2020)
Keywords: optimization
deep learning

Abstract: State-of-the-art models are now trained with billions of parameters, reaching hardware limits in terms of memory consumption. This has created a recent demand for memory-efficient optimizers. To this end, we investigate the limits and performance tradeoffs of memory-efficient adaptively preconditioned gradient methods. We propose \emph{extreme tensoring} for high-dimensional stochastic optimization, showing that an optimizer needs very little memory to benefit from adaptive preconditioning. Our technique applies to arbitrary models (not necessarily with tensor-shaped parameters), and is accompanied by regret and convergence guarantees, which shed light on the tradeoffs between preconditioner quality and expressivity. On a large-scale NLP model, we reduce the optimizer memory overhead by three orders of magnitude, without degrading performance.

Published as a conference paper at ICLR 2020
EXTREME TENSORING FOR LOW-MEMORY PRECON- DITIONING
Xinyi Chen
Naman Agarwal
Google AI Princeton Princcton, NJ xinyicegoogle.
Google Al Princeton
Princeton, NJ
namanagarwal hawpgggawwwww@gggglooo @gw.glg. com
Elad Hazan
Cyril Zhang
Princeton University & Google Al Princeton Princeton, NJ
Princeton University Google Al Princeton
Princeton, NJ cyril. www.ggprinceton.ncooomoo edu
Raazaness.prncces edu
Yi Zhang
Princeton University
Princeton, NJ y 2.azgaagssss.co wrwp:ettosriicneto..c..
ABSTRACT
State-of-the-art models are now trained with billions of parameters, reaching hard- ware limits in terms memory consumption. This has created a recent demand for memory-efficient optimizers. To this end, we investigate the limits and per- formance tradeoffs memory-efficiet adaptively preconditioned gradient meth- ods. We propose extreme tensoring for nigh-dimensional stochastic optimization, showing that an optimizer needs very little memory to benefit from adaptive pre- conditioning. Our technique applies arbitrary models (not necessarily with tensor-shaped parameters), and is accompanied by regret and convergence guaran- tees, which shed light on the tradeoffs between preconditioner quality and expres- sivity. On a large-scale NLP model, we reduce the optimizer memory overhead by three orders magnitude, without degrading performance.
INTRODUCTION
Among the most influential and important optimization techniques in machine learning are adaptive learning-rate methods, otherwise known diagonal-matrix adaptive preconditioning. Essentially all of the most-commonly used incarnations of adaptive preconditioning (AdaGrad, Adam, RMSprop, Adadelta, etc.) accumulate recond-moment estimators of each coordinate of the gradient, then scale the parameter updates by the square roots of these accumulators. These methods come with an overhead memory cost storing these accumulators, thereby doubling the memory consumption. In the regime where model size encroaches upon the same order of magnitude as the total amount of RAM, need has arisen to view memory as limited resource in large-scale optimization.
We address the question of whether the benefits of adaptive preconditioning can be attained with- out significant memory overhead. To this end, we introduce extreme tensoring, family of generic modifications to any conditioner which takes the form of tensor product arbitrary order, with simple updates, without necessarily requiring the parameters to be tensor-shaped. In regret analysis, quantify how extreme tensoring competes provably with full-memory AdaGrad in the online convex optimization framework, with multiplicative data-dependent constant that can be measured empirically.
based adaptive optimizer. Our method uses compressed pre-
In a large-scale language modeling setting, we demonstrate that an optimizer requires very little additional memory benefit from adaptive preconditioning. Furthermore, the inherent flexibility

---

Published as a conference paper at ICLR 2020
of our method enables us to conduct, to the first of our knowledge, the first empirical study of the tradeoff between training convergence and memory in the optimizer.
1.1 RELATED WORK
The most widely-adopted form of adaptive preconditioning are second-momente moment-basedt for example, AdaGrad (Duchi et al., 201 I), Adam (Kingma and Ba, 2014), RMSprop (Tieleman and Hinton, 2012), and Adadelta (Zeiler, 2012). Some recent preconditioning methods are not based on second moments (Bello et al., 2017; Chen and Gu, 2018; Bernstein et al., 2018), and fall beyond our scope.
fentoref preconditioners. Many works (Martens and Grosse, 2015; Gupta et al., 2018; Martens et al., 2018) investigate lensor-factorized preconditioners. These are presented in the view of restricted full-matrix preconditioning (vs. diagonal for us) of tensor parameters (vs. general- purpose for us). In parameter count regimes relevant to this paper, these full-matrix adaptive meth- ods suffer from prohibitive time and memory overhead issues. Extreme tensoring can be viewed as the diagonal version of Shampoo. The matrix version of diagonal Shampoo was discussed in the appendix of (Gupta et al., 20; in this paper we extend the algorithm to arbitrary tensor factor- izations, and focus on the tradeoff between memory consumption and optimizer performance. The theoretical part of our work follows proof techniques seen in (Gupta et al., 2017) and (Gupta et al., 2018), although our diagonal restriction results in distinct updates and incomparable regret bounds.
Optimizer memory reduction in deep learning. Perhaps the most closely related work to ours is Adafactor (Shazeer and Stern, 2018), an empirical work which achieves tive regularization by restricting preconditioners on matrix-shaped gradients to "row" and "column" learning rates, with similar update rule. Similar algorithms have appeared in prior work (Gupta et al., 2014; Shazeer et al., 2017). As special case of our regret analysis, we provide some theory method. Con-
memory adap-
for this line of work, while proposing a more general and versatile memory-reducte current to this work, Anil al. (2019) are also motivated by memory efficiency, and the algorithm uses a different approximation to the preconditioner than a tensor product. Their algorithms' index partitions compatible with our tensor indices, and provide comparison in Appendix C.
Optimizer memory reduction outside of deep learning. Older than our work is the classic L- BFGS algorithm (Liu and Nocedal, 1989), quasi-Newton method with tunable memory window hyperparameter. This does not immediately apply to our setting, as the tunable memory overhead is (window size)x, rather than (1 + =)x. Furthermore, second-order methods turn out to be brittle un- der stochastic gradients; it remains an open empirical problem to robustly exploit Hessian curvature of the loss surface in modern deep learning settings.
Intrinsic dimensionality in deep learning, More broadly, there are various empirical attempts locate and characterize low-dimensional phenomena hiding within dibbeeimnnionall parameter spaces in deep learning (Li et al., 2018a;b) 8a;b). Our comparison provides a lens on measuring the intrinsic dimensionality of the preconditioner. Another (Agarwal al., 2020), a different way to impose parameter bottlenecks on adaptive preconditioners; however, the proposed method does not reduce memory consumption.
Our contributions:
Using the user-selectable degree of memory reduction, we conduct an empirical study of training convergence vs. preconditioner memory usage. We find, in a realistic large-scale language modeling setting, that reallocating memory savings from the optimizer to the model itself can pay off in terms of end-to-end performance.
We propose extreme tensoring as a modification to AdaGrad, Adam, etc. for reducing the overhead memory cost in adaptive preconditioning, accompanying algorithms such as (Shazeer and Stern, 2018; Anil et al., 2019).
We derive a regret bound for extreme-tensored AdaGrad, and provide the first numerical comparisons of regret bounds in deep learning. Though they obviously do not provide convergence guarantees in the non-convex setting, they act as proxy for understanding the AdaGrad accumulator overestimation induced by memory constraints. Via this diagnostic lens, we see that the regret competes with AdaGrad within a multiplicative constant, which we measure be small in practice.
2

---

Published as a conference paper at ICLR 2020
ceshage
x(T() x(6
Figure I: Illustration of tensor index I: the vector x is reshaped into tensor x of dimensions, and the bijection I maps the ith coordinate of x, z[i], to x(I(i)).
PRELIMINARIES
2.1 STOCHASTIC OPTIMIZATION AND ADAPTIVE METHODS
We will state our algorithmic contribution in the usual framework of stochastic optimization of differentiable function f(.), equipped with an unbiased stochastic gradient oracle V.((). At each of T iterations, the algorithm queries stochastic gradient 9: := (zz). In practice, our target is the large-scale ERM setting, where is the population loss training data, and Vf is mini-batch stochastic gradient with respect to the model parameters.
However, since the introduction of AdaGrad (Duchi et al., 2011), the standard analysis of adaptive regularization has been the online convex optimization (OCO) framework (see, e.g. Gupta et al. (2017: 2018)). this language, online decisions chosen from a convex set K, and the as a sequence of gradients of adversarially chosen convex loss functions ft. We are interested minimizing the regret, defined
view the æ; queried by the learning algorithm
sequence of
T
S((e)) min Σ f(( aEK
This is a generalization of stochastic convex optimization; for a broad survey, see Hazan (2016). A framework for reducing the stochastic non-convex setting the online convex setting is intro- duced in Agarwal et al. (2018), translating a sublinear regret bound to convergence to approximate stationary points.
2.2 TENSORS AND INDEXING CONVENTIONS
We introduce some notation and terminology for clarity and convenience.
By [n] we denote the index set {l,.... n2. Ia and O, refer to the d x d identity matrix and d- dimensional zero vector, respectively. We use denote the trace inner product (Hadamard prod- uct) between two matrices, and to denote the element-wise product of two vectors, or vector and a scalar.
As is standard in optimization contexts, tensor of order is an element of R41x... p- dimensional table of real numbers, indexed by a p-tuple / (I1...-4p 4p) of integers, where 1, E [d.]. Our methods will require the user to specify a relabeling of gradient vectors, which we will call tensor index:
Definition 2.1. Let di....,dp be positive integers whose product is d. A tensor index bijection I: [d] X [d.] between indices for Rd and pd.xx..xd,
We will refer to the "reshape" conversion between vectors Id and tensors x Rddx.... Xd, specified by the index bijection Z. For this, we will use the shorthand notation x := I(x) and 1x). Throughout this paper, will square brackets to refer to vector and tensor indices: for example, the previous definition can be restated as enforcing xii = x(I(i)). See
Figure for an illustration of such conversion.
Although we do not need them to state the simple algorithm, it will be useful in the analysis to introduce few other pieces of notation. We denote the tensor (or Kronecker) product of matrices
3

---

Published as a conference paper at ICLR 2020
B, by A B. Given positive definite matrix A Raxd for any R define the matrix norm ||>||A- Vx Ax and its dual ||=||A V A-1 For square matrix M, (iag(M) refers the diagonal matrix of the same dimensions and diagonal entries as M.
PRECONDITIONING BY EXTREME TENSORING
We begin by presenting our proposal for black-box memory reduction in adaptive preconditioning. We call this extreme tensoring, it takes the view of restricting high-dimensional preconditioner to an arbitrary-order tensor product.
The algorithm implicitly maintains a preconditioner which is rank-one tensor product of the same dimension as given tensor index. A second-moment accumulator is maintained for sums of squared gradient magnitudes across each (p- - -)-dimensional tensor slice; the adaptive step size for each coordinate is the inverse square root of the geometric mean of its corresponding slice sums. The formal specification of the base algorithm (AdaGrad with extreme tensoring) given in Algorithm
Algorithm AdaGrad with extreme tensoring
1: Input: Initializer z1, learning rate schedule {mc}. tensor index I with dimensions (d1,..., dp),
2: Initialize (S(,),..., (0a, (0d,..... ...., Oa,,)
3: for = 1,...,''s do
4: Receive stochastic gradient g:
5: Reshape: g: := -(g)))
6: Accumulate slice sums:
Vi E (p),
[ - 5
Σ
g: [7
7: Get step sizes: 6,[4] = (5 + I|: 1= [1))
8: Update: Tc+1 f- 1/: I (6.) 9t
9:
end for
We make few important remarks:
AdaGrad is special case of Algorithm 1, with 1,d, d. The analogues for Adam, RMSprop, Adadelta, etc. obtained straightforwardly by decaying the accumulator (S [j] B2 [ + +.... Extreme tensoring is compatible with first-moment esti- mation (i.e. momentum), although the memory savings disappear. In the setting for our main empirical result, removing momentum did degrade performance, corroborating the findings of Shazeer and Stern (2018).
If the tensor dimensions (di....,dp) dp) in the decomposition are close equal, the memory overhead scales Opp d!/").
As is standard in matrix and tensor optimization methods (Gupta et al., 2018; Shazeer and Stern, 2018; Martens al., 2018), independent copies of Algorithm can be run on each tensor-shaped parameter group; optimizer APIs in standard deep learning packages pro- mote this convention. This can viewed as maintaining tensor sum of preconditioners, each of which is p-wise tensor product.
The accumulator update step can be implemented concisely and efficiently using re- indexing procedures (typically called reshape view) in standard deep learning and numerical linear algebra packages.
Lemma 4.3 in the analysis shows that the per-coordinate learning rates underestimates of those prescribed by AdaGrad. This interpretation is used in our proof, and serves as the basis of the empirical measurements in Section 5.3.
4

---

Published as a conference paper at ICLR 2020
4 REGRET ANALYSIS
In this section, we justify the choice of update rules for the compressed preconditioner using the regret-minimizing adaptive regularization framework introduced the original analysis Ada- Grad (Duchi et al., 2011). Throughout the analysis, we will adopt the conventions of online convex optimization, as discussed in Section 2.1. This gives rise to an intuitively appealing and inter- pretable regret bound, which we discuss in this section and revisit in the empirical measurements of Section 5.3.
4.1 STATEMENT AND DISCUSSION OF REGRET BOUND
First, we define and prove the main regret bound. As is standard in the regret analysis of diagonal AdaGrad (see Duchi et (2011)), this most naturally stated with diameter term on 3;; D = max,, Ie x" oo. This can be constrained using a projection onto an foo ball in the algorithm, but this is seldom seen in practice. It will also be helpful to define G to be the diagonal matrix with values
G = Σ g:1112
1:1,
on the diagonal.
Theorem 4.1. Define HT ('L, + ET= G and HT diag(=I + Σ.= 9:9 ) Then, there exists a choice of constant learning rate schedule m = _... = 7) := 17 and > such that the {xt) chosen by Algorithm I satisfy the regret bound
T Σ - (H 1=1
where x² = min, El ft(x)) is the loss-minimizing decision in hindsight.
As a direct consequence, we recover the AdaGrad regret bound as a special case where = 1. Noting further that Hy Hr, we restate this well-known result in our notation, for clarity and comparisons:
Corollary 4.2. setting of Theorem 4.1, when 1, the {w+} chosen by Algorithm satisfy
E f(a)) - f((x))
Thus, AdaGrad with extreme tensoring satisfies a regret bound most V (r((Hr)/T(((r)) times
than that of AdaGrad. This quantity can be as large as Ω ((Vd); this the worst-case price of the dramatic memory savings in the regularizer. However, in the presence of sparsity, this ratio can be much smaller; include a discussion and empirical measurements in Section 5.3.
4.2 PROOFS
The derivation of the regret bound uses the proof framework introduced by Gupta et al. (2017). Our first important lemma states that the regularizer in Algorithm 1 is spectral (per-coordinate) upper bound for the diagonal AdaGrad preconditioner. In other words, this lemma establishes that the per-coordinate learning rates under extreme tensoring underestimates of those dictated by AdaGrad.
gr are tensors of dimension d1 x... x dp, and let = I- for all t, where I is any tensor index whose image is X"=]Idi.. Then for t E [T).
Lemma 4.3. Suppose
diag(eI +
Σ gag. 1/22 - ( + Σ G/P
$=1
5

---

Published as a conference paper at ICLR 2020
Proof. Let d = d1d2. .. Let [d] and let d1....,, ap be such that I(j) = [a1, (01,02," 0g,
appp-
diag(=I + Σ geg. ) |.. (E Σ Be a1, 02. ...
, .p13)pp = IP E
Σ
g 8,1113).
==1 1:l, ma;
Taking the pth root both sides,
diag(eI+ Σ Iag. )i. mP Σ
Σ
g. B.(1)3)//= -T ( Σ
G2
[as u;))
8=1 I:1,=0.
8=1
- + Σ G G)) [j.
8=1
Taking square of the above inequality yields the lemma.
4.3 REGRET BOUND
The following lemma bounds the regret in terms of the quadratic norms in the regularizer and its dual:
Lemma 4.4. Denote by H₃ the d-by-d diagonal matrix whose entries are (5)) Take m ?) == n to be the constant learning rate, and let the loss minimizing choice hindsight.
Then, the regret of Algorithm is bounded by
T
7 (1 x² 21,, |04+1 x² 1171) +++
E
(llaelli,)2
2n
2
Proof. By the definition of Algorithm for any x²,
Tt+1 x=' xg x² HH yt, and H2 - -x)) H1 -x)) -99-
Taking the inner product of the above vectors.
(Ie+1 - x²)² H.( - x²) = (Iz x²) He - x²) - 22(2x - x²)² gt + 12990 Hr gt.
Rearranging and dividing by 2n, we have
1
9e (x+ - x²) = (1 x 112,, - |1 - x" 11, ) (11 1l,ellix))
2T
Using convexity of the loss functions and summing over yields the lemma.
Next, we present lemma that later used to bound the second term in the regret.
Lemma 4.5. (Gupta et al., 2018) Let U1.92.... " be sequence of vectors, and let M, diag(2 9.9.)) ). Let S+ denote the set of diagonal positive definite matrices. Given func- tion &b over St, define A, = argmin sES+ {M, A + +(A)) and assume that the minimum is attained for all Then
T
T
2
(1g. ( + 4(Ar) -(AA))
1=1
We proceed to prove Theorem 4.1.
Proof. (Theorem 4.1) Recall the definition of H, (E + + /2P, bound the first term in Lemma 4.4 as follows.
T
(1 x" 12, - || x'(lir))
<
[1) - x x²²² (H, (1 x²) + ||ll x² 11,,
1=2
<
D2
Tr(H,-H,,1 - D2 Tr(Hi) D2 Tr(HT)
1=2

---

Published as a conference paper at ICLR 2020
The last inequality is due to the fact that H He- Recall H, diag(=I+ [=I+ E )12 and by Lemma 4.3, H. < Ht. Now we use Lemma 4.5 with function $(A) = Tr(A) + (TT(A-)).
Let M, = diag E faAle ) and let S+ denote the set of diagonal positive definite matrices, we have argmin {M, A + ((A)) = argmin,est Tr(Hf A-1 + A) = Hg. The last equality can be derived by minimizing each diagonal entry of H individually.
By Lemma 4.5,
T (lloel17 1
T
T
Σ (1lg Il
o ((H)) - ((Ho) 1 (1l9el (l ft. 1,, + 1 ((H,)
Σ
diag Σ n:9
Hrp Tr(Hr) + + -
We can take the diagonal of Ei 9:96 in the first equality since H1, is a diagonal matrix as well. We proceed to bound the second term in Lemma 4.4. Since il. 1 He and both He and il, are (Tr((r)) Summing up (I1g=ll At )² full-rank, H ž H Therefore E/= (llee 12, E
conclude that the regret of Algorithm I is bounded /Tr(HT)
the two terms and taking 17 Dro (2Tx(Hr)
1 Do Tr(HT) + Tr(r))))) = Do (Tr(HT))TrHa) by 24
EXPERIMENTS
In this section, we provide several empirical studies on extreme tensoring. Our main experiment interpolates between the memorylessness of SGD and the full memory consumption of AdaGrad on a large-scale language model; we additionally isolate the effects of preconditioner expressivity in synthetic experiment of the same form, and provide parallel CIFAR-10 experiment in the appendix.
To accompany the main experimental result, we provide empirical measurements of the competitive ratio in the regret bound, as well as a quick comparison in which the memory savings are used to train larger model.
5.1 MEMORY-PERFORMANTE TRADFOFF IN T.ARGE-SCALE NLP
Our main empirical study focuses on large-scale language modeling with the Transformer architec- ture (Vaswani al., 2017) on the Google Billion Words (GBW) dataset (Chelba al., 2013), and the results are shown in Figure 2. We use the pipeline from the open-source Tensor2Tensor package (Vaswani et al., 2018) for preprocessing the GBW dataset, and use the base Transformer architec- ture in the same repository as our base architecture. The decoder-only model has 6 identical layers with hidden dimension dmodel = 512, and feedforward dimension d, = 2048. addition, the weights are shared between the embedding and softmax layers, and the model has total of ~ 35M parameters.
learning schedule of ?e min(10-6.. t, 3x) and is hyperparameter we tune for each experiment. The learning rate schedule is the same as the one used in (Shazeer and Stern, 2018): 8): linear warmup stage, followed by inverse square root decay. We train each model for 500K steps on GPUs, with a max sequence length of 256 tokens, and a max number of 4096 tokens in a batch. Global learning rates are selected by hyperparameter search.
For experiments,
Except in our comparison with Adam, we do not use momentum our extreme tensoring experi- ments. Momentum incurs an overhead memory cost linear in the model dimension. In addition, we empirically tested the benefit of exponentially decaying the second moment estimator (82 < 1 in the language of Adam and Adafactor). We found that in language modeling experiments, decaying the second moment estimator did not contribute to better performance; on the other hand, the vision experiments in the appendix do use this decay (32 0.99).

---

Published as a conference paper at ICLR 2020
52
50
550
Optimizer
Parameter count
Final ppl.
48
andeleepppinn uvelperplaityy
AdaGrad
3.5 10'
41.18
66
ET3
ET1
1.2 x 10"
39,84
64-
ET2
10
42.10
62-
ET2
ET3
5.0 x 103
45.19
ETxg
90
47.94
46
SGD
49.80
Adam
7.0 x 10"
38.47
107
10a
10
107
Adafactor
1.2 x 10"
38.86
Parameter count
Figure 2: peroory-erformmnce tradeoff for GBW language modeling with Transformer network. Cyan points: performance after dou- bling model size using optimizer RAM savings.
Table 1: Performance and memory comparison
between adaptive optimizers for GBW language modeling with a Transformer network.
Extreme tensoring gives us class of optimizers which interpolate smoothly between AdaGrad and SGD. Between these endpoints, we choose three levels of extreme tensoring (denoted by 2,3)), with tensor index dimensions fully specified in the appendix. While extreme ten- soring beyond level 2 have limited memory savings in most scenarios, this family of optimizers allow us to have additional interpolating points to investigate the memory tradeoff. The dark blue dots in Figure 2 records the training performance for these interpolating algorithms. Note that the horizontal scale is logarithmic, so extreme tensoring saves memory by orders magnitude. We also provide training results for Adam (which consumes and Adafactor (similar to ETI but with a different step size scaling) in Table 1, for completeness; however, these are not part of the interpolation study.
memory for first-moment estimation)
To provide an additional interpolating point, we include in the comparison a closely related algo- rithm, which selects single learning per parameter group (e.g. embedding, bias, convolutional layer); thus, the preconditioner is a tensor sum of scalar multiples of the identity matrix. Each such scalar is chosen be the inverse square root of the sum of squared Eg norms the parameter group over time; it is easily seen that this achieves the same regret as online gradient descent (Zinkevich, 2003), so we include as the least granular adaptive optimizer, and call ET Note that this algorithm is known, and Ward et al. (2018) in particular provide fine- tne-grained analysis.
Finally, in Appendix 6, we re-run the ET 2019), finding the same monotonic trend. In addition to model performance, we also provide in the appendix wall clock comparisons for the algorithms we have benchmarked.
experiments with the SM3 update rule (Anil al.,
5.2 DOUBLING THE MEMORY ALLOWANCE
As an extension of the above comparison, we argue that the memory consumption freed up by choosing a compressed adaptive optimizer can be usefully reallocated to training a larger model. We double the number of layers of the Transformer network used in the previous section, keeping all embedding sizes the same; this results in network with 56M parameters.
Results are shown in Figure and Table 2. where final perplexities are given with the same running time allowance as the corresponding main experiment (middle column), as well as the same iteration count (500K steps; right column). The light blue dots in Figure 2 plots the final evaluation perplexity of the larger model in the third column of Table 2 for a more intuitive visualization. Holding the memory consumption constant, we find that it pays off to use larger model rather than adaptive optimizer that stores all of the second- moment accumulators. Training the larger model with extreme tensoring consumes essentially the same amount of memory training the original model with Adam AdaGrad, but the performance improves significantly. Even holding the running time constant, the larger models are competitive with the fully-converged smaller ones.
Table 2 reports performance from the larger model,
8

---

Published as a conference paper at ICLR 2020
Optimizer
Final ppl. (time)
Final ppl. (iters)
ETI
39.25
36.23
ET2
43.81
40.04
ET3
44.70
40.45
ET
42.95
41.68
Trace of regulanterr
10a
Trondl
20'
Table 2: Comparison of memory-efficient op-
30*
timizers on double-sized model, so that to- tal memory consumption is lower than that of AdaGrad/Adam on the smaller model. Final perplexities are given with the same running time allowance as the corresponding main ex- periment (middle column), well as the same iteration count (500K steps; right column),
301
ET2
AdaOrad
Figure Comparison of quantities in the nu- merical regret bounds. Since the vertical scale logarithmic, the multiplicative regret bound gap compared to AdaGrad is half the height dif- ference between blue and red bars.
SGD
ET2
Final population loss
Population Ingistic
AdaGrad
SGD
ET1
ET2
10a
ET1
10a
10a
200
400
600
800
1000
10a
102
10*
Parameter count
Figure 4: Training curves and final loss comparison for a convex problem with synthetic data. Left: Loss curve for each optimizer. Right: Final loss vs. optimizer parameter count.
5.3 EMPIRICAL MEASUREMENTS OF THE REGRET BOUND
An appealing part of our regret analysis is the interpretability of the trace quantities in Theorem 4. Figure 3 shows comparison of traces Tr ((( and Tr(Hr) of the final regularizers in the language modeling experiment. Then, by Theorem 4.1 and Corollary 4.2, the upper bound for AdaGrad's regret scales with the latter quantity, while that of extreme tensoring scales with the geometric mean of the two traces. Note that the vertical scale of the figure is logarithmic, so the multiplicative regret bound gap compared AdaGrad is half the height difference between blue and red bars.
Intriguingly, the multiplicative gap, which depends on the loss function and training trajectory via the sequence of gradients, appears to be very small in practice. In the case of one-level extreme
tensoring, we have / (TTHHHr)TT((EE =e 5.7.
5.4 COMPARISON ON SYNTHETIC DATA
In this section, we exhibit simple experiment with a convex optimization problem, in which there clear tradeoff between preconditioner quality and expressivity. We generated Gaussian data ar E 12512 with large condition number (~ 104), and Gaussian matrix W E IR10x512 Labels are generated according to the log-linear model Pr(y: = j] x exp((Wa.),)) Then, the optimization problem in W is to minimize the empirical negative pog-probability under the same model. Our findings are robust with respect to the batch size; we use the full gradient on 10 samples in our plots for clarity.
As in the large-scale experiment, we use successively deeper tensor factorizations of the precondi- tioner, along the feature dimension of the matrix W. For depths 1, 2, and 3, choose tensor indices of dimensions (10, 512), (10, 16, 32), and (10, 8, 8, 8), respectively. Global learning rates are tuned individually. Results shown in Figure 4.

---

Published as a conference paper at ICLR 2020
Algorithm Wall-clock time (hrs)
Adafactor 20.64
Adam 18.26
AdaGrad 17.84
ET1
ET2
EТ3 24.82
ET 19.00
SGD 17.50
18.61
20.17
Table 3: Wall clock comparisons for optimizing the base Transformer model.
5.5 IMPLEMENTATION DETAILS
For completeness, we provide some details on our implementation of the algorithm. Table 3 gives the wall clock comparisons for optimizing the base Transformer model. We point out that the ET algorithms incur some computational overhead due additional tensor operations. We have not optimized our implementations, and expect that the time efficiency gap could be reduced with architecture-aware optimizations. We trained the model on one V100 GPU, with parallel hyperpa- rameter tuning on the global learning rate multiplier. We used the Tensor2Tensor (Vaswani et al., 2018) framework for training, and Vizier (Golovin et al., 2017) for hyperparameter tuning.
6 CONCLUSION
We have introduced extreme lensoring, a modification to any second-moment-based adaptive opti- mizer which drastically shrinks memory overhead. Our experiments characterize performance- memory tradeoff in the optimizer, and demonstrate the possibility of negligible memory overhead without degrading convergence. Our regret analysis provides a competitive ratio with uncompressed adaptive methods, giving an additional empirical lens on data-dependent preconditioner quality.
REFERENCES
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. The case for full-matrix adaptive regularization. arXiv preprint arXiv: 02958, 2018.
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint 2xXiv:2002./8803 2020.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory-efficient adaptive optimiza- tion for scage learning, 2019.
Irwan Bello, Barret Zoph, Vijay Vasudevan, and Quoc V Le. Neural optimizer search with rein- forcement learning. arXiv preprint arXiv:100.00411,,00
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar: signsgd: compressed optimisation for non-convex problems. arXiv preprint arXiv:1802.04444, 2018.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint 131..333 2013.
Jinghui Chen and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. arXiv preprint wrXiv:1806.06763. 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal Machine Learning Research, 12:2121-2159, 2011.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D Sculley. Google vizier: service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 1487-1495. ACM, 2017.
Maya R Gupta, Samy Bengio, and Jason Weston. Training highly multiclass classifiers. The Journal of Machine Learning Research, 15()):1661-149) 2014.
Vineet Gupta, Tomer Koren, and Yoram Singer. unified approach to adaptive regularization in online and stochastic optimization. arXiv preprint arXiv: / 706.066669 2017.
10

---

Published as a conference paper at ICLR 2020
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In Interaational Conference on Machine Learning, 2018.
Elad Hazan. Introduction to online convex optimization. Foundations and Trends in Optimization, 2(3-4):157-325, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1r12v6:412.66880 2014.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. arXiv preprint arXiv:/804.088.1 2018a.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. In Advances in Neural Information Processing Systems, pages 6389-6399, 2018b.
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical Programming, 45(1-3):503-528, 1989.
James Martens and Roger Grosse. Optimizing neural networks with kronecker factoree approximate curvature. In International Conference on Machine Learning, pages 2408-2411 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kroneeker-factored curvature approximations for recurrent neural networks. 2018.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
In International Conference on Machine Learning, 2018.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gan gated mixture-of-experts layer. arXiv preprint 1rxiv:1701.065:0 2017.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26~ 31, 2012.
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 Infor- mation Processing Systems, pages 5998-6008, 2017.
Ashish Vaswani, Samy Bengio, Fugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Eukasz. Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. CoRR, abs/1800.00446. 2018. URL http:///arxx...rrr//////113 100.///1113.11166...
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes, from any initialization. arXiv preprint 18x1::1806.018/// 2018.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701, 2012.
Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In
Proceedings of the 20th International Conference Machine Learning (ICML-03), pages 928- 936, 2003.
11

---

Published as a conference paper at ICLR 2020
A CIFAR-10 10 EXPERIMENT
In this section, we evaluate the memory-performance trade-off of our proposed algorithm on the CIFAR-10 dataset (Krizhevsky, 2009). Specifically, we compare the test accuracy of a ResNet-18 model (He et al., 2016) trained with SGD, Adam and levels of extreme tensoring.
A.1 SETUP
With each optimizer, the model is trained for 150 epochs with optimally tuned constant learning rate. To prevent a 2X memory overhead, no momentum is used for any of the optimizers (i.e. 31 = 0 in Adam). We batch size 128, weight decay 5 10-1 for all the experiments in this section.
A.2 TENSOR INDICES
We show tensor decomposition scheme in the table below. Note that bias parameters not shown in each layer in the table since are treated as vectors and thus not decomposed at all.
Parameter Shape
Shape (ETT)
Shape (ET2)
Shape (ET3)
(64, 3, 3,3) (64, 3, 9)
(8, 3, 9)
(8, 8, 3.9. 9)
(64, 64, 3,
(64, 64, 9)
88,8,8,8, 8,
(8,8, 8, 8, 9) (8,4,4,,,8,9
(128, 6,33
(128, 64, 9)
(8, 1,,88,899
(128, 128, 3,3)
(128, 128, 9)
(8,16, 16,9)
8,4.4.8,4.4.9)
(256, 128, 3,3)
(256, 128,9 9)
(16, 16, 16,9
4,8,4.448,44441
(256, 256, 3,33
(256,, 256, 9)
(16, 16, 16, 16,9)
4444,444,,44444440
(512, 251,,266,3 3,
(512, 256, 9)
(32, 16, 16, 16,9)
4.4.4444.4444444499
(512, 512,3,3
(512,512,
(32, 16, 32, 9)
(8, 8,,,444,8,4,44444
(128, 64, 1)
(128, 64)
(16, 8, 8,
(4,4,8.8,8)
(256, 128, 1, 1)
(256, 128)
[16, 16, 16, 8)
44444444.44444.8)
(512, 128, 1)
(512, 128)
(32, 16, 16, 8)
4,4,44444444.
Table 4: Tensor indices used for different levels of extreme tensoring for the ResNet- 18 model on Cifar-10.
A.3 RESULTS
In this section, we report the best test errors seen in the first 150 epoch along with the parameter count in the optimizer. A similar trend of trade-offs between performance and parameter count is observed.
Optimizer
Parameter count
Final test error
Adam(8
0)
1.1 x 10"
8.40
ET1
2.3 x 10a
7.22
ET2
1.6 x 104
8.49
ET3
1.5 x 10*
8.52
ETоo
62
8.57
SGD
9.27
Table 5: Performance and memory comparison between adaptive optimizers for Cifar- 10 classifica- tion with ResNet-18 network.
12

---

Published as a conference paper at ICLR 2020
20
3875
3.70
Figure 5: semorryperformance comparison for CIFAR-10 10 classification with an 18-layer ResNet: final test error vs. optimizer parameter count. Note that the horizontal scale is logarithmic.
B LANGUAGE MODELING EXPERIMENT DETAILS
In this section, we provide more details on the tensor decomposition in the main empirical result. The following table records parameter shapes for the Transformer model in main experiments, and the tensor shapes of the reshaped parameters.
Parameter
Parameter Shape
Shape (ETI)
Shape (ET2)
Shape (ET3)
Q, K, V, position-wise FF
(512, 512)
(512, 512)
(16,32, 16, 32)
44,4,4,8, 4,4,4,,,
Embedding Weights
(2000, 512)
(2000, 512)
(40, 50, 16, 32)
(5,8, 5, 10,,4,4,44
Layer Norm
(512.)
(512.) )
(16,32 (6,32)
(4,4,
FC Weights
(512, 2048)
(512, 2048)
(16, 16,32, 32, 64)
8,4,,,884,888888
FC Bias
(2048,) )
(2048,)
(32, 64)
(4,8,
FC Weights
(2048, 512)
(2048, 512)
(32, 64, 16, 32)
(4,8, 8,
PC Bias
(512,)
(512.)
(16,32)
(4, (4,4,4,8) 4,8)
B.1 WALL CLOCK COMPARISONS
In this subsection we provide the wall clock comparisons for optimizing the base Transformer model. We trained the model on one V100 GPU, with parallel hyperparameter tuning on the global learning rate multiplier. We used the Tensor2Tensor (Vaswani et al., 2018) framework for training, and Vizier (Golovin et al., 2017) for hyperparameter tuning.
Algorithm
Adafactor
Adam
AdaGrad
ET1
ET2
ET3
ET
SGD
Wall-clock time (hrs)
20.64
18.26
17.84
18.61
20.17
24.82
19.00
17.50
C SM3 COMPARISON
An alternate way to modulate the memory- perforyner tradeoff is to use the user-selectable parti- tions in SM3 (Anil et al., 2019), which provide a alternate scheme for overestimating the AdaGrad preconditioners. These are compatible with the tensor factorizations in our work, so we provide Table 6 the final test perplexities for the same Transformer experiments on the ET (T(1,22 2, de- compositions (denoted as 2,1,233) 3}). We used SM3-II reported in the pseudocode of Anil et al. (2019). Note that we have not attempted to retune any hyperparameters other than the global scalar learning rate multiplier. As such, these experiments should not be viewed comparison between these families of optimizers; rather, these experiments corroborate a loss of adaptivity as the precon- ditioner is constrained to extremely low dimension.
13

---

Published as a conference paper at ICLR 2020
Optimizer
Parameter count
Final ppl.
AdaGrad
3.5 107
41.18
S1
1.2 x 10
41.59
S2
1.0 x 104
43.25
S3
5.0 x 10
49.45
SGD
1
49.80
Table 6: Results of Transformer experiments with SM3 preconditioner estimates. The same mono- tonic behavior is exhibited as when the extreme tensoring granularity is varied.
14


Reviews:
[{"review_id": "r1lCAy65cH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "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": "============= Update after rebuttal\n\nThanks for the clarifications; I have updated by score and recommend acceptance.\n\nPlease do not forget to implement the promise changes in the camera ready version (e.g. about L-BFGS tradeoff work; as well as gap between OCO and non-convex framework).\n\n\n==============\n\nMotivated from NLP applications where models with billions of parameters are used, this paper proposes a *memory efficient* variant of Adagrad by maintaining a rank-one tensor approximation of the \"second-moment\" accumulator normally used in Adagrad. From the theoretical side, a simple regret bound is proved which provides an intuitive quantity quantifying the convergence guarantee loss for memory-efficient version (vs. Adagrad), and which is empirically evaluated to be small on a large-scale language modeling task (Section 5.3). From the empirical side, the memory vs. generalization performance tradeoff is evaluated on this language modeling task, showing that similar performance to Adagrad can be obtained with much a smaller (optimization overhead) memory footprint. A convex toy task also indicate a similar result.\n\nI like this paper, I am leaning towards acceptance. The write-up can be improved in a few places (see detailed comment below); but overall, I find the idea refreshing for optimization and the proof is simple and elegant.\n\nThe main motivation for the paper is that these large models used in NLP are often making us hit the memory limitation of the hardware just to store the model, and so one then has to tradeoff the size of the model with the memory requirement of the *optimization* algorithm (e.g. one step-size accumulator per dimension for Adagrad). Trying to get the gains of an adaptive preconditioned gradient method but with lower memory footprint thus seems valuable (and Section 5.2 which compares doubling the size of the model + using their memory efficient method vs. original model + Adagrad highlights the gains one can get).\n\nOn the other hand, I wish the paper provided a bit more intuition on when the rank-one tensor approximation structure will give good results (e.g. Section 5.3 gives an empirical measure; but what about some simple theoretical examples which give low values, to provide more intuitions?). In particular, the main inequality relating the Adagrad step-size with the \"extreme indexing\" step-size is arising from the first equation in the proof of Lemma 4.3 on p.5 -- basically one replace one entry of the squared gradient to store with the whole sum over all entries of the squared gradient except a slice. This inequality can be very loose (O(d) in the worst-case), so it is surprising that one could get good approximation ratios much better than the O(sqrt(d)) worst-case (a square-root is taken afterwards), and I wonder what structure yields this.\n\nI also note that Anil et al. (2019) tackles a similar problem with a different approach, and so ideally a more detailed comparison (theoretically and empirically) would be provided (rather than just one sentence as in the current submission), especially since this appeared on arXiv 9 months before the ICLR deadline. But the authors presented it as concurrent work (and it does not appear published anywhere yet) [and a Google search shows the authors had a first version on arXiv only a few weeks after this one, so this seems righs], and so I decided to not hold it against this submission. Disclaimer: I made my general opinion about the paper before doing these Google searches and finding the paper on arXiv.\n\nI am not very familiar with the modern transformer architecture experiments, so cannot evaluate their quality, but they seem sensible from an outsider's perspective.\n\n== Detailed comments ==\n\n- p.1 \"the first empirical study of the tradeoff between training convergence and memory in the optimizer.\" This should be properly qualified by \"in the Adagrad setup\" or something like that; I am pretty sure it is false *for general optimizers*. For example, I guess there are several papers which empirically studied the memory - performance tradeoff for L-BFGS optimizers...\n\n- Important: I think the paper should be more transparent about the *big gap* between the online *convex* optimization framework and stochastic optimization for a *non-convex* loss. There is only one sentence at the end of Section 2.1 mentioning vaguely that Agarwal et al. (2018) provide a reduction from stochastic non-convex setting to the online convex setting, but this is buried in the appendix of the said paper, with several caveats, and for a modified algorithm (i.e. as far as I understand, no convergence guarantee would be given for Algorithm 1 to find a stationary point on a non-convex loss; but one could apply Algorithm 1 on a series of convex problems to obtain overall guarantees on the non-convex loss [but this is a different algorithm!]). I suggest that either the argument of Agarwal et al. (2018) is summarized in a few sentences at the end of Section 2.1 to clarify the real link; or give more caveats [also just before 4.1] (e.g., my current guess is that we use the convex analysis setup to gain insights on the behavior in a controlled setup; and then just hope that some of it applies in the non-convex setup, even though no guarantees is provided whatsoever). \n\n- Some undefined notation: line 8 of Algorithm 1, it seems they use the dot for the Hadamard product between two vectors. The big dot in Lemma 4.5 is most likely the dot product between matrices (but please mention explicitly for clarity). \n\n- Typos in Section 5.1 & 5.2: several references seem wrong. E.g. Figure 5 is probably Figure 2; Table 4 is Table 1, etc... Please correct!\n\n- Clarity: I suggest to put much more description in the captions of Figure 2, Table 2, etc. (for example it was much clearer in their arXiv version). I was very confused by the extra light blue dot in Figure 2 which is only explained in 5.2; I suggest that it is already mentioned (with forward pointer) in the caption. Use more of p.9 for the clarity sake...\n\n- Appendix B.1: please provide hardware information when you give wall-clock comparison.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "HJx3ljfTFB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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============================== Update after rebuttal =======================================================\n\nI did not have any major concerns about the paper in my initial review, only some suggestions for improving the presentation. The authors have addressed most of these issues in their revision. I would like to keep my score as it is. The work seems simple and sound, but somewhat incremental. To have a more meaningful impact, I strongly encourage the authors to make an optimized implementation publicly available as an easy-to-use, plug-and-play type optimizer.\n\n========================================================================================================\n\nThis paper proposes a new memory-efficient pre-conditioning scheme for stochastic optimizers. The basic idea is to store a coarse-grained pre-conditioner, expressed as a rank-one tensor product, instead of the full-dimensional pre-conditioner typically used in algorithms like AdaGrad. I am not very familiar with prior work in this literature, but the proposed approach seems simple and sound. \n\nA regret bound is provided for the proposed method, however the analysis here seems to be a straightforward application of the results and techniques from a few prior works. So, I was a bit surprised to see so much space devoted to the proofs. These can be safely moved to the appendix in my opinion. Moreover, the bound does not seem to be very useful in practice. For example, in simulations in Figure 3, the proposed ET1 performs better than AdaGrad, but the bound is not able to capture this at all. \n\nThe presentation of the experimental results in section 5 can be improved in my opinion. The authors keep referring to “Figure 5” in this section, but I think this is a typo and these should be “Figure 2” instead. In Figure 2, please indicate on the figure itself what dark blue and light blue colors correspond to (smaller and larger models, respectively). \n\nIn Appendix B, wall clock results are presented for different algorithms. These show that the proposed approach is slower than standard algorithms like Adam or AdaGrad. Please explicitly mention this result in the main text (last paragraph of section 5.1) and discuss why this is the case (is this because of the extra reshaping operations required in the updates?).", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "H1grZ3oJ9r", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "This paper considers the problem of the need for memory-efficient optimizers given the increase in model complexity. They study memory-efficient adaptively preconditioned gradient methods, and see the trade-offs among expressivity and preconditioner quality. They show results on a large-scale NLP model, and show that the memory overhead can be reduced by 3 orders of magnitude, without sacrificing performance. \n\n+ show tradeoff among training convergence and memory in the optimizer.\n+ introduce extreme tensoring -- a modification that can be applied to *any* 2nd moment-based adaptive optimizer. It uses a compressed preconditioner. \n+ extend diagonal Shampoo to tensor factorization \n+ nicely written Related Work.\n+ applied their idea to widely used optimizers, such as AdaGrad, Adam, etc.\n+ the derived regret bound is only a multiplicative constant from the regret bound of AdaGrad.\n+ interesting experiments showing tradeoff among lack of memory (SGD) and full memory (AdaGrad), applied on a real-world machine learning setting of large-scale natural language modeling with Transformers. They show how their extreme tensoring modification can achieve an intermediate ground among lack of memory and full memory.\n- I am missing a plot similar to Figure 4 but instead of being on the synthetic data, being on the application of NLP.\n\nOverall, although I do not have a lot of experience in this area, it appears to be that the introduced family of algorithms can be promising as a nice interpolation among SGD-type algorithms and AdaGrad ones. I would personally have preferred to see a wider range of experiments, considering other application scenarios besides NLP, just to see the wide success of the proposed approach. Also, I would urge the authors to spend a bit more text on explaining the algorithm, so that it is easy for practitioners to try it out.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Post author rebuttal the score of this paper increased.
Discussions with reviewers were substantive and the AC recommends acceptance.
Decision: Accept (Poster)