Learning Rate Grafting: Transferability of Optimizer Tuning
Authors: Naman Agarwal
Rohan Anil
Elad Hazan
Tomer Koren
Cyril Zhang
Venue: ICLR 2022 Submitted (2022)
Keywords: Optimization
Learning Rate Schedules
BERT

Abstract: In the empirical science of training large neural networks, the learning rate schedule is a notoriously challenging-to-tune hyperparameter, which can depend on all other properties (architecture, optimizer, batch size, dataset, regularization, ...) of the problem. In this work, we probe the entanglements between the optimizer and the learning rate schedule. We propose the technique of optimizer grafting, which allows for the transfer of the overall implicit step size schedule from a tuned optimizer to a new optimizer, preserving empirical performance. This provides a robust plug-and-play baseline for optimizer comparisons, leading to reductions to the computational cost of optimizer hyperparameter search. Using grafting, we discover a non-adaptive learning rate correction to SGD which allows it to train a BERT model to state-of-the-art performance. Besides providing a resource-saving tool for practitioners, the invariances discovered via grafting shed light on the successes and failure modes of optimizers in deep learning.

Under review as a conference paper at ICLR 2022
LEARNING RATE GRAFTING: TRANSFERABILITY OF OPTIMIZER TUNING
Anonymous authors Paper under double-blind review
ABSTRACT
In the empirical science of training large neural networks, the learning rate sched- ule is notoriously challenging-to-tune hyperparameter, which depend all other properties (architecture, optimizer, batch size, dataset, regularization, ...) the problem. In this work, we probe the entanglements between the optimizer and the learning rate schedule. We propose the technique of optimizer grafting, which allows for the transfer of the overall implicit step size schedule from a tuned op- timizer to a new optimizer, preserving empirical performance. This provides a robust plug-and-play baseline for optimizer comparisons, leading to reductions the computational cost of optimizer hyperparameter search. Using grafting, we discover non- -daptiive learning rate correction SGD which allows BERT model to state-of-the-arr performance. Besides providing a resource-saving tool for practitioners, the invariances discovered via grafting shed light on the suc- cesses and failure modes of optimizers in deep learning.
train a
INTRODUCTION
Adaptive gradient methods and learning rate schedules are cornerstones of optimization for deep learning. The pursuit of faster and more robust convergence in training deep neural networks has led to a preponderance of optimizer update rules, annealing heuristics, and hyperparameters, with no clear or principled way to select between them. The difficulty makes tuning state-of-the-an deep learning pipeline very opaque and expensive.
Consequently, large body of work aims to understand the structure of these search spaces, ana- lyzing the properties optimizers and their components in idealized models predict directional trends, provide rules of thumb, and propose mechanisms for emergent phenomena (like general- ization and transfer). Other works study the intrinsic tradeoffs between competing objectives, such generalization. The continued prevalence of black-box hyperparameter
as faster convergence search points to a consensus that a quantitatively predictive understanding of training dynamics cur- rently eludes
Towards forming more robust beliefs this space, we provide a recipe that enables a wide class of ablation studies -tate-of-the-art models. We consider an unnatural-looking meta-algorithm, which combines the update rules of two different optimizers, forming single grafied update rule. Our findings serve not only sanity check on explanatory theory, but also way simplify the optimizer search space in practical settings.
1.1 OUR CONTRIBUTIONS
We introduce learning rate grafting, meta- alge--aag which blends the steps of two optimizers by combining the step magnitudes of one (M) with the normalized directions of the other (D). Applied to state-of-the-art deep models, grafting leads to the following:
Transfer of optimizer performance. We discover that in many deep learning settings, the per- formance of an optimizer can be explained by the implicit step size schedules induces. Grafting provides an adaptive way transfer these schedules between algorithms, consistently closing per- formance gaps between different optimizers used to train the same model architecture. As a result, grafting provides strong baseline for learning rate schedule tuning: given tuned baseline M,

---

Under review as a conference paper at ICLR 2022
D
M#D
Figure 1:
The optimizer grafting operation introduced and studied in this work: meta-optimizer MDD uses the update magnitude of optimizer .,, and the direction of DD
grafting can enable different optimizer D match (and sometimes exceed) the performance of M. without tuning the learning rate schedule for D. We demonstrate this phenomenon in multiple settings: across vision and language models, and different batch sizes.
Schedule discovery. We show how to grafting for the exploratory analysis of optimization trajectories and implicit learning rate schedules. Using grafting as a diagnostic tool, we discover per-layer learning rate correction that allows SGD with momentum (without adaptive precondi- tioning) to pretrain a BERT model to state-of-the-art accuracy for the first time. This correction is non-adaptive, and does not require the parallel execution of the original algorithm from which sched- ule is derived. We show analogous results for vision models, finding novel and effective global (as opposed per-layer) schedules for AdaGrad. We discuss the implications of this for contempo- rary theory and engineering practices in large-scale settings: the distinctions between widely-used optimizers may boil down to just a few implicit step size schedules, not the usual understanding of per-coordinate "preconditioned directions".
1.2 RELATED WORK
Learning rate schedules. Choosing learning rate schedules is a difficult empirical problem, for which many heuristics exist; the design and search of these schedules can result in dramatic perfor- mance differences (Gotmare et al. 2018: Smith 2018 |Smith & Topin 2019 Loshchilov & Hutter) 2016). In large-scale language modeling, source of particularly resource-intensive training prob- lems, the predominant empirical practice is to tune a warmup and decay alongside an adaptive optimizer (Popel & Bojar| 2018 You et al. 2019). Beyond classical minimax convergence rates for stochastic optimization, some theory has been established to explain the benefits of certain learn- ing rate schedules. Ge et al. (2019) provide a fine-grained theoretical account for quadratic losses. Agarwal et al. (2021) show that optimized non-adaptive learning rate schedules can induce stable and usable Nesterov-like acceleration. Li & Arora (2019) and Arora et al. (2018) study the inter- action of learning rates with batch normalization (loffe & Szegedy) 2013), while Li et al. (2019) provide a theoretical mechanism by which initial large learning rates improve generalization. The interaction of learning rates with batch size is explored in (Krizhevsky 2014 Goyal et al.) 2017]
Bottou et al.] 2018 Shallue et al. 2019). Anil et al. 2020 use the grafting method to stabilize an exotic full-matrix optimizer.
Adaptive optimizers in deep learning. Adaptive methods have turned out to be extremely robust in training deep neural networks. receiving tens thousands of citations for this reason. In partic- ular, Adam has been the de facio standard in fields such as NLP (e.g. (Devlin et al.] 2018 Yang
al. 2019; Liu et al.] (009)), deep generative modeling (e.g. (Karras et al. 2017 Brock et al.)
2018 Kingma & Dhariwal (2018), and deep RL (c.g. (Haarnoja et al. (018)). Adaptive methods
have seen adoption in extremely large-scale settings, necessitating modifications to reduce resource
consumption (Shazeer & Stem 2018 Anil al. 2019
Chen et al. 2019).
An important discussion was sparked by Wilson et al. (2017). ), who presented empirical and theoreti- cal examples where adaptive methods generalize poorly. Building on this premise, Keskar & Socher (2017) suggest switching from Adam to SGD during training. A variety of theoretical mechanisms have been proposed for the behavior of adaptive methods; for example, learning rate schedule auto-
2

---

Under review as a conference paper at ICLR 2022
tuning (Ward et al. 2019). escaping saddle points (Staib al.) 2019), and mitigating long-tailed distributions (Zhang et al. 2019). Reddi et al. (2018) construct, then mitigate, pathological set- tings where Adam fails to converge. However, in the vast majority of cases, out-of-the- adaptive methods are perfectly suitable for practitioners. Large-scale empirical studies (Choi et al.] 2019) Schmidt et al., 2020 Nado al. (202T) confirm that hyperparameter search, especially learning rate schedules, are essential to state-of-the-art performance, and undertuned baselines can often lead to faulty optimizer comparisons.
Layer-wise adaptive methods. Various algorithms have been proposed recently that focused on adaptive layer-wise scalings on top of base algorithms. Adaptive layer-wise scalar obtained via weight norms have been derived in LARS/LAMB You et al.,, 2017 2019) and successfully applied
to both Momentum and Adam. Several works (Shazeer & Stern 2018 Anil et al., 2019 Chen
al., 2019) propose compressed preconditioners to reduce the space of per-coordinate adaptive multipliers (as maintained by algorithms such as AdaGrad/Adam) to factored space. Novograd (Ginsburg al. 2019) reduces the per-coordinate scalars to per layer scalars via addition. LARS- like scalars on top of normalized SGD with momentum was used to train BERT in Cutkosky & Mehta 2020). To the best of our knowledge, non-adaptive per-layer learning rate corrections, such as those presented in Section 5 have not been explored.
PRELIMINARIES
2.1 OPTIMIZERS AS GENERIC STATE MACHINES
There are various abstractions of optimization in deep learning: for example, we could reason about stochastic first-order oracles V F(w) (w) which are unbiased estimators of some gradient of population objective F(w). However, this fails to capture various popular deviations from the standard empiri- cal risk minimization model, such as sequential batching in NLP (Merity et al. 2018 2017) deep reinforcement learning, where optimization deviates from its theoretical abstractions in many ways
[llys al. 2019 Engstrom et 2020).
For this reason, we take the maximally agnostic approach for our notation and definitions in this paper. We choose the software abstraction of optimization: we view a gradient-based optimizer state machine which takes a sequence of gradients Jr E IR and outputs update rule for the weights wt Idd. This abstraction is agnostic to properties or theoretical guarantees of the optimizer; in supervised learning, the gradients are usually understood to be averaged over minibatch, but we do not use this property either.
Thus, an algorithm A takes the current iterate and gradient (We, 9t), and outputs the next iterate W this update can depend on the algorithm's 's internal state, hyperparameters, side information, or randomness. Then, vanilla SGD with learning rate schedule 1): is the algorithm A specified by
A(wi, 9:) = w/ T)-tt
(1)
Then, in the usual process of iterative training, the model weights are iteratively updated, as in 1 A(w,, 9:). Even in simple theoretical settings like least-squares regression (Jain et al., 2018 Ge et al. 2019), the choice of the time-varying scalar hyperparameter {nc} already presents subtleties and bias-variance tradeoffs.
2.2 ADAPTIVE METHODS
We now bring in the second ubiquitous object of study: adaptive gradient methods. Broadly, they refer to the technique of maintaining a sequence of linear transformations He, dependent on the observed information (w.,ge so far, modify the (stochastic) gradient descent steps:
A(ws, gt) == w. -HHg
Despite their extreme popularity and theoretically principled origins, there is considerable uncer- tainty regarding the dynamics adaptive methods. They were first motivated by regret bounds dependent the geometry of the data (MeMahan & Streeter 2010 Duchi al. 2011 hence
3

---

Under review as a conference paper at ICLR 2022
the terminology "adaptive regularization". The view of loss curvature estimation leads to the term "adaptive preconditioning" where H, serves as an adaptive scaling factor to make the problem less ill-conditioned (Tieleman & Hinton 2012 Kingma & Ba 2014).
Although more exotic adaptive optimizers exist (Martens & Grosse 2015 Gupta et al. 2018 Agar-
wal al., 2019), focus on diagonal second-moment preconditioning methods, which have yet to be dethroned as the de facto standard. This family of optimizers accumulates the squares of the entries of gt, and sets H, be the diagonal matrix of entrywise inverse square roots of the accumu- lators.
Entanglements between preconditioners and learning rates. Notice that in Equation 2.2
scalar multiplier on the preconditioner H, can be absorbed into the learning rate schedule T)- In other words, the dynamics of He can induce an implicit step size schedule over the course of train- ing, and different choices of optimizers (or configurations of the same optimizer) lead to different implicit schedules. This entanglement motivates the introduction of our grafting meta-algorithm. Specifically, we design experiment to answer the following question: does the empirical success of an adaptive optimizer arise from the true preconditioner H,, or can be equivalently described by learning rale schedule Tte? By combining the step magnitudes of optimizer M and step directions of an optimizer D, we can isolate and disentangle adaptive preconditioners He from their induced step size schedules.
THE GRAFTING META-ALGORITHE
We begin by outlining the grafting procedure, which maintains two child optimizers M,, D and outputs a grafted update step; we call the entire algorithm MMDD. Before developing the meta- algorithm in full generality, walk through the execution of the simplest case:
1. At each training iteration, MDD feeds the same input (wi, 9t) both children M, D. which manage their states independently and produce outputs WM, WD.
2. We do not update W++1 yet. Instead, we compute I|w well and ||w - well. the norms of the steps the child optimizers would have taken.
3. The update step that grafting outputs is
Wt +
|WM (WD we)
WD K Wt otherwise
IeD we
which combines M's update magnitude with D's update direction.
We remark on a few elementary properties of this binary operation on optimizers:
Grafting is idempotent. If A. Z1, It are deterministic, then A#A= = A (meaning that these two optimizers have the same trajectory under identical sequences of inputs (w, Ic)).
If M, D differ only by learning rate schedule (i.e. WMM w; is always positive scalar multiple of wp Wt, under identical inputs (Wr, 9:)) then M#D = M.
Disregarding the special case where A outputs zero step, we have
(M#A)#D = M M# M#D.
Grafting is not commutative. M#D DHM.
Grafting is not necessarily a descent method. Even if M,D are guaranteed to make progress on a deterministic objective (i.e. when is VF(w)), M#D is not. For in- stance, if M,, are gradient descent with different preconditioners, there no guarantee that M converges, or ever makes progress on F(w).
Grafting can be used as a diagnostic tool to answer the question posed at the end of Section 2 by letting observe the optimization trajectory and end-to-end performance when M and D are forced to have the same overall learning rate schedule. This gives a controlled way to study whether H, or me is more important in determining an optimizer's empirical success.
4

---

Under review as a conference paper at ICLR 2022
Algorithm Grafted meta-optimizer MDD
init: optimizers M, D: hyperparameters for M, D; partition P of model parameters.
2: Initialize M, D with their hyperparameters.
3: for each iteration do
4: Receive input (Wz,gr).
5: Query steps from imner optimizers: WM f MMw,, 9t), WpD + D(wr, gt).
6: for each parameter group E P do
7:
Compute grafted update:
M (w)
(0) w
t ww(
w
W
m1
(Or, if (p) w!? let w (p)
(P)))
8: Output grafted weight update:
(M#D)) == w..
3.1 GRANULARITY OF GRAFTING
Before proceeding the experiments, will introduce more dimension the grafting proce- dure: its granularity. Formally, let P be partition of {1.... d), supplied as hyperparameter. For v E Rd P E P, we denote by e(e) E Rlel the restriction of v to the coordinates indexed at Then, transfer norms separately for each set of coordinates in the partition; this is the complete method
summarized by Algorithm Let us note some special cases of P, and their associated properties and interpretations:
((1, . d then schedule from M to D. This schedule is not in general a static sequence of scalars }mc), and depends on the trajectory of inputs (w.,gt). We call this instantiation global grafting.
think of MD transferring overall learning rate
If we choose instead the most granular partition P {{I)...,, {d}), grafting transfers one learning rate schedule per parameter. If the optimizers differ only by positive diagonal preconditioner (like SGD and AdaGrad), so that the steps have the same sign pattern, then it holds that M = M..
Interpolating between these two extremes, we can set the elements of P to be the (usually tensor-shaped) parameter groups in the model's implementation. This a natural partition provided by the optimizer interfaces of all commonly-used deep learning frameworks. We call this variant layer-wise grafting, which transfers one learning rate schedule per param- eter group from M to D.
This provides generalization of the schedule transfer methodology: for any choice of partition, we can force M and to have the same learning rate schedule simultaneously over the parameter groups. By making this partition more granular, obtain more expressive families of learning rate corrections.
4 GRAFTING FOR IMPLICIT SCHEDULE TRANSFER
We first present an empirical study on the transfer of implicit step size schedules between optimizers, on ofate-of-the-art training benchmarks. Starting with well-tuned baseline M, our protocol com- pares two optimizer hyperparameter searches with the same computational budget: an optimizer 'D which is less frequently used and has worse performance than M on the corresponding architecture, and MWD, fixing the tuned hyperparameters of M. Note that the set of hyperparameters tuned in both cases are also the same i.e. the hyperparameters of DD Under the same hyperparameter tuning protocol and budget consistently found across architecturesttaks and batch sizes, that grafting induced positive transfer of end-to-end model performance, i.c. the performance of D M closes the gap (sometimes fully) between the performance D and M.
5

---

Under review as a conference paper at ICLR 2022
BERT (batch size 8192)
BERT Ibatch size 32768)
10
0.7
0.0-
0.4
0.3
0.3
1.2
1.5
0.1
0.5
2.0
toering cape
taniggtips stops
Figure 2: BERT experiments for implicit hyperparameter transfer, comparing hyperparameter search for SGD (with momentum) vs. grafting with M Adam. Adam#SGD outperformed all pure SGD runs significantly; most hyperparameter settings for SGD caused training to diverge. Trials on batch sizes 8192 and 32768 are shown.
BERT pretraining. As a timely choice of a resource-intensing training task, we ran this hyper- parameter transfer experiment for BERT (Devlin et al.) 2018) model pretraining setup. We make two choices of M, Adam and L.MMB (You et al. 2019) commonly reported as the state-of-the-art optimizer for this problem (also see Nado al. (2021). For D we choose SGD with momentum, which has been widely reported to be unsuccessful in reaching mate-of-theear performance. In par- ticular, You et al. (2019) explicitly point out in their abstract that their LAMB algorithm is motivated by the fact that LARS You et al. 2017). its poorly on self-attention models. We repeated the experiment at batch sizes 8192 (measuring final performance 56000 training steps) and 32768 (14063 steps). Further details of the hyperparameter search space as well as the full hyperparameter tables are supplied in the appendix.
pwapiiiveelyyrrconiittinneed predecessor, performs
The model contains 24 transformer layers with 1024 hidden dimensions and 16 attention heads with a total of 340M parameters. It was trained on the combined Wikipedia and Books corpus (Zhu et al.] 2015) datasets (2.5B and 800M words, respectively). We used sequence length of 128. We follow the tuning protocol from Nado et al. (202T) that uses a quasi-random search (Bousquet et al. 2017) with a simple search space. Hyperparameters included are learning rate 7). the moment parameters 31. 32, the polynomial power for the learning rate warmup Pwarmup. and weight decay X. We fixed the e in Adam 1e 11 for all BERT experiments. We selected the best trial using the masked language model accuracy over 10k examples from the training set. Further details are provided in the appendix.
Figure Table I our findings. We were able to corroborate the reported difficulty of training Transformers with SGD, with SGD leading to significantly worse performance as compared to Adam or LAMB. On the other hand, we found that Adam#SGD (at the layer-wise granularity) was able to closely match the performance of Adam, even exceeding it at the smaller batch size. This establishes that the performance gap between the untuned optimizers in this setting arises not from Adam's 's per-coordinate preconditioning, but at the granularity of layer-wise implicit step size schedules. Zhang et al.] (2019) propose theoretical mechanism towards explaining the apparent necessity of adaptive methods for training self-attention models; here, grafting serves as an ablation tool, showing successful convergence with less granular adaptive preconditioning rule.
summarize
BERT MLM top-1 accuracy
Algorithm
Batch size 8192
Batch size 32768
Adam
69.5
70.9
SGD
63.3
48.3
Adam#SGD (Layer-wise)
70.1
67.6
LAMB
70.4
LAMB#SGD (Layer-wise)
71.0
Table I: Final Top-1 masked language model validation accuracies for BERT pretraining with Adam, Lamb at 8k and 32K batch size. The global version of grafting were found to perform significantly worse and have been omitted. All the results were repeated multiple times with random seeds and the results are consistent upto a 1 % deviation

---

Under review as a conference paper at ICLR 2022
ImagoNee//RisNeet555 50
Restlet11CCIBa11
0.8
3.5
40
3.50
3.0-
2.85
2.5
2.00
0.8-
2.25
1.5-
3.70
1.8
0.65
05
120
152
20
200
130
50
302
150
200
3.60
50
300
250
200
to6niing oporhs
traiing
thaiing wOOeTs
Figure 3: ResNet experiments for implicit hyperparameter transfer, comparing hyperparameter search for AdaGrad vs. grafting with a well-tuned baseline M SGD. SGD#AdaGrad outper- formed all pure AdaGrad runs. The batch size for the ImageNet run was 8192 and the CIFAR-10 run was 128. Further details can be found in the appendix.
ImageNet/CIFAR-I classification with ResNets. As alternative setting for grafting, ex- amine the behavior of grafting on convolutional architectures for image classification tasks: Ima- geNet ResNet-50 with 26M parameters) and CIFAR- parameters). In this setting, we make two choices of M. SGD (with momentum) well Adam, with tuned baseline setups derived from Goyal et al. (2017) (ImageNet) and Loshchilov & Hutter (016) (CIFAR-10). A full description of the hyperparameters is deferred to the appendix. For DD' we choose AdaGrad, which is commonly believed to lead to low performance on these models. For these models observe that the global version of grafting is effective as the local version and both lead to performance improvements over the Adagrad baseline for these models closing the gap between performance. The results summarized in Figure and the ImageNet results summa- rized in Table 2 We conducted these experiments at multiple regimes of batch sizes for both models and observed consistent trends. The precise numbers for the CIFAR-10 study can be found in the appendix.
ResNet-18 model with about 11M
has been noted many times (Zeiler 2012 Wilson et 2017 Bottou et 2018) that the AdaGrad optimizer, with commonly used hyperparameter search spaces, fails to train -tate-of-the-art vision models to the same performance as other, closely related optimizers. These experiments show that the gap can be partially closed with a global learning rate schedule transferred from successful run of SGD/Adam.
ImageNet Top-1
Accuracy Batch size 8192
Algorithm
Batch size 1024
SGD
76.85
76.68
AdaGrad
72.93
70.63
SGD#AdaGraa (global)
74.43
72.53
SGD#AdaGrad (layer-wise)
73.85
72.59
Adam
76.47
76.27
Adam#AdaGrad (global)
73.61
72.09
Adar#d (layer-wise)
73.42
72.45
Table 2: Top-1 accuracy at 150 training epochs for ImageNet experiments across batch sizes. The accuracy stabilized at 90 epochs. For all algorithms we performed an extensive tuning over hyper- parameters and used the step-decay schedule. SGD#AdaGrad and Adami Adan#Adaaaad are the grafted versions. All the results were repeated multiple times with random seeds and the results are consis- tent upto 0.3% deviation.
5 GRAFTING FOR EXPLICIT SCHEDULE DISCOVERY
A natural question, in light of the results in Section 4 is whether the benefits of grafting can be distilled into a non-adaptive correction to the algorithm D, eliminating the need to run M in par- allel. In this section, demonstrate that this pipeline can be completed in both of the large-scale setups from the previous section. We extract the step size ratios | - w.ll ||w - well from a successful execution of grafting, and using these sequences learning rate schedules for D in the

---

Under review as a conference paper at ICLR 2022
101
20-2
20-2
102 120
Lrawingg sprc7
Figure 4: A learning rate schedule for AdaGrad discovered using grafting. Before the first learning rate decay (imposed by the tuned schedule for .M=GDD), grafting discovers implicit polynomial warmup when transferring the performance to D=AdaGrad.
traditional non-adaptive. Applying this form of transfer from global grafting results in the discovery of novel learning rate schedules in the usual sense: global, time-dependent, non-adaptive multipliers for the learning rate πt- Applied more granularly layer-wise grafting, this allows for the discovery of higher-dimensional correction.
Training a ResNet with AdaGrad. Using the global variant of Algorithm I with (M, D)= (SGD, AdaGrad) and keeping track of the norms of the steps produced by M and D periodically at every epoch, we arrive at a polynomial warmup-like correction overlaid on the ing (Goyal et al., 2017) schedule employed by SGD. Figure A displays this discovered learning rate schedule compared to the one used on SGD. The final training performance was top-1 accuracy was 72.46%.
constant anneal-
Training BERT with SGD. We demonstrate that the protocol of learning rate schedule discovery can be done with per-layer grafting well. Strikingly, this enables to discover a simple per-layer step size correction that allows the standard SGD(with momentum) to train Transformer model without adaptive preconditioning for the first time to the best of our knowledge. In contrast to global grafting which transferred a non-adaptive global learning rate schedule, herein we only transfer a single per-layer scalar correction which is held constant through training.
In particular, we performed the protocol of layer-size grafting as discussed the previous section, and recorded the per-layer ratio of the norms. Figure 5 provides collection of some representative ratios. Table [5] and Figure X in the appendix provide all the ratios as well as additional visualizations finer granularity.
We computed the median of these individual corrections for the first 2000 steps and applied this layer-wise multiplier along with SGD. We note that these corrections effectively rescale the relative learning rate between the layers. We find that this offline transfer is sufficient to push the test accu- racy of SGD to the state-of-the-art Adam baseline achieving eval accuracy 69.5. Full experimental details with can be found in the supplementary material.
Simplifying the discovered schedule. The above discussion in particular highlights that standard SGD with momentum can be used train BERT state-of-the-an performance, as long as an ap- propriate per-layer learning rate correction is provided. In this section we explore whether the space of these corrections can be further simplified. This highlights the robustness of our transfer ap- proach, showing that performance is preserved as long as relative per-layer scales are approximately preserved and provides proof of concept that search space for these corrections is significantly smaller than the number of parameter groups. We only consider the case of batch-size 8192 here, similar results are obtained for batch size 32K and the presentation deferred to the appendix.
To provide a proof-of-concept. motivated by trends observed in Figure we tie the per-layer scalar corrections for all the parameter groups the attention layers to a fixed value 3, and simply dis- cretize the obtained medians for the other parameter groups/ layers (eg, softmax and embeddings) to the nearest power of 10 which is smaller. This results in a very simple scalar correction scheme which has values in the set D3 (0.1. 1. 10, 100, B}. We provide the achieved accuracy results via the above corrections in Table 6b over some choices of B. We find that performance tends to be relatively stable small changes in choices of B. We stress that our aim is to demonstrate
8

---

Under review as a conference paper at ICLR 2022
self aalf-attention
feadforward
atlfnattenion
softmax
Adam/SGD norm ratoo duina graftins
10a
attentttnnnn
emboddings
self self-attention
micc (LayerNorm, ....
10
10a
10=
10a
10a
500
750
1000
1250
1500
1750
2000
Training
Figure 5: Per-layer learning rate corrections during the first 2000 iterations of BERT pretraining with #dGD enabling the discovery of the learning rate correction to SGD. More visualizations are provided in the supplementary material.
Parameter group
Median
Rounded (D50o)
Algorithm
Eval accuracy
Position embeddings
18.8
10
Adam
69.5
Word embeddings
312.9
100
SGD
63.3
Layer Q
525.5
500
Adam#SGD
70.1
Layer K
466.6
500
SGD (Medians)
69.5
Layer V
312.9
500
SGD (Do00)
70.0
Layer 24
1090.1
500
SGD (D250)
68.9
Layer 24
1081.3
500
SGD (Dso)
67.8
Layer 24 V
193.9
500
Softmax
62.8
10
(a) Final top-1 validation accuracies for BERT pretraining with discovered per-
(b) Some representative per-layer learning rate corrections covered using grafting, then simplified manually. These correc- tions, which can be viewed 168 (i8-dimensional hyperparam- enable SGD pretrain BERT with performance petitive with adaptive methods.
layer learning rate corrections.
proof-of-concept that simplified correction scheme (in this case composed of only 5 distinct val- ues) works and thus shed light on the extent of the role of adaptivity training these models. A careful study of the space of all possible discretizations is out of the scope and if left for future work.
CONCLUSION
We have introduced learning rate grafting, binary operation which blends the behavior of two op- timization algorithms, towards investigating the entanglements between widely-used adaptive pre- conditioning rules and learning rate schedules. We have presented an empirical study with popular optimizers for ttate-of-the-art deep architectures, discovering that well well-performing optimizer M can transfer performance to D via grafting its sequence of implicit step size schedules. Fur- thermore, have shown that grafting can be used to extract standalone learning rate corrections, enabling us to train a Transformer using SGD (with momentum) for the first time. We hope that this finding will stimulate further empirical research the power of simple per-layer learning rate schedules.
The empirical phenomena examined in this work seem be unexplained by current theory. We hope that the experiments enabled by grafting will aid in developing more robust beliefs about both adaptive methods and learning rate schedules, and guide future theoretical inquiry.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning, pp. 102-110. PMLR, 2019.
Naman Agarwal, Surbhi Goel, and Cyril Zhang. Acceleration via fractal learning rate schedules. arXiv preprint arXiv:2103.01338 2021.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory--fficient adaptive optimiza- tion for large-scale learning. arXiv preprint arXiv.1901.11150 2019.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. arXiv e-prints, art. arXiv:2002.09018 February 2020.
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. arXiv preprint arXiv: 2018.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-31) -311,, 2018.
Olivier Bousquet, Sylvain Gelly, Karol Kurach, Olivier Teytaud, and Damien Vincent. Critical hyper-parameters: No random, no cry. arXiv, 2017. URL https://azxxt org/abs/ 1706. 03200
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1096, 2018.
Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. arXiv preprint arXiv:1902.046200 2019.
Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint
arXiv:1910.05446. 2019.
Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd. International Confer- ence on Machine Learning, pp. 2260-2268. PMLR, 2020.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810.0488 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep policy gradients: A case study on ppo and trpo. arXiv preprint arXiv:2005. 12729, 2020.
Rong Ge, Sham M Kakade, Rahul Kidambi, and Pranceth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares. In Advances in Neural Information Processing Systems 32, pp. 14951-14962. 2019.
Boris Ginsburg, Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M Cohen. Stochastic gradi- methods with layer-wise adaptive moments for training of deep networks. arXiv preprint arXiv: 1286, 2019.
Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. arXiv preprint OrXiv:/8/00 13243, 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 170.00777 2017.
10

---

Under review as a conference paper at ICLR 2022
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprini arXiv:1801.01299, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings deep residual networks. In European Conference Computer Vision, pp. 630-645. Springer, 2016.
Andrew Ilyas, Logan Engstrom. Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. A closer look at deep policy gradients. In International Con- ference on Learning Reprexentations, 2019.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:/502.03167, 2015.
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerat- ing stochastic gradient descent for least squares regression. In Conference On Learning Theory, pp. 545-604. PMLR, 2018.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for im- proved quality, stability, and variation. arXiv preprint arXiv: 10196, 2017.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv: 1rXiv:171,.07628, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14X2::412.66960 2014.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems, pp. 10215-10224, 2018.
Alex Krizhevsky, One weird trick for paralellizngg convolutional neural networks. arXiv preprint 1rXiv:1404.5997 2014.
kuangliu. Pytorch-cifar.
2017.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint arXiv:1907.04595, 2019.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. arXiv preprint arXiv:1910.074555 2019.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 19Xiv:1907.11699 2019.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: /608.0093,, 2016.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factoree approximate curvature. In International conference machine learning, pp. 2408-2417. PMLR, 2015.
Peter Mattson, Christine Cheng, Cody Coleman, Greg Diamos, Paulius Micikevicius, David Patter- son, Hanlin Tang, Gu- Ycon Wei, Peter Bailis, Victor Bittorf, et al. MIperf training benchmark. arXiv preprint 100.v:1910 2019.
Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex opti- mization. arXiv preprint 4900 2010.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and Optimizing LSTM Language Models. arXiv preprint 1708./7708.02 182, 2017.

---

Under review as a conference paper at ICLR 2022
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. An Analysis of Neural Language Mod- eling at Multiple Scales. arXiv preprint wrxiv:/803.08240, 2018.
Zachary Nado, Justin M. Gilmer, Christopher J. Shallue, Rohan Anil, and George E. Dahl. A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes. arXiv e-prints, art. arXiv:2102.06356 February 2021.
Zachary Nado, Justin M Gilmer, Christopher J Shallue, Rohan Anil, and George E Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. arXiv preprint (02.:::20000066566, 2021.
Martin Popel and Ondřej Bojar. Training tips for the transformer model. arXiv preprint arXiv:1804.00220 00247, 2018.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of Adam and beyond. In International Conference Learning Representations, 2018.
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through a crowded valley- benchmarking deep learning optimizers. arXiv preprint arXiv:2007.01547, 2020.
Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl, Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(112): 2019. URL httt::/ wttp:cpprppppp/////2/..
18 789. html
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprint arviv: 0888.0042355 2018.
Leslie N Smith. A disciplined approach to neural network hyper-parameters: Part 1 l-learning rate, batch size, momentum, and weight decay. arXiv preprint arXiv: 2018.
Leslie N Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using Operations
large learning rates. In Artificial Intelligence and Machine Learning for Applications, volume 1 11006, pp. 1100612. International Society for Optics and Photonics, 2019.
Matthew Staib, Sashank Reddi, Satyen Kale, Sanjiv Kumar, and Suvrit Sra. Escaping saddle points with adaptive gradient methods. In International Conference on Machine Learning, pp. 5956 5965. PMLR, 2019.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26- 31, 2012.
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. In International Conference on Machine Learning, pp. 6677-6686. PMLR, 2019.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stem, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pp. 4151-4161, 2017.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv: 1oXiv:/906.00 09099996622333 2019.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 038888 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint 1rXiv:190.000960 2019.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.570; 2012.
12

---

Under review as a conference paper at ICLR 2022
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? arXiv preprint 11X0:11912.03194,2000
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watch- ing movies and reading books. Proceedings of the 2015 IEEE International Conference on Computer Vision (ICCV), ICCV '15, pp. 19-27, USA, 2015. IEEE Computer Society. ISBN 9781 178146738391.. doi: 10.1109/CCV22215.11. URL https: doi http::/:///o.....0010.11110.1000 110//.0/ 2015.11
A ADDITIONAL EXPERIMENTS AND DETAILS
We provide additional experiments and all the implementation details in this section. In the next section we provide full parameee search tables for all our experiments to aid reproduction of experiments.
A.1 BERT EXPERIMENTS
A.1.1 IMPLEMENTATION DETAILS
We used the same experimental setup as the official BERT codebasq and the standard train/test split from the previous literature. We trained on Google TPUs, using TPUv3-128 / TPUv3-256 for the 8K and 32K batch size experiments.
We trained the two pretraining objectives on the combined Wikipedia and Books corpus (Zhu et al. 2015) datasets (2.5B and 800M words, respectively). We used sequence lengths of 128.
Our experimentation protocol is identical to tyat described in Nado et al. (2021): : we tuned hyperpa- rameters using quasi-random search (Bousquet al. 2017) with fixed budget in a simple search space described in Table We selected the best trial using the masked language model accuracy over 10k examples from the training set. The experiments were run for a total of 56352 steps for batch size 8K and 14063 steps for batch size 32K.
Adam. Hyperparameters included learning rate 7, exponential window parameters 31, 32, the polynomial power for the learning rate warmup and decay p. and weight decay We fixed the e in Adam to 10-11 for all BERT experiments. Polynomial power was searched over (1, for batch size 8K and set to for 32K. The entire hyperparameter search and associated performance can be found in Tables 7 and IT for batch sizes 8K and 32K respectively.
SGD with momentum. Hyperparameters included learning rate 7, $1, the polynomial power for the learning rate warmup and decay p. and weight decay X. Polynomial power was searched over 1,2 for batch size 8K and set to for 32K. The entire hyperparameter search and associated performance can be found in Tables 8 and 12 for batch sizes 8K and 32K respectively.
Grafting: Ada We use the above described implementations of Adam and SGD with mo- mentum as M, D respectively. The Adam hyperparameters were fixed the optimal ones found. The SGD hyperparameters were tuned with the same protocol as for SGD alone. The overall learn- ing rate scalar was also tuned. Tables 9 and [3 contain the hyperparameter search and associated performance for batch sizes 8K and 32K respectively.
https://github.coom reseaarc
13

---

Under review as a conference paper at ICLR 2022
Hyperparameter
Range {1.2)
Scaling
discrete
10 10-9,1.0
logarithmic
31
10-5,0.5
logarithmic
B2
[10 10-$,0.5
logarithmic
10
logarithmic
Table The search space used to tune Adam/SGD on BERT. À refers to weight decay and p refers to the polynomial power in the learning rate schedule for both the warmup and decay phases.
LAMB We use the exact same settings as in the original paper You et al. 2019).
Grafting: LAMB#SGD. We use the above described implementations of LAMB and SGD with momentum M, D respectively. The LAMB hyperparameters were fixed to the optimal ones. The SGD hyperparameters were tuned with the same protocol as for SGD alone. The overall learning rate scalar also tuned. Table (0ble|T0come contains the hyperparameter search and associated performance for batch sizes 8K and 32K respectively.
A.1.2 TRAINING BERT SGD (BATCH SIZE 8K AND 32K)
Some these results are partly presented in the main paper. We provide a full exposition for completeness.
demonstrate that the protocol of learning rate schedule discovery can be done with per-layer grafting well. Strikingly, this enables to discover simple per-layer step size correction that allows the standard SGD(with momentum) to train a Transformer model without adaptive precondi- tioning for the first time the best of our knowledge. In contrast global grafting which transferred a non-adaptive global learning rate schedule, herein we only transfer a single per-layer scalar cor- rection which is held constant through training.
In particular, we performed the protocol of layer-size grafting as discussed in the previous section, and recorded the per-layer ratio of the norms. Figure7ssows these ratios for batch sizes 8K and 32K. Given the apparent instability of the ratios on the later part of the runs for the 8K, we zoom in on the initial part of the training particular, we computed the median of these individual corrections for the first 2000 (respectively 500) steps for batch size 8K (respectively 32K) and applied this layer-wise multiplier along with SGD (with momentum); see Figures and for visualizations of the schedule corrections, and Table 5 for the median values. We find that this offline transfer is sufficient to push the test accuracy of SGD much closer to the state-of-the-art Adam baseline (see Table
Masked LM Accuracy
Algorithm
Batch size 8K
Batch size 32K
Adam
69.5
70.8
SGD
63.3
49.0
Adam#SGD
70.1
67.6
SGD (Medians)
69.5
63.8
Table Final top-1 validation accuracies for BERT pretraining with Adam, SGD(with momentum), Adam#SGD and the discovered per-layer learning rate corrections via medians.
Despite these exponentially large correction ratios, the grafted optimizer converges successfully. We leave investigation of this curious phenomenon for future work.
14

---

Under review as a conference paper at ICLR 2022
10
altwriion
Westeward
eltonion
enteddigg
3:8 alloncnn.c CAA
ma. (Layshoom....
101
101
303
260
400
500
Traiing trniiore
(a) Batch size 8K
(b) Batch size 32K
Figure 6: Per-layer learning rate corrections during the first 2000 (resp. 500) iterations of BERT pretraining with SdD with batch size 8K (resp. 32K), enabling the discovery of the learning rate correction to SGD. median filter of width applied to each sequence, for clarity.
BAdhoasdd selimee
sct ullnt
TEL (Lapertkmm..
32000
2000C
30000
40000
59300
2000
4000
6000
5000
30500
12000
14003
mmiliig
(a) Batch size 8K
(b) Batch size 32K
Figure 7: Per-layerleamme rate corrections for all training iterations, in the BERT pretraining setups with Adam#SGD. Some of these corrections increase exponentially; note that despite performing these large corrections, the grafted optimizer converges to nearly ffate-ofthhe-ee accuracy.
15

---

Under review as a conference paper at ICLR 2022
Parameter group
Median correction (8K)
Median correction (32K)
Embedding LN B
9.10
0.54
Embedding LN
16.34
2.26
Position embeddings
18.83
2.48
Token type
emhed-
0.45
0.03
dings
Word embeddings
312.94
7.94
Attention Q bias
378.07
38.56
Attention K bias
154101.58
25119.54
Attention bias
37.69
2.01
Attention out bias
21.92
1.16
Attention FC1 bias
81.30
4.57
Attention FC2 bias
23.87
1.24
Attention LNI B
21.60
1.13
Attention LN1 Y
41.99
4.08
Attention LN2 3
22.05
1.12
Attention LN2 7
45.42
3.83
Pooler bias
32.69
0.64
Pooler weights
48.73
1.02
MLM output bias
479.54
8.74
MLM IN |
21.78
0.57
MLM LN y
17.19
0.82
MLM pNe-LN bias
48.04
0.71
MLM Pre-LN weights
62.77
1.06
NSP bias
0.66
0.01
NSP weights
2.00
0.04
Attention Q weights
(525.50, 575.42, 714.81, 794.62,
900.46,
{39.00,
80.19, 103.98, 104.84. 146.52,
(layers 1-24)
904.97,
1150.11,
1290.34.
1411.08.
202.22,
187.58. 229.54, 279.68. 236.02,
1867.67,
2074.47,
1975.08,
859.45,
267.71,
272.56, 255.77, 227.56. 216.35,
2026.33, 1739.56,
1549.03,
1602.61,
207.21,
194.48.
172.97, 150.94, 145.22.
1516.13,
1291.10.
1101.66,
1124.54,
108.78,
84.44,
90.25, 69.70)
850.50, 993.93, 1090.10)
Attention K weights
(466.59, 503.57, 662.77, 679.56,
764.49,
[34.70,
66.95, 98.20. 90.28, 123.40.
(layers 1-24)
714.15, 958.21,
1086.80,
1266.25.
185.35.
180.04. 227.60, 265.37, 206.99.
1731.75, 1980.33, 1842.08,
595.65,
262.27,
276.03, 253.69, 231.60, 227.34.
1949.40,
1669.08,
1335.04,
1495.53,
211.37,
205.17, 179.20, 155.37, 142.24,
1452.81,
1098.97, 1027.76.
1096.50,
101.69, 72.56, 77.83, 64.13)
709.72, 838.97, 1081.26)
Attention FCI weights
(82.02, 108.41, 131.46, 151.14, 169.54,
{7.17, 11.03, 14.30, 16.96, 18.66, 19.93,
(layers 1-24)
186.14, 198.20, 208.18, 18, 217,43, 228.15.
20.77, 21.24, 21.95, 22.31, 22.27, 21.70,
230.22, 228.92, 237.99, 239.71, 236.57,
21.41. 21.04, 20.27, 19.43, 18.38, 17.47,
235.85, 231.04, 224.10, 219.76, 214.12,
16.04. 14.29, 12.70, 10.49, 8.66, 6.24)
206.27, 203.64, 207.25, 216.04)
Attention
weights
(48.61, 66.94, 85,83, 96.88, 112.34,
{2.58. 4.56, 7.23, 9.69, 10.97. 12.42,
(layers 1-24)
129.41, 135.96, 150.30, 167.87, 184.03,
13.39, 13.25, 14.18, 14.48, 14.72. 14.27.
203.72, 202.35, 178.28, 214.61, 210.79,
14.25, 14.16, 13.65, 12.89, 12.26. 11.58.
209.75, 216.17, 226.42, 221.07, 220.83.
10.93, 9.85, 8.55, 7.40, 6.32, 4.71)
219.30, 197.48, 187.25, 193.95)
Attention weights (layers 1-24)
(56.17, 74.65, 97.39, 105.54, 120.98,
[2.92, 5.59, 8.01, 10.51, 12.55, 13.42.
137.54, 145.35, 167.26, 186.15, 204.79.
13.61, 14.52, 14.51, 14.45, 14.40, 14.44
212.83, 212.14, 202.61, 232.44, 222.43,
14.49, 14.41, 13.88, 13.13, 12.61, 11.76.
228.23. 228.55, 238.57, 234,49, 232.85,
10.89, 9.91, 8.84, 7.40, 6.35, 4.65)
241.87, 229.22, 222.86. 213.22)
Attention FC2 weights
(71.84, 88.45, 103.35, 115.39, 126.44,
{5.56. 7.32, 8.64, 9.54, 10.34, 11.12.
(layers 1-24)
137.14, 148.70, 159.31, 167.05, 179.80,
12.05, 12.71, 13.44, 13.89, 14.37, 14.74,
181.04, 183.60, 192.79, 200.97, 201.98,
14.83. 14.89, 15.00. 14.69, 14.14, 13.83.
207.50, 208.11, 206.15, 208.88, 210.20,
13.19, 12.05, 10.81, 9.16. 7.38, 5.26)
207.35, 209.14, 216.43, 218.14)
Table Corrections for each parameter group for the BERT schedule discovery experiments, found using grafting but executed indepedently as a confection on SGD.

---

Under review as a conference paper at ICLR 2022
BERT. batch size 8192
Attention Q
Attention
Attention V
Attention out
30a
10a
10a
10a
10"
500
1000
1500
2000
500
1000
1500
2000
0
5.00
1000
1500
2000
500
1000
1500
2000
Attention FC1
Actention
Embeddings
Others
107
10
10*
301
10f
1500 2000
500
1000
1500
500
2000
1500
2000
500
1000
1500
2000
BERT, batch size 32768
Attention Q
Attention K
Attention on V
Attentian aut
107
102
0
200
400
200
400
200
400
200
400
207
10-
10'
10"1
10a
400
400
Figure 8: Per-layer learning rate schedule corrections from Figure 6 displayed on separate plots by type of parameter group.
17

---

Under review as a conference paper at ICLR 2022
A.2 IMAGENET EXPERIMENTS
A.2.1 IMPLEMENTATION DETAILS
For the image classification task with Imagenet we use the standard Resnet-50 model He et al. (016)). The model code, training pipeline and other parameters are identical to the MLPerf training benchmark for ResNet-50 v1.5 on ImageNet Mattson et al. (2019). All the training experiments carried out up to 150 epochs (however all training stabilizes around 90 epochs). We report the achieved test top-1 accuracy. For all the cases we implement a distributed version of batch-norm akin to the one used in Goyal et al. (2017). Each core locally computes and maintains batch-norm statistics over sub-batches of 64 (akin to the n parameter in Goyal et al. (2017)) The maintained batch norm statisties are synchronized at the end of every epoch. The batch norm momentum is set to 0.9 for batch size 8192 and 0.99 for batch size 1024 and epsilon is set to 0.00001. The experiments are performed on a 16 TPU core v3 (for batch size 1024) and a 64 TPU core v3 (for batch size 8 8192). The timings for single run are approximately 1.5 hours and 35 minutes for batch size 1024 and 8192 respectively.
SGD with momentum. We use the standard implementation of SGD with Heavy Ball Momen- tum with the momentum parameter set to 0.9 and the weight decay set to 0.0001. We use the warmup+step decay schedule proposed in Goyal et al. (2017). The schedule warms up the learning rate linearly to the base learning rate for epochs and then the learning rate is held constant until cut by factor of 0.1 at epoch 30,50,80. Table 15 contains the hyperparameter tuning information where the search was performed over the base learning rate and across two batch sizes.
Adam. We use the standard implementation of Adam and tune the momentum parameter and the weight decay. We fixed Br to the default 0.999. We use the warmuptstep decay schedule proposed in Goyal et al. (2017). The schedule warms up the learning rate linearly the base learning rate for 5 epochs and then the learning rate is held constant until cut by a factor of 0.1 at epoch 30,50,80. Tables [6 and [7 contain the hyperparameter tuning information where the search was performed over the base learning rate and across two batch sizes.
AdaGrad. We use the standard implementation of AdaGrad (present in TensorFlow). We use all the default hyperparameters except the learning rate which we search over. We use the same learning rate schedule as reported SGD with Momentum (step-decay). We highlight that this schedule worked significantly better than no schedule at all which is traditionally what has been reported for AdaGrad baselines leading to the stronger baselines for AdaGrad baselines we present. Table 14 contains the hyperparameter tuning information where the search was performed over the base learning rate and across two batch sizes.
Grafting: AGD#Aaaarad/AddatttAA.o
We use the above described implementations of SGD with cooentuumAAam and AdaGrad as M, D respectively. Since the only hyper-paramedt for Adagrad is learning rate tune the base learning rate supplied to SGD (as the M algorithm governs the overall learning rate of the grafted algorithm) Table [8] contains the hyperparameter tuning information where the search was performed over the base learning rate and across two batch sizes.
A.3 CIFAR- 10 EXPERIMENTS
For the image classfication task with CIFAR-10 dataset we use a standard PreActivation Resnet- 18 model (He et al. 2016). The model code is implemented identically in the following repository (kuangliu) 2017). .. All the training experiments are carried out upto 200 epochs. We report the achieved test accuracy. The experiments are performed on single Tesla V100-SXM2 GPU. The timings for a single run are approximately 90 minutes.
SGD with Momentum. We use the standard implementation of SGD with Heavy Ball Momentum with the momentum parameter set to 0.9 and the weight decay set to 0.0005. We use the single cycle cosine decay learning rate schedule suggested by Loshchilov & Hutter 2016). Formally the learning
18

---

Under review as a conference paper at ICLR 2022
rate at the start of epoch t given the base learning rate 33 is given by
B cos( cos(t/2(t/200)
Table 20 contains the hyperparameter tuning information where the search was performed over the base learning rate and across three batch sizes.
AdaGrad. We use the standard implementation of AdaGrad (present in TensorFlow). We use all the default hyperparameters except the learning rate which we search over. We use the same cosine schedule as reported in SGD with Momentum as we found it to be better than applying no schedule all. Table [9 contains the hyperparameter tuning information where the search was performed over the base learning rate and across three batch sizes.
Grafting: SGD#AdaGrad. We use the above described implementations of SGD with momen- tum and AdaGrad M, D respectively. Further tune the base learning rate supplied SGD (as the M algorithm governs the overall learning rate of the grafted algorithm) Table 21 contains the hyperparameter tuning information where the search and across three batch sizes. It is worth noting that while the additional learning rate tuning leads to small benefits, running grafting with the tuned SGD learning rate already provides strong baseline.
performed over the base learning rate
Table contains the comparisons between the tuned results.
CIFAR-10 Top-1 I Accuracy
Algorithm
Batch size 128
Batch size 512
Batch size 2048
SGD AdaGrad SGD# AdaGGad
94.93
94.39
93.37
89.81
88.25
85.71
93.42
92.36
89.81
Table 6: Top-1 accuracy at 200 training epochs for CIFAR-I0 experiments across batch sizes.
19

---

Under review as a conference paper at ICLR 2022
B FULL HYPERPARAMETER SEARCH TABLES
B.1
BERT EXPERIMENTS
B.1.1
BERT (BATCH SIZE 8K)
Poly Power (p)
1 B1
1- 32
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
2
0.014533
0.305531
3.469230
0.022672
0.000000
1
0.019056
0.068170
0.174876
0.027856
0.000000
1
0.043181
0.418378
1.930781
0,000571
0.000000
1
0.012960
0.067712
1.286336
0.002391
0.000000
2
0.396269
0.019148
0.003123
0.087673
0.000000
1
0.104732
0.49316 163
0.049373
0.000430
0.000000
2
0.439164
0.126745
0.009961
0,010594
0.000000
1
0.012784
0.077563
0.005028
0,002035
0.000000
1
0.242170
0.097739
0.653928
0.002798
0.0337 10
1
0.014564
0.054001
0.002970
0,000214
0.0337 10
1
0.116967
0.286304
4.401929
0.006926
0.0337 I0
2
0.406820
0.072237
0.008125
0.015782
0.048026
1
0.118934
0.012215
0.473520
0.003294
0.048026
0.124374
0.010643
4.738276
0,002856
0.048646
1
0.114625
0.129364
0.762753
0.026747
0.057222
2
0.016942
0.010706
0.010191
0.000042
0.596373
2
0.018620
0.479633
0.280896
0.000030
0.600682
2
0.056560
0.025435
0.545407
0.000050
0.602233
1
0.277735
0.316840
0,003889
0.000037
0.614416
1
0.051057
0.055620
0,001138
0.000051
0.634374
1
0.099721
0.473195
1.174743
0,000412
0.667841
2
0.039101
0.133753
1,472112
0.000956
0.672648
2
0.235072
0.186723
0.828703
0,000751
0.693680
1
0.047688
0.246138
0.075992
0.000151
0.696613
Table 7: Hyperparameter search space for Adam, BERT, batch size 8K.
20

---

Under review as a conference paper at ICLR 2022
Poly Power (p)
1 1 81
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
2
0.015823
0.036840
0.054733
0.000000
1
0.096770
0.426730
0.056828
0.000912
2
0.112533
0.382084
0.024687
0,057222
0.153500
0,105432
0.026578
0.057222
1
0.235137
4.153353
0.056153
0.057222
1
0.204847
0.104750
0.036520
0.057222
1
0.011275
4.299596
0.003757
0,057222
0.224600 2
1.709058
0.052954
0.057222
1
0.084911
0.002031
0.010545
0.057222
0.046596 2
2.510824
0,003590
0.057222
0.030383 2
1.391879
0.002720
0.057222
1
0.174070
0,033287
0,005125
0.057299
0.017656 2
4.237251
0.000495
0.087236
1
0.277091
0,003713
0,000025
0.107089
2
0.2591 132
1.357934
0.000693
0.107096
1
0.255655
1.054776
0.000041
0.107335
2
0.104964
1.613354
0,000066
0,111124
0.276399 1
0.011414
0.000054
0.119279
0.022524 2
3.054689
0.000019
0.121293
0.052927 2
3.238563
0.000501
0.125202
2
0.177523
0,002682
0,000 126
0.137419
2
0.311772
0.009484
0,000209
0.138753
0.034004 1
0,001888
0,00003
0.146332
1
0.131274
0.942633
0.001947 1441
0,281 134
2
0.105729
0.153407
0,000419 19
0,328938
0.180943
0.001704
0.000379
0.335963
1
0.013722
0.222883
0.000048
0.343942
0.165404
0.017183
0.000395
0,345 149
0.045047
0.948157
0,000550
0.358314
1
0.085101
0.060706
0.000309
0.377145
1
0.157212
0.015188
0.000701
0.431824
2
0.082335
0.030545
0.000618
0,449213
1
0.183162
0.024755
0.003892
0.465375
2
0.011765
0.738450
0.000429
0,491509
2
0.020559
0.007351
0.001888
0.632697
Table 8: Hyperparameter search space for SGD, BERT, batch size 8K.
21

---

Under review as a conference paper at ICLR 2022
Poly Power (p)
1 81
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
2
0.008817
1.523654e-03
0.247838
0,000000
2
0.025831
7.554433e-09
0.010177
0,000000
2
0.004666
7.457963e-05
0.041305
0.000000
2
0.042243
5.682363e-08
0.161833
0.000000
2
0.001035
2.020054e-10
0.002286
0.000000
2
0.034458
7.282971c-10
0.082530
0,002035
2
0.022501
1.747059c-06
0.005476
0,032267
2
0.002092
4.521310c-03
0.003641
0.032267
2
0.017431
1.669821e-2 0-
0.009074
0.048646
2
0.005098
7.112310e-03
0.142092
0.057222
2
0.016812
2.641299e-10
0.052162
0.057222
2
0.009421
3.166860c-08
0.006412
0.057222
2
0.001396
4.255182e-1 05
0.120496
0,057222
2
0.001615
8.588940c-07
0.001236
0,057222
2
0.013455
5.000904e-07
0.201592
0.057222
2
0.016216
8.084465e-03
0.034562
0.057222
2
0.003546
5.341552e-08
0.002101
0.057222
2
0,048880
7.500002e-01
0,018438
0. 0.100170
2
0.072832
3.817323e-01
0.033588
0.110570
2
0.190460
8.796291e-02
0.464158
0,357879
2
0.349429
2.119862e-06
0.002255
0.672795
2
0.070881
1.240433e-08
0.001685
0.685996
2
0.042185
7.085360e-04
0.054075
0.687168
2
0.120310
3.566263c-07
0.038694
0,691329
2
0.075832
4.531688c-05
0.015860
0,701189
Table 9: Graft Bert 8K: Adam#SGD
22

---

Under review as a conference paper at ICLR 2022
Learning Rate
- 81
Weight Decay Rate (A)
Masked LM Accuracy (as fraction)
0.557378
0.052662
0.008343
0.000000
0.000051
0.010000
0.000100
0.000000
0.010168
0.104408
0.118838
0.057684
0,000001
0,250000
1,000000
0.082367
0.000001
0.010000
0.000100
0.108002
0.000002
0.228890
0,046535
0.109349
0.000339
0.250000
1.000000
0. 13145516
0.001 165
0.250000
1,000000
0.140487
0.001274
0.010000
0.005104
0.296781
0.000014
0.024513
0.000535
0.309587
0.000051
0.099820
0.072073
0.346744
0.001021
0.184613
0.167212
0.415491
0.000052
0.105862
0.003519
0.431919
0.001568
0.037527
1.000000
0.439019
0.000514
0.040944
0.254123
0.503474
0.000501
0.212621
0.017731
0.529111 17
0,001231
0.209148
0,008710
0.599985
0.000909
0.017010
0.218740
0.601494
0.000159
0.010000
0.000100
0.631 107
0.000997
0.250000
0.000100
0.642831
0.000865
0.101639
0.004944
0.647802
0.001007
0.149186
0.000757
0.657476
0.001000
0.050000
0.010000
0.669588
0.000962
0.049001
0.007475
0.676035
0.000810
0.010000
0.016944
0.688638
0.001808
0.031094
0.004114
0.710352
Table 10: Hyperparameter search space for LAMB#SGD, BERT, batch size 8K.
23

---

Under review as a conference paper at ICLR 2022
B.1.2
BERT (BATCH SIZE 32K)
Poly Power (p)
1 - P1
1 - 32
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
1
0.446131
0.025372
0.107242
0.065235
0.000000
0.010053
0.056364
0,171185
0,000041
0.000000
0.260209
0.311713
0.010859
0,003011
0,000000
1
0.043480
0.018414
0.003984
0.045078
0.000000
1
0.1 16116
0.374498
0.338193
0.057663
0.000000
1
0.01 169
0.063804
0,003009
0.000230
0.000000
1
0.010612
0.379281
0.010484
0.000680
0.000000
1
0.049212
0.021061
0.154526
0.005826
0.000000
0.011473
0.010763
0.023060
0,004987
0.000000
1
0.094829
0.034333
0,007080
0,003053
0.00 0.000000
0.182225
0.116353
0.002412
0.005812
0.00 0.000000
0.035019
0.2971 140
0.523223
0.009278
0.00 0.000000
0.156544
0.277251
0.662085
0.059544
0.000000
0.016542
0.193700
0.707171
0.021833
0.000000
0.150077
0.286466
0.299879
0.093842
0.000000
0.012764
0.012950
0.692721
0,001412
0.000029
0.016431
0.026711
0.012575
0.001505
0.000057
0.072118 18
0.017038
2.161319
0.059588
0.029695
0.242228
0.014281
7.335423
0.000740
0.057367
0.014321
0.032597
7,421835
0,000292
0.057367
0.044166
0.023843
0,047968
0.000014
0.429771
0.107729
0.013147
2.653585
0.000018
0.453644
0.334832
0.069228
0,040715
0.000043
0.549747
0.079745
0.019922
0.007470
0.000053
0.564666
1
0.054556
0.035450
0.005083
0.000059
0.575131
1
0.209573
0.011596
3.090723
0.000189
0.601642
1
0.085371 1
0.020573
0.396168
0,000114
0,627102
1
0.147607
0.012472
0.420754
0,000135
0.639349
1
0.187887
0.045957
0.008888
0,000262
0.679573
1
0.073969
0.207516
0.010205
0.000782
0.704228
0.095339
0.030376
0.080869
0,000589
0.707279
1
0.065729
0.010705
0.314657
0.000594
0.708929
Table 11: Hyperparameter search space for Adam, BERT, batch size 32K.
24

---

Under review as a conference paper at ICLR 2022
Poly Power (p)
1- B1
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
0.130375
0.00001
0.008209
0,057367
0.130338
0.000298
0.005776
0.057367
1
0.145380
0.000295
0.005523
0,057367
1
0.161747
0.031063
0.006454
0,057367
0.211265
0.00001
0.003328
0.057367
1
0.222030
0.667330
0.000016
0.099826
0.107412
0.000204
0,00001
0,102016
I
0.248541
0.002621
0,000027
0,102487
0.241575
0.232839
0.000025
0,102706
0.206147
0.234432
0.000026
0.104275
0.190062
0.003350
0.000059
0,115585
0.129872
0.434452
0.000051
0.1 0.118843
0.225584
0.000080
0.000092
0.121976
0.218541
0.013667
0.000113
0,126746
0.173291
0.000074
0.000146
0,137803
0.133697
0.000275
0.000199
0.154843
0.204741
0.793410
0.002174
0.279284
0.159334
0.055297
0.000434
0.2891 12
0.222569
0.060583
0.003277
0.349611
0.102489
0.000315
0.000507
0.350370
1
0.204039
0,002175
0,000906
0,351956
0.249884
0.017145
0.001747
0,408409
1
0.166861
0.000449
0.001317
0.414829
0.199344
0.013357
0.001739
0.437989
1
1.100030
0.000124
0.002287
0.490207
Table 12: Hyperparameter search space for SGD, BERT, batch size 32K.
Poly Power (p)
1 - 81
Weight Decay Rate (A)
Learning Rate
Masked LM Accuracy (as fraction)
0.054284
0.000185
0.000383
0.032528
1
0.062899
0.000077
0.000248
0.032528
1
0.057734
0.000225
0,000325
0,057367
1
0.144769
0,000206
0,000473
0,057367
0.093987 1
0.000179
0,000401
0,057367
1
0.114372
0.000224
0.000495
0.057367
0.058382
0.000064
0.000463
0.057367
1
0. 0.149345
0.000213
0.000294
0,061835
0.071465
0.000186
0.000253
0.240233
1
0.142337
0.000096
0.000113
0.580597
I
0.134273
0.000208
0,000216
0.629878
0.114651
0.000124
0.000290
0,654084
0.063694
0.000224
0.000291
0.676234
Table 13: Hyperparameter search space for AdamSGD, BERT, batch size 8K.
25

---

Under review as a conference paper at ICLR 2022
B.2 IMAGENET EXPERIMENTS
Batch size 1024
Batch size 8192
learning rate
accuracy
learning-rate
accuracy
0.000100
11.908
0.000100
37.788
0.000158
16.586
0,000127
42.612
0.000200
19.972
0.000394
60,916
0.000316
28.038
0,001220
68.168
0.000398
32.874
0.002028
69,590
0.000501
37.754
0.002972
70.116
0.001 000
52.020
0.002999
70.134
0.001995
61.424
0.003277
70.088
0.002512
63.574
0,003442
70,488
0.003981
66.834
0.004383
70.432
0.005012
68.158
0.004410
70.356
0.007943
70.272
0.004417
70.314
0.010000
70.760
0.006360
69.872
0,012589
71.520
0.007823
69.672
0.015849
71.716
0.009804
69.538
0.025119
72.554
0,010000
69.848
0.031623
72.664
0,010554
69.502
0.039811
72.562
0.014409
69.068
0.050119
72.656
0.021884
68.306
0.063096
72.926
0.095870
25.560
0.079433
72.558
0.309611
7.162
0.100000
72.658
1,000000
0.122
Table 14: AdaGrad: ImageNet
Batch size 1024
Batch size 8192
learning rate
accuracy
learning rate
accuracy
0.001000
28.636
0.005480
65.374
0.001627
43.516
0.017648
73.556
0.003376
59,428
0.029705
75.218
0.005493
65.994
0.048524
75.688
0.008936
70.644
0,050000
75.996
0.014539
73.312
0.054569
76.108
0.023654
74.638
0.059523
76.294
0.038483
75.686
0.066660
76.218
0.079860
76.394
0.082167
76.278
0. 0.929
76.612
0.108805
76.670
0.165723
76.850
0.109256
76.646
0.269620
76,512
0.146364
76,442
0.343903
76.232
0.182960
76.504
0.438653
75.714
0.278865
76.428
0.559508
75.168
0.500000
75.356
Table 15: SGD: ImageNet
26

---

Under review as a conference paper at ICLR 2022
learning_rate
momentum
weight_decay
objective_valut
0,000010
0.886563
0.100137
0.100
0.000882
0.999000
0.088058
0.100
0.007832
0.886460
0.030099
0.100
0.084027
0.867351
0.114223
0.100
0.012791
0.871 190
0.762656
0.100
0.001059
0.779209
0.001784
0.100
0.003035
0.951726
0.001060
0.100
0.002368
0.920451
0.006413
0.100
0.043405
0.649227
0.379924
0.100
0. 0.000064
0.987255
0.586063
0.100
0,001212
0.600000
1.000000
49.392
0,028276
0.621864
0.007048
49.550
0,000010
0.600000
0.001000
50.098
0.000010
0.831437
O. 0.103619
50.924
0,000018
0.700728
0.008100 109
59.668
0.000058
0.999000
0.005061
66.848
0,000036
0.837616
0.001 366
66.908
0,001036
0.999000
1.000000
70.878
0,000144
0.889598
0.009050
72.672
0.0001 15
0.600000
0.036989
72.816
0.000158
0.675443
0.001000
72.840
0,000628
0.606453
0.001624
73.446
0.000292
0.779432
0.008273
73.740
0.000485
0.626856
0.004095
73.744
0.001786
0.877224
0.198602
73.944
0.000315
0.600000
1.000000
74.306
0.000258
0.999000
0.116977
74.562
0.000572
0.600000
0.014438
74.766
0.000232
0.622821
0.192101
75.740
0,000614
0.747656
0.040759
75.968
0.001000
0.799500
0.031623
76.146
0,001600
0.931992
0.058883
76.378
0.001030
0.803263
0.119977
76.398
0.001
0.848430
0.105139
76.474
0.000927
0.600000
0.101738
76.580
Table 16: Imagenet Batch Size 1024
27

---

Under review as a conference paper at ICLR 2022
learning_rate
momentum
weight_decay
objective_value
0.100000
0.638615
0.000010
0.098
10000000
0.828515
0.000828
0.100
100000
0.911213
0.000494
0.102
0,007128
0.714390
0.000112
3.004
0.017291
0.950635
0,000310
29.582
0.000010
0.999000
0.005794
35.810
0,000010
0.999000
0.010000
36.060
0. 0.000209
0.999000
0.000010
59.206
O. 0.000328
0.999000
0.000010
59.268
0. 0.000010
0.690038
0.000010
62.290
0,000010
0.746399
0.010000
62.454
0,000348
0,999000
0.010000
62.476
3.000013
0.797380
0.002920
64.452
0.000025
0.867802
0.000010
69.180
0.000049
0.984536
0.000378
70.680
0.002565
0.991471
0.0001 46
70.786
0,000052
0.850277
0.010000
71.704
0.000055
0.628945
0.000010
71,826
0,001567
0.697594
0.000473
71.870
0.003220
0.871350
0.000041
71.912
0.001823
0.800658
0.000302
72.114
0,001373
0.601300
0.002455
72.138
0,000973
0.803842
0.000208
72.436
0,000093
0.875276
0.000058
72.690
0.007858
0.964910
0.010000
72.694
0.000097
0.714558
0.000029
72.874
0.000761
0.819951
0.000481
72.914
0.000979
0.802539
0.002260
72.930
0.000233
0.780565
0,000107
72.958
0.000274
0.816199
0.000225
72.978
0.000197
0.773215
0.000843
72.988
0,000235
0.792918
0.000015
73.070
0.000208
0.896992
0.000010
73.150
0.000427
0.796783
0.000815 15
73.184
0,001725
0,809302
0.004951
73.408
0,000322
0.734479
0.010000
73.538
0,000647
0,851409
0.006530
73.594
0.000727
0.969539
0.010000
73.602
0,000263
0.869540
0.010000
73.658
0.001207
0.844760
0,010000
73,800
0.004101
0.934473
0.010000
73.944
0.002380
0.969588
0,010000
73,984
0,000771
0.856323
0.010000
74.030
0,001655
0.966147
0,010000
74.042
0.002968
0.961813
0.010000
74.196
0.002168
0.952292
0.010000
74.244
0,002870
0.960954
0.010000
74.298
0,000961
0.916564
0.207196
74.380
0.002261
0.869219
0.062905
74.862
0,000374
0.911343
0.376673
75.268
0.0001 00
0.870343
1.000000
75.430
0.001082
0.944840
O. 105851
75.692
0.001048
0.902476
0.093750
75.788
0.000411
0.950000
0.08 0.081495
75.856
0.001008
0.881673
0.100641
75.926
0,000967
0.885877
0.097476
75.964
0.000803
0.850000
0.073826
76.048
0.000905
0.850000
0.092752
76.066
0.001000
0.900000
0.100000
76.068
0.000921
0.850000
0.098263
76.096
0.000913
0.850000
1.1366
76.276
Table 17: Imagenet.Adam Batch Size 8192

---

Under review as a conference paper at ICLR 2022
Batch size 1024
Batch size 8192
learning_rate rate
accuracy
learning-rai
accuracy
0.003162
14.846
0.001283
6.694
0.003981
18.744
0,004114
17.484
0.006310
31.468
0,013190
43.236
0.007943
37.240
0.068529
63,408
0.010000
44.284
0,075832
64.202
0.015849
54.534
0,100000
65,428
0.019953
58.262
0.103259
65.768
0.025119 19
60.640
0.160894
67.610
0,03 163
63.074
0,185812
68. 104
0.039811
64.830
0.200426
68.268
0.063096
67.668
0.204970
68.280
0.079433
68.638
0.299255
69.638
0. 0.100000
69.144
0.303184
69.588
0.125893
69.928
0.414487
70.664
0. 158489
70.544
0.438243
70.708
0.199526
71.102
0,524030
71.312
0.251188 189
71.664
0.526898
71.426
0.398107
72.452
0.596139
71.438
0.501187
72.864
0.626390
71.736
0.630957
72.874
0.706443
72.086
0.794328
73,456
0.721966
72.044
1,000000
73,502
0.726457
72.214
0.876767
73.664
0.727370
72.160
0.960900
73.728
0.756109
72.164
1.053106
73.742
0.760034
72.388
1.264911
73.832
0.775082
72.284
1,386290
73.956
0.783008
72.274
1.519316
74.256
0.796600
72.150
1.824887
74.396
1.568864
0.100
2,000000
74.434
10.000000
0.100
Table 18: Graft SGD#Adagrad (Global): ImageNet
29

---

Under review as a conference paper at ICLR 2022
B.3
CIFAR-10 10 EXPERIMENTS
Batch size 128
Batch size 512
Batch size 2048
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
0.00075
85.418
0.00075
84.778
0.00075
82.032
0.0015
88.394
0.0015
86.778
0.0015
85.710
0.003
89.814
0,003
88.254
0.003
84.502
0.006
88.616
0.006
87.676
0.006
82.138
0.0125
87.270
0.0125
84.070
0.0125
73.304
0.025
84.822
0.025
77.164
Table 19: AdaGrad CIFAR-10
Batch size 128
Batch size 512
Batch size 2048
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
0.006
94.206
0.006
92.032
0.006
88.734
0.0125
94.564
0.0125
92.998
0.0125
90.724
0.025
94.872
0.025
93.786
0,025
91.974
0.05
94.932
0.05
94.386
0.05
92.820
0.075
94.784
0,075
93.638
0,075
93.372
0.1
94.772
0.1
93.412
0.1
88.818
0.15
10.000
0.15
10.000
Table 20: SGD (with momentum 0.9 and weight decay 0.0005) CIFAR--0
Batch size 128
Batch size 512
Batch size 2048
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
Learning Rate
Test Accuracy
0.0125
91.610
0.0125
86.998
0.025
85.332
0.025
92.878
0.025
89.764
0.05
88. 136
0.05
93.306
0.05
91.268
0.075
89.074
0.075
93.296
0.075
92.248
0.1
89.818
0.1
93,420
0.1
92.360
0.2
84.770
0.2
10.000
0.2
83.226
0.4
46.332
0.4
10.000
0.4
65.310
Table 21: SGD#AdaGrad CIFAR-10
30


Reviews:
[{"review_id": "oO9VWLfDTaa", "reviewer": "Reviewer_B67d", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "As a note, I reviewed a prior submission of this paper. The bulk of my review is the same as in the prior submission (with modifications corresponding to the modifications in the paper).\n\n## Originality\n\nThe paper is original, performing an experiment that I have not seen in the literature, and satisfactorily discusses background and related work.\n\n## Quality\n\nThe experiments performed within the paper are of high quality, satisfactorily demonstrating the main claims of the paper, that groupwise grafting results in performance equivalent to the better-tuned optimizer that the learning rates are derived from.\n\nMy main qualm with the experiments is the lack of motivation for the specific experimental configurations. Instead, the results of several disparate experiments in disparate settings are reported, making it hard to reason about the generality of the results. Specifically:\n- There are no reported results for global grafting on BERT, other than a statement that it performs \"significantly worse\". It would be helpful to see these results (in the main body or appendix) to understand how much worse and the tradeoffs between global and layer wise grafting.\n- The \"simplifying the discovered schedule\" section also feels arbitrary. The correction scheme is described as \"simple\", which it is, but it is also a bespoke correction for a single experiment which is not validated to transfer to other settings. It's not obvious what the contribution of this section is beyond that of the rest of the paper.\n\nI would also appreciate further validation that the differences in directions between M and D are significant on the layerwise and global scale: given that M # D == M when using per-weight groups, it seems possible that when using smaller groups that are e.g. the size of layers, M # D \\approx M. It would be worthwhile to show that the group-wise directions of M and D are significantly different. (my expectation is that they are essentially orthogonal, but this is worth validating).\n\n## Clarity\n\nOn the whole, the paper is very clear, with enough details to easily understand and reproduce the experiments reported in the paper. However, I do have the following questions and suggestions for improvements:\n- An explicit statement of the precise hypothesis being tested would significantly improve the paper. As-is, the results are possible to interpret in many different ways (e.g., that grafting either gives the performance of M, gives the better between M and D, or something else); and explicit statement of the hypothesis along with additional experiments to precisely test the bounds of that hypothesis would result in a much more precise takeaway from the paper.\n- Section 2.1: how does this deal with stochasticity in the gradient estimate? Based on Eq. 1, it seems like $g_t$ is the gradient for a minibatch, but the algorithm does not specify minibatch selection. \n- The notation of M#D and granularity could be improved; I kept having to scroll up and try to figure out which refers to magnitude and which refers to direction, and the granularity is specified separately in text.\n- The plots (specifically, Fig 2,3) are very small and hard to read.\n\n## Significance\n\nThough I am not an expert in optimizers, the results seem like a useful step towards empirically understanding the contrast and resultant difference in real-world performance between different optimizers.\n\n # Update after author response\n\nThanks to the authors for the response, especially the explicit statement of the research question. I do understand that the two experiments that I raised slight concerns with (\"simplifying the discovered schedule\" and global grafting for BERT) are not intended to be reusable methodologies for training networks in the future, but my concern is that without clearly stated methodologies/justifications for these one-off experiments or broader replication, the results may just be the result of chance (having continued to try variants of the approach until one worked) rather than supporting the broader claim that the (effective) learning rate schedule is the primary driver of accuracy.\n\nRegardless, my critiques are ultimately minor, and I do still think that the paper should be accepted. My main point of disagreement from the other reviewers is that I believe this paper is a significant contribution even if nobody ever uses the technique to train a network to deploy to end-users, because the paper's contribution is to the empirical understanding of optimizers rather than to the actual usage of optimizers. Because of that, I think that a slightly lower methodological bar is reasonable -- the paper is presenting a new finding moreso than a new technique, and we should not require that the approach works in all settings or even that the authors describe methodologies for finding similar results in new settings (e.g., with different hyper parameters or choices of M and D). \n", "summary_of_the_paper": "The authors propose learning rate grafting as a method to explore the power and dynamics of optimizers. Learning rate grafting partitions the parameters of the networks into groups, and for each group takes the direction of the weight update from one optimizer and the magnitude from another optimizer. The paper then shows that grafting allows for achieving the performance of a tuned optimizer using that tuned optimizer's group-wise magnitudes along with an untuned optimizer's group-wise directions.", "summary_of_the_review": "Accept. While I am not entirely satisfied with the motivation behind the choices of reported experiments and the missing statement of a falsifiable hypothesis to test, the paper proposes an interesting experiment and provides a technically correct evaluation of the claims in the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "a1KBb1Sh3AT", "reviewer": "Reviewer_4k79", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "The Algorithm and the method are quite simple. Magnitudes are taken from one, and directions are taken from \nanother. \nWhen you take, D and M#D, in the results, M and M#D are comparable in performance. So there is no clear benefit\nto using M#D, when M#D is comparable to M. A consequence of using M#D, they claim, is transferring the \nschedules? \nBut do schedules really generalize well?  Generalize to what? Suppose I change the data and architecture, \ncan I use this new rate scheduler?  This to be is the question they should address thoroughly.\n\nI think what the authors are proposing is this: If the change is in the optimizer only; how should I change the \nlearning rate schedule? But, there are additional questions.\n\nQ1. What is the impact of this method on generalization error? If yes, with what confidence?\nQ2. Does this new schedule improve the order of convergence? If yes, how?\nQ3. Can you discuss general properties of the grafting operator? \nQ4. What's the intuition behind transferring schedules (either theoretical or empirical) \nand usefulness of grafting?\nQ5. why should we expect A#B to perform better than B#A when there are two optimizers A and B, say?", "summary_of_the_paper": "The authors investigate the entanglements between the optimizer and the learning rate schedule and propose \nthe technique of optimizer grafting, which allows for the transfer of the overall implicit step size schedule\nfrom a tuned optimizer to a new optimizer, preserving empirical performance. This provides a robust \nplug-and-play baseline for optimizer comparisons, leading to reductions to the computational cost of \noptimizer hyperparameter search. Using grafting, they discover a non-adaptive learning rate correction to \nSGD which allows it to train a BERT model to state-of-the-art performance.", "summary_of_the_review": "The paper, although is interesting, lacks the technical/empirical novelty to merit publication", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XMidHaweD3", "reviewer": "Reviewer_JLsZ", "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": "The paper presents an interesting technique of grafting for the problem of step size hyperparameter tuning, and opens up questions as to the power of simple per-learning rate schedules. The strengths of the paper include performing analysis on state of the art benchmarks (ImageNet, CIFAR-10 for image classification and Wikipedia, Books for BERT pertaining). Another strength is assessing the context of transferring implicit step size schedules to another optimizer, for assessing global and per-layer variants, and for assessing simple learning rate discovery. \n\nHowever, empirical evidence was not sufficient to be conclusive on the methods presented. For example, the results were on 2 tasks  (image classification with ResNets and BERT pretraining) with 2 datasets each, and 2 batch sizes (8192 and 32768). Further, (for example) , Figure 2 displays results for the “best trial” performance, but does not include error bars. Adding error bars and the empirical computational cost reductions would be useful in better supporting the claims.\n\nOn several steps, potentially concerning results were identified but not adequately justified, and without more sufficient empirical results, the reader is uncertain if these are indications of more serious flaws in the approach. (e.g., for BERT (p.6), the global version of grafting were signficantly worse and “have been omitted” and (p. 14) despite these exponentially large correction ratios, the grated optimizer converges, but this “curious phenomenon” was left to future work).\n\nThe study also does not present theoretical underpinnings for the technique that would be helpful for understanding if indeed the results could be more widely applicable. \n\nThe study presents the approach M#D, but does not give guidance on how M or D might be selected more generally for other tasks, or how one might assess tasks to decide on M and D. There are also decisions for parameters that are not explained and it is not clear now sensitive results may be to these decisions. For example for learning rate discovery (p.8),  authors describe choosing the layer-wise multiplier as the median of individual corrections for the first 2000 iterations, or discretizing to the nearest power of ten. Providing additional guidance and/or theoretical underpinnings for such choices would be useful.", "summary_of_the_paper": "The authors report on a technique to address learning rate hyperparameter tuning for deep learning referred to as optimizer grafting. Specifically, the paper proposes a meta-algorithm (referred to as M#D) that blends the steps of two optimizers by combining the step magnitude of one (M) with the direction of the other (D). The technique of optimizer grafting allows for the transfer of the overall implicit step size schedule to a new optimizer, resulting in reductions in computational cost of optimizer hyper parameter search. The second primary result is leveraging the technique to identify a non-adaptive per-layer learning rate correction to SGD which allows it to train a BERT model to state-of-the-art performance. Analogous results are presented for vision models for global (non-per-layer) schedules for AdaGrad.\n\nThe authors describe grafting meta algorithm  (M#D) as, at each iteration, M#D feeds the same input (w_t, g_t) to both M and D which manage their states independently and produce w_M, w_D. Then the norms of the steps each would have taken is computed, and used to combine M’s magnitude update with D’s direction update. Partitioning is managed to implement global versus per-layer grafting.2 optimizer hyperparameter searches with the same computational budget, but different performances.\n\nThe authors present an empirical study on the transfer of implicit step size schedules between optimizers, comparing SGD and Adam to Adam#SGD for task of BERT pre-training. They show that Adam#SGD is able to achieve performance at/near Adam. \n\nThe paper also presents results for image classification for ImageNet and CIFAR-10), for AdaGrad, SGD, and SGD#AdaGrad, showing  SGD#AdaGrad outperforms SGD and AdaGrad. However, without error bars, it is hard assess the actual results.\n\nFinally, the paper shows results for grafting distilling a non-adaptive correction to D, eliminating the need to run M in parallel - that is, transferring a global, time-dependent non-adaptive multipliers for the learning rate. The results show for the global variant for ResNET (SGD, AdaGrad), the discovered learning rate is comparable to the one used on SGD and achieves a top1 accuracy of 72.46. For the per-layer variant, learning rate schedule enables a simple per-layer step size correction without adaptive preconditioning. The authors present proof-of-concept results for simplifying the discovered schedule as a way to support the robustness of their transfer approach.\n", "summary_of_the_review": "Overall, the paper presents an interesting technique and opens interesting questions, but without further empirical results or theoretical underpinnings, the results presented are insufficient to assess how generalizable the findings might be.", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "GrpHofCjiuF", "reviewer": "Reviewer_9GnN", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "S1. A large amount of experiment is conducted and plenty of result is shown in appendix.\n\nS2. A novel optimizing mode of grafting two different optimizers is proposed.\n\nW1. The paper structure is strange. I recommend to read some published proceedings to try to make this paper more clearly.\n\t\nW2. Some format maybe not legal. Such as the caption of table and content in Page 16.\n\t\nW3. The theory is not reasonable. In other word, you just tell me you do it like this but not why it’s reasonable. Actually, I don’t think ADAM#SGD will be better than ADAM. ADAM calculates the update direction according to loss function. In a multi-dimensional space, this direction is composed of the value of each gradient and positive or negative(symbol) of each gradient. However, you change the symbol of some parameters’ gradient according to SGD. This is why? I’m confused. In my view, this method is more like a SGD with multiplying a large const to its gradient.\n\t\nW4. I have a question, how to compute the norms (||w_m-w_t ||)/(||w_D-w_t ||). Is ||w_m-w_t || calculated with all the parameters in neural network? If not, I think Figure~1 is a wrong example, cause M#D will step to different direction with D in multi-dimensional space.\n\t\nW5. The results shown in tables are not strong enough.", "summary_of_the_paper": "This paper proposes a method called optimizer grafting. It uses two optimizers in one training session. One is to decide the update direction of parameters, and the other is to decide the update stride of parameters. This paper proposes a new optimizing mode and take a large amount of experiment exploration.\n", "summary_of_the_review": "Even though some idea is interesting, the theoretical work in this paper is insufficinet. ", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposed Trained ML oracles to find the decent direction and step size in optimization. The process they call grafting. Reviewers raised several concerns about the reliability of ML oracles in general settings which is valid. The rebuttal could not convince the reviewers to change their opinion.  Ideally for an empirical only paper with heavy reliability on ML for critical decisions, to meet the high bar of ICLR there must be several experiments (5-10 datasets or more) on diverse datasets and settings. Also, there should be discussions on when and how the method fails and related discussions. In that sense the paper does not meet the bar for publication.
Decision: Reject