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

2024-05-22

arXiv.org

AI Summary

Tuning weight decay (WD) in AdamW for large-scale models and datasets remains challenging due to its opaque interaction with model and data scale. Method: We introduce the exponential moving average (EMA) timescale—measured in epochs—as a physically interpretable core variable, decoupling and reformulating WD accordingly. Contribution/Results: Through theoretical analysis and extensive experiments across architectures (ResNet-18, ViT, NanoGPT) and datasets (CIFAR-10, ImageNet, OpenWebText), we find that the optimal EMA timescale is approximately invariant in epoch units. This yields an inverse scaling law: WD should increase with model size but decrease with dataset size. This law coheres with muP learning rate scaling; neglecting WD scaling while only scaling learning rate causes significant performance degradation in mid-to-late training. Our work establishes the first interpretable, principled scaling framework for AdamW weight decay, enhancing training stability and generalization in large-model regimes.

Technology Category

• Machine Learning: Scalability of ML Systems
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Responsible Web: Ethical and legal aspects of web-scale data analysis, uses and collection practices
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

The scaling of the optimal AdamW weight decay hyperparameter with model and dataset size is critical as we seek to build larger models, but is poorly understood. We show that weights learned by AdamW can be understood as an exponential moving average (EMA) of recent updates. This gives critical insights for how to set the weight decay in AdamW, and how the weight decay should scale with model and dataset size. In particular, the key hyperparameter for an exponential moving average is the EMA timescale. Intuitively, the EMA timescale can be understood as the number of recent iterations the EMA averages over. We find that the optimal timescale, measured in epochs, is roughly constant as we change model and dataset size. Moreover, given a learning rate, there is a one-to-one mapping from the EMA timescale to the weight decay hyperparameter. Thus, if the optimal EMA timescale is constant, that implies that as the dataset size increases, the optimal weight decay should fall and as the model size increases, the optimal weight decay should increase (if we follow the muP recommendation for scaling the learning rate). We validate these scaling rules on ResNet-18 and Vision Transformers trained on CIFAR-10 and ImageNet, and on NanoGPT pre-training on OpenWebText. Finally, we found that as training progresses, muP's learning rate scaling breaks down for AdamW unless weight decay is scaled appropriately.

Problem

• AdamW Optimization
• Model Scaling
• Weight Decay Parameter

Innovation

• Exponential Moving Average
• Weight Decay Adjustment
• muP Learning Rate Scaling