TLDR: A new research paper introduces Bias-Corrected Exponential Moving Average (BEMA), an enhancement to the popular EMA technique for fine-tuning large language models. BEMA eliminates the optimization lag caused by traditional EMA, leading to significantly faster convergence and improved performance on various language tasks. It is theoretically motivated and practically simple to implement, outperforming existing stabilization methods.
Fine-tuning large language models (LMs) is a crucial step to adapt them for specific tasks and improve their performance. However, this process often faces a significant challenge: training instability. This instability typically arises from using small batch sizes, which are common in fine-tuning to maximize the information extracted from limited high-quality data. Small batch sizes lead to increased variance in the training gradients, causing large oscillations in the model’s generation quality and making training difficult to stabilize.
A widely adopted technique to combat this instability is the Exponential Moving Average (EMA) of model weights. EMA works by averaging the model’s weights over time, which effectively smooths out the training process and reduces the impact of stochasticity. While EMA is successful in stabilizing training and improving final model performance, it introduces a drawback: a ‘lag’ in optimization. This lag occurs because EMA incorporates older weight values, which can bias the optimization process and slow down convergence compared to training without any stabilization.
A new research paper, EMA Without the Lag: Bias-Corrected Iterate Averaging Schemes, introduces a novel solution called Bias-Corrected Exponential Moving Average (BEMA). BEMA is designed to overcome the limitations of traditional EMA by retaining its variance-reduction benefits while completely eliminating the optimization lag caused by bias from old training steps. This means BEMA offers the best of both worlds: stable training and accelerated convergence.
The motivation behind BEMA is rooted in a simple theoretical model. The researchers demonstrate that BEMA can provably accelerate optimization compared to both standard EMA and even vanilla training (without any stabilization). This theoretical backing provides a strong foundation for its practical application.
From a practical standpoint, BEMA is remarkably simple to implement. It requires only a two-line change to existing EMA implementations, making it an easy ‘drop-in’ replacement for practitioners. The algorithm involves a bias-correction update that adjusts the EMA to account for the lag, ensuring a more direct path to optimal performance.
Extensive experiments on various language models, including Qwen2.5-1.5B, Gemma3-1B, and Llama3.2-1B, and standard LM benchmarks like BoolQ, GSM8K, and MMLU-HS, confirm BEMA’s effectiveness. The results show that BEMA consistently leads to significantly improved convergence rates and better final performance compared to both EMA and vanilla training. This improvement is robust across different optimizer settings, including learning rate decay schedules and varying batch sizes.
The paper also compares BEMA to other stabilization methods like OUEMA and Double EMA (DEMA). While OUEMA and DEMA show improvements over standard EMA, BEMA consistently outperforms them in terms of acceleration and overall performance on generation tasks. This highlights BEMA as a superior and theoretically motivated intervention for achieving more stable and efficient fine-tuning of language models.
Also Read:
- Enhancing Trustworthiness in Language Models: A Deep Dive into Calibration and Label Smoothing
- Smart Routing for AI at the Edge: Boosting LLM Performance
Future Directions
The researchers suggest several exciting avenues for future work. These include investigating optimal choices for initial model weights in BEMA, exploring how BEMA can be integrated with more sophisticated adaptive optimizers like AdamW, and considering the co-design of optimizers and stabilizers as a stochastic control problem. Additionally, applying BEMA to other training paradigms, such as Reinforcement Learning from Human Feedback (RLHF), could yield further benefits.
