Relative Generalization Invariance of LLM Pretraining

Fengzhuo Zhang; Shuche Wang; Shenggui Li; Tianyu Ruan; Jianliang He; Ivor Tsang; Tianyu Pang; Chao Du

cs.LGSep 26, 2026

Relative Generalization Invariance of LLM Pretraining

Authors: Fengzhuo Zhang, Shuche Wang, Shenggui Li, Tianyu Ruan, Jianliang He, Ivor Tsang, Tianyu Pang, Chao Du,

Organizations: Yale University · National University of Singapore · Nanyang Technological University · A*STAR ·

Abstract

Large Language Model (LLM) pretraining performance is jointly shaped by three components of the training triplet: the optimizer, model architecture, and training data stream. However, how these components influence performance in distinct ways remains unclear. We take a first step toward isolating their effects by studying relative generalization. We introduce Relative Generalization Invariance (RGI), the invariance of the validation-loss difference between any two tokens across models. We show that RGI approximately holds across a wide range of optimizers and moderate architectural variations, suggesting that these choices induce an approximately uniform shift in token-wise losses. In contrast, changing the training data stream can substantially alter relative generalization. We further show that RGI cannot be explained by the neural tangent kernel or mean-field regimes alone and prove that it can emerge in an overparameterized quadratic model. Overall, our work identifies RGI as a new phenomenon in LLM pretraining that helps distinguish the effects of optimizers and architectures from those of training data.

Figures & tables

Panel of Figure 1 — Figure 1: rgi emerges across optimizers and architectures and is not explained by the ntk or mf regimes alone. Panel (a) shows that training-loss curves across optimizers and architectures exhibit similar shapes after $2{,}000$ steps. Panel (b) shows that rgi persists under substantial feature learning and at reduced model widths, providing evidence beyond ntk - and mf -based explanations.

Panel of Figure 2 — Figure 2: Empirical evidence of rgi among optimizers. Panel (a) plots the validation losses of the $124$ M GPT models trained with Adam and Muon at the $10{,}000$ -step checkpoint, evaluated on $20{,}480$ validation batches. Panel (b) reports $r_{\mathrm{CCC}}(\text{Adam},\cdot)$ $(\uparrow)$ for the other five optimizers over training, all of which rise above $0.95$ after $2{,}000$ steps. Panel (c) reports $r_{\mathrm{diff}}(\text{Adam},\cdot)$ $(\downarrow)$ for the other five optimizers over training, all of which drop below $0.05$ after $2{,}000$ steps. These results show that rgi holds across all considered optimizers after sufficient training when they are trained from the same initialization and on the same training data stream.

Panel of Figure 3 — Figure 3: rgi across optimizers under different hyperparameter settings. Panels (a) and (b) plot $r_{\mathrm{CCC}}$ under different learning-rate multipliers, while panels (c) and (d) plot $r_{\mathrm{CCC}}$ under different momentum values. rgi persists across all tested learning-rate and momentum settings.

Panel of Figure 4 — Figure 4: rgi across optimizers, initializations, and model sizes. Panel (a) plots $r_{\mathrm{CCC}}$ across optimizers with different initializations. Panels (b) and (c) plot $r_{\mathrm{CCC}}$ across models with hidden-state dimensions $768$ , $384$ , and $64$ . rgi continues to hold across different initializations and moderate changes in model size.

Panel of Figure 5 — Figure 5: rgi across optimizers and models with different attention and ffn modules. Panels (a)–(c) plot $r_{\mathrm{CCC}}$ between models using RoPE fa trained with $\mathsf{Adam}$ and models using NoPE fa , ALiBi fa , and RoPE swa , each trained with all optimizers, respectively. Panel (d) plots $r_{\mathrm{CCC}}$ between models using GELU and SReLU activations. rgi continues to hold across different attention and ffn designs.

Figure 6: Values of

|\alpha_{t}^{\mathfrak{T}}|

and

|\beta_{t}^{\mathfrak{T}}|

. It shows that

|\beta_{t}^{\mathfrak{T}}|

is much larger than

|\alpha_{t}^{\mathfrak{T}}|

Panel of Figure 7 — Figure 7: rgi across optimizers trained on different data streams. The panels plot $r_{\mathrm{CCC}}$ when Adam is trained on FineWeb and the compared optimizer is trained on C4 (Panel (a)), ArXiv (Panel (b)), or CodeParrot (Panel (c)). The results show that rgi is strongly affected by the training data stream: changing the data stream can substantially weaken, or even eliminate, rgi across optimizers.

Panel of Figure 8 — Figure 8: rgi across optimizers and architectures evaluated on different validation datasets. The panels plot $r_{\mathrm{CCC}}$ evaluated on C4, ArXiv, CodeParrot, and Random. These results show that rgi can transfer to related validation distributions, but becomes weaker on more distributionally distinct validation datasets.

Panel of Figure 9 — Figure 9: Within- and cross-optimizer $\mathrm{cos\text{-}sim}(\uparrow)$ for parameters and hidden states. Panels (a)–(b) compare $\theta_{0}^{\mathfrak{T}}$ with $\theta_{t}^{\mathfrak{T}}$ , while Panels (c)–(d) compare $\theta_{t}^{\mathfrak{T}}$ with $\theta_{t}^{\mathfrak{T}^{\prime}}$ . Hidden-state similarities are averaged over layers $l\in[L]$ . Within-optimizer comparisons show that all optimizers learn nontrivial features throughout training, whereas cross-optimizer comparisons show that different optimizers learn distinct features.

Panel of Figure 10 — Figure 10: rgi across models of different widths. Panels (a)–(c) plot $r_{\mathrm{CCC}}$ for models with different widths, showing that rgi holds across all tested widths. Panel (d) plots the validation loss of Adam at $10{,}000$ steps across learning rates and widths, showing that the width- $192$ model exhibits behavior inconsistent with the mf regime.

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 11 — Figure 11: Validation-loss linearity across Adam and other optimizers. Panels (a)–(d) plot the validation-loss relationship between Adam and Adam-mini, Lion, Soap, and Scion, respectively, at the $10{,}000$ -step checkpoint. The near-perfect fits, with $R^{2}\geq 0.99$ , support approximate rgi across optimizers.

Panel of Figure 12 — Figure 12: Baseline-optimizer ablation for $r_{\mathrm{diff}}$ . Panels (a)–(e) plot $r_{\mathrm{diff}}$ when the baseline optimizer is Adam-mini, Lion, Muon, Soap, and Scion, respectively. Approximate rgi is insensitive to which optimizer is used as the evaluation baseline.

Panel of Figure 13 — Figure 13: Baseline-optimizer ablation for $r_{\mathrm{CCC}}$ . Panels (a)–(e) plot $r_{\mathrm{CCC}}$ when the baseline optimizer is Adam-mini, Lion, Muon, Soap, and Scion, respectively. Approximate rgi is insensitive to which optimizer is used as the evaluation baseline.

Panel of Figure 14 — Figure 14: Per-token evaluation of rgi . Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, when each validation batch contains a single token, i.e., $|D_{\mathrm{val},k}|=1$ . The small $r_{\mathrm{diff}}$ and large $r_{\mathrm{CCC}}$ values show that rgi also holds at token-level granularity.

Panel of Figure 15 — Figure 15: Baseline rgi for $0.7$ B models. Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, between Adam and Muon when they share the same initial parameters and training data stream. The results show that approximate rgi persists at the $0.7$ B scale.

Panel of Figure 16 — Figure 16: Baseline rgi for $4.5$ B models. Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, between Adam and Muon when they share the same initial parameters and training data stream. The results show that approximate rgi persists at the $4.5$ B scale.

Panel of Figure 17 — Figure 17: Initialization ablation for $0.7$ B models. Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, for Adam and Muon runs started from different initial parameters. Approximate rgi still emerges at the $0.7$ B scale, showing that initialization differences do not prevent the phenomenon.

Panel of Figure 18 — Figure 18: Learning-rate ablation for $0.7$ B models. Panels (a) and (b) plot $r_{\mathrm{diff}}$ for $0.25\eta^{\mathsf{A}}$ and $1.25\eta^{\mathsf{A}}$ , while Panels (c) and (d) plot $r_{\mathrm{CCC}}$ for the same two learning-rate multipliers. Approximate rgi remains robust to these learning-rate changes at the $0.7$ B scale.

Panel of Figure 19 — Figure 19: Momentum ablation for $0.7$ B models. Panels (a)–(c) plot $r_{\mathrm{diff}}$ for $\gamma=0.85$ , $0.9$ , and $0.95$ , while Panels (d)–(f) plot $r_{\mathrm{CCC}}$ for the same momentum values. Approximate rgi remains robust to the tested momentum settings at the $0.7$ B scale.

Panel of Figure 20 — Figure 20: Warmup-step ablation for $0.7$ B models. Panels (a) and (c) plot $r_{\mathrm{diff}}$ for $100$ and $1{,}400$ warmup steps, while Panels (b) and (d) plot $r_{\mathrm{CCC}}$ for the same warmup settings. Approximate rgi remains robust to the tested scheduler warmup lengths at the $0.7$ B scale.

Panel of Figure 21 — Figure 21: Training data-stream ablation for $0.7$ B models. Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, when Adam is trained on FineWeb and Muon is trained on CodeParrot. The weak agreement shows that changing the training data stream strongly disrupts rgi at the $0.7$ B scale.

Panel of Figure 22 — Figure 22: Validation-distribution ablation for $0.7$ B models. Panels (a)–(c) plot $r_{\mathrm{diff}}$ evaluated on C4, ArXiv, and CodeParrot, while Panels (d)–(f) plot $r_{\mathrm{CCC}}$ on the same validation datasets. rgi remains clearest on C4 and weakens on the more distributionally distinct ArXiv and CodeParrot validation sets.

Panel of Figure 23 — Figure 23: Learning-rate ablations for $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ . Panels (a)–(d) plot $r_{\mathrm{diff}}$ for learning-rate multipliers $0.25\eta^{\mathsf{A}}$ , $0.5\eta^{\mathsf{A}}$ , $0.75\eta^{\mathsf{A}}$ , and $1.25\eta^{\mathsf{A}}$ , respectively. Panels (e)–(h) plot $r_{\mathrm{CCC}}$ for them. The results show that approximate rgi holds robustly across the tested learning rates.

Panel of Figure 24 — Figure 24: Momentum ablations for $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ . Panels (a)–(c) plot $r_{\mathrm{diff}}$ for momentumns $\gamma=0.85.0.9,0.95$ , respectively. Panels (d)–(f) plot $r_{\mathrm{CCC}}$ for them. The results show that approximate rgi holds robustly across the tested momentum values.

Panel of Figure 25 — Figure 25: Warmup-step ablation. Panels (a) and (b) plot $r_{\mathrm{diff}}$ for $100$ and $1{,}400$ warmup steps, while Panels (c) and (d) plot $r_{\mathrm{CCC}}$ for the same warmup settings. Approximate rgi holds robustly across the tested scheduler warmup lengths.

Panel of Figure 26 — Figure 26: Weight-decay ablation. Panels (a) and (b) plot $r_{\mathrm{diff}}$ for $0.05$ and $0.1$ weight decay, while Panels (c) and (d) plot $r_{\mathrm{CCC}}$ for the same weight decay settings. Approximate rgi holds robustly across the tested weight decay values.

Panel of Figure 27 — Figure 27: rgi across optimizers, initializations, and model sizes. Panel (a) plots $r_{\mathrm{diff}}$ across optimizers with different initializations. Panels (b) and (c) plot $r_{\mathrm{diff}}$ across models with hidden-state dimensions $768$ , $384$ , and $64$ . rgi continues to hold across different initializations and moderate changes in model size.

Panel of Figure 28 — Figure 28: rgi across optimizers and models with different attention and ffn designs. Panels (a)–(c) plot $r_{\mathrm{diff}}$ between the RoPE fa baseline trained with $\mathsf{Adam}$ and NoPE fa , ALiBi fa , and RoPE swa models, respectively, each trained with all optimizers. Panel (d) plots $r_{\mathrm{diff}}$ between models using GELU and SReLU activations. rgi persists across different attention and ffn designs.

Panel of Figure 29 — Figure 29: rgi across optimizers and models with different ffn designs. Panels (a) and (b) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, between the GELU baseline trained with Adam and SiLU models trained with other optimizers. Panels (c) and (d) plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ , respectively, between the non-gated ffn baseline trained with Adam and gated ffn models trained with other optimizers. rgi persists across different ffn designs.

Panel of Figure 30 — Figure 30: rgi across optimizers trained on different data streams. The panels plot $r_{\mathrm{diff}}$ when Adam is trained on FineWeb and the compared optimizer is trained on C4 (Panel (a)), ArXiv (Panel (b)), or CodeParrot (Panel (c)). The results show that rgi is strongly affected by the training data stream: changing the data stream can substantially weaken, or even eliminate, rgi across optimizers.

Panel of Figure 31 — Figure 31: rgi across optimizers and fa / swa models evaluated on different validation datasets. The panels plot $r_{\mathrm{diff}}$ evaluated on C4, ArXiv, CodeParrot, and Random. These results show that rgi can transfer to related validation distributions, but becomes weaker on more distributionally distinct validation datasets.

Panel of Figure 32 — Figure 32: rgi across optimizers and gated/non-gated models evaluated on different validation datasets. The panels plot $r_{\mathrm{diff}}$ and $r_{\mathrm{CCC}}$ evaluated on C4, ArXiv, CodeParrot, and Random. These results show that rgi can transfer to related validation distributions, but becomes weaker on more distributionally distinct validation datasets.

Panel of Figure 33 — Figure 33: Movement from initialization in parameter and hidden-state space. Panels (a) and (b) plot $r_{\ell_{2}}(\downarrow)$ between initialization and later checkpoints for parameters and hidden states, respectively. The large relative distances show substantial movement during training, providing evidence against an ntk -style lazy-training explanation.

Panel of Figure 34 — Figure 34: Cross-optimizer distances in parameter and hidden-state space. Panels (a) and (b) plot $r_{\ell_{2}}(\downarrow)$ between Adam and the other optimizers at matched training steps for parameters and hidden states, respectively. The large relative distances show that different optimizers need not share similar parameters or hidden states even when their validation-loss differences exhibit rgi .

Panel of Figure 35 — Figure 35: Model-width ablations for $r_{\mathrm{diff}}$ . Panels (a)–(c) plot $r_{\mathrm{diff}}$ for model widths $576$ , $384$ , and $192$ , respectively. rgi holds across all tested widths, supporting the conclusion that the phenomenon is not restricted to the standard model width.

Metric	Adam-mini	Adam	Lion	Muon	Scion	Soap
Average $\widehat{H}(d)$	0.001	0.001	0.001	0.001	0.001	0.001
Average $L(\theta_{10000},d)$	3.500	3.457	3.554	3.409	3.459	3.403
Average $\widehat{H}(d)/L(\theta_{10000},d)$	0.003	0.004	0.003	0.004	0.004	0.004

Appendix

Table 1: Comparison of

H(d)

and

L(\theta_{t},d)

across optimizers. The results show that

H(d)

is negligible relative to

L(\theta_{t},d)

Adam step	Validation loss	Adam-mini	Lion	Muon	Scion	Shampoo	SOAP
0	11.023	1.000	1.000	1.000	1.000	1.000	1.000
1000	4.261	0.773	0.753	0.641	0.612	0.396	0.559
2000	3.815	0.658	0.604	0.475	0.469	0.186	0.400
3000	3.654	0.580	0.520	0.425	0.411	0.150	0.303
4000	3.588	0.526	0.465	0.377	0.383	0.130	0.257

Appendix

Table 2: Parameter cosine similarity across optimizers with aligned validation loss.

Step	Adam-mini	Lion	Muon	SOAP	Scion
0	0.000	0.000	0.000	0.000	0.000
1000	0.321	0.738	-0.275	-0.341	-0.226
2000	0.129	0.276	-0.138	-0.164	-0.085
3000	0.092	0.181	-0.099	-0.117	-0.046
4000	0.073	0.130	-0.089	-0.100	-0.033
5000	0.067	0.107	-0.072	-0.083	-0.018

Appendix

Table 3: Loss difference relative to Adam, where

C_{A,\mathrm{Adam}}(t)=L_{A}(t)-L_{\mathrm{Adam}}(t)

Panel of Figure 36 — Figure 36: Empirical assessment of Assumptions 7.1 and 7.2 . Panel (a) shows that the columns of the targeted linear mappings are nearly orthogonal, while panel (b) shows that the hidden states of different tokens remain nearly orthogonal throughout training across all optimizers. These results provide empirical support for Assumptions 7.1 and 7.2 in practical llm training.

Language Model Pretraining

Language model pretraining teaches models language patterns and knowledge from large text corpora before they are adapted to specific tasks.

See all 422 papers

RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training
Jun 2, 2026Rachit Bansal, Clara Mohri, Tian Qin +2Language Model Pretraining Supervised Fine-Tuning
Harvard University
Generalization Is Stability, Not Accuracy: Multi-Axis Evaluation of LLMs
Oct 1, 2026Nagham Omar, Mahmoud Jabarin, Maya Rozenshtein +3Prompt Sensitivity LLM Evaluation
Technion – Israel Institute of Technology
Why Does Reinforcement Learning Generalize? A Feature-Level Mechanistic Study of Post-Training in Large Language Models
Apr 27, 2026Dan Shi, Zhuowen Han, Simon Ostermann +3Reinforcement Learning RL for Language Model Reasoning
TJUNLP Lab, School of Computer Science and Technology, Tianjin University, China · German Research Center for Artificial Intelligence (DFKI), Saarbrücken, Germany · Saarland University, Saarbrücken, Germany

Jun 2, 2026J/K move · Enter open · S save

RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training

Rachit Bansal, Clara Mohri, Tian Qin, David Alvarez-Melis+1

Harvard University

Language Model Pretraining Supervised Fine-Tuning RL for Language Models Reinforcement Learning RL Post-Training

Open paper AskPDF arXiv

The standard LLM training pipeline applies reinforcement learning (RL) only after pre-training and supervised fine-tuning (SFT). We question this status quo by training a LLM from scratch and applying RL, SFT, and SFT followed by RL directly to intermediate pre-training checkpoints. We find that RL is effective very early, and often matches the full SFT

\to

RL pipeline early as well. Through experiments on harder problems, we find that targeted pre-training data composition is a strong lever for RL effectiveness, even more so than model scale. Beyond reasoning accuracy, applying RL directly to base checkpoints expands the model's distribution; the sharpening effect reported in recent work arises only when RL follows SFT. The general capabilities of the model remain essentially unchanged by RL, while they degrade following SFT. Finally, we merge RL and SFT objectives by parallel averaging, which outperforms across all other training methods discussed, across metrics, while preserving general capabilities. Together, these results suggest that LLM training might benefit from an expanded use of RL.

Abstract

Language Model Pretraining

Explore similar work

RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training

Generalization Is Stability, Not Accuracy: Multi-Axis Evaluation of LLMs

Why Does Reinforcement Learning Generalize? A Feature-Level Mechanistic Study of Post-Training in Large Language Models