Learning Rate Grafting: Transferability of Optimizer Tuning

1 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 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 . 1.1 OUR CONTRIBUTIONS . 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 ” . 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 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 . 2 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 ∇̃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 . 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 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 .

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