Early Weight Averaging meets High Learning Rates for LLM Pre-training (2306.03241v2)

Abstract: Training LLMs incurs significant cost; hence, any strategy that accelerates model convergence is helpful. In this paper, we investigate the ability of a simple idea checkpoint averaging along the trajectory of a training run to improve both convergence and generalization quite early on during training. Here we show that models trained with high learning rates observe higher gains due to checkpoint averaging. Furthermore, these gains are amplified when checkpoints are sampled with considerable spacing in training steps. Our training recipe outperforms conventional training and popular checkpoint averaging baselines such as exponential moving average (EMA) and stochastic moving average (SWA). We evaluate our training recipe by pre-training LLMs, where high learning rates are inherently preferred due to extremely large batch sizes. Specifically, we pre-trained nanoGPT-2 models of varying sizes, small (125M), medium (335M), and large (770M)on the OpenWebText dataset, comprised of 9B tokens. Additionally, we present results for publicly available Pythia LLMs, ranging from 1B to 12B, which were trained on the PILE-deduped dataset containing 207B tokens.

• The paper introduces LAWA, a method leveraging early checkpoint averaging as a surrogate for learning rate decay to boost convergence in LLM pre-training.
• The study shows that LAWA outperforms traditional techniques like EMA and SWA, delivering faster training on nanoGPT-2 and Pythia models.
• The approach enhances zero-shot task performance and scalability, reducing computational costs while improving model generalization.

Introduction

The paper "Early Weight Averaging meets High Learning Rates for LLM Pre-training" presents a novel approach to improve the convergence and generalization of LLMs during pre-training by employing early checkpoint averaging. The methodology is particularly effective when combined with high learning rates, which are inherently preferred due to the large batch sizes involved in LLM pre-training. The paper evaluates various pre-training techniques on nanoGPT-2 and Pythia models, demonstrating substantial improvements over conventional training methods and popular averaging baselines.

Methodology

Optimization and Diversity Insight

The paper introduces a strategy leveraging checkpoint averaging early during training, proposing that this technique can act as a surrogate for learning rate decay. By averaging weights post hoc from high learning rate trajectories, the approach mitigates oscillations in sensitive weight dimensions while enhancing generalization efficiently.

Additionally, the paper incorporates insights from model ensembling literature, suggesting that averaging distant checkpoints in a training trajectory induces model diversity. Increased diversity in checkpoints correlates with improved ensemble performance, thus enhancing the robustness and generalization capabilities of the final model.

LAWA Algorithm

The Latest Weight Averaging (LAWA) algorithm is adapted for LLM pre-training by integrating regular sampling of distant checkpoints and maintaining a sliding window of the latest checkpoints for averaging. The approach avoids restarting training with new schedulers and does not require intricate adaptations for batch normalization, simplifying integration into large-scale training regimes.

(Pseudocode for LAWA is provided in Algorithm 1, demonstrating its implementation in a Pytorch-style environment.)

Experimental Setup

The paper details the controlled experimental setup conducted across nanoGPT-2 models (125M, 335M, 770M parameters) and larger Pythia LLMs (up to 12B parameters). The models are trained on substantial datasets like OpenWebText and PILE-deduped, with evaluations conducted on held-out test sets and through zero-shot performance on downstream tasks such as Lambada OpenAI and SciQ.

Validation of the LAWA approach is characterized by analyzing log perplexity and zero-shot task performance, leveraging a range of checkpoints sampled throughout the pre-training trajectory.

(Figure 1, the paper's experimental results, show enhanced validation loss improvements with LAWA compared to EMA and SWA across different nanoGPT-2 model scales.)

Results

Early and Efficient Convergence

The experimental results affirm that models trained with higher learning rates exhibit substantial improvements and faster convergence when LAWA is applied. This gain is notably pronounced in early training stages and diminishes as the training progresses, due to inherent learning rate scheduling decays in traditional schemes.

The paper highlights that LAWA consistently outperforms traditional training paradigms and baseline averaging techniques, such as Exponential Moving Average (EMA) and Stochastic Weight Averaging (SWA). Notably, SWA applied earlier in training diverged, underscoring LAWA's versatility and robustness.

(Figure 2 demonstrates the performance trajectory of the nanoGPT-2 models using LAWA against original training, showcasing observable improvements in convergence rates and performance metrics.)

Scalability to Large Models

Extensive testing on Pythia models indicates that LAWA improves generalization across different scales, with accelerated convergence saving significant computational resources and training costs. The performance gains for larger models suggest the necessity of employing diverse checkpoint strategies to accommodate increased learning dynamics inherent to large-scale learning rates.

(Figure 3 highlights the substantial convergence savings for Pythia-2.8B and 6.9B models through LAWA-derived checkpoints compared to traditional training checkpoints.)

Improved Zero-Shot Performance

LAWA enhances zero-shot performance across a spectrum of downstream tasks, affirming the correlation between reduced training perplexity and improved zero-shot accuracy. These improvements are persistent across various checkpoints, offering significant computational advantages and early stopping opportunities in compute-constrained environments.

(Figure 4 illustrates the augmented zero-shot task performance with LAWA, providing empirical evidence of LAWA's effectiveness in optimizing zero-shot accuracy against traditional training checkpoints.)

Conclusion

The paper concludes that early weight averaging via LAWA offers substantial advantages for LLM pre-training, effectively accelerating convergence and improving generalization without increased computational overhead. The technique highlights several extensions for future research, including federated fine-tuning and continual intermediate checkpoint training, suggesting impactful implications for realizing efficient and scalable AI systems.