How to set AdamW's weight decay as you scale model and dataset size

Overview

• The paper shows that the weights learned by the AdamW optimization algorithm can be understood as an exponential moving average (EMA) of recent updates.
• This insight provides critical guidelines for setting the weight decay hyperparameter in AdamW, which should scale with the model and dataset size.
• The optimal EMA timescale, which determines the weight decay, should be on the order of the number of training epochs and not change much as the model and dataset scale.

Plain English Explanation

The paper examines the AdamW optimization algorithm, a popular method for training large machine learning models. The key insight is that the weights learned by AdamW can be viewed as an exponential moving average (EMA) of the recent updates made during training.

This EMA interpretation provides important guidelines for setting the weight decay hyperparameter in AdamW. The weight decay controls how quickly the model "forgets" past updates and is a critical setting that impacts the model's performance. The authors show that the optimal weight decay should scale with the model and dataset size.

Specifically, the EMA timescale - the number of recent iterations the EMA averages over - is a key parameter that determines the weight decay. The authors argue that the EMA timescale should be on the order of the total number of training epochs, but not much smaller or larger. This is because we want the EMA to average over all the training data, but also to forget about very early updates.

Importantly, this implies that as the dataset size increases, the optimal weight decay should decrease. And as the model size increases, the optimal weight decay should increase (following the recommendation to also scale the learning rate with model size).

The paper validates these guidelines empirically and shows they are consistent with the hyperparameter choices made in recent large-scale language model pretraining efforts like Llama and Stable LM. This provides confidence that the insights around AdamW and weight decay are important for successfully scaling up machine learning models.

Technical Explanation

The paper provides a theoretical analysis showing that the weights learned by the AdamW optimization algorithm can be understood as an exponential moving average (EMA) of the recent parameter updates. This EMA interpretation leads to critical insights about how to set the weight decay hyperparameter in AdamW.

Specifically, the authors show that there is a one-to-one mapping between the EMA timescale (the number of recent iterations the EMA averages over) and the usual weight decay hyperparameter, given a fixed learning rate. Intuitively, the EMA timescale can be thought of as the "memory" of the optimization algorithm - how many recent updates it remembers.

The authors argue that the EMA timescale should be on the order of the total number of training epochs, but not much smaller or larger. This is because we want the EMA to average over all the training data, but also to forget about very early updates that may no longer be relevant. Empirically, the authors find that optimal EMA timescales are consistent with these guidelines.

Critically, this implies that as the dataset size increases, the optimal weight decay should decrease, since the same number of epochs will cover more data. And as the model size increases, the optimal weight decay should increase, following the recommendation to also scale up the learning rate.

The authors validate these insights through experiments and show that the hyperparameter choices made in recent large-scale language model pretraining runs, such as Llama and Stable LM, are consistent with the guidelines derived from the EMA interpretation of AdamW.

Critical Analysis

The paper provides a compelling theoretical analysis and empirical validation of the relationship between the EMA timescale, weight decay, and model and dataset scaling in the context of the AdamW optimization algorithm. The insights offered are potentially quite impactful, as they provide clear guidelines for tuning a critical hyperparameter in large-scale machine learning models.

One potential limitation of the work is that it focuses exclusively on the AdamW algorithm, leaving open the question of whether the insights generalize to other optimization methods. Additionally, the paper does not delve into the potential reasons why the EMA timescale guidelines would hold, beyond the intuitive explanations provided. A deeper exploration of the underlying dynamics and theoretical justifications could further strengthen the conclusions.

It would also be valuable to see the guidelines validated on a broader range of model architectures and tasks, beyond just the language models discussed. Applying the insights to other domains, such as computer vision or reinforcement learning, could help establish their broader applicability.

Finally, the paper does not address potential issues that could arise from the proposed scaling of weight decay, such as interactions with other hyperparameters or potential negative impacts on model performance or generalization. A more comprehensive analysis of the practical implications and potential drawbacks would be a useful addition.

Overall, the paper presents an important and insightful contribution to the understanding of optimization algorithms and their role in scaling up machine learning models. The clear guidelines and empirical validation make it a valuable resource for practitioners working on large-scale model development.

Conclusion

This paper provides a critical analysis of the AdamW optimization algorithm, showing that the weights it learns can be understood as an exponential moving average (EMA) of recent updates. This EMA interpretation leads to important guidelines for setting the weight decay hyperparameter in AdamW, which should scale with the model and dataset size.

Specifically, the authors show that the optimal EMA timescale - the number of recent iterations the EMA averages over - should be on the order of the total number of training epochs, but not much smaller or larger. This implies that as the dataset size increases, the optimal weight decay should decrease, while as the model size increases, the optimal weight decay should increase.

These insights are validated empirically and shown to be consistent with the hyperparameter choices made in recent large-scale language model pretraining efforts. The work offers valuable guidance for practitioners working on scaling up machine learning models, and highlights the importance of understanding the underlying dynamics of optimization algorithms.