Progressive Residual Warmup for Language Model Pretraining
Authors: Tianhao Chen, Xin Xu, Lu Yin, Hao Chen, Yang Wang, Shizhe Diao, Can Yang
Organizations: The Hong Kong University of Science and Technology, Hong Kong, China · University of Surrey, Guildford, UK · The University of Hong Kong, Hong Kong, China · NVIDIA, Santa Clara, US
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.
Figures & tables
Figure 1: Evaluation perplexity ( ↓ ) of different methods as model depth increases. See Section 4.4 for details.
Method
ProRes
Forward Equation
Pre-LN
✘
xl+1=xl+F(Norm(xl))
✔
xl+1=xl+α(l,t)⋅F(Norm(xl))
Post-LN
✘
xl+1=Norm(xl+F(xl))
✔
xl+1=Norm(xl+α(l,t)⋅F(xl))
Sandwich-LN
✘
xl+1=xl+Norm(F(Norm(xl)))
✔
xl+1=xl+α(l,t)⋅Norm(F(Norm(xl)))
Table 1: Forward equations of Transformer variants, with (✔) and without (✘) ProRes.
Method
ProRes
130M
350M
1.3B
Pre-LN
✘
14.67
12.36
10.32
✔
14.30
11.74
9.86
Pre-LN (DS-Init)
✘
14.62
12.24
10.32
✔
14.29
11.73
9.85
Pre-LN (Scaled Init)
✘
14.63
12.29
10.30
✔
14.28
11.72
9.84
Table 2: Perplexity ( ↓ ) on C4-en test set across model scales. Cells with ProRes are shaded by improvement magnitude.
Method
ProRes
PIQA
SIQA
Hella
Wino
ARC-e
ARC-c
OBQA
RACE
Lambada
MMLU
Avg
Pre-LN
✘
72.85
41.86
52.45
55.41
53.54
26.45
33.40
31.29
39.30
27.69
43.42
✔
73.34
41.04
56.39
57.54
54.84
28.24
34.20
33.21
42.71
28.44
45.00
Pre-LN (DS-Init)
✘
72.58
41.66
52.92
56.35
53.79
27.39
34.20
32.25
39.69
27.79
43.86
✔
73.45
40.74
56.02
57.77
56.52
27.13
35.60
32.73
42.34
27.80
45.01
Pre-LN (Scaled Init)
✘
72.47
39.97
53.65
54.54
52.86
26.11
32.00
32.15
37.78
27.30
42.88
✔
73.83
41.10
56.48
56.83
56.78
28.07
35.20
31.10
42.34
28.08
44.98
Table 3: Zero-shot accuracy ( ↑ ) on reasoning benchmarks for 1.3B models.
Method
ProRes
WikiText
Lambada
Pre-LN
✘
25.61
20.46
✔
23.67
15.05
Pre-LN (DS-Init)
✘
25.77
19.83
✔
23.98
15.57
Pre-LN (Scaled Init)
✘
25.63
20.93
✔
23.62
15.42
Table 4: Perplexity ( ↓ ) on WikiText and Lambada for 1.3B models. Final row shows mean perplexity reduction (baseline minus ProRes).
Schedule
Warmup
α(l,t)
linear
T⋅L
min(T×lt,1)
linear-sqrt
T⋅L
(min(T×lt,1))1/2
linear-square
T⋅L
(min(T×lt,1))2
equal
T
min(t/T,1)
reverse
T⋅L
min(T×(L−l+1)t,1)
stagewise- 0
T⋅L
clip(Tt−T×(l−1),0,1)
Table 5: ProRes warmup schedules. Here l∈{1,…,L} denotes the layer index, t the training step, and T the warmup length for the first layer. “Warmup” denotes the number of training steps required for all layers to complete residual warmup. See Appendix D for visualization.
Schedule
Warmup
Pre-LN
Post-LN
Sandwich-LN
–
–
15.21
16.83
14.76
linear
1k
14.60
15.98
14.63
linear
3k
14.49
15.65
14.56
linear
6k
14.44
15.33
14.51
linear
12k
14.41
15.15
14.45
linear
24k
14.43
15.12
14.46
Table 6: Ablation of ProRes warmup schedules. Perplexity ( ↓ ) on C4 for 350M models trained with 6B tokens. “Warmup” denotes the number of training steps required for all layers to complete residual warmup. Blue indicates improvement while red indicates degrading performance relative to no-ProRes baseline. Perplexities higher than 1000 are marked as diverged . Additional baselines: DeepNorm (15.71) and LNS (14.63), trained under the same configuration without ProRes.
Figure 2: Spike score ( ↓ ) of loss and gradient norm as model depth increases.
Figure 3: Layerwise activation norm of Pre-LN, with and without ProRes.
Figure 4: Cosine similarity of layerwise residuals between intermediate (first 10% steps) and final checkpoint.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Loss gap across training steps of 7B models. Left: training loss gap (0.99 EMA smoothed) between Pre-LN and its ProRes variant. Right: evaluation loss gap calculated every 1000 steps.
Method
ProRes
130M
350M
1.3B
Pre-LN
✘
11.11
9.43
8.38
✔
10.82
9.13
8.09
Appendix
Table 7: Perplexity ( ↓ ) on ClimbMix test split across model scales.
Figure 6: Visualization of various schedules in Table 5 . We use T=1000 , L=12 for visualizing α(l,t) . Darker lines indicate shallow layers while brighter lines indicate deeper layers.
The Chinese University of Hong Kong, Shenzhen, China · Google LLC, Mountain View, CA, United States · Shenzhen International Center for Industrial and Applied Mathematics, Shenzhen2026 Research Institute of Big Data +1