EMA Without the Lag: Bias-Corrected Iterate Averaging Schemes

TL;DR

The paper introduces BEMA, a bias-corrected exponential moving average method that improves stability and convergence in language model fine-tuning by eliminating bias inherent in traditional EMA.

Contribution

BEMA extends EMA by removing bias, offering provable acceleration and improved performance in language model fine-tuning.

Findings

BEMA outperforms standard EMA and vanilla training in convergence speed.

BEMA achieves higher final performance on language model benchmarks.

Theoretically demonstrates acceleration over existing methods.

Abstract

Stochasticity in language model fine-tuning, often caused by the small batch sizes typically used in this regime, can destabilize training by introducing large oscillations in generation quality. A popular approach to mitigating this instability is to take an Exponential moving average (EMA) of weights throughout training. While EMA reduces stochasticity, thereby smoothing training, the introduction of bias from old iterates often creates a lag in optimization relative to vanilla training. In this work, we propose the Bias-Corrected Exponential Moving Average (BEMA), a simple and practical augmentation of EMA that retains variance-reduction benefits while eliminating bias. BEMA is motivated by a simple theoretical model wherein we demonstrate provable acceleration of BEMA over both a standard EMA and vanilla training. Through an extensive suite of experiments on Language Models, we show that BEMA leads to significantly improved convergence rates and final performance over both EMA and vanilla training in a variety of standard LM benchmarks, making BEMA a practical and theoretically motivated intervention for more stable and efficient fine-tuning.