Learning Rate Grafting: Transferability of Optimizer Tuning

Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, Cyril Zhang

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.

One-sentence Summary: We show implicit learning rate schedules can be transferred between optimization algorithms closing gaps in performance, overall leading to training of a BERT model to SOTA by SGD with no adaptivity.