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.
As the cost of pretraining large language models grows, there is continued interest in strategies to improve learning efficiency during this core training stage. Motivated by cognitive development, where humans gradually build knowledge as their brains mature, we propose Curriculum-Guided Layer Scaling (CGLS), a framework for compute-efficient pretraining that synchronizes increasing data difficulty with model growth through progressive layer stacking (i.e. gradually adding layers during training). At the 100M parameter scale, using a curriculum transitioning from synthetic short stories to general web data, CGLS outperforms baseline methods on the question-answering benchmarks PIQA and ARC. Pretraining at the 1.2B scale, we stratify the DataComp-LM corpus with a DistilBERT-based classifier and progress from general text to highly technical or specialized content. Our results show that progressively increasing model depth alongside sample difficulty leads to better generalization and zero-shot performance on various downstream benchmarks. Altogether, our findings demonstrate that CGLS unlocks the potential of progressive stacking, offering a simple yet effective strategy for improving generalization on knowledge-intensive and reasoning tasks.
We ask whether language-model pre-training can be decomposed into smaller, independently trainable jobs that can later be recomposed into a coherent larger model. We introduce Mixture of Training (MoT), a scaffolded modular pre-training procedure that partitions a target Transformer into contiguous layer blocks, trains each block inside a frozen pretrained aligner scaffold, and then recomposes the trained blocks with an optional short end-to-end adaptation pass. On a 1.3B-parameter Gemma-style model trained on C4, MoT provides a small-scale proof of mechanism: independently trained depth slices can be recomposed into a usable language model, and a quality-parity schedule reaches the same reported perplexity as the monolithic baseline. This parity setting processes more aggregate tokens and has a shorter idealized layer-equivalent critical path after aligner preparation; its effective compute advantage depends on reusing the aligner across runs. We therefore present MoT not as a general replacement for monolithic pre-training, but as a small-scale framework for studying whether scaffolded sub-runs can act as reusable training units.
Mohammed Sabry, Sean Augenstein, Keith Rush +1
Google, Mountain View, CA, USA · School of Computing, Dublin City University, Dublin, Ireland.
We propose a preconditioning (PC) layer, a weight parameterization via polynomial preconditioner that ensures stable weight conditioning throughout LLM training. The PC module reshapes the singular-value spectrum of weight matrices via low-degree polynomial preconditioning. After training, the preconditioned weights can be merged back into the original architecture, incurring no inference overhead. We demonstrate the advantage of the proposed PC layer over standard transformers in Llama-1B pre-training, for both the AdamW and Muon optimizers. Theoretically, we justify this spectrum-control principle by proving that uniformly bounding each layer's singular values ensures geometric convergence of gradient descent to global minima, for certain deep linear networks. Our code is available at https://github.com/Empath-aln/PC-layer.
Senmiao Wang, Tiantian Fang, Haoran Zhang +4
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