Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Progressive Residual Warmup for Language Model Pretraining

Tianhao Chen, Xin Xu, Lu Yin, Hao Chen, Yang Wang, Shizhe Diao, Can Yang

TL;DR

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.

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.

Progressive Residual Warmup for Language Model Pretraining

TL;DR

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.

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.
Paper Structure (39 sections, 3 equations, 10 figures, 7 tables)

This paper contains 39 sections, 3 equations, 10 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Transformer initialization and normalization.
  4. Progressive growing and freezing layers.
  5. Method
  6. Overview
  7. Motivation of ProRes
  8. Principle 1: Identity behavior at initialization.
  9. Principle 2: Bounded model update with respect to depth and time.
  10. Principle 3: Respecting sequential learning and contribution order.
  11. Experiments
  12. Experiment Setup
  13. Model architecture.
  14. Baselines and ProRes schedule.
  15. Pretraining data.
  16. ...and 24 more sections

Figures (10)

  • Figure 1: Evaluation perplexity ($\downarrow$) of different methods as model depth increases. See \ref{['sec:depth_scaling_exp']} for details.
  • Figure 2: Spike score ($\downarrow$) of loss and gradient norm as model depth increases.
  • Figure 3: Pre-LN without ProRes
  • Figure 4: Pre-LN with ProRes
  • Figure 5: Layerwise activation norm of Pre-LN, with and without ProRes.
  • ...and 5 more figures