Relative Generalization Invariance of LLM Pretraining

1 Introduction

• RGI is the invariance of the validation-loss difference between any two tokens across models, equivalent to a token-independent constant gap in validation loss.The paper extends the definition to training triplets through the models they produce.Source: §1 · p. 2
• RGI holds across a wide range of optimizers and moderate architectural variations, but can break when models are trained on data streams from different sources.Optimizer and architecture choices induce approximately uniform shifts in token-wise losses, whereas data-stream changes can alter relative loss structure.Source: §1 · p. 2
• RGI cannot be explained by NTK or MF regimes alone; it persists when models depart from initialization and when model width is reduced.The findings challenge an NTK-style lazy-learning explanation and provide evidence beyond the MF regime.Source: Figure 1 · p. 1, §1 · p. 2
• A stylized quadratic model exhibits RGI across GD, Adam, and Muon, different initializations, and learning rates despite low parameter similarity.The result identifies a setting for RGI beyond NTK and MF regimes.Source: §1 · p. 2
• The paper uses RGI to distinguish the effects of optimizers, architectures, and data through its definition, empirical demonstration, and theoretical characterization.Pretraining performance is shaped jointly by these three components.Source: §1 · p. 2

2 Related Work

• Earlier optimizer research proposes new optimization methods or improves their performance, while this paper studies shared relative-generalization behavior across existing optimizers.The cited optimizer literature includes coordinate-wise methods such as Adam and matrix-based methods such as Muon.Source: §2 · p. 3
• Research on pretraining factors examines architecture, optimization, and data mixtures through scaling laws, performance, efficiency, and data-mixture construction.This paper instead studies the distinct roles of optimizers and data in shaping relative generalization.Source: §2 · p. 3
• NTK analysis studies training with an effectively fixed kernel, while MF analysis represents neurons collectively through a distribution and allows an evolving kernel.The paper’s relative-generalization behavior is not explained by either perspective.Source: §2 · pp. 3–4

3 Preliminaries

• A training triplet consists of an optimizer, an initialization specifying its model architecture, and a data stream sampled from the corpus.The data stream may use a time-dependent distribution conditioned on prior batches and independent of the initialization.Source: §3 · p. 4
• The training map generates parameters from a triplet and time step, with the map induced by the optimizer’s update rule.The setup defines average cross-entropy loss on a batch and its corresponding gradient.Source: §3 · p. 4
• Coordinate-wise optimizers such as Adam use entry-wise gradient statistics, whereas matrix-based optimizers such as Muon exploit the matrix structure of parameter blocks.Adam maintains first- and second-moment states; the idealized Muon update uses an SVD-based transformation.Source: §3 · p. 4
• The analysis focuses on deterministic optimizers; randomized optimizers can be incorporated by treating algorithmic randomness as an additional input to the training map.The passage also points readers to further discussion of the optimizers.Source: §3 · p. 4

4 Relative Generalization Invariance

• RGI means that the validation-loss difference between any two tokens is the same across models, equivalently that their token-wise losses differ by a token-independent constant.For training triplets, this property must hold almost surely over the joint randomness in initialization and data streams; asymptotic RGI holds from some time onward.Source: §4 · p. 5
• rCCC measures agreement between centered validation losses across subsets, with rCCC = 1 indicating that the losses differ only by a constant shift.The subsets partition the validation set.Source: §4 · p. 5
• rdiff measures normalized discrepancy between corresponding validation-loss differences; with single-token subsets, rdiff = 0 implies exact RGI.Smaller rdiff indicates closer alignment with RGI.Source: §4 · p. 5
• After 2,000 steps, all five compared optimizers exceed rCCC(Adam, ·) = 0.95 and fall below rdiff(Adam, ·) = 0.05.The comparison uses models trained from the same initialization and on the same training data stream.Source: Figure 2 · p. 6

5.1 RGI Holds Across Optimizers

• With Adam as baseline, rdiff < 0.05 and rCCC > 0.95 show approximate RGI across optimizers after 2,000 steps; tested learning rates and momentum values preserve it eventually.Hyperparameters can change when RGI emerges, while the tested learning-rate, momentum, scheduler, and weight-decay variations leave the overall finding intact.Source: §5.1 · p. 7, Figure 3 · p. 7

5.2 RGI Holds Across Moderately Different Architectures

• RGI persists across attention and FFN designs and moderate model-size changes, but fails at hidden-state dimension 64.It also emerges after about 2,000 steps across independent initializations; changing architecture or optimizer can affect its emergence time.Source: §5.2 · pp. 7–8, Figure 4 · p. 8, Figure 5 · p. 8
• RGI helps explain shared training-loss curve shapes across optimizers and architectures: the batch-to-batch term is about 50 times larger than the update-related term.The dominant term describes loss differences for the same parameters across adjacent batches.Source: §5.2 · pp. 8–9, Figure 6 · p. 9

5.3 RGI is Strongly Influenced by Pretraining Data Streams

• Across FineWeb, ArXiv, and CodeParrot training streams, changing the data source can weaken or prevent RGI, even when the optimizer remains Adam.For ArXiv and CodeParrot, rCCC stays below 0.95 throughout training; C4 shows a different trajectory.Source: §5.3 · pp. 9–10, Figure 7 · p. 9

5.4 RGI Persists under Moderate Validation Shifts

• RGI can transfer to related validation distributions, but weakens on ArXiv and CodeParrot and shows no clear emergence on Random.The experiments compare validation rCCC on C4, ArXiv, CodeParrot, and an extreme OOD Random dataset.Source: §5.4 · p. 10, Figure 8 · p. 10

6.1 RGI Emerges Beyond the NTK Regime

• The NTK test uses cosine similarity and relative ℓ2 comparisons of parameters and hidden states, both within an optimizer over training and across optimizers at the same step.The main text reports cosine similarities; relative-ℓ2 results are deferred to Appendix B.Source: §6.1 · pp. 10–11
• After 5,000 steps, parameters and hidden states deviate substantially from initialization, while cross-optimizer similarities decrease, so RGI is not explained by NTK-style lazy learning or shared representations.The figure compares within-optimizer similarities to initialization and cross-optimizer similarities at the same training step.Source: §6.1 · p. 11, Figure 9 · p. 11

6.2 RGI Emerges Beyond the MF Regime

• RGI persists at width 192, one quarter of the baseline width, although its validation-loss–learning-rate profile differs qualitatively from wider models, suggesting departure from the same MF-style approximation.The width-192 profile is compared with widths 768, 576, and 384.Source: Figure 10 · p. 11, §6.2 · p. 12

7 Case Study of Over-Parameterized Quadratic Models

• The model is a two-layer overparameterized factorization fA,B(x) = ABx trained to fit the linear target y(x) = Hx with squared loss.The setup requires q ≥ 3r and s ≥ 3r.Source: §7 · p. 12
• Assumption 7.1 requires H to have orthogonal equal-norm columns and equal-magnitude nonzero entries, conditions met by one-hot or Hadamard bases.Source: §7 · p. 12
• Assumption 7.2 requires training covariates to cover the feature space approximately isotropically.The authors call these assumptions a stylized setting and leave relaxing them for future work.Source: §7 · p. 12
• The analysis compares gradient descent, Adam’s coordinate-wise normalized gradient updates, and Muon’s spectrally normalized gradient updates.Source: §7 · p. 12
• RGI holds across GD, Adam, Muon, and different initializations under Assumptions 7.1 and 7.2, even as parameter cosine similarity decreases during training.The theorem also permits a broad range of learning rates, with upper-bound constraints for Adam and Muon.Source: §7 · p. 13

8 Conclusion

• RGI persists under joint optimizer and architecture variations, while changing the training data stream can substantially alter it.The paper also reports that neither NTK nor MF alone explains RGI, and identifies a quadratic setting where it emerges robustly.Source: §8 · p. 13

A.1 Concordance Correlation Coefficient

• The appendix presents the detailed definition of concordance correlation coefficient (CCC) among centered losses.Source: §A.1 · p. 20

A.2 Definition of rdiff

• The appendix presents the definition of rdiff and specifies that its expectation is over θ0 and D1:t.Source: §A.2 · p. 20

A.3 Experimental Setups

• Experiments use GPT-2-tokenized sequences of length 1024, batch size 524,288, and 124M and 0.7B GPT-style models.Main-text validation uses batches of 1,024 across 20,480 batches, approximately 21M validation tokens; single-token tests use 5M tokens.Source: §A.3 · p. 20

B.1 Additional Results of Validation Losses Between Optimizers

• At step 10,000, validation losses across optimizers are highly linearly correlated, supporting approximate RGI.Figure 11 reports near-perfect fits with R2 ≥ 0.99 for Adam versus Adam-mini, Lion, Soap, and Scion.Source: Figure 11 · p. 21, §B.1 · p. 21

B.2 Additional Results of rdiff and rCCC with Other Optimizers as Baselines

• Changing the evaluation-baseline optimizer does not qualitatively affect approximate RGI, consistent with transitivity up to accumulated error.Source: §B.2 · p. 21, Figure 12 · p. 22, Figure 13 · p. 22

B.3 Per-Token Results of Relative Generalization Invariance

• RGI among optimizers also holds when evaluated token by token, with each validation batch containing a single token.Source: §B.3 · p. 23, Figure 14 · p. 23

B.4 Results for 0.7B and 4.5B Models

• Approximate RGI among Adam and Muon also holds at 0.7B and 4.5B scales when initial parameters and training data are shared.Because of computational cost, the larger-scale comparisons focus on Adam and Muon.Source: §B.4 · p. 23, Figure 15 · p. 23, Figure 16 · p. 24
• The 0.7B ablations report approximate RGI across different initializations and tested learning rates, momentum settings, and warmup lengths.Changing the training stream from FineWeb to CodeParrot weakens agreement, and RGI is clearest on C4 but weaker on more distinct validation sets.Source: Figure 17 · p. 24, Figure 18 · p. 25, Figure 19 · p. 25

B.5.1 Additional Results for Section 5.1

• Approximate RGI remains consistently observed across tested warmup lengths and weight-decay values.The tested warmup lengths are 100 and 1,400 steps; weight decay is 0.05 and 0.1.Source: §B.5.1 · pp. 28–29

B.5.2 Additional Results for Section 5.2

• RGI persists across different initializations, tested model sizes, and attention and FFN designs.The reported designs include NoPE and ALiBi full attention, RoPE sliding-window attention, GELU and SReLU activations, and gated and non-gated FFNs.Source: Figure 27 · p. 30, Figure 28 · p. 30, Figure 29 · p. 30

B.5.3 Additional Results for Section 5.3

• Changing the training data stream can substantially weaken or eliminate RGI across optimizers.Figure 30 compares Adam trained on FineWeb with optimizers trained on C4, ArXiv, or CodeParrot.Source: Figure 30 · p. 31

B.5.4 Additional Results for Section 5.4

• RGI can transfer to related validation distributions but becomes weaker on more distributionally distinct datasets.The reported evaluations include C4, ArXiv, CodeParrot, and Random.Source: Figure 31 · p. 31, Figure 32 · p. 32

B.6.1 Additional Results for Section 6.1

• Large relative distances from initialization and across optimizers in parameter and hidden-state spaces argue against an NTK-style lazy-training explanation.Source: Figure 33 · p. 32, Figure 34 · p. 33

B.6.2 Additional Results for Section 6.2

• RGI holds across all tested model widths, so the phenomenon is not restricted to the standard width.Figure 35 tests widths 576, 384, and 192.Source: Figure 35 · p. 33

C Discussions of Alternative Explanations for RGI

• About 0.4%: the model-independent component H(d) is negligible relative to the smallest observed loss, falsifying the intrinsic data-difficulty hypothesis.The paper estimates an upper bound on H(d) by optimizing model parameters for each data example.Source: C Discussions of Alternative… · p. 34, Table 1 · p. 34
• Similarity between models aligned to Adam’s validation loss decreases as training progresses, providing evidence against the different convergence-speed hypothesis.The comparison is limited to validation-loss levels of at least 3.55 because some optimizers do not reach lower losses.Source: C Discussions of Alternative… · p. 34, Table 2 · p. 34
• The loss difference between each optimizer and Adam at step 2,000, when RGI emerges, differs substantially from its value at convergence at step 10,000.This provides evidence against the explanation that optimizers produce the same learning trajectory at different speeds.Source: C Discussions of Alternative… · p. 34, Table 3 · p. 35

D Empirical Verification of Assumptions 7.1 and 7.2

• The empirical check estimates each layer’s target transformation as the best linear mapping between adjacent-layer hidden states after training.An average angle close to 90° between mapping columns supports Assumption 7.1; Assumption 7.2 concerns isotropy of layer inputs within their spanned subspace.Source: D Empirical Verification of… · pp. 35–36
• Figure 36 shows that targeted mapping columns are nearly orthogonal and hidden states of different tokens remain nearly orthogonal throughout training across optimizers.The average mapping-column angle is close to 90° at training’s end, while hidden-state angles remain close to 90° along training.Source: Figure 36 · p. 35, D Empirical Verification of… · pp. 35–36

E Proof of Proposition 4.2

• The proof establishes both directions of the equivalence between RGI and a constant generalization gap.For deterministic optimizers, the gap is expressed as a deterministic function of initialization and the training data streams.Source: E Proof of Proposition 4.2 · p. 36

Step 1: Derive closed-form expressions for the optimizer trajectories

• For GD, Adam, and Muon, the iterates remain in specified low-dimensional invariant subspaces for every t ≥ 0.This holds for positive step sizes under the specified tie-breaking rules and the stated initialization.Source: F Proof of Theorem 7.3 · p. 37

Step 2: Use these expressions to analyze the resulting losses

• On the validation sphere, GD, Adam, and Muon have exactly equal loss differences between any two points.The loss may vary by optimizer and time, but is constant across points on the sphere.Source: F Proof of Theorem 7.3 · p. 39

Step 3: Characterize the cosine similarity between the learned weights

• Proposition F.2 characterizes the cosine similarities between parameters learned by GD, Adam, and Muon under Proposition F.1’s conditions.A corresponding similarity is defined as zero when its denominator is zero.Source: F Proof of Theorem 7.3 · p. 40

Further trajectory bounds

• For Adam, the proof establishes that the positive quadrant is invariant and that a_t remains O(1) as t increases.The appendix also states that a_t and b_t are uniformly bounded while λ_t grows linearly for Adam and Muon, but this part provides only partial supporting results for those claims.Source: F Proof of Theorem 7.3 · pp. 41–42

• The proof uses induction to show that bMuon_t remains O(1) as t increases.Source: F Proof of Theorem 7.3 · p. 42
• Using boundedness, it derives parameter cosine-similarity bounds for GD and Adam, then applies an analogous argument for Muon.The stated inequalities use Cauchy–Schwarz and boundedness of aAdam_t; the proof then concludes Theorem 7.3.Source: F Proof of Theorem 7.3 · p. 42