E XTREME T ENSORING FOR L OW -M EMORY P RECONDITIONING
A BSTRACT
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-efﬁcient optimizers.
To this end, we investigate the limits and performance tradeoffs of memory-efﬁcient adaptively preconditioned gradient methods.
We propose extreme tensoring for high-dimensional stochastic optimization, showing that an optimizer needs very little memory to beneﬁt 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
I NTRODUCTION
Among the most inﬂuential and important optimization techniques in machine learning are adaptive learning-rate methods, otherwise known as diagonal-matrix adaptive preconditioning.
Essentially all of the most-commonly used incarnations of adaptive preconditioning (AdaGrad, Adam, RMSprop,
Adadelta, etc.) accumulate second-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 of 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, a need has arisen to view memory as a limited resource in large-scale optimization.
We address the question of whether the beneﬁts of adaptive preconditioning can be attained without signiﬁcant memory overhead.
To this end, we introduce extreme tensoring , a family of generic modiﬁcations to any second-moment-based adaptive optimizer.
Our method uses a compressed preconditioner which takes the form of a tensor product of arbitrary order, with simple updates, without necessarily requiring the parameters to be tensor-shaped.
In our regret analysis, we quantify how extreme tensoring competes provably with full-memory AdaGrad in the online convex optimization framework, with a multiplicative data-dependent constant that can be measured empirically.
In a large-scale language modeling setting, we demonstrate that an optimizer requires very little additional memory to beneﬁt from adaptive preconditioning.
Furthermore, the inherent ﬂexibility of our method enables us to conduct, to the ﬁrst of our knowledge, the ﬁrst empirical study of the tradeoff between training convergence and memory in the optimizer .
1. R ELATED WORK
The most widely-adopted form of adaptive preconditioning are second-moment-based: for example,
AdaGrad (Duchi et al., 2011), 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.
Tensor-factorized preconditioners.
Many works (Martens and Grosse, 2015; Gupta et al., 2018;
Martens et al., 2018) investigate tensor-factorized preconditioners.
These are presented in the view of restricted full-matrix preconditioning (vs. diagonal for us) of tensor parameters (vs. generalpurpose for us).
In parameter count regimes relevant to this paper, these full-matrix adaptive methods 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., 2018); in this paper we extend the algorithm to arbitrary tensor factorizations, 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 sublinear-memory adaptive regularization by restricting preconditioners on matrix-shaped gradients to “row” and “column” learning rates, with a similar update rule.
Similar algorithms have appeared in prior work (Gupta et al., 2014;
Shazeer et al., 2017).
As a special case of our regret analysis, we provide some theory for this line of work, while proposing a more general and versatile memory-reduction method.
Concurrent to this work, Anil et al.
(2019) are also motivated by memory efﬁciency, and the algorithm uses a different approximation to the preconditioner than a tensor product.
Their algorithms’ index partitions are compatible with our tensor indices, and we provide a comparison in Appendix C. Optimizer memory reduction outside of deep learning. Older than our work is the classic LBFGS algorithm (Liu and Nocedal, 1989), a quasi-Newton method with a tunable memory window hyperparameter. This does not immediately apply to our setting, as the tunable memory overhead is (window size) × , rather than (1 + ε ) × .
Furthermore, second-order methods turn out to be brittle under 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 to locate and characterize low-dimensional phenomena hiding within high-dimensional parameter spaces in deep learning (Li et al., 2018a;b).
Our comparison provides a lens on measuring the intrinsic dimensionality of the preconditioner .
Another is (Agarwal et 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 ﬁnd, 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 modiﬁcation 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 ﬁrst numerical comparisons of regret bounds in deep learning. Though they obviously do not provide convergence guarantees in the non-convex setting, they act as a 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 to be small in practice.
P RELIMINARIES 2. S TOCHASTIC OPTIMIZATION AND ADAPTIVE METHODS
We will state our algorithmic contribution in the usual framework of stochastic optimization of a differentiable function f ( · ) , equipped with an unbiased stochastic gradient oracle (cid:101) ∇ f ( · ) .
At each of T iterations, the algorithm queries a stochastic gradient g t := (cid:101) ∇ f ( x t ) .
In practice, our target is the large-scale ERM setting, where f is the population loss on training data, and (cid:101) ∇ f is a 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 in the online convex optimization (OCO) framework (see, e.g. Gupta et al.
In this language, we view the x t queried by the learning algorithm as a sequence of online decisions chosen from a convex set K , and the g t as a sequence of gradients of adversarially chosen convex loss functions f t .
We are interested in minimizing the regret, deﬁned as
[Equation]
This is a generalization of stochastic convex optimization; for a broad survey, see Hazan (2016).
A framework for reducing the stochastic non-convex setting to the online convex setting is introduced in Agarwal et al.
(2018), translating a sublinear regret bound to convergence to approximate stationary points.
2. T ENSORS AND INDEXING CONVENTIONS
We introduce some notation and terminology for clarity and convenience.
By [ n ] we denote the index set { 1 , . . . , n } .
I d and 0 d refer to the d × d identity matrix and ddimensional zero vector, respectively.
We use • to denote the trace inner product (Hadamard product) between two matrices, and · to denote the element-wise product of two vectors, or a vector and a scalar.
As is standard in optimization contexts, a tensor of order p is an element of R d 1 × ...
× d p : a pdimensional table of real numbers, indexed by a p -tuple I
=
( I 1 , · · · I p ) of integers, where I i ∈
[ d i ] .
Our methods will require the user to specify a relabeling of gradient vectors, which we will call a tensor index :
Deﬁnition 2.1.
Let d 1 , . . .
, d p be positive integers whose product is d .
A tensor index is a bijection
I : [ d ] → × p i =
[ d i ] between indices for R d and R d 1 × ...
× d p .
We will refer to the “reshape” conversion between vectors x ∈ R d and tensors x ∈ R d 1 × ...
× d p speciﬁed by the index bijection I .
For this, we will use the shorthand notation x := I ( x ) and x :=
I − 1 ( x ) .
Throughout this paper, we will use square brackets to refer to vector and tensor indices: for example, the previous deﬁnition can be restated as enforcing x
[ i ] = x
[ I ( i )] .
See
Figure 1 for an illustration of such a conversion.
Although we do not need them to state the simple algorithm, it will be useful in the analysis to introduce a few other pieces of notation.
We denote the tensor (or Kronecker) product of matrices
A, B by A ⊗ B .
Given a positive deﬁnite matrix A ∈ R d × d , for any x ∈ R d , deﬁne the matrix norm (cid:107) x (cid:107) A := √ x (cid:62) Ax and its dual (cid:107) x (cid:107) ∗ A := √ x (cid:62) A − 1 x .
For a square matrix M , diag( M ) refers to the diagonal matrix of the same dimensions and diagonal entries as M .
P RECONDITIONING BY EXTREME TENSORING
We begin by presenting our proposal for black-box memory reduction in adaptive preconditioning.
We call this extreme tensoring , as it takes the view of restricting a high-dimensional preconditioner to an arbitrary-order tensor product.
The algorithm implicitly maintains a preconditioner which is a rank-one tensor product of the same dimension as a given tensor index.
A second-moment accumulator is maintained for sums of squared gradient magnitudes across each ( p − 1) -dimensional tensor slice; the adaptive step size for each coordinate is the inverse square root of the geometric mean of its p corresponding slice sums.
The formal speciﬁcation of the base algorithm (AdaGrad with extreme tensoring) is given in Algorithm 1. − 12 p 8:
Update: x t +1 ← x t − η t · I − 1 ( δ t ) · g t 9: end for
We make a few important remarks:
• AdaGrad is a special case of Algorithm 1, with p = 1 , d 1 = d .
The analogues for Adam, RMSprop, Adadelta, etc. are obtained straightforwardly by decaying the accumulator ( S ( i )
[ j ] ← β 2 · S ( i )
[ j ] + . . . ).
Extreme tensoring is compatible with ﬁrst-moment estimation (i
e. momentum), although the memory savings disappear.
In the setting for our main empirical result, removing momentum did not degrade performance, corroborating the ﬁndings of Shazeer and Stern (2018).
• If the tensor dimensions ( d 1 , . . .
, d p ) in the decomposition are close to equal, the memory overhead scales as O ( p · d 1 /p ) .
• As is standard in matrix and tensor optimization methods (Gupta et al., 2018; Shazeer and Stern, 2018; Martens et al., 2018), independent copies of Algorithm 1 can be run on each tensor-shaped parameter group; optimizer APIs in standard deep learning packages promote this convention.
This can be viewed as maintaining a tensor sum of preconditioners, each of which is a p -wise tensor product.
• The accumulator update step can be implemented concisely and efﬁciently using reindexing procedures (typically called reshape or view ) in standard deep learning and numerical linear algebra packages.
Lemma 4.3 in the analysis shows that the per-coordinate learning rates are 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.
R EGRET ANALYSIS
In this section, we justify the choice of update rules for the compressed preconditioner using the regret-minimizing adaptive regularization framework introduced in the original analysis of AdaGrad (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 interpretable regret bound, which we discuss in this section and revisit in the empirical measurements of
Section 5.3.
4. S TATEMENT AND DISCUSSION OF REGRET BOUND
First, we deﬁne and prove the main regret bound.
As is standard in the regret analysis of diagonal
AdaGrad (see Duchi et al.
, this is most naturally stated with a diameter term on x t : D ∞ := max t,x ∗ (cid:107) x t − x ∗ (cid:107) ∞ .
This can be constrained using a projection onto an (cid:96) ∞ ball in the algorithm, but this is seldom seen in practice.
It will also be helpful to deﬁne G it to be the diagonal matrix with values
[Equation]
on the diagonal.
Theorem 4.1.
Deﬁne H T = ⊗ pi =
1 ( ε I d i + (cid:80) Tt =
1 G it ) 1 / 2 p , and ˆ H T = diag( ε I + (cid:80) Tt =1 g t g (cid:62) t ) 1 / 2 .
Then, there exists a choice of constant learning rate schedule η 1 = . . .
=
η t :=
η and ε > 0 such that the { x t } chosen by Algorithm 1 satisfy the regret bound
[Equation]
where x ∗ :=
min x (cid:80) Tt =1 f t ( x ) is the loss-minimizing decision in hindsight.
As a direct consequence, we recover the AdaGrad regret bound as a special case where p = 1 .
Noting further that H T
= ˆ H T , we restate this well-known result in our notation, for clarity and comparisons:
Corollary 4.2.
In setting of Theorem 4.1, when p = 1 , the { x t } chosen by Algorithm 1 satisfy
[Equation]
Thus, AdaGrad with extreme tensoring satisﬁes a regret bound at most (cid:113)
Tr ( H T ) /
Tr ( ˆ H T ) times than that of AdaGrad.
This quantity can be as large as Ω( √ d ) ; this is the worst-case price of the dramatic memory savings in the regularizer.
However, in the presence of sparsity, this ratio can be much smaller; we include a discussion and empirical measurements in Section 5.3.
4. P ROOFS
The derivation of the regret bound uses the proof framework introduced by Gupta et al.
Our ﬁrst important lemma states that the regularizer in Algorithm 1 is a 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 are underestimates of those dictated by
AdaGrad.
Lemma 4.3.
Suppose g 1 , g 2 , · · · , g T are tensors of dimension d 1 × · · · × d p , and let g t
= I − 1 ( g t ) for all t , where I is any tensor index whose image is × p
i
=1 [ d i ] .
Then for t ∈ [ T ] ,
[Equation]
Proof.
Let d = d 1 d 2 · · · d p .
Let j ∈ [ d ] and let a 1 , . .
. , a p be such that I ( j ) =
[ a 1 , a 2 , · · · , a p ] .
[Equation]
Taking the p th root on both sides,
[Equation]
Taking square roots of the above inequality yields the lemma.
4. R EGRET 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 t the d -by- d diagonal matrix whose entries are ( I − 1 ( δ t )) − 1 .
Take η 1 = . . .
=
η t := η to be the constant learning rate, and let x ∗ be the loss minimizing choice in hindsight.
Then, the regret of Algorithm 1 is bounded by
[Equation]
Proof.
By the deﬁnition of Algorithm 1, for any x ∗ ,
[Equation]
Taking the inner product of the above vectors,
[Equation]
Rearranging and dividing by 2 η , we have
[Equation]
Using convexity of the loss functions and summing over t yields the lemma.
Next, we present a lemma that is later used to bound the second term in the regret.
Lemma 4.5.
(Gupta et al., 2018) Let g 1 , g 2 , · · · , g T be a sequence of vectors, and let M t = diag( (cid:80) tx =1 g s g (cid:62) s ) .
Let S + denote the set of diagonal positive deﬁnite matrices.
Given a function Φ over S + , deﬁne A t = argmin
A ∈ S + { M t • A − 1 + Φ( A ) } and assume that the minimum is attained for all t .
Then
[Equation]
We proceed to prove Theorem 4. [Equation]
Proof.
(Theorem 4.1) Recall the deﬁnition of H t = ⊗ pi =
1 ( ε I d i + (cid:80) ts =1 G is ) 1 / 2 p , we bound the ﬁrst term in Lemma 4.4 as follows,
[Equation]
The last inequality is due to the fact that H t − 1 (cid:22) H t .
Recall ˆ H t = diag( ε I + (cid:80) ts =1 g s g (cid:62) s ) 1 / 2 , and by Lemma 4.3, ˆ H t (cid:22) H t .
Now we use Lemma 4.5 with function Φ( A ) =
Tr ( A ) + ε
Tr
( A − 1 ) .
Let M t = diag( (cid:80) ts =
g s g (cid:62) s ) and let S + denote the set of diagonal positive deﬁnite matrices, we have argmin A ∈ S
+ { M t • A − 1 + Φ( A ) } = argmin A ∈ S
+ Tr ( ˆ H 2 t
A − 1 + A ) =
ˆ H t .
The last equality can be derived by minimizing each diagonal entry of H individually.
By Lemma 4.5,
[Equation]
We can take the diagonal of (cid:80)
Tt =1 g t g (cid:62) t in the ﬁrst equality since ˆ H − 1 T is a diagonal matrix as well.
We proceed to bound the second term in Lemma 4.4.
Since ˆ H t (cid:22) H t and both H t and ˆ H t are full-rank, ˆ H − 1 t (cid:23)
H − 1 t .
Therefore (cid:80) T t =1 ( (cid:107) g t (cid:107) ∗ H t ) 2 ≤ (cid:80) Tt =1 ( (cid:107) g t (cid:107) ∗ ˆ H t ) 2 ≤ 2 Tr
( ˆ H T ) .
Summing up the two terms and taking η = D ∞ √
Tr ( H T ) √
Tr ( ˆ H T ) , we conclude that the regret of Algorithm 1 is bounded by 12 η D
2 ∞
Tr ( H T ) + η Tr ( ˆ H T ) =
D ∞ (cid:113)
Tr ( H T )
Tr ( ˆ H T ) .
E XPERIMENTS
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 a synthetic experiment of the same form, and provide a 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 a larger model.
5. M EMORY - PERFORMANCE TRADEOFF IN LARGE - SCALE NLP
Our main empirical study focuses on large-scale language modeling with the Transformer architecture (Vaswani et al., 2017) on the Google Billion Words (GBW) dataset (Chelba et 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 architecture in the same repository as our base architecture.
The decoder-only model has 6 identical layers with hidden dimension d model = 512 , and feedforward dimension d ff = 2048 .
In addition, the weights are shared between the embedding and softmax layers, and the model has a total of ∼ 35 M parameters.
For our experiments, we use a learning rate schedule of η t = c · min(10 − 6 · t, 1 √ t ) , and c is a hyperparameter we tune for each experiment.
The learning rate schedule is the same as the one used in (Shazeer and Stern, 2018): a 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.
Except in our comparison with Adam, we do not use momentum in our extreme tensoring experiments.
Momentum incurs an overhead memory cost linear in the model dimension.
In addition, we empirically tested the beneﬁt of exponentially decaying the second moment estimator ( β 2 < 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 ( β 2 = 0 . 99 ). [Figure] F i n a l e v a l p e r p l e x i t y
SGD
ET
ET ET ET AdaGrad
[Table]
Extreme tensoring gives us a class of optimizers which interpolate smoothly between AdaGrad and SGD.
Between these endpoints, we choose three levels of extreme tensoring (denoted by
ET { 1 , 2 , 3 } ), with tensor index dimensions fully speciﬁed in the appendix.
While extreme tensoring 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 of magnitude.
We also provide training results for Adam (which consumes more memory for ﬁrst-moment estimation) and Adafactor (similar to ET1 but with a different step size scaling) in Table 1, for completeness; however, these are not part of the interpolation study.
To provide an additional interpolating point, we include in the comparison a closely related algorithm, which selects a single learning rate 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 to be the inverse square root of the sum of squared (cid:96) 2 norms of the parameter group over time; it is easily seen that this achieves the same regret as online gradient descent (Zinkevich, 2003), so we include it as the least granular adaptive optimizer, and call it ET ∞ .
Note that this algorithm is known, and Ward et al.
(2018) in particular provide a ﬁne-grained analysis.
Finally, in Appendix 6, we re-run the ET { 1 , 2 , 3 } experiments with the SM3 update rule (Anil et al., 2019), ﬁnding the same monotonic trend.
In addition to model performance, we also provide in the appendix wall clock comparisons for the algorithms we have benchmarked.
5. D OUBLING 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 a network with ∼ 56 M parameters.
Results are shown in Figure 2 and Table 2.
Table 2 reports performance from the larger model, where ﬁnal 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 ﬁnal evaluation perplexity of the larger model in the third column of Table 2 for a more intuitive visualization.
Holding the memory consumption constant, we ﬁnd that it pays off to use a larger model rather than an adaptive optimizer that stores all of the secondmoment accumulators.
Training the larger model with extreme tensoring consumes essentially the same amount of memory as training the original model with Adam or AdaGrad, but the performance improves signiﬁcantly.
Even holding the running time constant, the larger models are competitive with the fully-converged smaller ones.
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Figure]
[Figure]
5. E MPIRICAL MEASUREMENTS OF THE REGRET BOUND
An appealing part of our regret analysis is the interpretability of the trace quantities in Theorem 4.1.
Figure 3 shows a comparison of traces
Tr ( H T ) and
Tr ( ˆ H T ) of the ﬁnal 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 ﬁgure is logarithmic, so the multiplicative regret bound gap compared to 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 (cid:113)
Tr ( H T ) /
Tr ( ˆ H T )
≈ 5 . 7 .
5. C OMPARISON ON SYNTHETIC DATA
In this section, we exhibit a simple experiment with a convex optimization problem, in which there is a clear tradeoff between preconditioner quality and expressivity.
We generated Gaussian data x i ∈ R 512 with a large condition number ( ∼ 10 4 ), and a Gaussian matrix W ∈ R 10 × 512 .
Labels are generated according to the log-linear model
Pr[ y i =
j ] ∝ exp(( Wx i ) j ) .
Then, the optimization problem in W is to minimize the empirical negative log-probability under the same model.
Our ﬁndings are robust with respect to the batch size; we use the full gradient on 10 4 samples in our plots for clarity.
As in the large-scale experiment, we use successively deeper tensor factorizations of the preconditioner, along the feature dimension of the matrix W .
For depths 1, 2, and 3, we choose tensor indices of dimensions (10 , 512) , (10 , 16 , 32) , and (10 , 8 , 8 , 8) , respectively.
Global learning rates are tuned individually.
Results are shown in Figure 4. [Table]
5. I MPLEMENTATION D ETAILS
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 to additional tensor operations.
We have not optimized our implementations, and expect that the time efﬁciency gap could be reduced with architecture-aware optimizations.
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.
C ONCLUSION
We have introduced extreme tensoring , a modiﬁcation to any second-moment-based adaptive optimizer which drastically shrinks memory overhead.
Our experiments characterize a performancememory 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.
A CIFAR-10 EXPERIMENT
In this section, we evaluate the memory-performance trade-off of our proposed algorithm on the
CIFAR-10 dataset (Krizhevsky, 2009).
Speciﬁcally, we compare the test accuracy of a ResNetmodel (He et al., 2016) trained with SGD, Adam and 3 levels of extreme tensoring.
A. S ETUP
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. β 1 = 0 in
Adam).
We use batch size 128 , weight decay 5 × 10 − 4 for all the experiments in this section.
A. T ENSOR INDICES
We show our tensor decomposition scheme in the table below.
Note that bias parameters are not shown in each layer in the table since are treated as vectors and thus not decomposed at all.
[Table]
Shape
[Table]
Shape (ET3) (64 , 3 , 3 , 3) (64 , 3 , 9) (8 , 8 , 3 , 9) (8 , 8 , 3 , 9) (64 , 64 , 3 , 3) (64 , 64 , 9) (8 , 8 , 8 , 8 , 9) (8 , 8 , 8 , 8 , 9) (128 , 64 , 3 , 3) (128 , 64 , 9) (8 , 16 , 8 , 8 , 9) (8 , 4 , 4 , 8 , 8 , 9) (128 , 128 , 3 , 3) (128 , 128 , 9) (8 , 16 , 8 , 16 , 9) (8 , 4 , 4 , 8 , 4 , 4 , 9) (256 , 128 , 3 , 3) (256 , 128 , 9) (16 , 16 , 8 , 16 , 9) (4 , 4 , 4 , 4 , 8 , 4 , 4 , 9) (256 , 256 , 3 , 3) (256 , 256 , 9) (16 , 16 , 16 , 16 , 9) (4 , 4 , 4 , 4 , 4 ,
[Table]
A. R ESULTS
In this section
[Table]
we report the best test errors seen in the ﬁrst 150 epoch along with the parameter count in the optimizer.
A similar trend of trade-offs between performance and parameter count is observed.
[Table]
[Figure]
F i n a l t e s t
e rr o r
SGD
ET ET
ET ET
Adam
B L ANGUAGE 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 our main experiments, and the tensor shapes of the reshaped parameters.
Parameter
Parameter Shape
Shape (ET1)
Shape (ET2)
Shape (ET3)
Q, K, V , position-wise FF (512 , 512) (512 , 512) (16 , 32 , 16 , 32) (4 , 4 , 4 , 8 , 4 , 4 , 4 , 8)
Embedding Weights (2000 , 512) (2000 , 512) (40 , 50 , 16 , 32) (5 , 8 , 5 , 10 , 4 , 4 , 4 , 8)
Layer Norm (512 , ) (512 , ) (16 , 32) (4 , 4 , 4 , 8)
FC Weights (512 , 2048) (512 , 2048) (16 , 32 , 32 , 64) (4 , 4 , 4 , 8 , 4 , 8 , 8 , 8)
FC Bias (2048 , ) (2048 , ) (32 , 64) (4 , 8 , 8 , 8)
FC Weights (2048 , 512) (2048 , 512) (32 , 64 , 16 , 32) (4 , 8 , 8 , 8 , 4 , 4 , 4 , 8)
FC Bias (512 , ) (512 , ) (16 , 32) (4 , 4 , 4 , 8) B [Table]
W ALL C LOCK C OMPARISONS
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.
[Table]
C SM3 C OMPARISON
An alternate way to modulate the memory-performance tradeoff is to use the user-selectable partitions 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 in Table 6 the ﬁnal test perplexities for the same Transformer experiments on the ET { 1 , 2 , 3 } decompositions (denoted as S { 1 , 2 , 3 } ).
We used SM3-II as reported in the pseudocode of Anil et al.
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 as a comparison between these families of optimizers; rather, these experiments corroborate a loss of adaptivity as the preconditioner is constrained to an extremely low dimension.
[Table]