Progressive Residual Warmup for Language Model Pretraining
This work improves the pretraining stability and convergence for researchers and practitioners working with Large Language Models, offering an incremental but effective optimization strategy.
This paper addresses the pretraining stability and convergence speed of Transformer-based Large Language Models by proposing Progressive Residual Warmup (ProRes). ProRes stabilizes pretraining, leads to faster convergence, stronger generalization, and better downstream performance by gradually warming up the residual connections of deeper layers.
Transformer architectures serve as the backbone for most modern Large Language Models, therefore their pretraining stability and convergence speed are of central concern. Motivated by the logical dependency of sequentially stacked layers, we propose Progressive Residual Warmup (ProRes) for language model pretraining. ProRes implements an "early layer learns first" philosophy by multiplying each layer's residual with a scalar that gradually warms up from 0 to 1, with deeper layers taking longer warmup steps. In this way, deeper layers wait for early layers to settle into a more stable regime before contributing to learning. We demonstrate the effectiveness of ProRes through pretraining experiments across various model scales, as well as normalization and initialization schemes. Comprehensive analysis shows that ProRes not only stabilizes pretraining but also introduces a unique optimization trajectory, leading to faster convergence, stronger generalization and better downstream performance. Our code is available at https://github.com/dandingsky/ProRes.