Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale
Authors: Ran Tian
Ankur P Parikh
Venue: Submitted to ICLR 2023 (2023)
Keywords: optimization
asymptotic behavior of stochastic optimization
learning-rate decay
weight decay
language model pre-training
Transformer pre-training

Abstract: We present Amos, a stochastic gradient-based optimizer designed for training deep neural networks. It can be viewed as an Adam optimizer with theoretically supported, adaptive learning-rate decay and weight decay. A key insight behind Amos is that it leverages model-specific information to determine the initial learning-rate and decaying schedules. When used for pre-training BERT variants and T5, Amos consistently converges faster than the state-of-the-art settings of AdamW, achieving better validation loss within <=70% training steps and time, while requiring <=51% memory for slot variables.

Under review as a conference paper at ICLR 2023
Amos: AN ADAM-STYLE OPTIMIZER WITH ADAPTIVE WEIGHT DECAY TOWARDS MODEL-ORIENTED SCALE
Anonymous authors Paper under double-blind review
ABSTRACT
We present Amos, stochastic gradient-based optimizer designed for training deep neural networks. It can viewed as an Adam optimizer with theoretically supported, adaptive learning-rate decay and weight decay. A key insight behind Amos that leverages model-specific information to determine the initial learning- rate and decaying schedules. When used for pre-training BERT variants and T5, Amos consistently converges faster than the state-of-the-ar settings of AdamW, achieving better validation loss within I 70% training steps and time, while requiring < 51% memory for slot variables. Our code is open-sourced at: https / //anonymous-uri
INTRODUCTION
The Adam (Kingma & Ba 2015) optimizer widely used for training deep neural networks, demonstrating fast convergence especially in the early stages of training. Although previous works have found issues regarding the theoretical convergence of Adam as the training proceeds (Reddi et al.) 2018), in practice it remedied by various learning-rate schedules and weight decay (Loshchilov & Hutter 2019). Specifically, Adam with linear learning-rate decay and constant weight decay is the standard setting for pre-training large language models such BERT (Devlin et al.) 2019 However, these decay settings are usually ad-hoc, increase the number of hyper-parameters, and may introduce additional complexities usage. For example, the linearly decaying learning-rate schedule requires knowing the number of training steps in advance, which makes it nontrivial to continuously train model after the learning-rate decays to O.
In this work, we present Amos, new optimizer with a theoretically supported and adaptive schedule for learning-rate and weight decay, which can significantly outperform the state-of-the-art Adam W settings for pre-training language models, provide guidance for hyper-paramete tuning, reduce the memory usage, and train continuously without having to specify the number of training steps priori.
A key insight behind Amos is hyper syper--arameter 7 to be provided by the model architecture, which indicates the expected scale of the trainable weights 0 of the model (8)2). i.e., theoretically assume that optimal point 0* exists within the 0 i diameter. Deep neural networks likely to satisfy such a constraint without degrading performance, because there exist many good local minima; and we show that an appropriate for Amos can improve generalization and accelerate convergence ($5.2). In this work, 7 is calculated in a consistent way from the input/output scale of neural network components, which is hinted by the model design (6) A.44. Given 7, Amos decides a learning-rate per variable, and its L2 regularization will lead the trained weights to the specified scale. The decay of the learning-rate rate is then determined by the L2 regularizer. Thus. Amos performs better because can utilize the model-oriented information v efficiently; the name Amos stands for "Adaptive weight decay towards Model-Oriented Scale".
Empirically, we focus the Transformer architecture Vaswani et al. 2017) since it is pre-dominant in pre-trained language models (Bommasani al., 2021). but add additional experiments on LSTM (Gers et al. 2000) and ResNet (He al., 2016). We apply Amos to the pre-training of 4 models: BERT (Devlin et al., 2019), two Transformer variants with relative position representations (Su et al. 2021 Shaw al. 2018), and the T5 model (Raffel al. (020); some with various model sizes and batch sizes. In all experiments, Amos consistently outperforms the state-of-the-a setting, achieving better validation loss within S 70% training steps and time ($ 5.1.. Compared Adam W,

---

Under review as a conference paper at ICLR 2023
the memory usage for slot variables is reduced to I 51% in Amos ($) A.8) In addition, Amos does not calculate learning-rat rate from a maximum number of training steps. so one can seamlessly continue training from any checkpoints, which is not trivial for Adam with linear learning-rate decay ($) (1...)..
2 THE ALGORITHM
For notation, we denote model weights by 0, and an online learning algorithm recursively calculates a sequence of weights, 01, 02.02.... from initial weights og and training examples 2, at each step 1 = 0, 0. An optimizer uses the gradient ÿ: VV(==;; O+) to compute a weight update 0.+1 <- 0, 8,, in order to minimize the loss function {2; 0). In neural network models, the model weights i is an array of trainable tensors (i.e. variables) collected from all model components; we view variable and its slices as subsets of the model weights (e.g. 0 C 0 is variable slice that functions in part of the model). We use a bold letter to denote an array (e.g. 0.. 0t), and the same normal letter to denote a scalar element of that array (e.g. 0z) for describing element-wise operations. We use tilde for information of the whole model (e.g. 0z), and drop tilde to indicate subsets (e.g. O₂)
To start, we recall the update rule of the RMSProp optimizer (Tieleman & Hinton| 2012). which computes the weight update by de +- Ver where is a scalar learning-rate and i, running average of the squared gradients g². Based on this, Adam (Kingma & Ba. 2015) replaces ge with its running average m, (i.e. momentum), and adopts bias correction my, Or for running averages. Further, AdamW (Loshchilov & Hutter (019)) allows schedule for learning-rate ar (depending on the step t) and adds a weight decay: oe < are vEr m, 76+). where 7/ is a constant hyper parameeer. For pre-training Transformer variants, the learning-rate schedule a; is set to linearly decay to 0 after warm-up. Therefore, maximum number of training steps before the learning-ratt decays 0 has to be set as a hyper-parameter Amos, with a similar construction, has the following update rule:
E1 Vor gI 2
1 )10))) where Tt Ca M2(gi)²
6 de
(1)
Here, M2(a)= = M= V at denotes the quadratic mean of entries of an array a E R The update rule consists of a gradient descent part (the term containing 9:) and an L2 regularization part (the similar to AdamW. The full Amos is shown in Algorithm | We explain several
term containing 0,
novel aspects below.
Model-oriented scale:
For each variable a C θ in the model weights, we specify the scale n(a) we expect a to converge to, i.c. M2(a*) ss n) for an optimal 0- 2 a'. Different variables may have different scale 1j's. For a common case of a linear transformation, = xW + u (W, u S
O, W m E RT), we calculate 77 (W) by assuming that æ is random Gaussian with standard deviation oz and y random Gaussian with standard deviation σy: so we have 7/ (W)- = ayy/ (0 in order to satisfy the input/output standard deviation (assuming entries of W to be Gaussian well). Additionally, we set 17(2) = 0,/2 to ensure that u has a slightly smaller magnitude than x wW. The input/output standard deviation can be hinted by other layers in the model; for example, the activation function GELU (Hendrycks & Gimpel (016) usually expects the inputs to have standard deviation
z eevin /m)
1, because its non-linearity mostly lies within that range; also the output standard deviation of LayerNormalizztion (Ba et al. 2016) is expected to be 1. For Transformer variants, we will discuss the input/output standard deviation of all types of non -oneeinear layers and derive 7 in AA.4
Factored initial learning-rate: In Amos, we use {1 as the initial learning raate where is the model- oriented scale specified for each variable, and & is global learning-rate shared across all variables. For online optimizers, the learning-rate is generally affected by both data and model; by factoring the initial learning-rate into & and 1). we disentangle the two to some extent: While & is tuned and may depend on the data, î is calculated from the model architecture
FFllowing Loshchilov & Hutter 2019). we decouple the gradient of an 1.2 regularization term (taking the form of weight decay) apart from the adaptive gradient normalization factor __ When an adaptive optimizer -
is used, Loshchilov & Hutter (2019) point out that the decoupled weight decay is not equivalent to the L2 regularization without explicit decoupling, and the former is more appropriate. In this work, we always treat 1.2 regularization decoupled weight decay, and the two terms interchangeably.
2

---

Under review as a conference paper at ICLR 2023
Adaptive L2 regularization: Unlike AdamW which uses constant y for weight decay, the Amos weight decay "/ is intended to control the scale of trained variables, rather than regularize the loss function; so 7/2 decays to 0 at t + oc to be less biased, and it is adaptive in the sense that % depends on Gr so that the variables not getting gradient updates are not regularized. Thus, the 1.2 regularization robust to sparse gradients, and it does not introduce any additional hyper-parameter. We will give heuristic derivation of the form of T/e in $
Decay factors: d,, ē are per-parameter decay factors such that dp co = and de, Cr monotonically decrease to 0 at -> 2o. We provide theoretical derivation of the asymptotic behavior of these factors in 8 ||A.2] together with a default form that works well empirically in all our experiments. The decay factors do not depend on a maximum number of training steps, thus enabling arbitrary continuous training.
Memory Reduction: Most previous optimizers operate element-wise, the slot variables (e.g. the rumning average De. m, in Adam) have the same shape as 0, which can be memory consuming. In Amos, two slot variables (ő,, b, Algorithm 1 are shared by certain slices in the model weights, reducing the memory usage of these slot variables. For example, if mR* e W C transformation, the corresponding V E <R1x is shared by the input dimension of W, reducing the memory usage by m times. As a result, in Equation I and Algorithm|] | Ut b,, C₁ and d, are reduced and become scalars, to be used and updated by vector-valued g: and 0,. In this work, we reduce the input dimension of linear transformations, the embed dimension of embedding matrix, and all dimensions for other variables by default. An ablative study with different settings is found in AA.
a linear
Algorithm The Amoy optimizer at step L.
Input g. = Vf(2;; 0.): The gradient of loss on a random example 24.
Input 0:: Trainable model weights at step f.
Input 1; b,: Slot variables of shape broadcastable to 0, initialized O.
Input (Optional) m. Slot variable of the same shape as Ó, initialized to 0 for momentum.
Hyper-parameter Global learning-rate.
Hyper-parameter îp: Expected scale for model weights θ.
Hyper-parameter ē; Decay factor for L2 regularization. Defaults to c₁ = (1+ IVE6.) Hyper-parameter d,: Decay factor for learning-rate. Defaults to de + IVEnb,
Hyper-parameter B E [0, 1): Exponential decay rate for running average 8t
max 19.1) 9r. 2: v + 30,- + (1 3) M2(gi)²
I: (Optional)
X
> Gradient clipping with hyper-paue parameter X > 0.
Running average squared gradients.
3: 0 + 01/(1 -
(1-B²) 3²)
> Bias correction.
4: 7 a 0r
2
Adaptive L2 regularization strength.
5: 8, de gt 7(0))
{n
Amos update rule.
V
6: 6,+1 + bt + 7/ (1 b²)
Decay factor update.
7: (Optional) 5, + m4-1 + pm, + (1 11)02
> Momentum with hyper-parameter E [0. 1).
Output: Updated model weights l.+1 +- 0, -
Output: Updated slot variables Fr, b and optional m
Hyper-parameter Tuning The running average U, Amos a low-cost estimator for (g )2), where the expectation is taken over the example 2, randomly drawn from the training data. It is similar to V/ in Adam except that the mean square (2(gt) 2 is used instead of element-wise g² due to the memory reduction. Thus, the hyper-parameter behaves similarly to 3o in Adam: Since the estimator mostly depends on the previous 1/(1 33 steps, / should be close enough to 1 to make the estimator accurate, but not too large that the model weights in the previous 1/(1 111--- (3) steps differ too much from the current step. We set 3 = 0.999 by default (the same as 3g in Adam), and it is found that & should be smaller with larger batch size (Shazeer & Stern 2018 Liu et al.) 2019).
3

---

Under review as a conference paper at ICLR 2023
The global learning- rrate can depend the step to follow a warm-up schedule the beginning of training. but schedule with learning-rate decay is not necessary, since the decay factor d, is already included in Algorithm[] | most of the time & remains a constant. While this constant the major hypeer-arameter be tuned, good rule of thumb is to set 5 to the same order of magnitude as 1/v N/V,, where N is the number of independent batches in the training set (see $ for justification). This value is usually larger than the typical learning-rates used for Adam. It also implies that & should be proportion to the square-root of the batch size, which we observe in practice as well (3A..5).
In addition, Algorithm | includes optional gradient clipping and momentum. Momentum in Amos is applied after the main update rule (unlike Adam which applies it before). can improve performance for pre-training Transformer variants, but consume memory because the slot variable m. must have the same shape as O. When momentum is applied, its decay rate μ is typically set to 0.9.
3 RELATED WORK
Besides RMSProp TTilleman & Hinton 2012), Adam (Kingma & Ba 2013) and AdamW (Loshchilov & Hutter 2019), the number of previous works on optimization is vast, so we focus on some directly related alternatives below. Also, we note that Amos a stochastic first-order optimizer, in contrast to recent progress in the second-order optimization methods (Gupta et al. 2018). The convergence of stochastic optimizers has been studied terms of stochastic approximation (Bottou) 1998), regret (Hazan 2019), or nonconvex stochastic programming (Ghadimi & Lan 2013). In particular, Reddi cet al. 2018) observed cases of non-convergence of Adam (with constant learning-rate) and proposed fix. In our work, we analyze the behavior of Amos in an intuitive and heuristic manner, but leave rigorous convergence proof (e.g. based on regret) to future work.
The update rule of AdaGrad (Duchi et al. 2011 is 8, + vh, -,, St. where 6.+1 b, + oe is similar to the by in Algorithm [ in the sense that both AdaGrad and Amos use (weighted) sum of squared gradients decay learning-rates. Such decay is "adaptive" because the learning-rate will decay more for parameters getting more updates, which is suitable for sparse gradients. On the other hand, conventional wisdom is that the learning-rate in AdaGrad might decay "too fast" in some cases, which makes the convergence slow, and Adam mitigates this issue by using running average of squared gradients instead of the decay factor. However, the AdamW W setting suggests that the normalization factor by running average of squared gradients is not a replacement for learning- rate decay; one still needs a linearly decaying learning-rat rate schedule for better convergence. Thus, Amos integrates both the Adam-style gradient normalization and the AdaGrad style learning- rate decay; with gradient normalization, the learning-rate can actually decay faster and it converges faster.
AdaGrad:
SGD with L2 Regularization: For the classic Stochastic Gradient Descent (SGD) algorithm, is recommended decay the learning-rate by the factor 1 of the Hessian (Murata 1998). Although X is generally unknown, adopting an L.2 regularizer of strength X' guarantees that X X, so can set the learning-rate to x7 we adopt a similar idea to heuristically derive the learning-ras rate decay (see A for more detailed discussion). by connecting the decaying speed with the strength of L2 regularization (i.e., the L2 strength ?/2 in Algorithm also appears in the update of by). Unlike SGD, both the learning-rate and L.2 regularization Amos decay adaptively. The adaptive L.2 regularization, in particular, a novel component unseen in previous optimizers.
Xe where À is the smallest eigen-value
- (Bottou 2012). In Amos,
The LAMB optimizer Youu et al. 2020) and its origin LARS (You et al. 2017) share several similar aspects with Amos. The idea of layer-wise learning-rate in LAMB and LARS is similar to the per-variable learning-rate 7 in Amos; they all normalize the gradients in some way; and they all imply scaling up the learning-rate as the batch size increases. In our experiments, scaling the global learning-rate of Amos in proportion to the square-root of the batch size indeed works AA...) although we leave a systematic study of scaling-up up to extremely large batch sizes and comparing with LAMB and LARS to future work.
LAMB:
AdaFactor: In Adam, the slot variable ē, for maintaining running average of squared gradients requires the same amount of memory the model weights 0. In order to reduce the memory usage, AdaFactor (Shazeer& & Stern/ 2018) proposes to use nonnegative matrix factorization to decompose any matrix into two vectors. In contrast, Amos reduces memory usage by simply reducing some axes
4

---

Under review as a conference paper at ICLR 2023
of the slot variables and broadcasting to the shape of model weights. This reduction is more efficient than AdaFactor, and our experiments suggest that it will not degrade performance (8 A.8)
4 DERIVATION OF AMOS
In this section, heuristically derive the Amos update rule (Equation I. form of the weight update for a given variable O,
We start from general
0 = 0, a:gt where g: V[(2,;
and gradually pin down to the specific form of Equation I Here, the step size at > 0 a scalar (due to our memory reduction mechanism in 8 and is shared across the elements of the vector-valued gt, 0g E R We are focusing on a subset of model parameters, but furthermore note that ag may differ for different variables.
Then, the following Descent Lemma Murata 1998) provides a sanity check for a wide range of possible forms of at, while also suggests constraints. Its proof can be found in AA.I]
Lemma 4.1 (Descent Lemma). If ax does not depend 2,, then there exists & >> 0 such that
E, ( 0e+1)l] < EE (f(=z;;)) for any 0t < ct.
where [x[²] denotes the expectation taken over the random example 2 drawn from the training data at step t, while conditioned on 2t-1....;20 20 of the previous steps.
In light of Lemma 4.1 we require (I) a does not depend on 2 (but may differ for different variables), and (II) decays to 0 at t + oo, so the step-size can be sufficiently small that the Descent Lemma applies and Equation 2 will always make progress on average.
In the Amos update rule, a d, /e, and ún depends on 2t, which seems to violate requirement (I) above. However, é should be regarded as approximation of [8[M2 (g.) where B[}) denotes the expectation taken over examples randomly drawn from the training data, which is 2, independent.
Next, we add an L2-regularization term Equation 2
0++1 = (arg: p.O,)
(3)
where P: 2 0 can depend on g. (hence "adaptive"), but we require (III) Bip:] does not depend on gi- The intuition behind is that an 1.2-regularization should have the same strength across all variables, rather than be affected by the typical gradient magnitude on each variable. It is the same intuition that motivates the weight decay decoupled from gradient adaptive factors (Loshchilov & Hutter 2019).
The first challenge for Amos is to keep a balance between a, and Pt. so that M2(0.) will converge to the pre-specified, per-variable hyper-parameter 7). In order to achieve this, we will declare some intuitions on the largeness of gt. E [g]]] and prO,, as a guide for our heuristic derivation. For deep neural networks, gt's 's upon different 2, 's appear to be randomly noisy, so they will cancel out when being averaged to E'gg.]: which means that (E[g.) is usually much smaller than M2(g.). On the other hand, 9, does not depend 22. and it changes slowly between different steps, so the update by p.O, is easier to accumulate than 0, This means that the magnitude of p.O, can be kept smaller M2(g.) decays faster than than ogg: while still compete with [(ge]] In Amos, Pt d ct b, ar (due the extra decay factor c2), which we assume will make p.0, small enough compared to a,gr, when is large.
8²
Quantitatively, we consider the error ē; = 0, - Ō", where θ" is a local minimum. Equation implies
Mz(EL41) = M2(E,)²
(aig: + p.O,) E, M2 + p.O.)
k
~ M2(E,) (a.ge + p.O.) E. af M2(g.)
where we investigate a time point large enough that the model nearly converges. At this point, p.O, is small compared o.gi: it can be approximately omitted in the third term. And should have
5

---

Under review as a conference paper at ICLR 2023
Eig]] 0 and M2(0,) 17 if the trained weights converge to scale 77. So taking [[]] of Equation we should get
2 - 0 0, E, + a(
[[Mg(E(+1) ~ M2(E,)²
Furthermore, in order for the model to converge, we should have E[M2(E++1) S M2(e,)² from the above. Hence, we should have
(g) 2 0l E, S 2 M2(0,) M2(e)) = 2Ilp
necessary condition for the trained weights to converge scale n. By setting Pt to the smallest possible, we get
2pm M2(:) at M2..)
which an important relation connecting Pt to aq, We require (IV) Equation 5 to be satisfied throughout the course of training, and use it ubiquitously in our derivation. It is out of the scope of this work to prove whether Equation actually makes M2(0:) converge to 77: but the requirements so far already determine a basic form of the Amos update rule (as shown in Lemma 4.2 below), and our experiments suggest that Amos indeed brings the trained weights to the specific scale (85.2).
Lemma 4.2 (Basic Form of Amos). Assume Equation requiring that a, does not depend on 2t and E[pe] does not depend on gt- Then, we huve
M2(g.² (g. and X 2
0x x
The proof is found in A.1 It noteworthy that the Adam-style gradient normalization naturally occurs in 0y Based on Lemma 4.2 Amos is derived by specifying the initial learning-rai rate and decay schedule. For that, we need the following assumption to quantify the largeness of g: and E[gt).
(2(((()))
Assumption 1. A scalar g > 0 exists such that
2 & for all t and across all variables.
E M
This assumption formalizes two intuitions, i.e. randomly noisy ge will cancel out when being averaged to E Elg.] (so & has a small value), and as the training proceeds, M2(g:) will decrease along with M2 (E(g.) (so the ratio remains larger than a constant 0). Assumption 1 is verified by
experiments (8) (...6)).
The value of & is related the global learning-rate in Amos (as shown in Lemma 4.3 below), which tuned as a
in practice. However, we also provide an intuitive estimation of g, which is usually good start for hyper-parameter tuning. The intuition is view the canceling out of gr averaged to E[gg] as similar to the average of N i.i.d. samples drawn from a distribution of mean 0. According to the Law of Large Numbers, the variance of the average (i.e. M2 ().)))))))))))))))))) is about 1/N of the variance of the distribution (i.e. (g, ) so & = V In reality, the gradients of deep neural networks, computed same order of magnitude as the number of independent batches in the training data.
mini-batches, appear be highly random. So N is usually of the
derive the optimal initial learning-rate as below, under an ideal condition that gu points to the same direction as En. The proof is found in $A..1
Now,
Lemma 4.3 (Initial Learning-rate Assume Equation that go points to the same direction as En. Then,
Assumption
/E and and
B[M2(e1) < M2(eo)² 20g M2(E0) + M2(50) = En-
and the RHS achieves minimum
As the training proceeds, [Bgg]] will converge to = 0, so Assumption | related to the observation that, for highly expressive models, g. Vl(2,; Ø*) can get close to for every 2t the training data (Ma al. 2018). However, Assumption only requires that Mz (g:) decreases as fast as BBgg]. which is empirically verified 6).. Whether Ma (ge] actually converges to is not guaranteed (because the training may stop early, or E8 [ggr]] not get exactly 0 due to I. 1.2-regularization, etc.) and not used in our theory, On the other hand, M2(g:) always large compared to Elge), hecause E small value.

---

Under review as a conference paper at ICLR 2023
Lemma 4.3 suggests the initial learning-rate do [9o
En
Then, we get Po = 2 Ma(ga)"
from Equation 5 By adding the decay factors, we reveal the Amos update rule (Equation D
8e + aggt + prOr, where a₁ = de En and Pt di de -e - cc
M2 (g
(6)
(
(g.)
Here, de and ct monotonically decrease to and do = cp = 1. In particular, oa decaying to 0 ensures that Pr decays to faster than a₁, so prO, can be sufficiently small compared to oggt for large t, which justifies the approximation of Equation
In A.2 we will further derive that ct and d, = (1+ + qb.) where q are constants, together with the update rule of b, The specific p IVE and q VEn are found through experiments and work well in practice.
+ pb:)
EXPERIMENTS
We focus the Transformer model /Vaswani al. 2017); 7, and pre-train several variants below.
BERT: A Transformer Encoder model with learned position embeddings (Devlin et al. 2019). We experiment with the base (12-layer 768-hidden) and large (24-layer 1024-hidden) model sizes.
RoPE: A Transformer Encoder variant with the Rotary Position Encoding (Su et al. 202T). RoPE integrated in some recent large-scale language models (Chowdhery et al. 2022). It encodes relative positions but the encoding is not Iearned. We experiment with the base (12-layer 768-hidden) and large (24-layer 1024-hidden) model sizes.
Relative Position Embeddings (RPE): A Transformer Encoder variant with learned relative po- sition embeddings (Shaw ct al. 2018). It achieves better performance but the pre-training is more costly TPU (Tian et al.] 2021p. We experiment with the base (12-layer 768-hidden) model size.
T5 Encoder- (T5): A Transformer Sicoodr--eecodes model implemented by Raffel et al. (2020). We experiment with the large (24-layer 1024-hidden) model size.
For encoder only models, pre-train with the Masked Language Modeling loss (Deliin et al. 2019) on Wikipedia and the Books Corpus (Zhu et al. 2015). Following Liu et al. 2019), we use batch size 1024 for base-sized models, and pre-train 200k or 300k steps. For BERTT-lare,, 4096 and pre-train 250k steps. For RoPF-llarge, due to memory limitations we have to use batch size 1024 and pre-train IM steps. For T5, the batch size is 4096 and we pre-train with the Span Corruption loss on the C4 corpus (Raffel et al.) 2020), for 250k steps. More detailed settings are found in A.5
use batch size
As an additional evaluation, also applied Amos to the ResNet model (He al. ImageNet (Deng et al. 2009) dataset. The experiment settings and results are shown in
2016) on the
A.9
BERT-basc
(b) RoPE-basc
RPE
Adami 200k
Adomt/12000
Adaniti-30k
Amoo
Amoo--SScale
28
130
135
2006
2008
100
AdamW-200s
Mamm
Amos
AoantW-Cont
100
120
128
110
100
200
200
Figure 1: Pre-training 3 models of the base (12-layer 768-hidden) size: (a) BERT, (b) RoPE and (c) RPE. We show training loss on the top and validation loss on the bottom.
hhttps://en wi kikeii .rgg/wt...... wiki /Ma+ i h_Page
7

---

Under review as a conference paper at ICLR 2023
3 tuavirnn toz
22
(a) BERT-larg
(b) RoPE-large
(c) TS-large
AdammY
Adamiy
AdamtW
Amoo
Ames
Amoss "Scle
AddFactor
00ec.rt ord www. pw. bwg.eeeeer
13
1/
Figure 2: Pre-training models of the large (24-layer 1024-hidden) size: (a) BERT, (b) RoPE and (c) T5. We show training loss on the top and validation loss on the bottom. For T5, the training loss is different from the eross-entropy loss due to an extra regularization term. See A.5 for details.
5.1 LEARNING CURVE OF PRE-TRAINING TRANSFORMER VARIANTS
In TigureTand Figure 2 we show training and validation loss of pretrrininng the Transformer variants. In all experiments, across different model architectures, model sizes, batch sizes, datasets and loss functions, Amos (pink curve) outperforms the state-of-the-ar Adam setting, with the loss always significantly lower beyond 30% of the training procedure and the validation loss achieving the final value of Adam W-300k within < 70% training steps or time 5 For BERT-base (Figure [1), Amos achieves the same within only 145k steps (< 50%), and the Amos checkpoint at 150k outperforms the final checkpoint of Adam W-300k in fine-tuning on MNLI (Williams et al. 2018) as well ($..7)
In Figure Da. also tried starting from the final checkpoint of AdamW-2000 and resetting the learning-raie rate as if it linearly decaying to max training step 300k CAdamW-Coni The loss spikes higher and does not go further lower than the value 200k, suggesting that the hyper-parameter of max training steps has to be set a priori, and continuous training is not trivial with AdamW. In addition, tried learning-rate schedule (AdamW-rsqrl that takes the same value step 10k but adopts decay in proportion to t 1/2 (where t is the step) beyond. Although this setting does not require max training steps, converges slower than both Adam W-200k and Adam Adam W-300k.
smaller value (Amos-* *Scale, see A for more details), and found significant impact especially on the validation loss. Similar results are observed when we change 17 for a certain type of layers in the BERT-large model (Figure 2h, Amos-*Scale, see A.4). It suggests that the model-specific
For the RPE model (Figure Te), tried setting of the relative position embeddings
(a) BERT-base
(b) BERT-base
(c) RPE
Scale of TokenEmbed
Scale LLL Dew.Le/ee/eeBBBBee
Scale of Layer//ePPPoo
AdamiW 200k
Adant/t-2004
AdamWW-300
233
Adomt//-300k
Amos
Amrs
121
11.0
6/4
Figure 3: Plots of the quadratic mean of entries of variables over pre-trained steps.
'We have tried different learning-rates in preliminary experiments and the hest chosen. A learning-rate search for BERT-base presented in AA.5
'In our JAX (Bradbury et al. (018)) implementation, the running time per training step for all optimizers (Adam W. Amos and AdaFactor) almost the same.
8

---

Under review as a conference paper at ICLR 2023
Train Loss
Validation Perplexity
Scale of Kernel
Maamii
Adien
40
84
120
DN0 120
10%
120%
105
108
300
Figure 4: Training a single layer LSTM on the PTB corpus.
information 7 indeed contributes to the performance of Amos, which according to previous work (Kaplan et al. 2020) is unlikely achieved by tuning the learning-rate schedule alone.
5.2 SCALES OF TRAINED VARIABLES
In Figure we show how the scale of entries some variables evolve as the training proceeds. With AdamW, both the token embeddings and the bias converge to similar scales (Figure Jab); while with Amos the token embeddings converge to vI/d (where d is the hidden size) and the bias 0.5, as specified by the hyper-parameter fj. It shows that the algorithm of Amos can lead variables to converge drastically different scales, which is unlikely with Adam In Figure 3r, comparing Amos and Amos-* *Scale, the relative position embeddings in typical layer of the RPE model converge to different scales, which shows that the scale is indeed controlled by the hyper-parameter rj. Recall that Figure shows this has impact on the performance.
In order further illustrate the relation among the optimizer, validation performance and the scale of variables, train single layer LSTM on the Penn Tree Bank (PTB) corpus (Marcus et al. [193). The model size is 256 for hidden states and 1024 for memory. We set dropout rate 0.55 for hidden states (which is important for training on PTB) and 0.1 for memory. Sequence length and batch size are set to 64. We compare Amos, Adam W, and Adam (without weight decay). For Amos, the global learning- rate is set to 0.01 and 77 for the LSTM kernel is set to /32 (calculated from input scale 1. input dimension 512 and output scale 1). For Adam W and Adam, the learning-rate is set 0.0015 (about the same as Amos for the LSTM kernel), and the weight decay is set to 0.01 for AdamW.
The results are shown in Figure Without weight decay, the scale of the LSTM kernel trained by Adam can keep increasing; so Adam is better than Adam on training loss but worse on validation perplexity (i.e. the model trained by Adam generalizes worse). On the other hand, Amos achieves the same training loss Adam, while keeping the scale of the kernel as specified. It results in much Overall, we conclude that controlling
better validation perplexity which matches the state-of-the-as the scale of trained variables help the generalization performance of deep neural networks, and the model-specific information from 77 enables Amos to do this.
CONCLUSION
We have presented the Amos optimizer, which uses an adaptive L2 regularizer control learning- -ate decay and guides trained weights towards a specified model-oriented scale. It demonstrates faster convergence than the ofate-of-the-art in pre-training language models, where the training process long and decaying schedule is crucial. On the other hand, its ability to control the scale of trained weights also brings better generalization to small models such as a single layer LSTM.
Besides pre-training, expect Amos to have advantages in fine-tuning as well, especially for multi-modal models that combine heterogeneous components of varied scales and/or pre-trained with different recipes. Hopefully, the model- model-specific information 7 can help fine-tune such models that were previously difficult with other optimizers (Liang et al. 2022 Kumar et al.) 2022).
"See Melis et al. 2020 for a setting that achieves the state-of the-art performance for single layer LSTM on PTB, It uses RMSProp and dynamically decays the learning matting-rate by watching the performance on the validation set. To our knowledge, no previous work has been able to achieve the state-of-the-art with straightforward setting of the optimizer we do with Amos.

---

Under review as a conference paper at ICLR 2023
Ethics Statement This work includes pre- language models, which have the potential risk of inherited bias from the training data. Our empirical contribution is on accelerating the pre-train are-training process and thus does not focus on addressing such risk. For fair comparison, the pre-training data we have used are the same as previous works, and consequently the models we trained to evaluate our approach are similar to those already open-sourced. We refer to Bommasani et al. (2021) for discussion of the risks of pre-trained language models.
Reproducibility Statement Proof of lemmas in $ is given in $ A.1 Following the derivation of the Amos update rule, a heuristic derivation of the asymptotic behavior of the Amos decay factors is found in 8[A.2] and its connection with SGD is discussed in 8 AA.3 Assumption in our derivation
verified by experiments A/A.6] We explain the calculation of rj for the Transformer models in
$A.44 For the pre-training experiments in $5.1 we describe detailed settings in | and present
learning-rate search for BERT-base as well. Fine-tuning experiments on MNLI are shown in AA.7
Furthermore, an ablation test for the memory reduction settings of Amos is found in A Additional
experiment settings and results of training the ResNet50 model on ImageNet are found in A.9 Our
code is open-sourced at: https:: haap:::/aanoymouuus //a///nnnoouuuuaaeaaaaa url
REFERENCES
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. abs/1607.06450 2016. URL [ttL::t///axxi..oo www./r////////o//////////..eeee/ 1600/66006000 160700647070000
CoRR,
Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri S. Chatterji, Annie S. Chen, Kathleen Creel, Jared Quincy Davis, Dorottya Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah D. Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel Ho. Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark S. Krass, Ranjay Krishna, Rohith Kuditipudi, and et al. On the opportunities and risks of foundation models. CoRR, aass/208..02588, 2021. URL https::/arxii: org/ awwp////////////oaeooo////s/sa// 2108.100000 07258
Léon Bottou. On-line learning and stochastic approximations. In David Saad (ed.), On-line Learning in Neural Networks, Publications of the Newton Institute, pp. 9-42. Cambridge University Press, 1998. URL htttps:// /ww.. bottou. .rg/ wawbliiaaiiss/pppoooooo /ondf//n// iineinne-i 1998. pd
Léon Bottou. Stochastic gradient tricks. In Grégoire Montavon, Genevieve B. Orr, and Robert Müller (eds.), Neural Networks, Tricks of the Trade, Reloaded, Lecture Notes in Computer Science (LNCS 7700), pp. 430-445. Springer, 2012. BtLh:t/s::/// http:////een..bttttuoo ort.t/ooooooooooooooooo wrg.papers/ papers/
bottou-tricks- 2012
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
http: pettp://gh /ttp://ggttt..tt// cww/gooole///////oooooooo
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann. Parker Schuh, Kensen Shi, Sasha Tsyyascchekko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay. Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Er- ica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern,
10

---

Under review as a conference paper at ICLR 2023
Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language model-
ing with pathways. CoRR, abs/2204.02311, 2022. doi: 10.48550aaXXi..204400111 URL
httpss:/dolt http///o.....g...// www.oo.goooo...... 48550/ arxiv. 22044.v...204... 02311
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA. pp. 248-255, IEEE URL https:tttti....ooo
Computer Society, 2009, doi: 10. 11..1///VV..2000902200 1109/CVP3 2009.52068442.
Jacob Devlin, Ming-Way Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171- doi: 10.18653//11119-1444 1423. URL
4186. Association for Computational Linguistics, 2019.
https: http:///ooo.og//.....
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):11112255, 2011. URL
pectps://mmlrs www.oo/gggpprr...
/vl v22 httpp//////r..rr//ppppeerrss duchilla html
Felix Gers, Jürgen Schmidhuber, and Fred Cummins. Learning to forget: Continual prediction with
LSTM. Neural computation, 12:2451-71, 10 2000. doi: 10. 1622009997600000000055
Saeed Ghadimi and Guanghui Lan. Stochastic first- and zeroth-order methods for nonconvex
stochastic programming. SIAM J. Optim., 23(4)2341-2366 2013. doi: 10.1133//20888811. URL
pettps://doi.orgee 1117001111171101111111
Vincet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. CoRR, 000//882200968. 2018. URL
09568
Elad Hazan. Introduction to online convex optimization. CoRR, abs/1909 2019. URL
05207
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778. IEEE Computer Society, 2016. doi: 10.1 2.11099VCR..200...99
сттоr linear units. CoRR, htR.tttss:/ http///oooog/.//./////oo//////// 1tt.9/.....11/////.. 2016.
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian 08415, 2016. URL
/006
08415
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 448 456. JMLR.org, 2015. URL http::/proceed/ http::/pprcceedlng/ www.mrr...mmmmm............ wwwpsrrssssee.e. wwt.////1/oo..f..m/ .hmmlml1.mmmm..mmmmm
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, 0as/20011.836/1 2020. htLhttps: arxiv. wwtp:/wwwwwgaaa/////aaaaaaaa
08361
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. BRRLhttts:///rrxt org/ abs. awtpaab/.aaaaaeaa 412. 6980
Ananya Kumar, Aditi Raghumathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform distriiiibitiinnn In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, OpenRe- view.net, 2022. URL htRlpttssst htttps:://ppereviiiii net w htt.:///www......eeooeeeeeee//en

---

Under review as a conference paper at ICLR 2023
Weixin Liang Yuhui Zhang, Yongchan Kwon, Serena Yeung, and James Zou. Mind the gap: Understanding the modality gap in multi-modal contrastive representation learning. CoRR. abx/2203.020533, 2022. doi: 10448500aXX1..220.00055.. URL /atpa:///ol.org/11 48550/arXiv. 1ttp:0/aaxiv........o. 2203.0003333
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeutlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692, 2019. URL
abs 119p//11907.2222
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL BtLLittsss//peeeeieeee net www.mmmiiiiiaiee.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern over-parametrized learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 3325-3334, PMLR, 10-15 Jul 2018. URL
https://proceecte wwwpeiinnnssm.....mmmm Tivees//v0//nee ptw.a/ppeee////mmmaaa...aa. htmml
Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building large annotated corpus English: The Penn Treebank. Computational Linguistics, 19())311-330, 1993. URL
https: http:::://aclaanthhoe wog..hooooy..oo ogy. org/ (a0g//933200
Gábor Melis, Tomáš Kočiský, and Phil Blunsom. Mogrifier Istm. In International Confer- ence on Learning Representations, 2020. URL https::/opperevii ew net. forum?
5Je5P6EYvS
Noboru Murata. A statistical study on on-line learning. In David Saad (ed.), On-Line Learning in Neural Networks, Publications of the Newton Institute, pp. 63-92. Cambridge University Press, 1998. URL httthh:t///www/ www.eeecheaaae....eee. .eteeennee/eee hwt.://ubb.ctttto//6666666 Statistical Statistical_Study. on l--linn__Learinn
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Johannes Fürmkranz and Thorsten Joachims (eds.), Proceedings of the 27th International Conference on Machine Learning (ICML-10), June 21-24, 2010, Haifa, Israel, pp. 807-814 pdf
Omnipress, 2010. URL https: /mmm..mmm..am Cwa.mnccc//oooie http:///.......ooeeeeeeee///////
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. J. Mach. Learn. Res., 21:140:1-140:61 2020. URL httt::/mmmrr.rgg
[papers/v21/20-00) htmll
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
met/foruu www.ooiidummmuuuuuiiua
RZ
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations.
In Marilyn A. Walker, Heng Ji, and Amanda Stent (eds.), Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT, New Orleans, Louisiana, USA, June 1-6, 2018, Volume (Short Papers), pp. 464-468. Association for Computational Linguistics. 2018. doi: 10.18653/v1/ n18-2074. URL (tRLhhtts:://cco /ttp://///oooooo//./o./o 1ww53//1//nnni 2074
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4603-4611 PMLL, 2018. URL mlr. prwss/ v80 http://////hhheeeeeeee swwzeerlseaee htm.llmml..mmmm...mmmmmmmmmmmm..
Charles Stein. Inadmissibility of the usual estimator for the mean of multivariate normal dis- tribution. In Berkeley Symposium on Mathematical Statistics and Probability, 1956. URL http::/proojet id. w...ddd..ooo.....oeeeooeooooooo httg:///wne....0mmmm//0000001110
12

---

Under review as a conference paper at ICLR 2023
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with
rotary position embedding. CoRR, abs/2104.09864, 2021. URL htt.ttaaa/aaxvv.o////////// 2104. 09864
Ran Tian, Joshua Maynez, and Ankur P. Parikh. Shatter: An efficient transformer encoder with
single-headed self-attention and relative sequence partitioning. CoRR, 0tbs/2108.130330 2021. URL
https://arxiv: org/ http//a///// 2108 13032
T. Tieleman and G. Hinton. Lecture 6.5 RMSProp, COURSERA: Neural networks for machine learning. Technical report, 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 Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 1998-6008, 2017. URL https: /tttps::/ppcceeeinngg/ .nw.eeeeeddnnn..nnpp............ neurips. cc h0ahhaahh/ Absstaaraaaaaaaaatatttatc html
Abstract
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long Papers), pp. 1112-111 1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653//1N11--111 1101. URL https 7/acl /accannho ogg.
Yang You, Igor Gitman, and Boris Ginsburg. Scaling SGD batch size to 32k for imagenet training. CoRR, 1008/100..00 03888, 2017. hRttp:////axx http://aaaii.ooaaa////////aa// abs/ /1708 03888
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. In International Conference on Learning Representations, 2020.
Yukun Zhu, Ryan Kiros, Richard S. Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In 20/5 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 19-27. IEEE Computer Society, 2015. doi: 10.1109/1CC22001.111. URL/tttp:: doi. org, /00 1101/ ICCV. 2015.
A APPENDIX
A.1 PROOF OF LEMMAS
Proof of Lemma 4. We have
E E 0.41) [f - ā, g.]]
EE+11((2441 0.) (ā. о g.) V((=411 0.)] + o(ār):
where @ denotes element-wise multiplication, o(á:)/ + 0 at || + O. and arrays are flattened vectors for the dot-product. Since 1++1 and 2, drawn from the same distribution,
have E,+1 [ 0+)] Eg [x(;; 0.)] and o.. - 0.)] = (.. (VEzz; 0.)].
Moreover, because og does not depend on 21, we have E, a @ gil ai C E. E.G. Thus,
[E4 0111) = E, [e(; 0,)] (ā. C E, g.
E, V2(2; 03) + o(ā)).
Now E, = E, (2) 0.)] by definition, so (ā, 3 E, (g.)) E, VV 0.) > 0. and the lemma follows by taking oz small enough that o(āx) can be omitted.
13

---

Under review as a conference paper at ICLR 2023
Proof of Lemma 4.2 In the LHS of Equation only Ps depend on 2,; while in the RHS, M2 (g:)² )² depends on 24 but a2 does not. In order to satisfy Equation on every 2t, it is necessary that pt has a (2)9)) factor: Pr M2,g) Moreover, require that Blpe] does not depend on gt. Pt should
be normalized by (g) ) Pa [[M2][9))]] This, substituted back into Equation implies that
Ma(o:)2
Og x
VE[Ma(g.)91
Proof of Lemma 4.3
Equation 2 implies E 8o 00go- Taking M of this equation,
have
B[M2(e1) M2(50)²
[ ED + Bg
=M2(E0)2 Elgol En +a? (())2
Since go and Eo point to the same direction, we have [ co Assumption we have M2 (Tlgg)))/ M2 2 &. Hence,
M2(Elgol) M2(50). By
B[M2(E1) S M2(E,)² - 2n5 M2(E,) +
The RHS above is quadratic function of a, which achieves minimum at M2((E)). Finally, since (usually) 0, is initialized close to 0, and M2(0") s 17, we have M2(50) ss 17.
A.2 HEURISTIC DERIVATION OF DECAY FACTORS
Substituting Equation 6 into Equation we get the following equivalent of Equation 5
M2(Et) = dz7)-
Without knowing any specific relation among gt. Or and εt. we found difficult to theoretically decide an optimal ce Given that es decreases to 0, we set to decrease according to M2(E,) in Amos, ie. C ~ r M2),,, where is a constant and denotes asymptotically equal t - oc. Thus, by Equation we have de ~ M2(E,)?. We will analyze the evolution of M2(5,) 2 to derive o and d,
y
Taking E of Equation 4 and applying Equation 6 and Equation
7 get
E Egg.
de El100,) -E + M2(e)?? (8)
[[M)(E(+1) x M2(e,)² coE M2(es)
2
As in the derivation of the initial learning-rate, we make optimistic estimation that Elg,] and e have the same direction. Then, applying Assumption [we have
1
Elg.]
(2(E(g (g.)
er
Ma(e:) M2(ee.),
K V M
and Equation iimplies
de E[v00 E, E,+ di M2(22)
B[ < 2 2 2
dl d. E|y 0,
< - -
where the last equation we have used the fact that e, S Now, order to estimate 0, Er, we assume that O, will be evenly distributed on the hypersphere of radius M2(ex) around 0 as the training proceeds. Then, if & 2 3, for most 0, from the distribution we will have 0* E, z 0. this case, we have 0
(E, + 0*) E, 2 M2(E,)². and "on average" it is safe assum that + E, 2 q 2 for some constant q 0. Then, Equation I becomes
E[M2(E+11) < 22 E[ de) M²(e.)²
(10)
'I not useful in this work to provide a rigorous definition of "on average". We only point out its deep connection with Stein's 's example (Stein) 1956) that if & 2 3, an estimator with 1.2 regularization can be hetter than the maximum likelihood estimator without 1.2.
14

---

Under review as a conference paper at ICLR 2023
where we have used the fact that E/y]] ci In light of Equation 10 consider the following asymptotic difference equation:
Et+1 ~ - - + deg)e
(11)
where e, is intended to follow the asymptotic behavior of M2(E:)². Since we have de ~ f M2(5:)² it is natural to assume de ccc. Then, we transform Equation IT as the following:
1
(1 (1(1 d(q))
+ rel 1
)
e(+1 6t 1 1 70 + d(g) Et
e,
et
which we have used the approximation 1/(1 x) ~ + x applied to x 7(l1 + deq). Thus, the
update rule of by in Algorithm can be revealed by setting be 17 1 qr e2
bit1 = be te + 1).
r And de ~ -e implies de ~ so we set de dg =1
1
to satisfy both the asymptotic behavior and
1]
qb:
1 qb.
Similarly, since c, ~ r (2 we have a
where p
9 r7)
So we set G. V1 + pb;
to
Vpb,
satisfy the asymptotic behavior and co = 1.
A.3 CONNECTION TO SGD
The derivation decay factors in Amos ($..22 is largely inspired by SGD (Murata 1998). this section, we recall the theory of learning-raie -rate schedule of SGD and discuss its relation with Amos.
The update rule of SGD is simply 8, f- 0r where ay is scalar learning-rate. It recommended to set the learning-ratee taatingg--tteeaaee schedule Og T+ where is the initial learning-rate and À is the smallest value of the Hessian (Bottou) 2012). This is based on the following discussion.
Lemma A.1. Assume 0, is in a neighborhood of a local minimum 0-, such that the gradient E [ approximated by Hé, via Taylor expansion. Here, H = [E[ 0*)] is the Hessian at 0* Let 0 < X be the smallest eigen-value of H. Then,
E{(2(E++1) < M2(E,)² - 2λo, M2(E,)² a R(22
(12)
and the minimum of RHS of Equation 12 is achieved by
M (
X M E[M2 )1
and $
(13)
Proof. Since 0* is local minimum, we have ViV((2 0*) positive definite. Given the smallest value of H, we have Elg.] ē, 2 E| Applying this to E (+)] of Equation? 2 we get Equation 12 Now the RHS is quadratic function of a, and it
0 and Eig ~ Hē, where H is
takes minimum Equation 13
the lemma follows.
Note that both Amos and SGD analyze the evolution of M2(e,)² by estimating oggr Et. For SGD this is achieved by approximating Elğ.) with the Hessian. For Amos, the other hand, we have to make Assumption I due to the gradient normalization factor 1/ E// (g.)). J²] In both cases, the learning-rate decay is derived by setting a, in terms of M2(e,) so that (e,), decreases fast, then solve the asymptotic behavior of M2(E,).
Heuristic derivation of a;: We assume lim (M 0. In light of Equation 13 we consider the following asymptotic difference equation:
2
Et+1 E
(14)
15

---

Under review as a conference paper at ICLR 2023
where ez is intended follow the asymptotic behavior of M2(E,)? We transform Equation 14 as:
1
1
1
~ (1
X²
1
X
c
c++1
Ct
Ct.
C
X 1. Now, since oz
A XM)(E,)
A
1 Xt
a
so we have
we have Or
v et ~
So Oz
(
1 a
satisfies both the asymptotic behavior and ag a.
In the above derivation, the assumption lim Bl = > states that E| (g.)² )a] will converge to some non-zero value and will not further decrease. This often described intuitively as "the stochastic noise of sampled gradients does not vanish", a characteristic feature in the theory of SGD. is in drastic contrast with Assumption] | We assume that 2 ( ) decreases along with (2 in Amos. Ma et al. (0118) 8) pointed out that the vanishing of (M(2(gg)) )²) might lead to faster convergence; but to our knowledge, Amos is the first work to use the vanishing of (M)2)
to actually develop optimizer that empirically converges faster.
For SGD, the hyper-parameter, À is generally unknown; but if we adopt an L2 regularizer of strength X, it is guaranteed that À Z X, so one can safely set the learning-rate to T+ax7 (Bottou 2012). In Amos, the strength of 1.2 regularization 5t takes similar role in controlling the speed of learning-rate decay. We expect this work to inspire more theoretical investigation into this principle.
A.4 THE CALCULATION OF 7
As explained in $2 for a linear transformation y æW + (W, C θ, W E R x E RT)). set (W) Fy (a (a=vm) and ry(t)) d,/2, where a, is the standard deviation of entries of x and by the standard deviation entries of y. The values of 0r and oy are constrained by connected layers, and non-linear layers usually expect entries of input/output tensors from some approximate range. In Table we show 3 types of non-linear layers that occur in Transformer, and specify their input/output range (i.e. expected standard deviation) used for calculating 7).
For activations, e.g. GELU (Hendrycks & Gimpel 2016) in the Multi-Layer Perceptron (MLP) block, the input range set to 1 because the non-linearity of the activation function mostly lies within that range; and the output range is set / 1/2 because the activation function, as similar to ReLU (Nair & Hinton, 2010), will map negative values (which account for 1 1/2 of the input dimension) to close to and approximately retain positive values.
For Softmax of n classes, the input range is set to because the derivative of exp(x) is close to 1 within the |oc| S range (so Softmax is most sensitive to values within this range); and the output range set to /I/n because the output is an n-dimension vector of L2 norm (so the quadratic mean of entries - vI/n).
For LayerNormalization (Ba et 2016), the input range is arbitrary because the input will be normalized. The output range is expected to be 1.
We will discuss the calculation of r for specific models in the next sub-sections.
A.4.1 BERT, ROPE AND RPE
For BERT, RoPE and RPE, the multi-headed attention layer receives the hidden state ær, and the linear transformations & (i.e. the query) and K (i.e. the key) are expected to have standard deviation
1 so that the dot-product v1/h(=Q) (æ K) (i.e. attention score) has standard deviation as well
Type of Non-linear Layer
Input Range
Ourtput Range
Activation in MLP
1
v
1/2
Softmax of classes
1
1/n
LayerNormalization
N/A
Table 1: The input/output range of non-linear layers we specify this work for calculating ń
16

---

Under review as a conference paper at ICLR 2023
Type of Variable
1
Remark
Bias in all Linears
0,5
Layealoomaiization Scale
Input Embeddings
1/d
d is the size of hidden states
MLP/Dense2/Kemel
V 2/m
is the size of intermediate activation in the MLP
Other Linear Kernels
V v1/d
d is the size of hidden states
Relative Position Embeddings
0.5
Table 2:
The 7 calculated for variables in BERT, RoPE and RPE. MLPP/Dens22/Kemel is the linear kernel for the output layer of the MLP block. Other linear kernels include e.g. query, key and value kernels in the multi-headed attention layer.
(and this is why there is the scaling factor / 1/hh where h is the size per-head), which is expected by the Softmax for calculating the attention probability. Therefore, the output ranges of Q and K are 1. For RoPE, the dot-produe is replaced by a bi-linear form which encodes relative positions, but this does not change the scale because the bi-linear form orthogonal.
For other linear transformations in the model, the outputs are either fed into the activation function of an MI.P (which requires input range 1), or serve as summand in Residual Connection where the residual part comes from LayerNormalizationn (which has range 1). So all the linear transformations have output range in these model architectures.
Thus, we set the y of bias in all linear transformations to 0.5, and the for kernels is categorized by the input range and dimension, as we show in T
The input embeddings (i.e. token embeddings, position embeddings and segment-type embeddings) are inputs to LayerNormalizatin so their scales are not constrained there: but the token embeddings also used the linear kernel for producing the logits of token generation, which expects input range (because it comes from LayerNormalization) and input dimension d (where d is the hidden size), so is set to VI/d.
For the linear kernel of the MLP output layer MMLP/Dense2/Kernell) the input range is /1/2 because comes from non-linear activation, and input dimension m the size of intermediate activation in the MIP, so is /2/m.
For all other linear kernels, the input range is because it comes from LayerNormalization, and input dimension is the hidden size d. So is VI/d.
The relative position embeddings in the RPE model used as input to the key and value transforma- tions at each layer, similar to the hidden state. We set 17 to 0.5 its scale is close to the hidden state (which has scale 1) but will not dominate
Experiments with Amos-*Scall In 5.1 we have experimented with pre-training RPE and BERT- large with different 7 (Amos-*Scale). For RPE (Figure [Tb, we tried setting 7) of the relative
Type of Variables
Remark
LayerNormalizationn Scale
Query Kernel
/ / (hd)
is the size per-head and d is the size of hidden states
Input Embeddings
MLP/wo/Kemnel
V 2/m
m is the size of intermediate activation in the MLP is the size of hidden states
Other Linear Kernels
V1/d
Relative Attention Bias
0.5
Table 3: The n calculated for variables in T5. MLP/wo/Kermel is the linear kernel for the output layer of the MI. block.
17

---

Under review as a conference paper at ICLR 2023
position embeddings to v1/d instead of 0.5. For BERT-large (Figure 24), tried setting 7) of MLP/Dense2/Kernel to / instead of /2/m. They both had impact on performance. Especially for BERTT-arge, V1/d and /2/m only differ by 2 factor (because m = 4d), still the performance gap is significant. It illustrates the importance of setting r appropriately.
A.4.2 T5
For the T5 model, is set in Table It is different from Table due to several differences between the T5 architecture and BERT, as discussed below.
1. Linear transformations do not have bias terms in T5.
2. Attention score is calculated by (aQ) (æ ) in T5, without the scaling factor. Instead, smaller scale vI/(hd), with extra /1/h factor
the query kernel Q is initialized compared to K. Thus, we accordingly set 7) of the query kernel to /1/(hd).
3. The token embeddings we set to 1, which the same as the scale for initialization.
longer re-used for producing logits of token generation. So
4. The MI.P activation function (i.e. gated-GELLU) used in T5 different from BERT. Still, for the linear kernel of the MLP output (MLP/wo/Kernel) is set the same.
5. We set 17 of the relative attention bias to 0.5 so its scale is close to the attention score (which has scale 1) but will not dominate
A.5 DETAILED EXPERIMENT SETTINGS AND LEARNING-RATE SEARCH
In this section, we discuss detailed settings of the pre-training experiments in 5..1 The hyper- parameters and required computation resources are shown Table For pre-training BERT with AdamW, follow the settings of Liu et al. (2019). For RPE, rre-training TPU is slow, use a different configuration with more TPU cores to train the base-sized model. For T5, we found that using / 0.98 for Amos and Adam W causes training instability,
decrease the value to 3 0.95. The settings of AdaFactor follow Raffel et al. 2020) and Shazeer Stern (2018).
For encoder-only models (i.e. BERT, RoPE and RPE) trained the Wikipedia+ Books corpus, use the Penn TreeBank corpus (Marcus et al. float32. Number of warm-up steps is set to 10k for AdamW and 20k for Amos.
1993) as the validation set. The training precision is
For T5, the training loss is cross-entropy with an extra regularization term, (log 2)² (where the normalization factor in Softmax), which makes the logits close to mean 0 and self-normalized. In
Batch Size
Optimizer
Learning-rate
#Steps
Resource
AdamW
0.98
2e-4
200k/300k
BERT-base
1024
Amos
0.98
0.01
300k
TPUv4 2x2x4
About 2 days
AdamW
0.98
2c-4
200k/300k
RoPE-base
1024
Amos
0.98
0.01
300k
AdamW
0.98
2e-4
200k/300k
TPUv3 8x8 About 4 days
RPE
1024
Amos
0.98
0.01
300k
Adam W
0.98
2e-4
250k
BERT-large
4096
Amos
0.98
0.01
250k
AdamW
0.99
le-4
1M
TPUv4 4x4x4
Rorg-
1024
Amos
0.99
5e-3
IM
About 4 days
AdamW
0.95
le-3
250k
T5-large
4096
Amos
0.95
0.01
250k
AdaFactor
0.8
0.01
250k
Table 4: Hyper-parament parameter settings and required computational resources. The Amos is corresponding to B2 Adam W and the (second moment) decay rate in AdaFactor.
/ in
18

---

Under review as a conference paper at ICLR 2023
(1) AdamW, Batch 1024
(b) Amos, Butch 1024
(c) Amos, Batch 256
Adsmi/ltte- te-4
Adaamii- te-4
Adomi/--o-
Amrs--O
Amos 0. 3,002
Amos-0002 02
Am1c--0.0
Arros 0.01
Arnos-O, 005
22
20
136
2206
3004
Figure 5: Validation loss for pre-training BERT-base. We compare different learning-rate for (a)
AdamW W with batch size 1024, (b) Amos with batch size 1024 and (c) Amos with batch size 256.
(a) Bayerr//Attemiinneeyybb
(b) Olytterl/tttenttinnnn Kemme
(c) Layer Of/
Aayee10/Atteemitee Attention VateenttooValueeiiiia
Amos-O
Amtze 02
144
1.0
Figure 6: Plot of the ratio M, for variables in the BERT-base model, over pre- steps.
(
we plot cross-entropy for validation loss instead of the loss used for training. The training precision of T5 is 6. Possibly because linear transformations in T5 do not have bias terms, we found the model easier to train than BERT, and Amos can be applied without warm-up of learning- rate. The number of warm-up steps is set to 10k for both AdamW W and AdaFaccoor Learning-rate decay is in proportion to t 112 (where t is the step) for AdaFactor and linear for Adam
Figure 2..
For pre-training BERT-base, present a learning-rate search in Figure 5 For AdamW (Figure 5p) a smaller learning-rate significantly slows down the convergence, while a larger one results in a bumpy validation loss but almost the same performance. On the other hand, both smaller larger learning-rate can degrade performance for Amos (Figure(3po) Comparing Fig and Figure also verify a theoretical prediction about the global learning-rate of Amos in $/4 i.e. the best learning-rate for Amos is in proportion to the square-root of the batch size: Training with 4x the batch size matches 2x the learning-rate.
5p)
A.6 VERIFICATION OF ASSUMPTION I
In Assumption we have assumed that (F(a))) 2 O for all i and across all variables. E[g.] (B[M,
and ([M2(g.) (g.) IT can be estimated by taking the running average of gi and M2(ge), respectively; so in Figure 6 we track the pre-training of the BERT-base model, calculate the running averages with exponential decay rate 0.98, and show some typical plots of the ratio. We note two characteristics of the plots: (1) the ratios are increasing the training proceeds, which suggests that taking global constant & to satisfy Assumption I is indeed possible; (2) starting points on the left of these plots are similar across different learning rates, which suggests that it is detectable in the early stage of training whether a learning-ran rate is too small or too large. In fact, in all plots for all variables we can see that
the ratio (.(g()) > 0.01; the appropriate global learning-rate can be read from these plots.
A.7 FINE-TUNING RESULTS
In Table 5 we show fine-tuning results on the MNLI (Williams et al. 2018) dataset. We compare checkpoints pre-trained for 150k and 300k steps with Amos, and the final checkpoints of AdamW- 200k and Adam 3-00k. We fine-tune all checkpoints using the Adam optimizer with learning-rate rate 5e-6, batch size 16, and evaluate by the best accuracy the MNLI dev set among every 1k of 200k training steps. We run each experiment 3 times and report the mean and standard deviation. The checkpoint pre-trained for 150k by Amos already outperforms the final checkpoint of Adam 3000.
19

---

Under review as a conference paper at ICLR 2023
MNI.Imatchhed
MNL.Imismatched
Amos 150k
84.15 #.40
84.17 #.37
Amos @ 300k
84.72 1.15
84.44 4.26
AdamW- W-200k
83.19 #.37
83.45 4.41
Adam W-300k
83.84 1.14
83.88 1.17
Table Fine-tuned accuracy on MNLI dev set. We show the mean and standard deviation of 3 runs.
Thus, the faster convergence by Amos in pre-training indeed transfers to better performance in fine-tuning; we can save 50% of the pre-training cost by using Amos instead of AdamW.
A.8 ABLATION OF MEMORY REDUCTION
In this section, we experiment with different settings of the memory reduction. We compare the current setting of reducing the input dimension for linear transformations (Reduce. Axis), to no memory reduction all Reduce, and the setting of reducing both axes for linear transformations (.educe.Dense). For embedding matrices, no axis is reduced in the No_ Reduce setting, and the embed dimension reduced for both Reduce_ LAxis and Dense. We have tried reducing both axes for embedding matrices as well, but found the training unstable in this setting. The comparison of memory usage slot variables is shown below.
AdaFactor (No Momentum) Reduce_Dense < -edis < AdamW < No Reduce.
Without memory reduction, Amos (eduReduce) consumes memory than AdamW because it has more slot variables (ö,, b. m. 66 m.). When memory reduction is applied, the memory usage of ēz, b. becomes negligible compared the momentum m,, so Amos (Reduce_ AAxis and Reduce Seduce_Dense) requires < 51% memory for slot variables than AdamW. The memory reduction method used by Amos more efficient than the matrix factorization used by AdaFactor, but in the pre- trainininng of T5 (FigurePp), AdaFactor achieved favorable performance (although slightly worse the end than Adam W with linear learning-rai rate decay) without using momentum, reducing the memory usage further. Whether Amos can achieve similar performance without using momentum unclear yet.
tran//air_Jsa
Fainnalld Joss
X
22
No Reduce
10
Reduce Dense
24
Reduce 1Axxs
22
1.
Figure 7: Pre-training BERT-base using Amos with different memory reduction settings.
show the training and validation loss of pre-training BERT-base by Amos with different memory reduction settings. Reduce Dense is slightly worse in training loss compared to No. RedRed but so much in validation loss. On the other hand, R is almost the same as No_ Reduee in training loss, and generalizes even slightly better in validation loss than the other two. So the current Reduce_ LAxis setting for Amos is favorable.
In Figure 7
A.9 TRAINING RESNET50 ON IMAGENET
In this section, apply Amos to the training of ResNet50 (He et al. 2016) the ImageNet dataset (Deng et al., 2009). ResNet50 is a deep Convolutional Neural Network of 50 layers, with Batch Normalization (loffe & Szegedy 2015) and Residual Connection. ImageNet 1000-class image
20

---

Under review as a conference paper at ICLR 2023
Type of Non-linear Layer
Input Range
Ourtput Range
ReLU Activation
1
11/2
BatchNormalization
N/A 1
Max-pooling on patch size n
1/v2lnn
Table 6: The input/output range of non-linear layers we use calculate 7j for ResNet.
classification task with 1 1.28M traning examples. We train with batch size 1024, on an 8-core TPU machine. The settings for Amos out-of-the-box: the syyper-paammete 3 set to 0.95, warmup steps 5k, and the global learning rate is set to 1 0.028, where N 1281167 11224 is the number of vN
batches in the traning data. We use the open-sourced init2wing codebase to run the experiments.
A.9.1 THE CALCULATION OF 7j FOR RESNET
In order to calculate the hyper-parameter 7 for ResNet, we specify the input/output range of 3 types of non-linear layers in Table 6 This is similar to Transformers, with the only specialty that the output range of Max-pooling layer is set to 1/v2lnn, where n is the patch size. This is because the maximum of normally distributed random variables has standard deviation of about 1 1/v 2 22
The calculated v for different types of variables ResNet is shown in Table
BatchNormalization is treated the same LayerNormalization in Transformer.
The projection kernel of the first residual block is scaled up by v2 In because of its previous max-pooling layer of patch size n.
The 2nd and 3rd convolution kernels in each residual block is scaled up by v2 because their inputs come from a ReLU activation.
The variables for bias and other linear kernels are treated the same as in Transformer.
Settings of Amos-"Scale We also tried an Amos-**Salll setting where the n for the projection kemel of the first residual block is set /1/d instead of (22ln (in ResNet50, x 3 9).
A.9.2 RESULTS
In Figure Sa, we show the validation error rate of Amos and Amos-*Scale, where the error rate for Amos (0.261 lowest) is slightly better than the Amos-*Scale setting (0.263 lowest). Furthermore, it is known that a strong L2 regularization is beneficial for many popular image classification tasks Loshchilov & Hutter 2019), but Amos does not have hyper-parameter parameter to adjust the strength of L2
Type of Variable
Remark
Bias
0.5
BatchNormalization scale Projection kernel of the first residual block
1
n is the patch size of the previous max max-pooling; d is the input size
The 2nd and 3rd convolution kernels in each residual block
2
d is the input size
Other linear kernels
Vidd
is the input size
Table The calculated for variables in ResNet. Other linear kernels include convolution kernels and the final linear classification kernel.
htttp:://ggttbb...ommgggg//////e hhttps://en.wwiiieddi// hheorendGumbelldiistnn
sher $8288082200000 9311 ppett k,2880 933 3.Gnekkko.
21

---

Under review as a conference paper at ICLR 2023
Validation Error Rute
(b) Amos vs. Adam
(c) Amos-Extra vs. Adam
Arnos- Scale
616
10-
335
Arnos-EBa Extra
610
m
42
320
334
69
Amoss
Figure 8: Training ResNet50 on ImageNet. We plot error rate of the validation set.
regularization; so we tried an ad hoc setting Amos-Extra, where the Amos update rule (Equation 1 is replaced by 6, + de Ve, + (17 + 0.001)0,) with everything else kept the same (we also tried other constants, but 0.001 was the best). As shown in Figure Amos-Extra (0.242 lowest error rate) significantly improves the performance on ImageNet.
In Figure (8, we compare the schedule of Adam is set to cosine decay with 5% warmup, and the number of training steps is set to 140k. The base learning-rate with 25 runs. Other outperforms all the 25 runs; the best 6 of the 25 are shown Figure Sp. As alternative settings for Amos, we have also tried = 0.98, 0.999, or & 0.02, even changed the decay factors to
Amos with Adam (no weight decay). The learning-rate
tuned by random search of log scale between 1e-5 and 1e-2, are set to the default (i.e. 81 = 0.9 and f3 = 0.999). Amos
(1+ in VEb.) and d, = (1 + 1 same validation error rate, sometimes with slightly slower convergence.
le VErjbe)
All the other settings converge to almost the
In Figure we compare Amos-Exte with the state-of fiite-of-theee ahe- settings of AdamW. The learning raae schedule of AdamW set to cosine decay with 5% warmup, and the number of training steps set to 187k. The base learning-rate, weight decay strength, and label smoothing rate (defaults to 0.1 for other experiments) are tuned by random search, of log scale between Ic-4 and le-2, log scale between le-2 and 1.0, and linear scale between 0.0 and 0.2, respectively, with 25 runs. Other hyper-parameten are set to the default (i.e. 31 0.9 and 8g = 0.999). Among the 25 runs, 9 of them outperform Amos-E Extra, which are shown in Figure Se. The best performing settings of AdamW gain their advantage close to the end training, which is probably due to the interaction between weight decay and cosine learning-rate schedule. On the other hand, Amos-Extra demonstrates faster and more stable convergence.
To conclude, when applied ResNet50 ImageNet, Amos can outperform Adam out-of-the-hoo and become comparable to the state-of-the-arr AdamW settings by adding small constant weight decay term. However, the extra weight decay term is ad hoc, cannot be covered by our current theory (because we have assumed that the L2 regularization is weak enough and decays to 0, not to bias the loss function but only constrain the scale of trained variables), and probably is not the optimal way to strengthen L2 regularization. It leaves the problem of searching for a more general working theory that enables stronger 1.2 future work.
22


Reviews:
[{"review_id": "846anjxQl3E", "reviewer": "Reviewer_Ebn1", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "The paper introduces Amos, a first-order DL optimizer with adaptive learning rate and decay. The proposed contributions are:\n\n* Outperforming AdamW for pre-training language models\n* Providing guidance for hyperparameter tuning\n* Reducing memory usage\n* Allowing continuous training and resuming from checkpoints\n\nThe proposed optimizer leverages model-specific information, by partitioning the model parameters and adding a per-partition norm constraint. This norm constraint, together with the gradients, is used to adapt the learning rate and weight decay.\nThe parameter update involves a series of newly introduced hyperparameters with different functions, and the paper provides a series of heuristic derivations and experiments to set their values.\n\nIn experiments pre-training current architectures for NLP tasks, Amos is shown to outperform AdamW with 2019 settings in terms of speed, performance and memory usage.", "summary_of_the_review": "My opinion is that this contribution clearly deserves attention from the community, but it needs more work: It presents an optimizer of increased complexity, involving a substantial amount (potentially arbitrary) heuristics. Experiments are very promising, but there are methodological concerns.\nThe paper would greatly benefit from improvements in clarity, both in formulation and presentation, especially given the increased amount of details needed to understand and tune this optimizer.\n\nFor those reasons I'm inclined to marginally reject, but I thank the authors for their contribution and encourage them to address/answer some of my points above in order to reconsider my evaluation. ", "strength_and_weaknesses": "The idea of using model information to replace the hyperparameters is great, we have good examples of similar strategies from the weight initialization literature (Xavier, He) and extending this to adaptive optimizers sounds exciting. Furthermore, this combines also very nicely with the interplay between learning rate and decoupled weight decay, which has been recently proven useful to solve issues with first-order adaptive DL optimizers like Adam. Finally, it also brings better performance using less parameters in a block-wise regularized fashion.\nCombining these ideas is a very promising and well-grounded line of work.\n\nStill, the question of why specifically this model, is not fully clear to me, for the following reasons:\n* Compared to its counterparts, the model introduces increased complexity, but the motivation (section 1) lists reducing complexity as a goal. This is expected to be achieved through the \"guidance for hyperparameter tuning\", but I'm not convinced of such guidance (see next points).\n* There is a clear effort to justify the choice for heuristics and hyperparametrizations, but the efforts lack clarity both in substance and presentation (see notes on clarity) and the claim that they are \"theoretically supported\" is in many cases not true: many crucial aspects are resolved through trial-and-error (see e.g. end of page 6 and Appendix 7). In many cases, I feel like an ablation study showing the contribution of the added components in an empirical way would be clearer and more compelling.\n* Even the more analytical hyperparameters don't give the impression of being \"easier\" to find. E.g. for \"eta\", Appendix A.3 shows that many factors must be taken into account in a non-automated way, and new architectures may require careful tuning. What if we e.g. have batchnorm instead of layernorm? How is the concept of \"range\" characterized? Will it hold under non-stationary, noisy and/or sparse gradients?\n* Regarding stability: While very promising, the experiments aren't comprehensive enough to convince that the proposed settings (theoretical or empirical) are not highly specific to the results reported, due to the amount of hyperparameters and tuning involved. One thing that could help with this is to analyze the loss landscape stability (see AdamW paper), and/or extend experiments to other tasks and architectures.\n* Regarding experiments, it is unclear how many settings were tried before finding the reported Amos hyperparametrizations. This makes difficult to compare the different optimizers in terms of their hyperparametrization budgets. Furthermore, error bars are not provided: they would be welcome (together with a more publication-friendly plotting mechanism), although the differences between Amos and the rest are very significant.\n\nA couple of thoughts regarding weight decay:\n* Please correct footnote in page 2: Loschilov&Hutter precisely mentions that L2 and weight decay are not equivalent for adaptive optimizers (beyond empirical success), the footnote seems to operate on the opposite idea.\n* We saw that decoupled weight decay helps overcoming convergence issues with Adam. Would re-introducing adaptive weight decay also re-introduce some of those issues?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7YzRdPMyyb", "reviewer": "Reviewer_f9Uo", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "This paper proposes Amos, an adaptive optimization method for deep learning. It combines various aspects, most notably an adaptive weight decay term and a step size involving an expected scale of model parameters, informed by structural knowledge about the model being trained. It is evaluated in extensive experiments on pretraining Transfomer models.", "summary_of_the_review": "I list several scientific shortcomings above, which are think are pretty fundamental. In light of that, the good experimental results for the proposed method on their own do not warrant acceptance in my opinion.", "strength_and_weaknesses": "### Strengths\n\n1. Improving optimization algorithms and/or estimating their hyperparameters using structural knowledge of the model is a very promising research direction. To my knowledge, this is one of only a few papers that uses model characteristics to inform the optimization hyperparameters (beyond layer-wise normalization and/or step sizes).\n\n2. The experiments on Transformers are extensive and consistently demonstrate significant improvements over state-of-the-art competitors. Given the importance and immense cost of such training tasks, that is a very significant improvement.\n\n### Weaknesses\n\n1. My main criticism is that the method combines a large number of aspects without proper ablations. If a method gives very good results, but we can't pin down which innovation really drives that improvement, the scientific value is greatly diminished. Comparing Eq. (1) to Adam or RMSProp, I can see at least 6 components/innovations that could be seen as largely independent and I would like to see investigated in isolation\n    - The factored learning rate involving the \"model-oriented scale\"\n    - The learning rate in front of the weight decay term being the square of the learning rate scaling the gradient ($\\xi^2 / \\hat{v}_t$ vs $\\xi/\\sqrt{\\hat{v}_t}$)\n    - The multiplier $M_2(g_t)$ in front of the weight decay term.\n    - The decay scheme for the global step size ($d_t$)\n    - The additional decay scheme for the weight decay term ($c_t$)\n    - Computing $\\hat{v}_t$ across model parameters instead of using running averages.\n\n2. The derivation of Amos is a series of approximations (some very crude) and heuristics/intuition. I could put some caveats on almost all the steps of the derivation. While the authors clearly state that this is a heuristic derivation, at some point we have to ask ourselves whether such a derivation is helpful or just obfuscates. Some steps that I found particularly problematic:\n    - Equation (6) suggests a particular _ratio_ of $\\alpha_t^2$ to $\\rho_t$. It is completely unclear to me how that should motivate $\\rho_t \\propto M_2(g_t)^2 / \\mathbf{E}[M_2(g_t)^2]$ and $\\alpha_t \\propto 1/\\sqrt{\\mathbf{E}[M_2[g_t]^2]]}$. This choice satisfies Eq. (6), but so does an infinite number of other choices.\n    - Assumption 1: This assumes $\\mathbf{E}[\\Vert g_t\\Vert^2] \\leq C \\cdot \\Vert \\mathbf{E}[g_t]\\Vert^2$, i.e., the variance of a stochastic gradient needs to go to zero as the true (expected) gradient goes to zero. That is the so-called interpolation regime: every data point is fit perfectly. This seems unlikely to be fulfilled in LLM pretraining tasks.\n    - Interpretation of $\\xi$. Sorry, but this is almost absurd. First, if we assume the per-example gradients to come from a zero-mean distribution this would imply $\\mathbf{E}[g_t] = 0$, in direct contradiction to the Equation below. Second, why should  $$\\left(\\frac{1}{N} \\sum_{i=1}^N x_i \\right)^2 = \\frac{1}{N^2}\\sum_{i=1}^N x_i^2$$ hold? That is simply not true and is not a meaningful approximation in any way. Third, the law of large numbers is completely irrelevant here.\n\n3. The experiments focus entirely on Transformer architectures. While this is certainly a model class of great practical importance, I would find it highly desirable to check the extent to which these findings transfer to other families of architectures and data modalities. Given the resources that went into the experiments involving Transformers, it would be \"cheap\" to throw in a ResNet on CIFAR-10/CIFAR-100/ImageNet. If Amos performs amicably there - great. If not, this would alert us of the fact that some of the proposed heuristics are specific to Transformers (which would be perfectly fine).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7V-WrBk7Ha", "reviewer": "Reviewer_Fga8", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "The submission presents a new optimizer, Amos, which can heuristically determine the initial learning rate and adjust its decaying schedules for LR and L2 regularizer during training. The experiments show that Amos converges faster than AdamW on language models.", "summary_of_the_review": "Due to the missing (1) experimental results on general tasks, (2) lack of convergence guarantee, and (3) too many confusing mathematic details in derivations, currently, I hold a negative score on the paper.", "strength_and_weaknesses": "#### Strength\n- The proposed optimizer can adaptively choose the decaying schedules for learning rate the L-2 regularize coefficient during training.\n- There are fewer sensitive hyper-parameters compared to the previous optimizers.\n\n#### Weaknesses\n-  Since the optimizer is not explicitly designed for the language model, the absent experimental results on other tasks, such as the vision or RL tasks, degenerate the contribution of this optimizer.\n- Some typos, e.g., step one in Algorithm 1, in which the $g_t$ should be bold. The bold and normal formats are used at will in the $\\alpha_t$ below Eq.(2).\n- Actually, $\\hat{v}_t$ is not nearly equal to $E[M_2(g_t)^2]$. In general, we have $$(\\hat{v}_t - E[M_2(g_t)^2])^2 \\leq \\beta (\\hat{v}_t - E[M_2(g_{t-1})^2])^2 + \\frac{\\beta^2}{1-\\beta}(E[M_2(g_t)^2] - E[M_2(g_{t-1})^2])^2$$ (see [1], Lemma 2), which is large until convergence. Hence it is strange to replace  $\\hat{v}_t$ with $E[M_2(g_t)^2]$.\n- The approximation in Eq.(4) is confusing. Since $\\alpha_t g_t$ can be small when we are close to convergence, but $\\rho_t \\theta_t$ is not. We can not ignore $\\rho_t \\theta_t$ directly.\n- Why $\\rho_t $  should be normalized by $E[M_2(g_t)^2]$? The stated reason ``the L2-regularizer should (on average) not depend on the norm of $g_t$'' seems weak and not solid.\n- Lack the comparison to some strong baselines, such as LAMB and Adam's variants.\n\n\n[1] Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject