Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

By Noam Shazeer, Mitchell Stern. 2018

1,323 citations recorded by the source corpus.

In several recently proposed stochastic optimization methods (e.g. RMSProp, Adam, Adadelta), parameter updates are scaled by the inverse square roots of exponential moving averages of squared past gradients. Maintaining these per-parameter second-moment estimators requires memory equal to the number of parameters. For the case of neural network weight matrices, we propose maintaining only the per-row and per-column sums of these moving averages, and estimating the per-parameter second moments based on these sums. We demonstrate empirically that this method produces similar results to the baseline. Secondly, we show that adaptive methods can produce larger-than-desired updates when the decay rate of the second moment accumulator is too slow. We propose update clipping and a gradually increasing decay rate scheme as remedies. Combining these methods and dropping momentum, we achieve comparable results to the published Adam regime in training the Transformer model on the WMT 2014 English-German machine translation task, while using very little auxiliary storage in the optimizer. Finally, we propose scaling the parameter updates based on the scale of the parameters themselves.

Subfields: Artificial Intelligence, Statistical and Nonlinear Physics, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Statistical and Nonlinear Physics

In Optimization & Training

Citations in the corpus

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (2019, methodology): “Apart from those changes mentioned above, we use the same training procedure and hyperparameters as our baseline (AdaFactor optimizer, inverse square root learning rate schedule for pre-training, constant learning rate for fine-tuning, dropout regularization, vocabulary, etc.).”
Scaling Laws for Neural Language Models (2020, methodology): “Due to memory constraints, our largest models (more than 1B parameters) were trained with Adafactor [SS18].”
PaLM: Scaling Language Modeling with Pathways (2022, methodology): “The model was trained with the Adafactor optimizer (Shazeer & Stern, 2018), without factorization.”
How Much Knowledge Can You Pack into the Parameters of a Language Model? (2020, methodology): “3 For ﬁne-tuning the T5 checkpoints, we follow the procedure used in Raffel et al. (2019) without any additional hyperparameter tuning: We use the AdaFactor optimizer (Shazeer and Stern, 2018) with a constant learning rate of 0 .”
Simple, Scalable Adaptation for Neural Machine Translation (2019, methodology): “All our models are trained with Adafactor (Shazeer and Stern, 2018) with momentum factorization, a learning rate schedule of (3.”
CTRL: A Conditional Transformer Language Model for Controllable Generation (2019, methodology): “We also experimented with explicit memory-saving optimizers including SM3 (Anil et al., 2019), Adafactor (Shazeer & Stern, 2018), and NovoGrad (Ginsburg et al., 2019) with mixed results.”
Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges (2019, methodology): “All our models are trained with Adafactor (Shazeer and Stern, 2018) with momentum factorization, a learning rate schedule of (3.0, 40k), 5 and a per-parameter norm clipping threshold of 1.0.”
Memory-Efficient Adaptive Optimization for Large-Scale Learning (2019, methodology): “One notable exception, which is directly related to our work, is the Adafactor algorithm [23] that was proposed as a way to reduce the memory costs of Adam, primarily for training large language models.”
fairseq: A Fast, Extensible Toolkit for Sequence Modeling (2019, methodology): “We provide wrappers around most PyTorch optimizers and an implementation of Adafactor (Shazeer and Stern, 2018), which is a memory-efficient variant of Adam.”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “Our optimizer is Adafactor (Shazeer & Stern, 2018) with first-moment decay β1 = 0, second-moment decay β2 = 0.99 with a 1 − t−0.8 decay schedule, update clipping threshold of 1.0, and factored second-moment estimation.”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “3) Memory Efficient Optimizer: [21], [22] focus on reducing memory consumption of adaptive optimization methods by maintaining coarser-grained statistics of model parameters and gradients, with potential impact on model convergence guarantees.”
Finetuned Language Models Are Zero-Shot Learners (2021, methodology): “…examples per dataset to 30k and follow the examplesproportional mixing scheme (Raffel et al., 2020) with a mixing rate maximum of 3k. 2 We finetune all models for 30k gradient steps with a batch size of 8,192 tokens using the Adafactor Optimizer (Shazeer & Stern, 2018) with a learning rate of 3e-5.”
Deduplicating Training Data Makes Language Models Better (2021, methodology): “Like T5, models were trained with the Adafactor optimizer (Shazeer and Stern, 2018).”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “We use 32-bit floating-point parameters, 16-bit floating-point activations, and Adafactor optimizer [35].”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, methodology): “We used the Adafactor (Shazeer & Stern, 2018) optimizer with a) factored second-moment estimation; b) ﬁrst moment decay β 1 = 0 .”
Efficient Content-Based Sparse Attention with Routing Transformers (2020, methodology): “Finally, we use the Adafactor optimizer (Shazeer and Stern, 2018) which is more memory efficient than Adam in training larger models.”
GLU Variants Improve Transformer (2020, methodology): “Similarly to [Raﬀel et al., 2019], we use the Adafactor optimizer [Shazeer and Stern, 2018] and an inverse-square-root learning-rate schedule.”
Towards a Human-like Open-Domain Chatbot (2020, methodology): “Additionally, to save memory, we chose the Adafactor optimizer (Shazeer and Stern, 2018) with 0.”
Reformer: The Efficient Transformer (2020, methodology): “We used the Adafactor optimizer (Shazeer & Stern, 2018) for training these models.”