Progressive Residual Warmup for Language Model Pretraining

TL;DR

ProRes is a novel warmup technique for language model pretraining that stabilizes training, accelerates convergence, and improves downstream performance by gradually enabling deeper layers.

Contribution

ProRes introduces a layer-wise residual warmup strategy that enhances stability and efficiency in large language model pretraining.

Findings

ProRes stabilizes pretraining across various models.

ProRes accelerates convergence speed.

ProRes improves downstream task performance.

Abstract

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.