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
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.
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 rCCC(Adam,⋅)(↑) for the other five optimizers over training, all of which rise above 0.95 after 2,000 steps. Panel (c) reports rdiff(Adam,⋅)(↓) 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.
Figure 3: rgi across optimizers under different hyperparameter settings. Panels (a) and (b) plot rCCC under different learning-rate multipliers, while panels (c) and (d) plot rCCC under different momentum values. rgi persists across all tested learning-rate and momentum settings.
Figure 4: rgi across optimizers, initializations, and model sizes. Panel (a) plots rCCC across optimizers with different initializations. Panels (b) and (c) plot rCCC across models with hidden-state dimensions 768 , 384 , and 64 . rgi continues to hold across different initializations and moderate changes in model size.
Figure 5: rgi across optimizers and models with different attention and ffn modules. Panels (a)–(c) plot rCCC between models using RoPE fa trained with Adam and models using NoPE fa , ALiBi fa , and RoPE swa , each trained with all optimizers, respectively. Panel (d) plots rCCC between models using GELU and SReLU activations. rgi continues to hold across different attention and ffn designs.
Figure 6: Values of ∣αtT∣ and ∣βtT∣ . It shows that ∣βtT∣ is much larger than ∣αtT∣
Figure 7: rgi across optimizers trained on different data streams. The panels plot rCCC 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.
Figure 8: rgi across optimizers and architectures evaluated on different validation datasets. The panels plot rCCC 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.
Figure 9: Within- and cross-optimizer cos-sim(↑) for parameters and hidden states. Panels (a)–(b) compare θ0T with θtT , while Panels (c)–(d) compare θtT with θtT′ . Hidden-state similarities are averaged over layers l∈[L] . Within-optimizer comparisons show that all optimizers learn nontrivial features throughout training, whereas cross-optimizer comparisons show that different optimizers learn distinct features.
Figure 10: rgi across models of different widths. Panels (a)–(c) plot rCCC 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
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 R2≥0.99 , support approximate rgi across optimizers.
Figure 12: Baseline-optimizer ablation for rdiff . Panels (a)–(e) plot rdiff 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.
Figure 13: Baseline-optimizer ablation for rCCC . Panels (a)–(e) plot rCCC 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.
Figure 14: Per-token evaluation of rgi . Panels (a) and (b) plot rdiff and rCCC , respectively, when each validation batch contains a single token, i.e., ∣Dval,k∣=1 . The small rdiff and large rCCC values show that rgi also holds at token-level granularity.
Figure 15: Baseline rgi for 0.7 B models. Panels (a) and (b) plot rdiff and rCCC , 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.
Figure 16: Baseline rgi for 4.5 B models. Panels (a) and (b) plot rdiff and rCCC , 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.
Figure 17: Initialization ablation for 0.7 B models. Panels (a) and (b) plot rdiff and rCCC , 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.
Figure 18: Learning-rate ablation for 0.7 B models. Panels (a) and (b) plot rdiff for 0.25ηA and 1.25ηA , while Panels (c) and (d) plot rCCC for the same two learning-rate multipliers. Approximate rgi remains robust to these learning-rate changes at the 0.7 B scale.
Figure 19: Momentum ablation for 0.7 B models. Panels (a)–(c) plot rdiff for γ=0.85 , 0.9 , and 0.95 , while Panels (d)–(f) plot rCCC for the same momentum values. Approximate rgi remains robust to the tested momentum settings at the 0.7 B scale.
Figure 20: Warmup-step ablation for 0.7 B models. Panels (a) and (c) plot rdiff for 100 and 1,400 warmup steps, while Panels (b) and (d) plot rCCC for the same warmup settings. Approximate rgi remains robust to the tested scheduler warmup lengths at the 0.7 B scale.
Figure 21: Training data-stream ablation for 0.7 B models. Panels (a) and (b) plot rdiff and rCCC , 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.
Figure 22: Validation-distribution ablation for 0.7 B models. Panels (a)–(c) plot rdiff evaluated on C4, ArXiv, and CodeParrot, while Panels (d)–(f) plot rCCC on the same validation datasets. rgi remains clearest on C4 and weakens on the more distributionally distinct ArXiv and CodeParrot validation sets.
Figure 23: Learning-rate ablations for rdiff and rCCC . Panels (a)–(d) plot rdiff for learning-rate multipliers 0.25ηA , 0.5ηA , 0.75ηA , and 1.25ηA , respectively. Panels (e)–(h) plot rCCC for them. The results show that approximate rgi holds robustly across the tested learning rates.
Figure 24: Momentum ablations for rdiff and rCCC . Panels (a)–(c) plot rdiff for momentumns γ=0.85.0.9,0.95 , respectively. Panels (d)–(f) plot rCCC for them. The results show that approximate rgi holds robustly across the tested momentum values.
Figure 25: Warmup-step ablation. Panels (a) and (b) plot rdiff for 100 and 1,400 warmup steps, while Panels (c) and (d) plot rCCC for the same warmup settings. Approximate rgi holds robustly across the tested scheduler warmup lengths.
Figure 26: Weight-decay ablation. Panels (a) and (b) plot rdiff for 0.05 and 0.1 weight decay, while Panels (c) and (d) plot rCCC for the same weight decay settings. Approximate rgi holds robustly across the tested weight decay values.
Figure 27: rgi across optimizers, initializations, and model sizes. Panel (a) plots rdiff across optimizers with different initializations. Panels (b) and (c) plot rdiff across models with hidden-state dimensions 768 , 384 , and 64 . rgi continues to hold across different initializations and moderate changes in model size.
Figure 28: rgi across optimizers and models with different attention and ffn designs. Panels (a)–(c) plot rdiff between the RoPE fa baseline trained with Adam and NoPE fa , ALiBi fa , and RoPE swa models, respectively, each trained with all optimizers. Panel (d) plots rdiff between models using GELU and SReLU activations. rgi persists across different attention and ffn designs.
Figure 29: rgi across optimizers and models with different ffn designs. Panels (a) and (b) plot rdiff and rCCC , respectively, between the GELU baseline trained with Adam and SiLU models trained with other optimizers. Panels (c) and (d) plot rdiff and rCCC , respectively, between the non-gated ffn baseline trained with Adam and gated ffn models trained with other optimizers. rgi persists across different ffn designs.
Figure 30: rgi across optimizers trained on different data streams. The panels plot rdiff 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.
Figure 31: rgi across optimizers and fa / swa models evaluated on different validation datasets. The panels plot rdiff 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.
Figure 32: rgi across optimizers and gated/non-gated models evaluated on different validation datasets. The panels plot rdiff and rCCC 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.
Figure 33: Movement from initialization in parameter and hidden-state space. Panels (a) and (b) plot rℓ2(↓) 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.
Figure 34: Cross-optimizer distances in parameter and hidden-state space. Panels (a) and (b) plot rℓ2(↓) 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 .
Figure 35: Model-width ablations for rdiff . Panels (a)–(c) plot rdiff 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 H(d)
0.001
0.001
0.001
0.001
0.001
0.001
Average L(θ10000,d)
3.500
3.457
3.554
3.409
3.459
3.403
Average H(d)/L(θ10000,d)
0.003
0.004
0.003
0.004
0.004
0.004
Appendix
Table 1: Comparison of H(d) and L(θt,d) across optimizers. The results show that H(d) is negligible relative to L(θ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 CA,Adam(t)=LA(t)−LAdam(t) .
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.
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→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.
Generalization in large language models (LLMs) is the ability to produce consistent and semantically stable outputs when the same input is expressed in different ways. Existing work typically evaluates generalization through aggregate accuracy on a single prompt format, task, or set of variations, which conflates robustness with overall benchmark performance. In this work, we show generalization evaluation at the level of individual examples, across multiple input variants, and across different aspects of model behavior, focusing on variability rather than reducing performance to a score that can be improved through narrow training or other ways that obfuscate generalization evaluation. Following this view, we introduce the Stability-Aware Generalization Objective (SAGO), a framework that measures how much model behavior changes for the same input under different variations and benchmarks, capturing variability across several dimensions including generation consistency, internal activations, confidence, and response mirroring. We show that many commonly used models exhibit statistically significant and consistent generalization instability: no model generalizes uniformly, behavioral axes capture independent failure modes, and cross-dataset variation can reverse model rankings.
Reinforcement learning (RL)-based post-training often improves the reasoning performance of large language models (LLMs) beyond the training domain, while supervised fine-tuning (SFT) frequently leads to general capabilities forgetting. However, the mechanisms underlying this contrast remain unclear. To bridge this gap, we present a feature-level mechanistic analysis methodology to probe RL generalization using a controlled experimental setup, where RL- and SFT-tuned models are trained from the same base model on identical data. Leveraging our interpretability framework, we align internal activations across models within a shared feature space and analyze how features evolve during post-training. We find that SFT rapidly introduces many highly specialized features that stabilize early in training, whereas RL induces more restrained and continually evolving feature changes that largely preserve base models' representations. Focusing on samples where RL succeeds but the base model fails, we identify a compact, task-agnostic set of features that directly mediate generalization across diverse tasks. Feature-level interventions confirm their causal role: disabling these features significantly degrades RL models' generalization performance, while amplifying them improves base models' performance. The code is available at https://github.com/danshi777/RL-generalization.
Dan Shi, Zhuowen Han, Simon Ostermann +3
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