Relative Generalization Invariance of LLM Pretraining
Fengzhuo Zhang, Shuche Wang, Shenggui Li, Tianyu Ruan, Jianliang He, Ivor Tsang, Tianyu Pang, Chao Du, Tianwei Zhang, Zhuoran Yang
cs.LG, cs.AI
2026-09-27
On 124M–4.5B GPTs, Adam, Muon and four other optimizers reach rCCC>0.95 after 2,000 steps, a near-constant token-loss shift; FineWeb→ArXiv or code breaks RGI.
Change the optimizer, the attention block, or the pretraining corpus, and validation loss moves. The unresolved question is whether those three knobs move the same structure. Soap, Muon, sliding-window attention, and dynamic data selection have all been used to drive pretraining loss down. What has been missing is a split of the training triplet that shows which layer of generalization each piece actually changes.
Yale, NUS, NTU, and Sea AI Lab make that testable: for any two tokens, is the validation-loss gap the same under two models? If yes, easy and hard tokens keep their order; the whole loss vector has only been shifted by a constant. They call this Relative Generalization Invariance (RGI).
RGI requires L(θ,d)−L(θ,d′) = L(θ′,d)−L(θ′,d′) for every validation pair. Equivalently, token-wise losses differ by a token-independent constant C. A run is a triplet T = (optimizer, init/architecture, data stream). RGI between two triplets at step t must hold almost surely over init and stream draws. If it holds for all large t, they call it asymptotic RGI (ARGI).
Two scores measure how close a pair of runs is. rCCC bins the validation set and takes the concordance correlation of centered bin losses; 1 is a pure shift. rdiff compares pairwise loss gaps; 0 is exact RGI.
The workhorse is a 124M Nano-GPT trained on FineWeb with a same-distribution held-out set. Six optimizers: Adam, Lion, Adam-mini on the coordinate side; Muon, Soap, Scion on the matrix side. Within a setting they share init and data, tune learning rate per optimizer, and use Warmup-Stable-Decay. Main-text figures are 124M; 0.7B and 4.5B are reported to agree. Across 10 seeds they quote the worst rCCC and rdiff.
Ablations twist one axis: learning-rate multipliers 0.25x to 1.25x, momentum 0.85/0.9/0.95, independent inits, hidden size 768 vs 384 vs 64, attention as NoPE, ALiBi, or RoPE SWA, FFN as SReLU or SiLU, streams C4 / ArXiv / CodeParrot, the same shifts at validation time, plus Random tokens drawn i.i.d. uniformly from the vocabulary.
Shared init and FineWeb, the optimizer mostly adds a constant. At 10,000 steps, on 20,480 validation batches, Muon versus Adam fits y=x-0.07 with R²=0.99. After 2,000 steps every other optimizer versus Adam has rCCC>0.95 and rdiff<0.05.
| Comparison | Metric | Result |
| Muon vs Adam, 10,000 steps, 20,480 batches | val-loss fit | y=x-0.07, R²=0.99 |
| Five optimizers vs Adam, ≥2,000 steps | rCCC / rdiff | >0.95 / <0.05 |
| Hidden size 768 vs 384 | RGI | holds |
| Hidden size 768 vs 64 | RGI | fails |
| FineWeb vs ArXiv or CodeParrot | rCCC | stays below 0.95 |
Independent inits start off RGI and return around 2,000 steps. Adam with RoPE full attention versus NoPE, ALiBi, or SWA, and GELU versus SReLU, still clear rCCC=0.95. The learning-rate and momentum grid does the same; it changes when the threshold is crossed, not whether. 1.5x learning rate is dropped because some optimizers diverge.
Training-loss curves look alike after 2,000 steps. The step-to-step change splits into an update term α and a batch-gap term β; |β| is about 50× |α|. β is the loss gap of the current parameters on two unseen batches, so approximate RGI pins β across optimizers and the curves share a shape.
The data stream can break this. Adam on FineWeb versus another optimizer on C4: rCCC drops, rises, then drops again; the middle rise may reflect shared Common Crawl ancestry. On ArXiv and CodeParrot, rCCC stays below 0.95 even when the second run also uses Adam. A mild validation shift to C4 still transfers RGI; ArXiv and code are weaker but still climb with training; Random tokens never show it.
NTK lazy learning does not cover the pattern. After 5,000 steps, parameters and hidden states have moved far from initialization, and cross-optimizer cosine similarity keeps falling. Mean-field is not enough either. At width 192, one quarter of the 768 baseline, RGI across optimizers at that width still holds, yet the validation-loss versus learning-rate curve no longer matches the 384/576/768 family. Comparing width 768 to 64, RGI fails. The edge sits between a moderate shrink and cutting the model too thin.
A quadratic case study fills the theory gap. For f=ABx, with an orthogonal equal-magnitude target and isotropic covariates, GD, momentum-free Adam (entrywise sign), and momentum-free Muon (spectral normalization) satisfy exact RGI on equal-norm inputs, while |cos-sim(GD,Adam)|=O(1/t) and |cos-sim(Adam,Muon)|=O(1/t²).
For people who pretrain, optimizer choice and moderate architecture tweaks look like a uniform offset on every token. Hard tokens stay hard. Change the corpus source and the relative structure can move.
If the goal is to change which tokens are learned well, move the data stream. Switching to Muon is a poor bet for reordering difficulty.
The observation that different optimizers produce similar loss curves now has a mechanism: shape is dominated by β, and RGI freezes β. The paper does not rank absolute loss, and it reports no downstream tasks. Before picking an optimizer from a loss curve, ask whether the axis you changed was allowed to rearrange relative generalization.
The authors restrict the analysis to autoregressive LLMs and leave diffusion models open. The quadratic assumptions are narrow, and the Adam/Muon updates drop momentum. Width 64 already breaks RGI, so moderate architecture change has a hard edge they do not map onto production-scale models.
All headline numbers are validation cross-entropy. There is no MMLU, code, or long-context score, so invariant relative loss is not invariant capability. The C4 rCCC rise-then-fall is explained only as distributional kinship, with no usable distance threshold. 4.5B appears in the appendix with the same qualitative claim; almost every number in the main text is 124M. Failure on Random validation tokens means RGI still needs semantic structure on the eval set.