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

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

summaries: 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.  The 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.  The 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.   The 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.  Finally, 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.

summaries: The authors investigate the entanglements between the optimizer and the learning rate schedule and 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, they discover a non-adaptive learning rate correction to  SGD which allows it to train a BERT model to state-of-the-art performance.

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

authors:
Naman Agarwal
Brian Bullins
Xinyi Chen
Elad Hazan
Karan Singh
Cyril Zhang
Yi Zhang

title: Efficient full-matrix adaptive regularization

authors:
Naman Agarwal
Surbhi Goel
Cyril Zhang

title: Acceleration via fractal learning rate schedules

venue: arXiv preprint arXiv:2103.01338,

authors:
Rohan Anil
Vineet Gupta
Tomer Koren
Yoram Singer

title: Memory-efficient adaptive optimization for large-scale learning

venue: arXiv preprint arXiv:1901.11150,

authors:
Rohan Anil
Vineet Gupta
Tomer Koren
Kevin Regan
Yoram Singer

title: Scalable Second Order Optimization for Deep Learning

venue: arXiv e-prints, art

authors:
Sanjeev Arora
Zhiyuan Li
Kaifeng Lyu

title: Theoretical analysis of auto rate-tuning by batch normalization

venue: arXiv preprint arXiv:1812.03981,

authors:
Léon Bottou
Frank E Curtis
Jorge Nocedal

title: Optimization methods for large-scale machine learning

venue: Siam Review,

authors:
Olivier Bousquet
Sylvain Gelly
Karol Kurach
Olivier Teytaud
Damien Vincent

title: URL https://arxiv.org/abs/ 1706.03200

venue: Critical hyper-parameters: No random,

authors:
Andrew Brock
Jeff Donahue
Karen Simonyan

title: Large scale gan training for high fidelity natural image synthesis

venue: arXiv preprint arXiv:1809.11096,

authors:
Xinyi Chen
Naman Agarwal
Elad Hazan
Cyril Zhang
Yi Zhang

title: Extreme tensoring for low-memory preconditioning

venue: arXiv preprint arXiv:1902.04620,

authors:
Dami Choi
Christopher J Shallue
Zachary Nado
Jaehoon Lee
Chris J Maddison
George E Dahl

title: On empirical comparisons of optimizers for deep learning

authors:
Ashok Cutkosky
Harsh Mehta

title: Momentum improves normalized sgd

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: pre-training of deep bidirectional transformers for language understanding

authors:
John Duchi
Elad Hazan
Yoram Singer

title: Adaptive subgradient methods for online learning and stochastic optimization

venue: Journal of machine learning research,

authors:
Logan Engstrom
Andrew Ilyas
Shibani Santurkar
Dimitris Tsipras
Firdaus Janoos
Larry Rudolph
Aleksander Madry

title: Implementation matters in deep policy gradients: A case study on ppo and trpo

authors:
Rong Ge
Sham M Kakade
Rahul Kidambi
Praneeth Netrapalli

title: The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares

venue: In Advances in Neural Information Processing Systems

authors:
Boris Ginsburg
Patrice Castonguay
Oleksii Hrinchuk
Oleksii Kuchaiev
Vitaly Lavrukhin
Ryan Leary
Jason Li
Huyen Nguyen
Yang Zhang
Jonathan M Cohen

title: Stochastic gradient methods with layer-wise adaptive moments for training of deep networks

venue: arXiv preprint arXiv:1905.11286,

authors:
Akhilesh Gotmare
Nitish Shirish Keskar
Caiming Xiong
Richard Socher

title: A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation

venue: arXiv preprint arXiv:1810.13243,

authors:
Priya Goyal
Piotr Dollár
Ross Girshick
Pieter Noordhuis
Lukasz Wesolowski
Aapo Kyrola
Andrew Tulloch
Yangqing Jia
Kaiming He

title: Accurate, large minibatch sgd: Training imagenet in 1 hour

venue: arXiv preprint arXiv:1706.02677,

authors:
Vineet Gupta
Tomer Koren
Yoram Singer

title: Shampoo: Preconditioned stochastic tensor optimization

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor

venue: arXiv preprint arXiv:1801.01290,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Identity mappings in deep residual networks

venue: In European Conference on Computer Vision,

authors:
Andrew Ilyas
Logan Engstrom
Shibani Santurkar
Dimitris Tsipras
Firdaus Janoos
Larry Rudolph
Aleksander Madry

title: A closer look at deep policy gradients

venue: In International Conference on Learning Representations,

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: arXiv preprint arXiv:1502.03167,

authors:
Prateek Jain
Sham M Kakade
Rahul Kidambi
Praneeth Netrapalli
Aaron Sidford

title: Accelerating stochastic gradient descent for least squares regression

venue: In Conference On Learning Theory,

authors:
Tero Karras
Timo Aila
Samuli Laine
Jaakko Lehtinen

title: Progressive growing of gans for improved quality, stability, and variation

venue: arXiv preprint arXiv:1710.10196,

authors:
Nitish Shirish Keskar
Richard Socher

title: Improving generalization performance by switching from adam to sgd

venue: arXiv preprint arXiv:1712.07628,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Durk P Kingma
Prafulla Dhariwal

title: Glow: Generative flow with invertible 1x1 convolutions

venue: In Advances in Neural Information Processing Systems,

title: One weird trick for parallelizing convolutional neural networks

venue: arXiv preprint arXiv:1404.5997,

authors:
Yuanzhi Li
Colin Wei
Tengyu Ma

title: Towards explaining the regularization effect of initial large learning rate in training neural networks

authors:
Zhiyuan Li
Sanjeev Arora

title: An exponential learning rate schedule for deep learning

venue: arXiv preprint arXiv:1910.07454,

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

authors:
Ilya Loshchilov
Frank Hutter

title: Sgdr: Stochastic gradient descent with warm restarts

venue: arXiv preprint arXiv:1608.03983,

authors:
James Martens
Roger Grosse

title: Optimizing neural networks with kronecker-factored approximate curvature

venue: In International conference on machine learning,

authors:
H Brendan McMahan
Matthew Streeter

title: Adaptive bound optimization for online convex optimization

venue: arXiv preprint arXiv:1002.4908,

authors:
Stephen Merity
Nitish Shirish Keskar
Richard Socher

title: Regularizing and Optimizing LSTM Language Models

venue: arXiv preprint arXiv:1708.02182,

title: An Analysis of Neural Language Modeling at Multiple Scales

venue: arXiv preprint arXiv:1803.08240,

authors:
Zachary Nado
Justin M. Gilmer
Christopher J. Shallue
Rohan Anil
George E. Dahl

title: A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes

authors:
Zachary Nado
Justin M Gilmer
Christopher J Shallue
Rohan Anil
George E Dahl

title: A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes

venue: arXiv preprint arXiv:2102.06356,

authors:
Martin Popel
Ondřej Bojar

title: Training tips for the transformer model

venue: arXiv preprint arXiv:1804.00247,

authors:
Sashank J Reddi
Satyen Kale
Sanjiv Kumar

title: On the convergence of Adam and beyond

authors:
Robin M Schmidt
Frank Schneider
Philipp Hennig

title: Descending through a crowded valley– benchmarking deep learning optimizers

venue: arXiv preprint arXiv:2007.01547,

authors:
Christopher J. Shallue
Jaehoon Lee
Joseph Antognini
Jascha Sohl-Dickstein
Roy Frostig
George E. Dahl

title: Measuring the effects of data parallelism on neural network training

venue: Journal of Machine Learning Research,

authors:
Noam Shazeer
Mitchell Stern

title: Adafactor: Adaptive learning rates with sublinear memory cost

venue: arXiv preprint arXiv:1804.04235,

title: A disciplined approach to neural network hyper-parameters: Part 1–learning rate, batch size, momentum, and weight decay

venue: arXiv preprint arXiv:1803.09820,

authors:
Leslie N Smith
Nicholay Topin

title: Super-convergence: Very fast training of neural networks using large learning rates

venue: In Artificial Intelligence and Machine Learning for Multi-Domain Operations Applications,

authors:
Matthew Staib
Sashank Reddi
Satyen Kale
Sanjiv Kumar
Suvrit Sra

title: Escaping saddle points with adaptive gradient methods

authors:
Tijmen Tieleman
Geoffrey Hinton

title: Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning

authors:
Rachel Ward
Xiaoxia Wu
Leon Bottou

title: Adagrad stepsizes: Sharp convergence over nonconvex landscapes

authors:
Ashia C Wilson
Rebecca Roelofs
Mitchell Stern
Nati Srebro
Benjamin Recht

title: The marginal value of adaptive gradient methods in machine learning

authors:
Zhilin Yang
Zihang Dai
Yiming Yang
Jaime Carbonell
Ruslan Salakhutdinov
Quoc V Le

title: Xlnet: Generalized autoregressive pretraining for language understanding

authors:
Yang You
Igor Gitman
Boris Ginsburg

title: Large batch training of convolutional networks

venue: arXiv preprint arXiv:1708.03888,

authors:
Yang You
Jing Li
Sashank Reddi
Jonathan Hseu
Sanjiv Kumar
Srinadh Bhojanapalli
Xiaodan Song
James Demmel
Kurt Keutzer
Cho-Jui Hsieh

title: Large batch optimization for deep learning: Training bert in 76

title: Adadelta: an adaptive learning rate method

venue: arXiv preprint arXiv:1212.5701,

heading: 1 INTRODUCTION

text:
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 a state-of-the-art deep learning pipeline very opaque and expensive.
Consequently, a large body of work aims to understand the structure of these search spaces, analyzing the properties of optimizers and their components in idealized models to predict directional trends, provide rules of thumb, and propose mechanisms for emergent phenomena (like generalization and transfer). Other works study the intrinsic tradeoffs between competing objectives, such as faster convergence vs. generalization. The continued prevalence of black-box hyperparameter search points to a consensus that a quantitatively predictive understanding of training dynamics currently eludes us.
Towards forming more robust beliefs in this space, we provide a recipe that enables a wide class of ablation studies on state-of-the-art models. We consider an unnatural-looking meta-algorithm, which combines the update rules of two different optimizers, forming a single grafted update rule. Our findings serve not only as a sanity check on explanatory theory, but also as a way to simplify the optimizer search space in practical settings.

heading: 1.1 OUR CONTRIBUTIONS

text:
We introduce learning rate grafting, a meta-algorithm 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 performance of an optimizer can be explained by the implicit step size schedules it induces. Grafting provides an adaptive way to transfer these schedules between algorithms, consistently closing performance gaps between different optimizers used to train the same model architecture. As a result, grafting provides a strong baseline for learning rate schedule tuning: given a tuned baseline M,
grafting can enable a different optimizer D to 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 use grafting for the exploratory analysis of optimization trajectories and implicit learning rate schedules. Using grafting as a diagnostic tool, we discover a per-layer learning rate correction that allows SGD with momentum (without adaptive preconditioning) 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 schedule is derived. We show analogous results for vision models, finding novel and effective global (as opposed to per-layer) schedules for AdaGrad. We discuss the implications of this for contemporary 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”.

heading: 1.2 RELATED WORK

text:
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 performance differences (Gotmare et al., 2018; Smith, 2018; Smith & Topin, 2019; Loshchilov & Hutter, 2016). In large-scale language modeling, a source of particularly resource-intensive training problems, 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 learning 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 interaction of learning rates with batch normalization (Ioffe & Szegedy, 2015), 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 of thousands of citations for this reason. In particular, Adam has been the de facto standard in fields such as NLP (e.g. (Devlin et al., 2018; Yang et al., 2019; Liu et al., 2019)), deep generative modeling (e.g. (Karras et al., 2017; Brock et al., 2018; Kingma & Dhariwal, 2018)), and deep RL (e.g. (Haarnoja et al., 2018)). Adaptive methods have seen adoption in extremely large-scale settings, necessitating modifications to reduce resource consumption (Shazeer & Stern, 2018; Anil et al., 2019; Chen et al., 2019).
An important discussion was sparked by Wilson et al. (2017), who presented empirical and theoretical 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-
tuning (Ward et al., 2019), escaping saddle points (Staib et al., 2019), and mitigating long-tailed distributions (Zhang et al., 2019). Reddi et al. (2018) construct, then mitigate, pathological settings where Adam fails to converge. However, in the vast majority of cases, out-of-the-box adaptive methods are perfectly suitable for practitioners. Large-scale empirical studies (Choi et al., 2019; Schmidt et al., 2020; Nado et al., 2021) 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 are 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 et al., 2019) propose compressed preconditioners to reduce the space of per-coordinate adaptive multipliers (as maintained by algorithms such as AdaGrad/Adam) to a factored space. Novograd (Ginsburg et al., 2019) reduces the per-coordinate scalars to per layer scalars via addition. LARSlike 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.

heading: 2 PRELIMINARIES

heading: 2.1 OPTIMIZERS AS GENERIC STATE MACHINES

text:
There are various abstractions of optimization in deep learning: for example, we could reason about stochastic first-order oracles ∇̃F (w) which are unbiased estimators of some gradient of a population objective F (w). However, this fails to capture various popular deviations from the standard empirical risk minimization model, such as sequential batching in NLP (Merity et al., 2018; 2017) or deep reinforcement learning, where optimization deviates from its theoretical abstractions in many ways (Ilyas et al., 2019; Engstrom et al., 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 as a state machine which takes a sequence of gradients gt ∈ Rd and outputs an update rule for the weights wt ∈ Rd. This abstraction is agnostic to properties or theoretical guarantees of the optimizer; in supervised learning, the gradients are usually understood to be averaged over a minibatch, but we do not use this property either.
Thus, an algorithm A takes the current iterate and gradient (wt, gt), and outputs the next iterate wt+1; this update can depend on the algorithm’s internal state, hyperparameters, side information, or randomness. Then, vanilla SGD with learning rate schedule ηt is the algorithm A specified by
A(wt, gt) := wt − ηtgt. (1)
Then, in the usual process of iterative training, the model weights are iteratively updated, as in wt+1 ← A(wt, gt). 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 {ηt} already presents subtleties and bias-variance tradeoffs.

heading: 2.2 ADAPTIVE METHODS

text:
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 Ht, dependent on the observed information (wt, gt) so far, to modify the (stochastic) gradient descent steps:
A(wt, gt) := wt − ηtHtgt.
Despite their extreme popularity and theoretically principled origins, there is considerable uncertainty regarding the dynamics of adaptive methods. They were first motivated by regret bounds dependent on the geometry of the data (McMahan & Streeter, 2010; Duchi et al., 2011), hence
the terminology “adaptive regularization”. The view of loss curvature estimation leads to the term “adaptive preconditioning”, where Ht 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; Agarwal et al., 2019), we 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 Ht to be the diagonal matrix of entrywise inverse square roots of the accumulators.
Entanglements between preconditioners and learning rates. Notice that in Equation 2.2, a scalar multiplier on the preconditioner Ht can be absorbed into the learning rate schedule ηt. In other words, the dynamics of Ht can induce an implicit step size schedule over the course of training, 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 an experiment to answer the following question: does the empirical success of an adaptive optimizer arise from the true preconditioner Ht, or can it be equivalently described by a learning rate schedule ηt? By combining the step magnitudes of an optimizer M and step directions of an optimizer D, we can isolate and disentangle adaptive preconditioners Ht from their induced step size schedules.

heading: 3 THE GRAFTING META-ALGORITHM

text:
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 M#D. Before developing the metaalgorithm in full generality, we walk through the execution of the simplest case:
1. At each training iteration, M#D feeds the same input (wt, gt) to both children M,D, which manage their states independently and produce outputs wM, wD.
2. We do not update wt+1 yet. Instead, we compute ‖wM − wt‖ and ‖wD − wt‖, the norms of the steps the child optimizers would have taken.
3. The update step that grafting outputs is{ wt +
‖wM−wt‖ ‖wD−wt‖ (wD − wt) wD 6= wt
wt otherwise ,
which combinesM’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, x1, gt are deterministic, then A#A ≡ A (meaning that these two optimizers have the same trajectory under identical sequences of inputs (wt, gt)).
• IfM,D differ only by a learning rate schedule (i.e. wM − wt is always a positive scalar multiple of wD − wt, under identical inputs (wt, gt)), thenM#D ≡M.
• Disregarding the special case where A outputs a zero step, we have (M#A)#D ≡M#(A#D) ≡M#D.
• Grafting is not commutative.M#D 6≡ D#M. • Grafting is not necessarily a descent method. Even if M,D are guaranteed to make
progress on a deterministic objective (i.e. when gt is ∇F (wt)), M#D is not. For instance, ifM,D are gradient descent with different preconditioners, there is no guarantee thatM#D 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 us observe the optimization trajectory and end-to-end performance whenM andD are forced to have the same overall learning rate schedule. This gives a controlled way to study whether Ht or ηt is more important in determining an optimizer’s empirical success.
Algorithm 1 Grafted meta-optimizerM#D 1: init: optimizersM,D; hyperparameters forM,D; partition P of model parameters. 2: InitializeM,D with their hyperparameters. 3: for each iteration t do 4: Receive input (wt, gt). 5: Query steps from inner optimizers:
wM ←M(wt, gt), wD ← D(wt, gt). 6: for each parameter group ρ ∈ P do 7: Compute grafted update:
w (ρ) # ← w (ρ) t + ∥∥∥w(ρ)M −w(ρ)t ∥∥∥∥∥∥w(ρ)D −w(ρ)t ∥∥∥ ( w (ρ) D − w (ρ) t ) .
(Or, if w(ρ)D = w (ρ) t , let w (ρ) # ← w (ρ) t .)
8: Output grafted weight update: (M#D)(wt, gt) := w#.

heading: 3.1 GRANULARITY OF GRAFTING

text:
Before proceeding to the experiments, we will introduce one more dimension to the grafting procedure: its granularity. Formally, let P be a partition of {1, . . . , d}, supplied as a hyperparameter. For v ∈ Rd, ρ ∈ P , we denote by v(ρ) ∈ R|ρ| the restriction of v to the coordinates indexed at ρ. Then, we transfer norms separately for each set of coordinates in the partition; this is the complete method summarized by Algorithm 1. Let us note some special cases of P , and their associated properties and interpretations:
• If P = {{1, . . . , d}}, then we can think ofM#D as transferring an overall learning rate schedule fromM to D. This schedule is not in general a static sequence of scalars {ηt}, and depends on the trajectory of inputs (wt, gt). We call this instantiation global grafting.
• If we choose instead the most granular partitionP = {{1}, . . . , {d}}, grafting transfers one learning rate schedule per parameter. If the optimizers differ only by a positive diagonal preconditioner (like SGD and AdaGrad), so that the steps have the same sign pattern, then it holds thatM#D ≡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 is 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 parameter group fromM to D.
This provides a generalization of the schedule transfer methodology: for any choice of partition, we can forceM and D to have the same learning rate schedule simultaneously over the parameter groups. By making this partition more granular, we obtain more expressive families of learning rate corrections.

heading: 4 GRAFTING FOR IMPLICIT SCHEDULE TRANSFER

text:
We first present an empirical study on the transfer of implicit step size schedules between optimizers, on state-of-the-art training benchmarks. Starting with a well-tuned baselineM, our protocol compares two optimizer hyperparameter searches with the same computational budget: an optimizer D which is less frequently used and has worse performance thanM on the corresponding architecture, andM#D, fixing the tuned hyperparameters ofM. Note that the set of hyperparameters tuned in both cases are also the same i.e. the hyperparameters of D. Under the same hyperparameter tuning protocol and budget we consistently found across architectures/tasks and batch sizes, that grafting induced positive transfer of end-to-end model performance, i.e. the performance of D#M closes the gap (sometimes fully) between the performance D andM.
BERT pretraining. As a timely choice of a resource-intensive training task, we ran this hyperparameter transfer experiment for a BERT (Devlin et al., 2018) model pretraining setup. We make two choices ofM, Adam and LAMB (You et al., 2019) commonly reported as the state-of-the-art optimizer for this problem (also see Nado et al. (2021). For D we choose SGD with momentum, which has been widely reported to be unsuccessful in reaching state-of-the-art performance. In particular, 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 non-adaptively-preconditioned predecessor, performs poorly on self-attention models. We repeated the experiment at batch sizes 8192 (measuring final performance at 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.
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. (2021) that uses a quasi-random search (Bousquet et al., 2017) with a simple search space. Hyperparameters included are learning rate η, the moment parameters β1, β2, the polynomial power for the learning rate warmup pwarmup, and weight decay λ. We fixed the in Adam to 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 2 and Table 1 summarize 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 per-coordinate preconditioning, but at the granularity of layer-wise implicit step size schedules. Zhang et al. (2019) propose a 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 a less granular adaptive preconditioning rule.
ImageNet/CIFAR-10 classification with ResNets. As an alternative setting for grafting, we examine the behavior of grafting on convolutional architectures for image classification tasks: ImageNet (a ResNet-50 with 26M parameters) and CIFAR-10 (a ResNet-18 model with about 11M parameters). In this setting, we make two choices ofM, SGD (with momentum) as well as Adam, with tuned baseline setups derived from Goyal et al. (2017) (ImageNet) and Loshchilov & Hutter (2016) (CIFAR-10). A full description of the hyperparameters is deferred to the appendix. For D we choose AdaGrad, which is commonly believed to lead to low performance on these models. For these models we observe that the global version of grafting is as 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 are summarized in Figure 3 and the ImageNet results are summarized 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.
It has been noted many times (Zeiler, 2012; Wilson et al., 2017; Bottou et al., 2018) that the AdaGrad optimizer, with commonly used hyperparameter search spaces, fails to train state-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 a successful run of SGD/Adam.

heading: 5 GRAFTING FOR EXPLICIT SCHEDULE DISCOVERY

text:
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 runM in parallel. In this section, we demonstrate that this pipeline can be completed in both of the large-scale setups from the previous section. We extract the step size ratios ‖wM − wt‖ / ‖wD − wt‖ from a successful execution of grafting, and using these sequences as learning rate schedules for D in the
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 to layer-wise grafting, this allows for the discovery of a higher-dimensional correction.
Training a ResNet with AdaGrad. Using the global variant of Algorithm 1 with (M,D) = (SGD, AdaGrad) and keeping track of the norms of the steps produced byM and D periodically at every epoch, we arrive at a polynomial warmup-like correction overlaid on the stagewise-constant annealing (Goyal et al., 2017) schedule employed by SGD. Figure 4 displays this discovered learning rate schedule compared to the one used on SGD. The final training performance was a top-1 accuracy was 72.46%.
Training BERT with SGD. We demonstrate that the protocol of learning rate schedule discovery can be done with per-layer grafting as well. Strikingly, this enables us to discover a simple per-layer step size correction that allows the standard SGD(with momentum) to train a 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 in the previous section, and recorded the per-layer ratio of the norms. Figure 5 provides a collection of some representative ratios. Table 5 and Figure 8 in the appendix provide all the ratios as well as additional visualizations at a 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 accuracy of SGD to the state-of-the-art Adam baseline achieving eval accuracy of 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 to train BERT to state-of-the-art performance, as long as an appropriate 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 approach, showing that performance is preserved as long as relative per-layer scales are approximately preserved and provides a 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 is deferred to the appendix.
To provide a proof-of-concept, motivated by trends observed in Figure 5, we tie the per-layer scalar corrections for all the parameter groups in the attention layers to a fixed value β, and simply discretize 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 Dβ = {0.1, 1, 10, 100, β}. We provide the achieved accuracy results via the above corrections in Table 6b over some choices of β. We find that performance tends to be relatively stable to small changes in choices of β. We stress that our aim is to demonstrate a
0 250 500 750 1000 1250 1500 1750 2000 Training iterations
100
101
102
103
104
105
106
Ad am
/S GD
n or
m ra
tio d
ur in
g gr
af tin
g
self-attention Q self-attention K self-attention V self-attention out
feedforward softmax embeddings misc (LayerNorm, ...)
proof-of-concept that a simplified correction scheme (in this case composed of only 5 distinct values) works and thus to shed a light on the extent of the role of adaptivity in training these models. A careful study of the space of all possible discretizations is out of the scope and if left for future work.

heading: 6 CONCLUSION

text:
We have introduced learning rate grafting, a binary operation which blends the behavior of two optimization algorithms, towards investigating the entanglements between widely-used adaptive preconditioning rules and learning rate schedules. We have presented an empirical study with popular optimizers for state-of-the-art deep architectures, discovering that a well-performing optimizerM can transfer its performance to D via grafting its sequence of implicit step size schedules. Furthermore, we 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 on the power of simple per-layer learning rate schedules.
The empirical phenomena examined in this work seem to 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.

heading: A ADDITIONAL EXPERIMENTS AND DETAILS

text: We provide additional experiments and all the implementation details in this section. In the next section we provide full hyper-parameter search tables for all our experiments to aid reproduction of our experiments.

heading: A.1 BERT EXPERIMENTS

text:
A.1.1 IMPLEMENTATION DETAILS
We used the same experimental setup as the official BERT codebase1 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 hyperparameters using quasi-random search (Bousquet et al., 2017) with a fixed budget in a simple search space described in Table A.1.1. 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 η, exponential window parameters β1, β2, the polynomial power for the learning rate warmup and decay p, and weight decay λ. We fixed the in Adam to 10−11 for all BERT experiments. Polynomial power was searched over {1, 2} for batch size 8K and set to 1 for 32K. The entire hyperparameter search and associated performance can be found in Tables 7 and 11 for batch sizes 8K and 32K respectively.
SGD with momentum. Hyperparameters included learning rate η, β1, the polynomial power for the learning rate warmup and decay p, and weight decay λ. Polynomial power was searched over 1,2 for batch size 8K and set to 1 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: Adam#SGD. We use the above described implementations of Adam and SGD with momentum asM,D respectively. The Adam hyperparameters were fixed to the optimal ones found. The SGD hyperparameters were tuned with the same protocol as for SGD alone. The overall learning rate scalar was also tuned. Tables 9 and 13 contain the hyperparameter search and associated performance for batch sizes 8K and 32K respectively.
1https://github.com/google-research/bert
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 asM,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 was also tuned. Table 10 contains the hyperparameter search and associated performance for batch sizes 8K and 32K respectively.

heading: A.1.2 TRAINING BERT WITH SGD (BATCH SIZE 8K AND 32K)

text:
Some of these results are partly presented in the main paper. We provide a full exposition for completeness.
We demonstrate that the protocol of learning rate schedule discovery can be done with per-layer grafting as well. Strikingly, this enables us to discover a simple per-layer step size correction that allows the standard SGD(with momentum) to train a 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 in the previous section, and recorded the per-layer ratio of the norms. Figure 7 shows 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.2 In 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 6 and 8 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 4).
2Despite these exponentially large correction ratios, the grafted optimizer converges successfully. We leave an investigation of this curious phenomenon for future work.
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. (2016). 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 are 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 statistics 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 8192). The timings for a 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 Momentum 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 5 epochs and then the learning rate is held constant until cut by a 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 β2 to the default 0.999. 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 5 epochs and then the learning rate is held constant until cut by a factor of 0.1 at epoch 30,50,80. Tables 16 and 17 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 in 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: SGD#AdaGrad/Adam#Adagrad. We use the above described implementations of SGD with momentum/Adam and AdaGrad asM,D respectively. Since the only hyper-parameter for Adagrad is a learning rate , we tune the base learning rate supplied to SGD (as theM algorithm governs the overall learning rate of the grafted algorithm) Table 18 contains the hyperparameter tuning information where the search was performed over the base learning rate and across two batch sizes.

heading: A.3 CIFAR-10 EXPERIMENTS

text:
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 as 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 a 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
rate at the start of epoch t given the base learning rate β is given by
β · cos(π/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 at all. Table 19 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 momentum and AdaGrad as M,D respectively. Further we tune the base learning rate supplied to SGD (as the M algorithm governs the overall learning rate of the grafted algorithm) Table 21 contains the hyperparameter tuning information where the search was performed over the base learning rate 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 a strong baseline.
Table 6 contains the comparisons between the tuned results.

heading: B.1.1 BERT (BATCH SIZE 8K)

heading: B.1 BERT EXPERIMENTS

heading: B FULL HYPERPARAMETER SEARCH TABLES

heading: B.1.2 BERT (BATCH SIZE 32K)

text:
B.2 IMAGENET EXPERIMENTS
Batch size 1024 Batch size 8192 learning rate accuracy
0.000100 11.908 0.000158 16.586 0.000200 19.972 0.000316 28.038 0.000398 32.874 0.000501 37.754 0.001000 52.020 0.001995 61.424 0.002512 63.574 0.003981 66.834 0.005012 68.158 0.007943 70.272 0.010000 70.760 0.012589 71.520 0.015849 71.716 0.025119 72.554 0.031623 72.664 0.039811 72.562 0.050119 72.656 0.063096 72.926 0.079433 72.558 0.100000 72.658
learning rate accuracy
0.000100 37.788 0.000127 42.612 0.000394 60.916 0.001220 68.168 0.002028 69.590 0.002972 70.116 0.002999 70.134 0.003277 70.088 0.003442 70.488 0.004383 70.432 0.004410 70.356 0.004417 70.314 0.006360 69.872 0.007823 69.672 0.009804 69.538 0.010000 69.848 0.010554 69.502 0.014409 69.068 0.021884 68.306 0.095870 25.560 0.309611 7.162 1.000000 0.122
Table 14: AdaGrad: ImageNet
Batch size 1024 Batch size 8192 learning rate accuracy
0.001000 28.636 0.001627 43.516 0.003376 59.428 0.005493 65.994 0.008936 70.644 0.014539 73.312 0.023654 74.638 0.038483 75.686 0.079860 76.394 0.129926 76.612 0.165723 76.850 0.269620 76.512 0.343903 76.232 0.438653 75.714 0.559508 75.168
learning rate accuracy
0.005480 65.374 0.017648 73.556 0.029705 75.218 0.048524 75.688 0.050000 75.996 0.054569 76.108 0.059523 76.294 0.066660 76.218 0.082167 76.278 0.108805 76.670 0.109256 76.646 0.146364 76.442 0.182960 76.504 0.278865 76.428 0.500000 75.356
Table 15: SGD: ImageNet

heading: B.3 CIFAR-10 EXPERIMENTS