PaperlayerView exact revision on arXiv

Source-linked AI summary

Progressive Residual Warmup for Language Model Pretraining

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

arXiv:2603.05369v1cs.CL

TL;DR

Deep Transformer pretraining must manage optimization challenges arising from sequentially dependent layers and simultaneous residual updates. ProRes schedules residual scaling from shallow to deep layers, delaying deeper contributions until earlier representations stabilize. Across scales and architectures, it improves perplexity, reasoning accuracy, generalization, and depth scaling.

  • Problem

    Scaling Transformers creates optimization challenges because deeper layers may contribute before upstream representations stabilize, while existing methods are generally not training-phase-aware.

  • Method

    ProRes multiplies each layer’s residual by a scalar that warms from 0 to 1, using longer warmup periods for deeper layers.

  • Results

    ProRes consistently improves performance and depth scaling across model sizes, initialization methods, and normalization schemes, including a 1.27% average reasoning-benchmark gain.

  • Takeaways & Limitations

    Training-phase-aware residual scheduling is an effective and practical direction for improving Transformer optimization.

Abstract

from arXiv · show

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.

1. Introduction

Transformer training becomes harder to stabilize and scale as layers are stacked, because deeper layers may contribute before upstream representations settle. ProRes addresses this by warming residual contributions progressively from shallow to deep layers, improving depth scaling and training behavior.

  • Motivation: Existing stabilization methods modify normalization, initialization, or residual connections but are generally not aware of training phase.These mechanisms are typically applied at initialization and then leave the optimizer to determine how learning proceeds.
  • Motivation: Transformer residual connections let all layers modify representations from initialization, potentially creating inefficient updates or conflicting learning signals.The paper frames this as a consequence of sequential layer dependencies and the absence of explicit coordination.
  • ProRes: ProRes assigns each layer’s residual a scalar that warms from 0 to 1, with deeper layers taking longer to activate than shallow layers.This prioritizes early-layer learning while delaying deeper contributions until upstream representations are more stable.
  • Results: ProRes consistently improves depth-scaling performance and lets deeper models achieve better perplexity without compromising training stability.The paper also analyzes how coordinating learning order across depth affects optimization dynamics and representation evolution.
  • Results: 71M-to-7B-parameter pretraining experiments show ProRes improves performance across model scales, initialization methods, and normalization schemes.The contribution summary presents this as evidence of broad effectiveness across pretraining settings.

2. Related Work

Related work stabilizes or accelerates Transformer training through normalization, initialization, progressive growth, freezing, and residual modifications. ProRes is positioned within this broader effort while emphasizing training-phase-aware coordination of layer contributions.

  • Normalization and initialization: Post-LN, Pre-LN, and related normalization schemes aim to stabilize activations and improve efficiency at scale.The discussion contrasts normalization placement and its role in controlling training behavior.
  • Normalization and initialization: Prior initialization methods bound model updates with respect to depth or width to improve deep-Transformer training stability.These approaches focus on controlling activation, weight, gradient, or update magnitudes.
  • Progressive training: Progressive-growth methods train deeper Transformers by stacking shallower models or training newly added layers, reducing overall training time.Related work also proposes progressive freezing because shallow layers converge earlier than deeper layers.

3. Method

ProRes modifies Transformer residual branches with predefined, layer- and time-dependent scaling factors. Its schedule starts residual contributions at zero and activates them progressively from shallow to deep layers, supporting identity behavior and controlled sequential learning.

  • Overview: In a Pre-LN Transformer, ProRes multiplies each residual branch by a predefined scalar α(l, t).x_l is the layer input representation, F is the attention or FFN module, and Norm is the normalization layer.
  • Overview: The linear ProRes schedule warms α(l, t) from 0 to 1, with warmup length increasing linearly with layer index.T denotes the warmup length for the first layer and L the total number of layers.
  • Overview: Under the linear schedule, residual branches activate sequentially from shallow to deep layers.This implements the paper’s intended learning and contribution order across depth.
  • Applicability: ProRes can be applied to Transformer variants that use stacked layers with residual connections.The paper summarizes the studied variants and their forward equations in Table 1.
  • Design principles: ProRes is motivated by identity behavior at initialization, bounded updates across depth and time, and sequential learning order.These principles target activation growth, gradient behavior, update magnitude, and interference from premature deeper-layer contributions.

4. Experiments

Experiments across model scales, architectures, and warmup schedules show that ProRes improves perplexity and reasoning performance while stabilizing training. Progressive shallow-to-deep activation is generally strongest, though optimal schedules depend on architecture and warmup length.

  • Experiment setup: ProRes was evaluated from 130M to 1.3B parameters, with additional scaling to 7B and depth experiments spanning 12 to 120 layers.The experiments used multiple Transformer variants, initialization methods, normalization schemes, and corpora.
  • Main results: ProRes consistently reduced pretraining-corpus perplexity across configurations and scales, with larger improvements for larger models.Post-LN benefited most, while combining ProRes with LNS produced smaller gains at larger scales.
  • Main results: 1.27% average gain in zero-shot reasoning accuracy was observed over corresponding baselines, with a 2.89% increase on LAMBADA.The top five average scores were all achieved by ProRes variants.
  • Main results: 4.86 average perplexity reduction on LAMBADA showed stronger out-of-distribution gains than the approximately 0.4 reduction on the pretraining corpus.The result aligned with the observed LAMBADA accuracy improvement.
  • Warmup schedule ablation: Shallow-to-deep residual activation consistently outperformed simultaneous or deep-to-shallow activation across Pre-LN, Post-LN, and Sandwich-LN.The authors connect this pattern to respecting the sequential dependency structure of stacked Transformer layers.
  • Warmup schedule ablation: The linear schedule was the most robust overall, but performance depended on architecture and warmup length; warmups below 1k or above 48k steps were unfavorable.For Post-LN, linear-square and stagewise-L outperformed linear, while fix schedules underperformed stagewise counterparts.
  • Depth scaling: Near-zero loss spikes as depth increased indicated that ProRes improved performance while stabilizing training.Spike scores were measured during the stable 10%–90% training phase.

5. Analysis

ProRes changes training dynamics by coordinating residual learning across depth. It mitigates activation growth and produces smoother, more stable representation evolution than vanilla Pre-LN.

  • Activation growth: ProRes mitigates the exponential activation growth observed in Pre-LN, producing gentler and more linear layerwise activation growth during early training.The comparison concerns 1.3B models during the first 20k training steps.
  • Training dynamics: Vanilla Pre-LN updates all layers concurrently from initialization, allowing deeper layers to operate on unstable upstream representations and fluctuating gradient signals.ProRes coordinates residual contributions to reduce these counterproductive updates.
  • Representation evolution: ProRes yields smoother representation evolution across layers, whereas vanilla Pre-LN shows frequent fluctuations in residual-output similarity during training.Both approaches show earlier convergence in shallow layers, but ProRes produces more stable updates across depth.

6. Conclusions

The paper introduces ProRes as a residual warmup scheme that coordinates layerwise learning during Transformer training. Across model sizes and training configurations, it improves performance and depth scaling, supporting training-phase-aware residual scheduling.

  • Method: ProRes progressively activates residual contributions across depth, prioritizing shallow layers before deeper layers engage after upstream representations stabilize.This explicitly coordinates layerwise learning over the course of training.
  • Results: Extensive pretraining experiments show consistent improvements across model sizes, initialization methods, and normalization schemes.The conclusion characterizes ProRes as simple, scalable, and effective across these settings.
  • Implication: Training-phase-aware residual scheduling is presented as an effective and practical direction for improving Transformer optimization.

Impact Statement

The paper aims to improve optimization and depth scaling for Transformer-based language models. It reports no new ethical concerns beyond those associated with existing large language models.

  • Research goal: The primary goal is to advance scalable training techniques for improving optimization and depth scaling in Transformer-based language models.
  • Potential impact: Improved training stability and efficiency may facilitate the development of larger and more capable models.
  • Ethical considerations: The broader societal impacts are described as consistent with those of existing large language models.The paper does not identify new ethical concerns or societal risks.

A. Learning Rate Tuning

The learning-rate tuning procedure accounts for scheduler behavior and model-specific stability. Choices vary by model scale and architecture, with lower rates selected for some Post-LN and DeepNorm configurations.

  • Tuning rationale: The final performance of the WSD scheduler is not highly sensitive to learning-rate choices when training remains stable.Smaller rates can lower loss during the constant-rate phase, while larger rates may catch up after decay.
  • 1.3B models: For 1.3B models, 6 × 10^-4 was selected for most baselines, while 3 × 10^-4 produced more stable training for Post-LN and DeepNorm.
  • 350M models: For 350M models, 1 × 10^-3 was adopted generally, except Post-LN, which performed best with 5 × 10^-4.
  • 130M models: For 130M models, selected rates were 1 × 10^-3 for Post-LN, 1.5 × 10^-3 for DeepNorm, and 2 × 10^-3 for remaining baselines.

B. Pretrain Experiments on 7B parameters

In 7B pretraining, ProRes uses a linear residual-warmup schedule and initially incurs higher loss than vanilla Pre-LN, but later achieves and maintains lower loss.

  • ProRes is evaluated against vanilla Pre-LN on 7B-parameter models using a learning rate of 3 × 10−4.Its linear schedule uses T = 1000, corresponding to 32,000 residual-warmup steps for 32 layers.
  • ProRes initially has higher loss because small α(l, t) values restrict model updates, but later achieves lower loss than the baseline.As training proceeds, residual warmup progressively unlocks the learning potential of each layer.
  • The loss gap continues widening after all residuals finish warming up, and ProRes maintains or slightly increases its advantage after learning-rate decay.Figure 5 reports both training-loss and evaluation-loss gaps across training steps.

C. Pretrain Experiments on Alternative Corpus

The paper tests ProRes generalization on the ClimbMix corpus and visualizes alternative residual-warmup schedules across layer depth and training time.

  • ProRes is evaluated on the ClimbMix dataset to test generalization beyond the C4 pretraining corpus.The experiments use the same preprocessing pipeline as Section 4.1 and report held-out-test evaluation perplexities in Table 7.
  • Held-out test-set evaluation perplexities are reported for the ClimbMix experiments.
  • Figure 6 visualizes the schedules in Table 5 using T = 1000 and L = 12, with darker lines for shallow layers and brighter lines for deeper layers.

E. Use of Compute Resources

The experiments ran on a single node with eight H800 GPUs, with training time varying substantially by model scale.

  • All experiments were conducted on a single node of 8×H800 GPUs.
  • Training required approximately 260h for 7B, 100h for 3B, and 50h for 1.3B models.
  • Training required approximately 20h for 350M models and 7h for 130M models.
Loading 2603.05369v1…