A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

# A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio
Source: [https://arxiv.org/html/2605.28975](https://arxiv.org/html/2605.28975)
Ali Shehper ali\.shehper1@gmail\.com &Ashish Vaswani11footnotemark:1 ashish\.vaswani@gmail\.com

###### Abstract

We study the log\-alignment ratio \(LAR\), a measure of parameter–activation alignment, introduced in parameterization theory\. We reformulate it as the overlap between a weight spectrumppof the normalized squared singular values of a matrix and an activation spectrumqqof the normalized squared projections of inputs onto its singular directions\. We show that unembedding LAR tracks the transition between memorization and generalization in two different settings by capturing the spread ofppandqqduring training\. In grokking, LAR predicts the effective dimension of the learned function:k≈n2​\(1−LAR\)k\\approx n^\{2\(1\-\\mathrm\{LAR\}\)\}, wherennis the input dimension of the matrix\. In 3B\-parameter language model pre\-training, its deviation from a non\-overfitting baseline tracks the generalization gap, and its rate of decline increases as overfitting approaches\. LAR is computable from quantities available during the forward pass with negligible computational overhead, and requires no held\-out validation data\.

## 1Introduction

Large\-scale pre\-training has driven the most visible recent advances in neural network capability\(Kaplanet al\.,[2020](https://arxiv.org/html/2605.28975#bib.bib28); Hoffmannet al\.,[2022](https://arxiv.org/html/2605.28975#bib.bib29)\), and as model and dataset sizes continue to grow, the cost of every training run grows with them\. This has made the study of training dynamics, and the design of cheap diagnostic tools for monitoring them, unusually valuable: a metric that can flag problems such as overfitting, instability, or saturation from training\-time quantities alone, without requiring expensive held\-out evaluation, can save compute and time for practitioners\.

A complementary line of work on neural network*parameterizations*, which prescribe how initialization, parameter multipliers, and learning rates should scale with width and depth, offers a principled approach to transferring hyperparameters from small to large models\.Yanget al\.\([2022](https://arxiv.org/html/2605.28975#bib.bib17)\)proposed the maximal update parametrization \(μ​P\\mu P\), which enables hyperparameter transfer under the assumption that updates at the readout layer are fully aligned with the incoming activations\.Everettet al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib2)\)showed that this alignment assumption is stronger than necessary and, by relaxing it, obtained a standard\-parameterization recipe that consistently outperformedμ​P\\mu Pacross tens of thousands of experiments at scales up to 26\.8B parameters\. As part of this analysis they introduced the*log\-alignment ratio*\(LAR\) as an empirical measure of how strongly a layer’s parameter matrix aligns with its input activations during training\.

WhileEverettet al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib2)\)used LAR as a proxy for parameter\-activation alignment, we discover in this paper that the metric carries a second, previously\-unexplored meaning: namely, it measures how diffuse or concentrated the layer’s activations and weight spectra are\. More precisely, we show that LAR measures the overlap between a weight spectrumppand an activation spectrumqq: it is large when both spectra are concentrated on the same small set of directions, and small when either is diffuse or when the two are misaligned\.

This insight makes LAR a natural candidate to measure*generalization*— the extent to which a model generalizes beyond its training data rather than overfitting to it\. Empirically, we observe that well\-generalizing networks tend to concentrate their computation on a small number of directions, while overfitting networks tend to spread it across many\. Most previously\-studied metrics for this behavior cover either the weight spectrum or the activation spectrum, but not both, and are computationally expensive as they require full SVD of the weight matrix or a PCA of the activations\. In contrast, LAR looks at both weights and activations and is computable from quantities already available in the model’s forward pass, with negligible computational overhead\.

We study the relationship between generalization and the LAR of a model’s unembedding matrix in two different settings\. First, on small algorithmic tasks exhibiting grokking\(Poweret al\.,[2022](https://arxiv.org/html/2605.28975#bib.bib1)\), we find that optimization concentratesppandqqon a shared set of directions, which causes LAR to grow\. As the numberkkof directions becomes small, grokking may occur, depending on factors such as task difficulty and training dataset size\. LAR predicts the dimensionality of the learned function,k≈n2​\(1−LAR\)k\\approx n^\{2\(1\-\\mathrm\{LAR\}\)\}, which we compare againstk95k\_\{95\}— the number of principal components explaining 95% of the unembedding input activation variance\. We find close agreement for final checkpoints across all experiments, including 12 grokking tasks and a sweep over training\-set fractions\.

Second, in 3B\-parameter language model pre\-training, we find that unembedding LAR decreases slowly and stabilizes when the model generalizes, but declines sharply when overfitting is approaching\. The deviation of LAR from a non\-overfitting baseline also closely tracks the generalization gap\. The accelerating rate of decline in LAR prior to overfitting could enable evaluation\-free methods to train generalizing models — a possibility we leave open for future investigation\.

Across both the grokking and large\-scale settings, we discover a consistent picture: lower LAR corresponds to more diffuse weight and activation distributions, which coincides with memorization, while higher LAR corresponds to more concentrated distributions and generalization\. In each case, we observe that the correlation betweenppandqqremains high throughout training, so LAR effectively tracks how far the weight and activation distributions are from uniform\.

#### Contributions\.

Our main contributions are as follows\.

- •Distributional reformulation of LAR\.We show that LAR can be written as the log\-overlap∑ipi​qi\\sum\_\{i\}p\_\{i\}q\_\{i\}between a weight spectrumppand an activation spectrumqq, together with equivalent covariance and correlation forms that make the metric’s behavior interpretable in terms of elementary statistical quantities\.
- •Effective dimension via LAR\.We show that whenWWconcentrates its energy onkksingular directions and activations lie along the same directions,LAR=1−12​logn⁡k\\mathrm\{LAR\}=1\-\\tfrac\{1\}\{2\}\\log\_\{n\}k\. We verify this empirically in the grokking setting by comparing the LAR\-predicted dimension tok95k\_\{95\}on 12 algorithmic tasks\.
- •LAR as an overfitting diagnostic at scale\.On 3B\-parameter Gemma\-2 models trained with Adam and Muon, we show that unembedding LAR stabilizes in the non\-overfitting regime and declines sharply under overfitting, and that its deviation from a non\-overfitting baseline closely tracks the generalization gap\.
- •LAR as a proxy for spread inppandqq\.We show that in practice, optimization drivesppandqqto be highly correlated, makingLAR\\mathrm\{LAR\}a proxy for the variance ofppandqq\. In pre\-training experiments, the top singular value of the unembedding matrix is much larger than the rest and is the primary driver of the spread inpp\. This makes the contribution from the top singular direction,L1=1\+12​logn⁡\(p1​q1\)L\_\{1\}=1\+\\tfrac\{1\}\{2\}\\log\_\{n\}\(p\_\{1\}q\_\{1\}\), an alternative proxy for the pre\-overfit signature of overfitting\.

## 2Related Work

#### Spectral measures of generalization\.

Several lines of work have connected spectral properties of weight matrices to generalization\.Bartlettet al\.\([2017](https://arxiv.org/html/2605.28975#bib.bib10)\)andNeyshaburet al\.\([2018](https://arxiv.org/html/2605.28975#bib.bib11)\)derived margin\-based generalization bounds in terms of products of spectral norms, andAroraet al\.\([2018](https://arxiv.org/html/2605.28975#bib.bib13)\)gave tighter compression\-based bounds using noise stability of trained networks\.Jianget al\.\([2020](https://arxiv.org/html/2605.28975#bib.bib12)\)later compared dozens of such measures in a large\-scale empirical study\. Empirically,Martin and Mahoney \([2021](https://arxiv.org/html/2605.28975#bib.bib8)\)propose that the spectral density of weight matrices develops heavy tails as training progresses, and use this as a data\-free signal for generalization and early stopping\.Yuniset al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib15)\)show that the normalized effective rank of weight matrices decreases broadly during training, and that the onset of low\-rank weights coincides with the transition from memorization to generalization in grokking\.

In contrast, LAR depends on both weights and activations and is computable from matrix norms alone\. It admits an explicit effective\-dimension formula verified in grokking experiments, and tracks the generalization gap at 3B scale\. Furthermore, effective rank only weakly distinguishes overfitting from generalization under Adam in 3B pre\-training experiments \([Section˜D\.1](https://arxiv.org/html/2605.28975#A4.SS1)\), while unembedding LAR separates them cleanly under both Adam and Muon\.

#### Neural network parameterizations\.

A complementary line of work studies neural network*parameterizations*as a principled approach to transferring hyperparameters from small to large models\. The maximal update parameterization \(μ\\muP\)\(Yang and Hu,[2021](https://arxiv.org/html/2605.28975#bib.bib18); Yanget al\.,[2022](https://arxiv.org/html/2605.28975#bib.bib17)\)enables width transfer for learning rate and related hyperparameters\. Subsequent work has extended this framework along new scaling axes:Deyet al\.\([2025](https://arxiv.org/html/2605.28975#bib.bib19)\)introduce CompleteP, which achieves hyperparameter transfer across depth while also ensuring non\-lazy learning in every layer, andMlodozeniecet al\.\([2025](https://arxiv.org/html/2605.28975#bib.bib20)\)further extend the framework to cover per\-module hyperparameters, batch size, and token horizon\.Everettet al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib2)\)revisited the alignment assumptions underlying the maximal update parameterization and showed that relaxing them yields per\-layer learning rate prescriptions that outperformμ\\muP at scales up to 26\.8B parameters\. LAR was introduced inEverettet al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib2)\)as an empirical measure of parameter–activation alignment as part of this analysis\. Our work repurposes LAR as a diagnostic for generalization: the same quantity that probes whether alignment assumptions hold during training also tracks overfitting dynamics, independent of its uses in parameterization theory\.

#### Grokking\.

Grokking, first identified byPoweret al\.\([2022](https://arxiv.org/html/2605.28975#bib.bib1)\), refers to the delayed transition from memorization to generalization observed long after training accuracy has saturated\. Several lines of work have offered explanations\.Baraket al\.\([2022](https://arxiv.org/html/2605.28975#bib.bib34)\)study hidden progress measures, i\.e\. scalar functions of the training state predictive of convergence time, and use sparse parity learning as a testbed in which grokking\-style phase transitions arise and are consistent with their hidden\-progress framework\.Nandaet al\.\([2023](https://arxiv.org/html/2605.28975#bib.bib3)\)mechanistically reverse\-engineer the generalizing solution for modular addition and define progress measures for its emergence;Liuet al\.\([2023](https://arxiv.org/html/2605.28975#bib.bib4)\)attribute grokking to a mismatch between the geometry of training and test loss landscapes;Davieset al\.\([2022](https://arxiv.org/html/2605.28975#bib.bib6)\)unify grokking and double descent as instances of the same dynamic; andVarmaet al\.\([2023](https://arxiv.org/html/2605.28975#bib.bib5)\)explain grokking as a competition between a memorizing circuit and a more parameter\-efficient generalizing circuit\. In the spectral\-dynamics literature,Yuniset al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib15)\)observe that the validation\-loss drop during grokking coincides with a sharp reduction in the effective rank of the weight matrices\.

Rather than propose a new theory, we use grokking as a testbed to study the behavior of unembedding LAR during the memorization\-to\-generalization transition\. We also find that LAR predicts the effective dimension of the learned function,k≈n2​\(1−LAR\)k\\approx n^\{2\(1\-\\mathrm\{LAR\}\)\}, which we verify empirically across 12 tasks\.

## 3Log\-Alignment Ratio

### 3\.1Definition

LetW∈ℝm×nW\\in\\mathbb\{R\}^\{m\\times n\}be a matrix,X∈ℝb×nX\\in\\mathbb\{R\}^\{b\\times n\}be a batch of input vectors, andW​X∈ℝb×mWX\\in\\mathbb\{R\}^\{b\\times m\}be the batch of corresponding outputs\. The log\-alignment ratio \(LAR\) of the matrix is defined byEverettet al\.\([2024](https://arxiv.org/html/2605.28975#bib.bib2)\)as:

LAR=logn⁡‖W​X‖RMS‖W‖RMS⋅‖X‖RMS\\text\{LAR\}=\\log\_\{n\}\\frac\{\\\|WX\\\|\_\{\\mathrm\{RMS\}\}\}\{\\\|W\\\|\_\{\\mathrm\{RMS\}\}\\cdot\\\|X\\\|\_\{\\mathrm\{RMS\}\}\}\(1\)where∥⋅∥RMS\\\|\\cdot\\\|\_\{\\mathrm\{RMS\}\}is the RMS\-norm of the matrix\.

LetW=∑i=1rui​si​vi⊤W=\\sum\\limits\_\{i=1\}^\{r\}u\_\{i\}s\_\{i\}v\_\{i\}^\{\\top\}be the SVD ofWW, wherer=rank⁡\(W\)r=\\operatorname\{rank\}\(W\)\. Each input vectorx∈ℝnx\\in\\mathbb\{R\}^\{n\}can be expanded in the basis of right\-singular vectors asx=∑i=1nxi​vix=\\sum\\limits\_\{i=1\}^\{n\}x\_\{i\}v\_\{i\}withxi=vi⊤​xx\_\{i\}=v\_\{i\}^\{\\top\}x, andvr\+1,…,vnv\_\{r\+1\},\\ldots,v\_\{n\}complete an orthonormal basis ofℝn\\mathbb\{R\}^\{n\}whenn\>rn\>r\.

We define the*weight distribution*ppand the*activation distribution*qqas the normalized squared singular values ofWWand the normalized squared projections ofxxonto its singular directions, respectively:

pi=si2∑j=1rsj2,qi=∑x∈B\(vi⊤​x\)2∑x∈B‖x‖2p\_\{i\}=\\frac\{s\_\{i\}^\{2\}\}\{\\sum\_\{j=1\}^\{r\}s\_\{j\}^\{2\}\},\\qquad q\_\{i\}=\\frac\{\\sum\_\{x\\in B\}\(v\_\{i\}^\{\\top\}x\)^\{2\}\}\{\\sum\_\{x\\in B\}\\\|x\\\|^\{2\}\}\(2\)
where∥⋅∥\\\|\\cdot\\\|denotes theL2L\_\{2\}\-norm of a vector\.ppis a distribution overrrelements capturing howWWdistributes its energy across singular directions, andqqis a distribution overnnelements capturing how the batch of activations distributes its variance along those directions and the null space\.111qiq\_\{i\}fori=r\+1,…,ni=r\+1,\\ldots,ndepends on the choice of basis vectors completing the null space\.

A direct calculation gives

LAR=1\+12​logn⁡\(∑i=1rpi​qi\)\.\\mathrm\{LAR\}=1\+\\frac\{1\}\{2\}\\log\_\{n\}\\left\(\\sum\_\{i=1\}^\{r\}p\_\{i\}\\,q\_\{i\}\\right\)\.\(3\)
To see this, write

‖W‖RMS2=1n​m​∑j=1rsj2,‖X‖RMS2=1\|B\|​n​∑x∈B‖x‖2,‖W​x‖RMS2=1\|B\|​m​∑x∈B∑i=1rsi2​xi2\.\\\|W\\\|\_\{\\mathrm\{RMS\}\}^\{2\}=\\frac\{1\}\{nm\}\\sum\\limits\_\{j=1\}^\{r\}s\_\{j\}^\{2\},\\qquad\\\|X\\\|\_\{\\mathrm\{RMS\}\}^\{2\}=\\frac\{1\}\{\|B\|\\,n\}\\sum\_\{x\\in B\}\\\|x\\\|^\{2\},\\qquad\\\|Wx\\\|\_\{\\mathrm\{RMS\}\}^\{2\}=\\frac\{1\}\{\|B\|\\,m\}\\sum\_\{x\\in B\}\\sum\\limits\_\{i=1\}^\{r\}s\_\{i\}^\{2\}\\,x\_\{i\}^\{2\}\.
The squared ratio becomes

‖W​x‖RMS2‖W‖RMS2​‖x‖RMS2\\displaystyle\\frac\{\\\|Wx\\\|\_\{\\mathrm\{RMS\}\}^\{2\}\}\{\\\|W\\\|\_\{\\mathrm\{RMS\}\}^\{2\}\\;\\\|x\\\|\_\{\\mathrm\{RMS\}\}^\{2\}\}=1\|B\|​m​∑x∈B∑isi2​xi21n​m​∑jsj2⋅1\|B\|​n​∑x∈B‖x‖2\\displaystyle=\\frac\{\\frac\{1\}\{\|B\|\\,m\}\\sum\_\{x\\in B\}\\sum\_\{i\}s\_\{i\}^\{2\}\\,x\_\{i\}^\{2\}\}\{\\frac\{1\}\{nm\}\\sum\_\{j\}s\_\{j\}^\{2\}\\;\\cdot\\;\\frac\{1\}\{\|B\|\\,n\}\\sum\_\{x\\in B\}\\\|x\\\|^\{2\}\}=n2⋅∑isi2​∑x∈Bxi2∑jsj2⋅∑x∈B‖x‖2\\displaystyle=n^\{2\}\\cdot\\frac\{\\sum\_\{i\}s\_\{i\}^\{2\}\\sum\_\{x\\in B\}x\_\{i\}^\{2\}\}\{\\sum\_\{j\}s\_\{j\}^\{2\}\\;\\cdot\\;\\sum\_\{x\\in B\}\\\|x\\\|^\{2\}\}=n2​∑i=1rsi2∑jsj2⋅∑x∈Bxi2∑x∈B‖x‖2\\displaystyle=n^\{2\}\\sum\_\{i=1\}^\{r\}\\frac\{s\_\{i\}^\{2\}\}\{\\sum\_\{j\}s\_\{j\}^\{2\}\}\\cdot\\frac\{\\sum\_\{x\\in B\}x\_\{i\}^\{2\}\}\{\\sum\_\{x\\in B\}\\\|x\\\|^\{2\}\}=n2​∑i=1rpi​qi\.\\displaystyle=n^\{2\}\\sum\_\{i=1\}^\{r\}p\_\{i\}\\,q\_\{i\}\.
Takinglogn\\log\_\{n\}of the square root gives the result in[Equation˜3](https://arxiv.org/html/2605.28975#S3.E3)\.

Treatingpip\_\{i\}andqiq\_\{i\}as twonn\-vectors of equally\-weighted data points \(withpj=0p\_\{j\}=0forj\>rj\>r\), we can writeLAR\\mathrm\{LAR\}in terms of their covariance,cov​\(p,q\)=1n​∑i=1rpi​qi−1n2\\text\{cov\}\(p,q\)=\\frac\{1\}\{n\}\\sum\\limits\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}\-\\frac\{1\}\{n^\{2\}\},

LAR=1\+12​logn⁡\(n​cov​\(p,q\)\+1n\)\.\\mathrm\{LAR\}=1\+\\frac\{1\}\{2\}\\log\_\{n\}\\left\(n\\,\\text\{cov\}\(p,q\)\+\\frac\{1\}\{n\}\\right\)\.\(4\)
Equivalently, in terms of the Pearson correlation coefficient,corr​\(p,q\)=cov​\(p,q\)σp​σq\\text\{corr\}\(p,q\)=\\frac\{\\text\{cov\}\(p,q\)\}\{\\sigma\_\{p\}\\,\\sigma\_\{q\}\}, whereσp\\sigma\_\{p\}andσq\\sigma\_\{q\}are the standard deviations ofppandqqwithσp2=1n​∑i=1npi2−1n2\\sigma\_\{p\}^\{2\}=\\frac\{1\}\{n\}\\sum\\limits\_\{i=1\}^\{n\}p\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\}andσq2=1n​∑i=1nqi2−1n2\\sigma\_\{q\}^\{2\}=\\frac\{1\}\{n\}\\sum\\limits\_\{i=1\}^\{n\}q\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\},

LAR=1\+12​logn⁡\(n​σp​σq​corr​\(p,q\)\+1n\)\.\\mathrm\{LAR\}=1\+\\frac\{1\}\{2\}\\log\_\{n\}\\left\(n\\,\\sigma\_\{p\}\\,\\sigma\_\{q\}\\,\\text\{corr\}\(p,q\)\+\\frac\{1\}\{n\}\\right\)\.\(5\)
Notably, smaller values ofσ2\\sigma^\{2\}correspond to more diffuse distributions \(see[Section˜B\.2](https://arxiv.org/html/2605.28975#A2.SS2)for a derivation\)\.

In practice, we compute LAR through[Equation˜1](https://arxiv.org/html/2605.28975#S3.E1)which has a small, negligible computational overhead ofO​\(\|B\|​m\+\|B\|​n\+m​n\)O\(\|B\|m\+\|B\|n\+mn\)\.[Equations˜3](https://arxiv.org/html/2605.28975#S3.E3),[4](https://arxiv.org/html/2605.28975#S3.E4)and[5](https://arxiv.org/html/2605.28975#S3.E5)provide alternative perspectives into the metric as overlap between two distributions\. While LAR is well\-defined for any matrixWW, we focus on the unembedding matrix in[Sections˜4](https://arxiv.org/html/2605.28975#S4)and[5](https://arxiv.org/html/2605.28975#S5)\. Unless otherwise stated, "LAR" refers to the unembedding LAR in these sections\. We discuss LAR for other matrices in[Section˜D\.5](https://arxiv.org/html/2605.28975#A4.SS5)\.

### 3\.2Properties

For any matrixWW, LAR satisfies the following properties\.

1. 1\.LAR∈\(−∞,1\]\\mathrm\{LAR\}\\in\(\-\\infty,1\], withLAR=−∞\\mathrm\{LAR\}=\-\\inftyif activations lie entirely in the null space ofWW, andLAR=1\\mathrm\{LAR\}=1ifWWis a rank\-1 matrix and activations lie entirely along its singular direction\. In practice, backpropagation in neural networks encourages alignment between weights and activations, and we observeLAR∈\(0\.3,0\.8\)\\mathrm\{LAR\}\\in\(0\.3,0\.8\)\.
2. 2\.Through[Equation˜4](https://arxiv.org/html/2605.28975#S3.E4),LAR≈12\\mathrm\{LAR\}\\approx\\tfrac\{1\}\{2\}whencov⁡\(p,q\)≈0\\operatorname\{cov\}\(p,q\)\\approx 0\. Equivalently,[Equation˜5](https://arxiv.org/html/2605.28975#S3.E5)givesLAR≈12\\mathrm\{LAR\}\\approx\\tfrac\{1\}\{2\}whenever any of the following holds: the weight spectrum is uniform \(σp≈0\\sigma\_\{p\}\\approx 0\), the activations are isotropic \(σq≈0\\sigma\_\{q\}\\approx 0\), orppandqqare uncorrelated \(corr⁡\(p,q\)≈0\\operatorname\{corr\}\(p,q\)\\approx 0\)\. Weights and activations of a neural network are uncorrelated at initialization and hence,LAR≈12\\mathrm\{LAR\}\\approx\\frac\{1\}\{2\}\.
3. 3\.IfWWhas its singular values distributed uniformly acrossk≤rk\\leq rdirections \(s1=⋯=sk\>0s\_\{1\}=\\cdots=s\_\{k\}\>0,si=0s\_\{i\}=0fori\>ki\>k\) and all activation variance is contained in thosekkdirections, then LAR=1−12​logn⁡\(k\)\.\\mathrm\{LAR\}=1\-\\tfrac\{1\}\{2\}\\log\_\{n\}\(k\)\.\(6\)To see this, notepi=1/kp\_\{i\}=1/kfori≤ki\\leq kandqi=0q\_\{i\}=0fori\>ki\>kimplies∑i=1rpi​qi=1k​∑i=1kqi=1/k\\sum\\limits\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}=\\tfrac\{1\}\{k\}\\sum\\limits\_\{i=1\}^\{k\}q\_\{i\}=1/k, and[Equation˜6](https://arxiv.org/html/2605.28975#S3.E6)follows\. In this setup, we say that LAR captures theeffective dimensionkkof the function learned by the neural network, withk=n2​\(1−LAR\)k=n^\{2\(1\-\\mathrm\{LAR\}\)\}\.

## 4Grokking Experiments

![Refer to caption](https://arxiv.org/html/2605.28975v1/x1.png)\(a\)12 binary operation tasks\.
![Refer to caption](https://arxiv.org/html/2605.28975v1/x2.png)\(b\)Modular addition with varying training dataset fraction\.

Figure 1:Unembedding LAR \(blue\) and validation accuracy \(orange\) vs\. epochs, with mean and ±1 std shaded bands across 5 seeds\. Red and green dashed vertical lines mark the mean memorization and grokking epochs\. The absence of a green line indicates that these tasks do not grok\.### 4\.1Setup

FollowingPoweret al\.\([2022](https://arxiv.org/html/2605.28975#bib.bib1)\), we train 2\-layer decoder\-only Transformer models \(with embedding dimension128128, feed\-forward dimension512512, and44attention heads\) on1212binary operation tasks \(cf\.[Section˜A\.1](https://arxiv.org/html/2605.28975#A1.SS1)\)\. We use AdamW\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2605.28975#bib.bib22)\)with learning rate10−310^\{\-3\}, weight decay=1\.0=1\.0, andβ2=0\.98\\beta\_\{2\}=0\.98, with batch size512512, and1010\-step linear warmup\. We train for up to50,00050,000epochs, terminating2,0002,000epochs after grokking if it occurs\.

We conduct two sets of experiments\. First, we train on all 12 tasks using a 50% train split\. Second, we focus on modular addition \(x\+y\(mod97\)x\+y\\pmod\{97\}\) and vary the training fraction across 5%, 10%, 20%, and 30% of the entire dataset to study how unembedding LAR behaves as a function of training set size\. Each configuration is trained with 5 seeds\. We compute validation accuracy every 20 epochs and check whether a model has memorized \(training accuracy\>99%\>99\\%, validation accuracy<50%<50\\%\) or grokked \(training accuracy\>99%\>99\\%, validation accuracy\>99%\>99\\%\)\. We also compute unembedding LAR at the same epochs using the entire training dataset rather than the current batch, as the dataset sizes are small in this setup\.

### 4\.2Experimental Results

For our two sets of experiments, we plot the mean unembedding LAR and validation accuracy across seeds as functions of training epochs in[Figure˜1\(a\)](https://arxiv.org/html/2605.28975#S4.F1.sf1)and[Figure˜1\(b\)](https://arxiv.org/html/2605.28975#S4.F1.sf2)\. We clip the range of x\-axis to 3 times the grokking epoch if grokking occurs, else to500500\. In each case, we plot vertical dashed lines in red and green to represent the mean memorization and grokking epochs respectively\. In 4 out of 12 cases in[Figure˜1\(a\)](https://arxiv.org/html/2605.28975#S4.F1.sf1)and in 2 out of 4 cases in[Figure˜1\(b\)](https://arxiv.org/html/2605.28975#S4.F1.sf2), the green line is missing as these models do not grok\.

In all cases, we make the following observations\.

1. 1\.LAR starts at∼0\.5\\sim 0\.5at initialization, as expected\. It then decreases briefly before increasing steadily and stabilizing in the range0\.60\.6–0\.720\.72at the end of training\.222The initial drop is shown more clearly in[Figure9\(a\)](https://arxiv.org/html/2605.28975#A2.F9.sf1)in[SectionB\.5](https://arxiv.org/html/2605.28975#A2.SS5), where we show that it stems from using a large batch size\.LAR at memorization is lower than at grokking, which is lower than at the final checkpoint\.
2. 2\.Grokking models achieve higher values \(LAR\>0\.65\\mathrm\{LAR\}\>0\.65\) than non\-grokking models \(LAR≤0\.65\\mathrm\{LAR\}\\leq 0\.65\)\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x3.png)\(a\)Modular addition,30%30\\%training data
![Refer to caption](https://arxiv.org/html/2605.28975v1/x4.png)\(b\)Modular addition,5%5\\%training data
![Refer to caption](https://arxiv.org/html/2605.28975v1/x5.png)\(c\)Modular Division
![Refer to caption](https://arxiv.org/html/2605.28975v1/x6.png)\(d\)x3\+x​y2\+y\(mod97\)x^\{3\}\+xy^\{2\}\+y\\pmod\{97\}

Figure 2:Weight distributionppand activation distributionqqat different points during training\. Memorization, grokking, and final points are shown for tasks that grok in \(a\) and \(c\)\. Only memorization and final points are shown for tasks that do not grok in \(b\) and \(d\)\. Subplot titles include epoch,LAR\\mathrm\{LAR\}, andcorr​\(p,q\)\\text\{corr\}\(p,q\)values\.
### 4\.3Analysis

Motivated by the interpretation of LAR as the overlap between the weight distributionppand activation distributionqq, we plotppandqqat memorization, grokking, and final checkpoints for four tasks in[Figure˜2](https://arxiv.org/html/2605.28975#S4.F2)\. In each case,ppandqqare diffuse at memorization, concentrated at grokking, and most concentrated at the final checkpoint\. Equivalently,σp\\sigma\_\{p\}andσq\\sigma\_\{q\}are smallest at memorization, larger at grokking, and largest at the final checkpoint\. Sincecorr⁡\(p,q\)\\operatorname\{corr\}\(p,q\)remains near11throughout,[Equation˜5](https://arxiv.org/html/2605.28975#S3.E5)then explains the evolution of LAR during training: it is smaller at memorization due to more diffuseppandqq, and larger at grokking and final checkpoints due to more concentrated distributions\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x7.png)\(a\)12 binary operation tasks\.
![Refer to caption](https://arxiv.org/html/2605.28975v1/x8.png)\(b\)Modular addition, varying train fraction\.

Figure 3:k95k\_\{95\}\(the number of principal components explaining95%95\\%of unembedding input activation variance\) vs\. LAR\-predicted dimensionalitykp=n2​\(1−LAR\)k\_\{p\}=n^\{2\(1\-\\mathrm\{LAR\}\)\}\. Each point corresponds to a memorization, grokking or final checkpoint\.We hypothesize that optimization forcesppandqqto concentrate on the same small subset of directions, across all tasks and seeds in our experiments\. We test this hypothesis by performing PCA on the input activations of the unembedding matrix and computing the number of principal components,k95k\_\{95\}, that explain95%95\\%of the variance\.k95k\_\{95\}is the number of directions along whichqqis mostly concentrated\. Ifppis concentrated along the samek95k\_\{95\}directions, property 3 from[Section˜3\.2](https://arxiv.org/html/2605.28975#S3.SS2)can provide a good approximation\.333Strictly, property 3 requiresppto beuniformacrosskkdirections, which does not hold exactly in our experiments\.In this case,LAR≈1−12​logn⁡\(k95\)\\mathrm\{LAR\}\\approx 1\-\\frac\{1\}\{2\}\\log\_\{n\}\(k\_\{95\}\), or equivalently,kpred:=n2​\(1−LAR\)k\_\{\\text\{pred\}\}:=n^\{2\(1\-\\mathrm\{LAR\}\)\}approximatesk95k\_\{95\}\.

We plotk95k\_\{95\}andkpredk\_\{\\text\{pred\}\}for grokking, memorization and final checkpoints of each individual experiment in[Figure˜3](https://arxiv.org/html/2605.28975#S4.F3)\. We observe that the approximation holds well for final and \(to some extent\) grokking checkpoints, butk95\>kpredk\_\{95\}\>k\_\{\\text\{pred\}\}invariably at memorization\. The root\-mean\-squared\-error values,‖k95−kpred‖\\\|k\_\{95\}\-k\_\{\\text\{pred\}\}\\\|, given in[Table˜1](https://arxiv.org/html/2605.28975#S4.T1)support this observation\. Larger values ofk95k\_\{95\}at memorization reflect more diffuseqq, while its smaller values and better approximationss tokpredk\_\{\\text\{pred\}\}at grokking and final checkpoints indicateppandqqbecoming concentrated on the same small subset of directions\.

Lastly, we note thatk95≈kpredk\_\{95\}\\approx k\_\{\\text\{pred\}\}takes larger values for final checkpoints of non\-grokking models as compared to grokking models\. Thus, a grokking model learns a more compressed representation than its non\-grokking counterpart, which is reflected in grokking models achieving higher values of LAR\.

Table 1:RMSE betweenkpk\_\{p\}andk95k\_\{95\}at memorized, grokked, and final checkpoints\. Values are RMSE±\\pmhalf the 95% percentile bootstrap CI \(10,000 resamples\)\.

## 5Large\-Scale Experiments

### 5\.1Experimental Setup

We train 3B\-parameter models with embedding dimension40964096and Gemma\-2 decoder blocks\(Gemma Team,[2024](https://arxiv.org/html/2605.28975#bib.bib26)\)using two optimizers, Muon\(Jordan,[2024](https://arxiv.org/html/2605.28975#bib.bib7)\)and Adam\(Kingma and Ba,[2015](https://arxiv.org/html/2605.28975#bib.bib21)\), across five dataset sizes: 200M, 400M, 800M, 1\.6B, and 100B tokens\. Each model is trained for 2B tokens, so the number of epochs varies with dataset size\. The smaller datasets are random subsets of the largest, ensuring similar distributions; in each case, 10% is held out for validation\. The training batch size is 24 sequences of length 8,192 \(196,608 tokens\), and we evaluate validation loss every 10 steps on a batch of 2\.6M tokens\. We sweep over learning rates on the 100B dataset to find the optimal rate for each optimizer, then use it for all other dataset sizes\. Further details are provided in[Section˜A\.2](https://arxiv.org/html/2605.28975#A1.SS2)\.

### 5\.2Results

![Refer to caption](https://arxiv.org/html/2605.28975v1/x9.png)Figure 4:Training loss, validation loss, and unembedding log\-alignment ratio \(LAR\) as functions of training steps\. Vertical dashed lines represent the training step with minimum validation loss\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x10.png)Figure 5:LAR differential, defined as the difference between the baseline LAR from the100​B100\\text\{B\}experiment and the LAR for a given dataset size, and generalization gap — the difference between validation loss and smoothed training loss as functions of training steps\.For each of the 10 models in our setup,[Figure˜4](https://arxiv.org/html/2605.28975#S5.F4)shows training loss, validation loss, and unembedding LAR, with the vertical dashed lines marking the step of minimum validation loss\.444We evaluated several criteria for detecting the onset of overfitting, aiming for one that triggered as early as possible without showing any false positives on the 100B baseline\. The best rule we found was simple: flag the first step at which the mean validation loss over a 5\-step window exceeds the running minimum by at least 2%\. In practice, this rule identifies essentially the same step as the minimum of the validation loss curve itself, which we attribute to our large validation batch \(2\.6M tokens\) producing a smooth loss trace with little step\-to\-step noise\.2 out of 5 models overfit with Adam, and 3 out of 5 models overfit with Muon\.

The case of the largest dataset size, 100B, establishes a baseline behavior of LAR for non\-overfitting models\. In this case, LAR starts at∼\\sim0\.5 at initialization, as expected, then rises to a peak of∼\\sim0\.676 with Adam and∼\\sim0\.688 with Muon before slowly declining throughout training\.555We suspect that the initial rise in LAR is due to the use of large optimization batch size\. We leave a verification of this hypothesis to future work\.In contrast, overfitting models exhibit larger decline in LAR towards the end of training\.

Table 2:Rate of decline in unembedding LAR in 1000\-step windows after the peak step, measured as negative slope per 1000 steps\. Higher values indicate steeper decline and are visualized as darker cell color\.Boldvalues indicate up to three windows ending at the detected overfitting step\. 100B and 1\.6B runs show slowdown in decline towards the end of training, while overfitting runs show sharp decline before the overfitting step\.In[Figure˜5](https://arxiv.org/html/2605.28975#S5.F5), we plot the LAR differential, defined as the difference between the baseline LAR \(from the non\-overfitting100​B100\\text\{B\}experiment\) and the LAR for a given dataset size, alongside the generalization gap — the difference between validation loss and smoothed training loss\.666We smooth training loss by averaging over a centered 14\-step sliding window, approximately matching the 2\.6M tokens used to compute validation loss\.We note that this differential has a similar shape to the generalization gap during training\. The LAR differential therefore serves as a signal of overfitting that requires no validation data\.

In[Table˜2](https://arxiv.org/html/2605.28975#S5.T2), we further give the rate of decline in LAR over periods of10001000steps in training\. Non\-overfitting models exhibit a steady rate of decline after the peak, followed by a stabilization phase towards the end of training\. In contrast, overfitting models show steeper declines as the overfitting step approaches\. Notably, Adam 800M maintains elevated decline rates throughout training without formally overfitting in our budget of 2B tokens, suggesting that it is approaching the overfitting step\. The rate of change in slope of LAR, i\.e\., its acceleration, can therefore be a useful metric to check for overfitting — a possibility we leave open for future investigation\.

### 5\.3Analysis

To understand why unembedding LAR tracks generalization, we computeppandqqexplicitly every10001000training steps during training\. Across all experiments, the top singular value is 3 times as large as the next value, sop1≈9​p2p\_\{1\}\\approx 9p\_\{2\}\. \(See[Appendix˜C](https://arxiv.org/html/2605.28975#A3)for detailed plots ofppandqqand related analysis\.\)

LAR\\mathrm\{LAR\}admits a natural decomposition,LAR=Lk\+Δk\\mathrm\{LAR\}=L\_\{k\}\+\\Delta\_\{k\}, into thecumulativecontribution from topkksingular directions,LkL\_\{k\}, and the residual,Δk\\Delta\_\{k\}:

Lk=1\+12logn\(∑i=1kpiqi\);Δk=−12logn\(∑i=1kpi​qi∑i=1rpi​qi\)L\_\{k\}=1\+\\frac\{1\}\{2\}\\log\_\{n\}\\left\(\\sum\\limits\_\{i=1\}^\{k\}p\_\{i\}q\_\{i\}\\right\)\\quad;\\quad\\Delta\_\{k\}=\-\\frac\{1\}\{2\}\\log\_\{n\}\\left\(\\frac\{\\sum\\limits\_\{i=1\}^\{k\}p\_\{i\}q\_\{i\}\}\{\\sum\\limits\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}\}\\right\)
We plot the evolution ofL1L\_\{1\}andΔ1\\Delta\_\{1\}in[Figure˜6](https://arxiv.org/html/2605.28975#S5.F6)\.L1L\_\{1\}diverges from the non\-overfitting baseline prior to the overfitting step, whileΔ1\\Delta\_\{1\}stays approximately constant through the overfitting step\. The pre\-overfit signature of overfitting is thus contained mostly inL1L\_\{1\}\. This is further confirmed in[Figure˜7](https://arxiv.org/html/2605.28975#S5.F7), where we plotLkL\_\{k\}fork=1,2,3,5,10k=1,2,3,5,10across 200M and 100B experiments\. The curves fork\>1k\>1differ fromL1L\_\{1\}only by a near\-constant offset, following the same trend otherwise\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x11.png)\(a\)L1L\_\{1\}
![Refer to caption](https://arxiv.org/html/2605.28975v1/x12.png)\(b\)Δ1\\Delta\_\{1\}

Figure 6:Contribution of the top singular direction,L1L\_\{1\}, toLAR\\mathrm\{LAR\}\(a\), and the residualΔ1=LAR−L1\\Delta\_\{1\}=\\mathrm\{LAR\}\-L\_\{1\}\(b\), as functions of training steps\. Different colors indicate dataset sizes; vertical dashed lines mark the onset of overfitting for the corresponding run\. Step 0 is omitted in \(b\) for clarity\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x13.png)Figure 7:Contribution of topkksingular directions,LkL\_\{k\}, to the unembedding LAR fork=1,2,3,5,10k=1,2,3,5,10in 200M and 100B experiments with Adam and Muon\. Step 0 is omitted for clarity in all cases\.Our observations are consistent withYanget al\.\([2023](https://arxiv.org/html/2605.28975#bib.bib33)\), who note that in the presence of feature learning, gradient descent induces alignment between a layer’s input activations and the top singular directions of its weight matrix\. The decline inL1L\_\{1\}before overfitting observed here reflects a weakening of this alignment\.

To further understand this weakening, we plotp1p\_\{1\}andq1q\_\{1\}separately in[Figure˜8](https://arxiv.org/html/2605.28975#S5.F8)\. We note that with Adam, the signature of overfitting is entirely inq1q\_\{1\}, withp1p\_\{1\}exhibiting growth across all dataset sizes\. With Muon, it is contained in bothp1p\_\{1\}andq1q\_\{1\}\. Notably, the unembedding layer is optimized with Adam in both cases\. We leave a mechanistic explanation of this asymmetry for future work\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x14.png)\(a\)p1p\_\{1\}
![Refer to caption](https://arxiv.org/html/2605.28975v1/x15.png)\(b\)q1q\_\{1\}

Figure 8:Weight\-spectrum componentp1p\_\{1\}\(a\) and activation\-spectrum componentq1q\_\{1\}\(b\) along the top singular direction, as functions of training steps\. Different colors indicate dataset sizes; vertical dashed lines mark the onset of overfitting for the corresponding run\.#### Interpretation in terms of spreads inppandqq\.

In grokking experiments of[Section˜4](https://arxiv.org/html/2605.28975#S4), we observed that memorization corresponds to more diffuse distributionsppandqq, while generalization corresponds to more concentrated distributions\. We make a similar observation here: whenp1p\_\{1\}andq1q\_\{1\}are much larger than their next counterparts, they dictate the spread in distributions asσp≈p1n\\sigma\_\{p\}\\approx\\frac\{p\_\{1\}\}\{\\sqrt\{n\}\}andσq≈q1n\\sigma\_\{q\}\\approx\\frac\{q\_\{1\}\}\{\\sqrt\{n\}\}\. With Adam, overfitting thus correlates with more diffuse activations and with Muon, it correlates with more diffuse activations and weight spectrum\. We discuss the behavior ofσp\\sigma\_\{p\}andσq\\sigma\_\{q\}in more detail in[Appendix˜C](https://arxiv.org/html/2605.28975#A3)\.

## 6Discussion

We have studied the unembedding log\-alignment ratio \(LAR\) in two settings\. On small algorithmic tasks, the final value of LAR tracks the effective dimensionalityk95k\_\{95\}used by the model, which can be interpreted as the dimensionality of the generalizing circuit when the model groks\. On 3B\-parameter language models, LAR decreases slowly and eventually stabilizes when the model generalizes, but declines sharply once overfitting sets in\. LAR is computed entirely from training\-time quantities with negligible computational overhead, and hence, does not depend on the evaluation data, making it a practically useful signal in situations where validation is expensive: for example, large\-scale pretraining with limited evaluation budgets\.

The decompositionLAR=1\+12​logn⁡\(n​σp​σq​corr⁡\(p,q\)\+1/n\)\\mathrm\{LAR\}=1\+\\tfrac\{1\}\{2\}\\log\_\{n\}\(n\\,\\sigma\_\{p\}\\,\\sigma\_\{q\}\\,\\operatorname\{corr\}\(p,q\)\+1/n\)explains*why*LAR declines during training: either because the weight or activation spectra become more diffuse, or because the two stop being concentrated on the same directions\. Empirically, we observe thatppandqqremain highly correlated in almost all cases, and the behavior ofLAR\\mathrm\{LAR\}is explained byσp​σq\\sigma\_\{p\}\\sigma\_\{q\}\. In pre\-training experiments, the top singular value of the unembedding is much larger than the rest, soσp​σq\\sigma\_\{p\}\\sigma\_\{q\}further reduces to the top\-direction overlapp1​q1p\_\{1\}q\_\{1\}, and the pre\-overfit signature lives almost entirely in the top\-component contributionL1=1\+12​logn⁡\(p1​q1\)L\_\{1\}=1\+\\tfrac\{1\}\{2\}\\log\_\{n\}\(p\_\{1\}q\_\{1\}\)\. The behavior depends on the optimizer: with Adam, overfitting is reflected almost entirely inσq\\sigma\_\{q\}, while with Muon both distributionsppandqqspread out\. Understanding this asymmetry, which is likely tied to how each optimizer propagates spectral structure through the rest of the network, is a natural direction for follow\-up work\.

#### Limitations\.

First, we focus only on the unembedding layer\. While Appendix[D\.5](https://arxiv.org/html/2605.28975#A4.SS5)shows that some other matrices exhibit related properties, a complete understanding of LAR for each matrix type and the information it provides about the training dynamics is a promising direction for future work\. Second, while we vary the dataset size and observe LAR behavior in both the overfitting \(D≪ND\\ll N\) and non\-overfitting \(D≫ND\\gg N\) regimes, our 100B baseline trains for onlyT≈2​BT\\approx 2\\text\{B\}tokens\. We therefore have not probed theD≫ND\\gg Nregime atT∼DT\\sim D\. We have no specific reason to expect LAR to break out of its plateau in this scenario, but collecting more empirical evidence will be quite useful\. Third, our evidence is correlational: LAR decline coincides with overfitting, but we have not tested whether directly regularizingσp\\sigma\_\{p\},σq\\sigma\_\{q\}, orcorr⁡\(p,q\)\\operatorname\{corr\}\(p,q\)during training would actually prevent overfitting\. Finally, we computedqqusing a batch of fixed size in all pre\-training experiments\. Studying the behavior of LAR as a function of batch size and its relationship with generalization is an important future direction\.

## Acknowledgments

We thank Michael Callahan and Philip Monk for help with training infrastructure for this project\. We are also grateful to Adarsh Chaluvaraju, Devaansh Gupta, Yash Jain, Noam Shazeer, Somanshu Singla, and Kalina Slavkova for useful discussions, and to Yash Jain for comments on the draft\.

## References

- S\. Arora, R\. Ge, B\. Neyshabur, and Y\. Zhang \(2018\)Stronger generalization bounds for deep nets via a compression approach\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Barak, B\. Edelman, S\. Goel, S\. Kakade, E\. Malach, and C\. Zhang \(2022\)Hidden progress in deep learning: sgd learns parities near the computational limit\.Advances in Neural Information Processing Systems35,pp\. 21750–21764\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- P\. L\. Bartlett, D\. J\. Foster, and M\. J\. Telgarsky \(2017\)Spectrally\-normalized margin bounds for neural networks\.InAdvances in Neural Information Processing Systems,Vol\.30\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1)\.
- X\. Davies, L\. Langosco, and D\. Krueger \(2022\)Unifying grokking and double descent\.InNeurIPS 2022 ML Safety Workshop,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- N\. Dey, B\. C\. Zhang, L\. Noci, M\. Li, B\. Bordelon, S\. Bergsma, C\. Pehlevan, B\. Hanin, and J\. Hestness \(2025\)Don’t be lazy: CompleteP enables compute\-efficient deep transformers\.InAdvances in Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px2.p1.2)\.
- K\. Everett, L\. Xiao, M\. Wortsman, A\. A\. Alemi, R\. Novak, P\. J\. Liu, I\. Gur, and J\. Sohl\-Dickstein \(2024\)Scaling exponents across parameterizations and optimizers\.arXiv preprint arXiv:2407\.05872\.Cited by:[§1](https://arxiv.org/html/2605.28975#S1.p2.2),[§1](https://arxiv.org/html/2605.28975#S1.p3.2),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px2.p1.2),[§3\.1](https://arxiv.org/html/2605.28975#S3.SS1.p1.3)\.
- Gemma Team \(2024\)Gemma 2: improving open language models at a practical size\.arXiv preprint arXiv:2408\.00118\.Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px1.p1.6),[§5\.1](https://arxiv.org/html/2605.28975#S5.SS1.p1.1)\.
- A\. Grattafiori, A\. Dubey,et al\.\(2024\)The Llama 3 herd of models\.arXiv preprint arXiv:2407\.21783\.Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px1.p1.6)\.
- J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. d\. L\. Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark,et al\.\(2022\)Training compute\-optimal large language models\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2605.28975#S1.p1.1)\.
- Y\. Jiang, B\. Neyshabur, H\. Mobahi, D\. Krishnan, and S\. Bengio \(2020\)Fantastic generalization measures and where to find them\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1)\.
- K\. Jordan \(2024\)Muon: an optimizer for hidden layers in neural networks\.Note:[https://kellerjordan\.github\.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)Cited by:[§5\.1](https://arxiv.org/html/2605.28975#S5.SS1.p1.1)\.
- J\. Kaplan, S\. McCandlish, T\. Henighan, T\. B\. Brown, B\. Chess, R\. Child, S\. Gray, A\. Radford, J\. Wu, and D\. Amodei \(2020\)Scaling laws for neural language models\.arXiv preprint arXiv:2001\.08361\.Cited by:[§1](https://arxiv.org/html/2605.28975#S1.p1.1)\.
- D\. P\. Kingma and J\. Ba \(2015\)Adam: a method for stochastic optimization\.InInternational Conference on Learning Representations,Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px2.p1.12),[§5\.1](https://arxiv.org/html/2605.28975#S5.SS1.p1.1)\.
- J\. Liu, J\. Su, X\. Yao, Z\. Jiang, G\. Lai, Y\. Du, Y\. Qin, W\. Xu, E\. Lu, J\. Yan,et al\.\(2025\)Muon is scalable for llm training\.arXiv preprint arXiv:2502\.16982\.Cited by:[footnote 9](https://arxiv.org/html/2605.28975#footnote9)\.
- Z\. Liu, E\. J\. Michaud, and M\. Tegmark \(2023\)Omnigrok: grokking beyond algorithmic data\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2017\)SGDR: stochastic gradient descent with warm restarts\.InInternational Conference on Learning Representations,Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px2.p1.12)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InInternational Conference on Learning Representations,Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px2.p1.12),[§D\.4](https://arxiv.org/html/2605.28975#A4.SS4.p1.1),[§4\.1](https://arxiv.org/html/2605.28975#S4.SS1.p1.11)\.
- C\. H\. Martin and M\. W\. Mahoney \(2021\)Implicit self\-regularization in deep neural networks: evidence from random matrix theory and implications for training\.Journal of Machine Learning Research22\(165\),pp\. 1–73\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Mlodozeniec, P\. Ablin, L\. Béthune, D\. Busbridge, M\. Klein, J\. Ramapuram, and M\. Cuturi \(2025\)Completed hyperparameter transfer across modules, width, depth, batch & duration\.arXiv preprint arXiv:2512\.22382\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px2.p1.2)\.
- N\. Nanda, L\. Chan, T\. Lieberum, J\. Smith, and J\. Steinhardt \(2023\)Progress measures for grokking via mechanistic interpretability\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- B\. Neyshabur, S\. Bhojanapalli, and N\. Srebro \(2018\)A PAC\-Bayesian approach to spectrally\-normalized margin bounds for neural networks\.InInternational Conference on Learning Representations,Cited by:[§D\.1](https://arxiv.org/html/2605.28975#A4.SS1.SSS0.Px1.p1.2),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1)\.
- A\. Power, Y\. Burda, H\. Edwards, I\. Babuschkin, and V\. Misra \(2022\)Grokking: generalization beyond overfitting on small algorithmic datasets\.InICLR 2022 Workshop on MATH\-AI,Cited by:[§A\.1](https://arxiv.org/html/2605.28975#A1.SS1.p1.1),[§1](https://arxiv.org/html/2605.28975#S1.p5.5),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1),[§4\.1](https://arxiv.org/html/2605.28975#S4.SS1.p1.11)\.
- O\. Press and L\. Wolf \(2017\)Using the output embedding to improve language models\.InProceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics,Cited by:[§D\.3](https://arxiv.org/html/2605.28975#A4.SS3.p1.3)\.
- O\. Roy and M\. Vetterli \(2007\)The effective rank: a measure of effective dimensionality\.In2007 15th European Signal Processing Conference \(EUSIPCO\),pp\. 1737–1741\.Cited by:[§D\.1](https://arxiv.org/html/2605.28975#A4.SS1.SSS0.Px2.p1.1)\.
- A\. Sanyal, P\. H\. S\. Torr, and P\. K\. Dokania \(2020\)Stable rank normalization for improved generalization in neural networks and GANs\.InInternational Conference on Learning Representations,Cited by:[§D\.1](https://arxiv.org/html/2605.28975#A4.SS1.SSS0.Px1.p1.2)\.
- N\. Shazeer \(2020\)GLU variants improve transformer\.arXiv preprint arXiv:2002\.05202\.Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px1.p1.6)\.
- V\. Varma, R\. Shah, Z\. Kenton, J\. Kramár, and R\. Kumar \(2023\)Explaining grokking through circuit efficiency\.arXiv preprint arXiv:2309\.02390\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- G\. Yang, E\. J\. Hu, I\. Babuschkin, S\. Sidor, X\. Liu, D\. Farhi, N\. Ryder, J\. Pachocki, W\. Chen, and J\. Gao \(2022\)Tensor programs V: tuning large neural networks via zero\-shot hyperparameter transfer\.arXiv preprint arXiv:2203\.03466\.Cited by:[§1](https://arxiv.org/html/2605.28975#S1.p2.2),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px2.p1.2)\.
- G\. Yang and E\. J\. Hu \(2021\)Tensor programs IV: feature learning in infinite\-width neural networks\.InInternational Conference on Machine Learning,pp\. 11727–11737\.Cited by:[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px2.p1.2)\.
- G\. Yang, J\. B\. Simon, and J\. Bernstein \(2023\)A spectral condition for feature learning\.arXiv preprint arXiv:2310\.17813\.Cited by:[§5\.3](https://arxiv.org/html/2605.28975#S5.SS3.p4.1)\.
- D\. Yunis, K\. K\. Patel, S\. Wheeler, P\. Savarese, G\. Vardi, K\. Livescu, M\. Maire, and M\. R\. Walter \(2024\)Approaching deep learning through the spectral dynamics of weights\.arXiv preprint arXiv:2408\.11804\.Cited by:[§D\.1](https://arxiv.org/html/2605.28975#A4.SS1.SSS0.Px2.p1.2),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2605.28975#S2.SS0.SSS0.Px3.p1.1)\.
- B\. Zhang and R\. Sennrich \(2019\)Root mean square layer normalization\.InAdvances in Neural Information Processing Systems,Cited by:[§A\.2](https://arxiv.org/html/2605.28975#A1.SS2.SSS0.Px1.p1.6)\.

## Appendix AAdditional Experimental Details

### A\.1Binary Operations For Grokking Experiments

The 12 tasks fromPoweret al\.\[[2022](https://arxiv.org/html/2605.28975#bib.bib1)\]are:

1. 1\.x\+y\(modp\)x\+y\\pmod\{p\}
2. 2\.x−y\(modp\)x\-y\\pmod\{p\}
3. 3\.x/y\(modp\)x/y\\pmod\{p\}fory≠0y\\neq 0
4. 4\.x/y\(modp\)x/y\\pmod\{p\}ifyyodd, elsex−y\(modp\)x\-y\\pmod\{p\}
5. 5\.x2\+y2\(modp\)x^\{2\}\+y^\{2\}\\pmod\{p\}
6. 6\.x2\+x​y\+y2\(modp\)x^\{2\}\+xy\+y^\{2\}\\pmod\{p\}
7. 7\.x2\+x​y\+y2\+x\(modp\)x^\{2\}\+xy\+y^\{2\}\+x\\pmod\{p\}
8. 8\.x3\+x​y\(modp\)x^\{3\}\+xy\\pmod\{p\}
9. 9\.x3\+x​y2\+y\(modp\)x^\{3\}\+xy^\{2\}\+y\\pmod\{p\}
10. 10\.x⋅yx\\cdot yforx,y∈S5x,y\\in S\_\{5\}
11. 11\.x⋅y⋅x−1x\\cdot y\\cdot x^\{\-1\}forx,y∈S5x,y\\in S\_\{5\}
12. 12\.x⋅y⋅xx\\cdot y\\cdot xforx,y∈S5x,y\\in S\_\{5\}

Tasks 1–5 and 10–12 grok reliably\. Tasks 6–9 do not grok with our experimental setup and our training budget \(50,000 epochs with 50% training fraction\)\.

### A\.2Experimental details of large\-scale experiments

#### Architecture\.

All models use Gemma\-2 decoder blocks\[Gemma Team,[2024](https://arxiv.org/html/2605.28975#bib.bib26)\]with88transformer layers, embedding dimensiondmodel=4096d\_\{\\text\{model\}\}=4096, MLP hidden dimensiondff=16,384d\_\{\\text\{ff\}\}=16\{,\}384with a gated GeLU activation\[Shazeer,[2020](https://arxiv.org/html/2605.28975#bib.bib25)\],3232attention heads of head dimension128128, and untied input/output embeddings over a vocabulary of128,256128\{,\}256tokens \(Llama\-3 tokenizer\[Grattafioriet al\.,[2024](https://arxiv.org/html/2605.28975#bib.bib27)\]\)\. Both attention sub\-layers within each Gemma\-2 block are global \(no sliding\-window attention\)\. Each attention and MLP sub\-block is wrapped by RMSNorm\[Zhang and Sennrich,[2019](https://arxiv.org/html/2605.28975#bib.bib24)\]both before*and*after the sub\-block \(i\.e\., both pre\-norm and post\-norm are enabled, following the original Gemma\-2 design\[Gemma Team,[2024](https://arxiv.org/html/2605.28975#bib.bib26)\]\); we do not apply QK\-norm and do not soft\-cap the final logits\.

#### Optimization\.

We train for10,00010\{,\}000steps with a cosine learning\-rate schedule\[Loshchilov and Hutter,[2017](https://arxiv.org/html/2605.28975#bib.bib23)\]preceded by linear warmup over the first4%4\\%of training \(400400steps\), and clip gradients at global norm1\.01\.0\. For Adam\[Kingma and Ba,[2015](https://arxiv.org/html/2605.28975#bib.bib21)\], we useβ1=0\.95\\beta\_\{1\}=0\.95,β2=0\.999\\beta\_\{2\}=0\.999,ϵ=10−8\\epsilon=10^\{\-8\}, and no weight decay \(we use the plain Adam optimizer rather than AdamW\[Loshchilov and Hutter,[2019](https://arxiv.org/html/2605.28975#bib.bib22)\]\)\. For Muon, we use momentum0\.950\.95with Nesterov updates,55Newton–Schulz iterations for the orthogonalization step, and decoupled weight decayλ=0\.1\\lambda=0\.1for all two\-dimensional parameters including the unembedding matrix\. We sweep the peak learning rate on the 100B\-token dataset \(in factor\-of\-two increments\) and pick the value with the lowest final validation loss for each optimizer; the resulting optima areηAdam=2\.44×10−4\\eta\_\{\\text\{Adam\}\}=2\.44\\times 10^\{\-4\}andηMuon=9\.77×10−4\\eta\_\{\\text\{Muon\}\}=9\.77\\times 10^\{\-4\}\. These same learning rates are reused for the 200M, 400M, 800M, and 1\.6B runs\.

#### Dataset\.

The pretraining mixture is an internal Essential AI corpus consisting predominantly of English web text with a smaller fraction of code\.

## Appendix BAdditional Results on LAR and Related Quantities

### B\.1Bounds on∑i=1rpi​qi\\sum\\limits\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}and LAR

Sinceppandqqare distributions \(overrrandnnelements respectively, with only the firstrrterms ofqqentering the sum\):

∑i=1rpi​qi\\displaystyle\\sum\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}≤maxj⁡qj⋅∑i=1rpi=maxj⁡qj≤1,\\displaystyle\\leq\\max\_\{j\}q\_\{j\}\\cdot\\sum\_\{i=1\}^\{r\}p\_\{i\}=\\max\_\{j\}q\_\{j\}\\leq 1,\(7\)∑i=1rpi​qi\\displaystyle\\sum\_\{i=1\}^\{r\}p\_\{i\}q\_\{i\}≥0,\\displaystyle\\geq 0,\(8\)with equality in the lower bound iffqi=0q\_\{i\}=0for alli≤ri\\leq r\(i\.e\.,xxis in the null space ofWW\)\. The upper bound is achieved whenppandqqare both point masses on the same index\.LAR\\mathrm\{LAR\}is−∞\-\\inftyat the lower bound andLAR=1\\mathrm\{LAR\}=1at the upper bound\.

### B\.2σp\\sigma\_\{p\}asL2L\_\{2\}distance to uniform distribution

We derive the claim thatppis more diffuse iffσp\\sigma\_\{p\}is small\. Letppbe a distribution overnnelements\. Expanding the squaredL2L\_\{2\}\-norm fromppto the uniform distributionui=1/nu\_\{i\}=1/n:

‖p−u‖2=∑i\(pi−1n\)2=∑ipi2−1n,\\\|p\-u\\\|^\{2\}=\\sum\_\{i\}\\\!\\left\(p\_\{i\}\-\\frac\{1\}\{n\}\\right\)^\{2\}=\\sum\_\{i\}p\_\{i\}^\{2\}\-\\frac\{1\}\{n\},and therefore,

σp2=1n​∑ipi2−1n2=1n​‖p−u‖2\.\\sigma\_\{p\}^\{2\}=\\frac\{1\}\{n\}\\sum\_\{i\}p\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\}=\\frac\{1\}\{n\}\\\|p\-u\\\|^\{2\}\.Soσp\\sigma\_\{p\}is proportional to theL2L\_\{2\}distance fromppto the uniform distribution, and is hence minimum whenppis uniform\.

### B\.3Bounds onσp\\sigma\_\{p\}

We can show thatσp\\sigma\_\{p\}is maximum whenppis concentrated on a single direction:pi=1p\_\{i\}=1for a fixedii\. To see this, note that transferring probability frompjp\_\{j\}topkp\_\{k\}increasesσp\\sigma\_\{p\}: suppose at least two componentspj,pk\>0p\_\{j\},p\_\{k\}\>0, then the change inσp2\\sigma\_\{p\}^\{2\}by transferringε\>0\\varepsilon\>0frompkp\_\{k\}topjp\_\{j\}gives:

Δ​\(∑ipi2\)=\(pj\+ε\)2\+\(pk−ε\)2−pj2−pk2=2​ε​\(pj−pk\)\+2​ε2\>0\\Delta\\\!\\left\(\\sum\_\{i\}p\_\{i\}^\{2\}\\right\)=\(p\_\{j\}\+\\varepsilon\)^\{2\}\+\(p\_\{k\}\-\\varepsilon\)^\{2\}\-p\_\{j\}^\{2\}\-p\_\{k\}^\{2\}=2\\varepsilon\(p\_\{j\}\-p\_\{k\}\)\+2\\varepsilon^\{2\}\>0Hence, transferring mass to a single component always increasesσp\\sigma\_\{p\}\. Withpj=1p\_\{j\}=1,

‖p−u‖2=\(1−1/n\)2\+\(n−1\)/n2=\(n−1\)/n,\\\|p\-u\\\|^\{2\}=\(1\-1/n\)^\{2\}\+\(n\-1\)/n^\{2\}=\(n\-1\)/n,soσp,max=\(n−1\)/n2=n−1/n\\sigma\_\{p,\\max\}=\\sqrt\{\(n\-1\)/n^\{2\}\}=\\sqrt\{n\-1\}/n\. In short,

σp∈\[0,n−1n\]\.\\sigma\_\{p\}\\in\\left\[0,\\;\\frac\{\\sqrt\{n\-1\}\}\{n\}\\right\]\.

### B\.4σp\\sigma\_\{p\}and Rényi 2\-entropy

The Rényi 2\-entropy ofppis defined asH2​\(p\)=−log2​∑ipi2H\_\{2\}\(p\)=\-\\log\_\{2\}\\sum\_\{i\}p\_\{i\}^\{2\}, so∑ipi2=2−H2​\(p\)\\sum\_\{i\}p\_\{i\}^\{2\}=2^\{\-H\_\{2\}\(p\)\}\. Substituting this intoσp2=1n​∑ipi2−1n2\\sigma\_\{p\}^\{2\}=\\frac\{1\}\{n\}\\sum\_\{i\}p\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\}gives

σp2=1n​\(2−H2​\(p\)−1n\)\.\\sigma\_\{p\}^\{2\}\\;=\\;\\frac\{1\}\{n\}\\\!\\left\(2^\{\-H\_\{2\}\(p\)\}\-\\frac\{1\}\{n\}\\right\)\.Hence,σp\\sigma\_\{p\}is small iffH2​\(p\)H\_\{2\}\(p\)is close to its maximum,log2⁡n\\log\_\{2\}n, which is attained whenppis uniform\. The two quantities are monotonic functions of∑ipi2\\sum\_\{i\}p\_\{i\}^\{2\}and capture the same notion of diffuseness\.

### B\.5LAR Dependence on Batch Size

![Refer to caption](https://arxiv.org/html/2605.28975v1/x16.png)\(a\)Unembedding LAR \(y\-axis\) vs\. training epochs \(x\-axis\) for grokking on modular addition \(mod97\\bmod 97\)\. Legend entries give batch size and the number of training steps per epoch\.
![Refer to caption](https://arxiv.org/html/2605.28975v1/x17.png)\(b\)Minimum unembedding LAR across first 10 epochs \(y\-axis\) and the ratio of batch size to training dataset size \(x\-axis\) across grokking experiments on modular addition tasks with varying primepp, training data fraction and batch size\.

Figure 9:Effect of the training batch size on unembedding LAR\.Batch size affects LAR both directly, throughqq\(which is computed over the batch\), and indirectly, through the optimization trajectory\. More precisely, the batch size used to computeqqmay be different from the optimization batch size; and we may study the effect of either of them on LAR\.

#### Dependence ofLAR\\mathrm\{LAR\}on batch size through optimization\.

In grokking experiments of[Section˜4](https://arxiv.org/html/2605.28975#S4), we computeLAR\\mathrm\{LAR\}over the entire training dataset due to the small dataset sizes\. Figure[9\(a\)](https://arxiv.org/html/2605.28975#A2.F9.sf1)plots unembedding LAR over the first 20 epochs of training on modular addition \(forp=97p=97\) for 5 choices of batch size\. We observe an early dip in LAR at large batch sizes, which disappears when the batch size is small\. We attribute this to the number of gradient updates per epoch: with few large updates, each step perturbs weights enough to push them out of alignment with activations\. This effect is transient, asLAR\\mathrm\{LAR\}rises above0\.50\.5after training for a few epochs\.

To study whether this effect persists over a range of training dataset sizes and batch size, we vary training dataset fraction, batch size and the primeppof modular addition in[Figure˜9\(b\)](https://arxiv.org/html/2605.28975#A2.F9.sf2)\. We note that the minimum unembedding LAR observed in the first 10 training epochs decreases as the ratio of batch size to dataset size increases\. This confirms that the presence of a few large gradient updates causes an early dip in unembedding LAR\. In each case, training for more epochs pushedLAR\\mathrm\{LAR\}above 0\.5, confirming that the effect is only transient\.

#### Dependence ofqqon batch size\.

The activation distributionqqdepends on the batchBBused to compute it\. If the inputsx∈Bx\\in Bare drawn i\.i\.d\. from a data distribution𝒟\\mathcal\{D\}, thenqqis an empirical estimate of the population quantity

qi∗=𝔼x∼𝒟​\[\(vi⊤​x\)2\]𝔼x∼𝒟​\[‖x‖2\]=vi⊤​C​vitr⁡\(C\),q\_\{i\}^\{\*\}\\;=\\;\\frac\{\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\\!\\left\[\(v\_\{i\}^\{\\top\}x\)^\{2\}\\right\]\}\{\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\\!\\left\[\\\|x\\\|^\{2\}\\right\]\}\\;=\\;\\frac\{v\_\{i\}^\{\\top\}C\\,v\_\{i\}\}\{\\operatorname\{tr\}\(C\)\},\(9\)whereC=𝔼x∼𝒟​\[x​x⊤\]C=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\[xx^\{\\top\}\]is the population uncentered second\-moment matrix\. The batch versionqi=∑x∈B\(vi⊤​x\)2/∑x∈B‖x‖2q\_\{i\}=\\sum\_\{x\\in B\}\(v\_\{i\}^\{\\top\}x\)^\{2\}/\\sum\_\{x\\in B\}\\\|x\\\|^\{2\}replaces each expectation with its empirical average overBB, andqi→qi∗q\_\{i\}\\to q\_\{i\}^\{\*\}as\|B\|→∞\|B\|\\to\\infty\.

## Appendix CDetailed Analysis ofppandqqDuring Pre\-Training

In[Section˜5](https://arxiv.org/html/2605.28975#S5)of the main text, we showed thatLAR\\mathrm\{LAR\}tracks the generalization gap in large\-scale pre\-training experiments\. We also showed that the pre\-overfit signature of overfitting lies mostly in the contribution from the top singular component,L1L\_\{1\}\. Here, we provide a more detailed analysis ofppandqqin[Section˜C\.1](https://arxiv.org/html/2605.28975#A3.SS1)and an alternative perspective on the relationship betweenLAR\\mathrm\{LAR\}and overfitting in[Section˜C\.2](https://arxiv.org/html/2605.28975#A3.SS2)\.

### C\.1Detailed plots ofppandqq

In[Figure˜10](https://arxiv.org/html/2605.28975#A3.F10)and[Figure˜11](https://arxiv.org/html/2605.28975#A3.F11), we give concrete examples ofppandqqfor Adam and Muon experiments respectively\. A row corresponds to a fixed dataset size: 100B, 800M, 400M or 200M, and a column represents a fixed training step:0,20002000,40004000,60006000,80008000, or99009900\.777We skip the experiments with dataset size of 1\.6B tokens as it is similar to the case of 100B dataset when training budget is 2B tokens\.Each subplot title contains Pearson correlation,corr​\(p,q\)\\text\{corr\}\(p,q\), and Spearman rank correlation,ρ​\(p,q\)\\rho\(p,q\), to capture linear and rank alignment betweenppandqq\. We also plot Pearson and Spearman correlations every 1000 training steps for all 5 dataset sizes in[Figure˜12\(a\)](https://arxiv.org/html/2605.28975#A3.F12.sf1)and[Figure˜12\(b\)](https://arxiv.org/html/2605.28975#A3.F12.sf2)respectively\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x18.png)Figure 10:Weight distributionpip\_\{i\}\(red\) and activation distributionqiq\_\{i\}\(blue\) against singular value indicesiiof the unembedding matrixWW, for 3B models trained with Adam on datasets of 100B, 800M, 400M and 200M tokens\. Each subplot title lists Pearson correlation,corr​\(p,q\)\\text\{corr\}\(p,q\), and Spearman’s rank correlation,ρ​\(p,q\)\\rho\(p,q\)\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x19.png)Figure 11:Weight distributionpip\_\{i\}\(red\) and activation distributionqiq\_\{i\}\(blue\) against singular value indicesiiof the unembedding matrixWW, for 3B models trained with Muon on datasets of 100B, 800M, 400M and 200M tokens\. Each subplot title lists Pearson correlation,corr​\(p,q\)\\text\{corr\}\(p,q\), and Spearman’s rank correlation,ρ​\(p,q\)\\rho\(p,q\)\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x20.png)\(a\)corr​\(p,q\)\\text\{corr\}\(p,q\)
![Refer to caption](https://arxiv.org/html/2605.28975v1/x21.png)\(b\)ρ\\rho
![Refer to caption](https://arxiv.org/html/2605.28975v1/x22.png)\(c\)σp\\sigma\_\{p\}
![Refer to caption](https://arxiv.org/html/2605.28975v1/x23.png)\(d\)σq\\sigma\_\{q\}

Figure 12:Pearson correlationcorr​\(p,q\)\\text\{corr\}\(p,q\)\(a\), Spearman rank correlationρ\\rho\(b\), and standard deviationsσp\\sigma\_\{p\}\(c\) andσq\\sigma\_\{q\}\(d\), as functions of training steps in 3B experiments\. Different colors indicate dataset sizes; vertical dashed lines mark the onset of overfitting for the corresponding run\.We make the following observations:

- •ppandqqat initialization\.At initialization,pproughly follows the prediction from the Marchenko–Pastur distribution of the asymptotic behavior of singular values of random matrix\.qqis isotropic withqi≈1/n=1/4096≈2×10−4q\_\{i\}\\approx 1/n=1/4096\\approx 2\\times 10^\{\-4\}\.
- •ppduring training\.After initialization,pip\_\{i\}follows a decaying distribution across all training steps\. The largest singular value is approximately 3 times larger than the next value, givingp1≈9​p2p\_\{1\}\\approx 9p\_\{2\}\. \(See[Table˜3](https://arxiv.org/html/2605.28975#A3.T3)\.\) Throughout training,p1≈0\.04−0\.07p\_\{1\}\\approx 0\.04\-0\.07, or equivalently,4−7%4\-7\\%of the weight’s total energy\. The smallest probability mass,p4096p\_\{4096\}, tends to be 1\-2 orders of magnitude smaller compared top2p\_\{2\}:p4096∼10−4p\_\{4096\}\\sim 10^\{\-4\}for Adam andp4096∼10−5p\_\{4096\}\\sim 10^\{\-5\}for Muon\.
- •qqduring training\.Similar topp,qqalso becomes more concentrated on the top singular directions but still remains more diffuse thanpp\. The ratioq1/q2q\_\{1\}/q\_\{2\}sits in the range1\.11\.1–4\.04\.0across post\-initialization checkpoints \(median≈2\.3\\approx 2\.3; see[Table˜4](https://arxiv.org/html/2605.28975#A3.T4)\)\.q1q\_\{1\}plateaus around0\.040\.04in the absence of overfitting, but overfitting causes a steady decline in its value\. \(See[Figure˜8\(b\)](https://arxiv.org/html/2605.28975#S5.F8.sf2)in the main text\.\)
- •Pearson Correlation betweenppandqq\.At initialization,corr​\(p,q\)≈0\.02\\text\{corr\}\(p,q\)\\approx 0\.02asqqis isotropic and is uncorrelated withpp\. It rises sharply to above0\.90\.9by step 1000 and remains there through training across all experiments, with the exception of 200M runs where prolonged training after overfitting causes a drop in the correlation\. We leave a detailed study of this exception for future work\.
- •Rank\-alignment betweenppandqq\.Spearman correlation \([Figure˜12\(b\)](https://arxiv.org/html/2605.28975#A3.F12.sf2)\) shows that activation distribution and weight distribution exhibit larger rank\-alignment with Muon as compared to Adam\. In[Figure˜10](https://arxiv.org/html/2605.28975#A3.F10)and[Figure˜11](https://arxiv.org/html/2605.28975#A3.F11), this reflects asqqtaking larger values along the tail ofppin the case of Adam, but not in the case of Muon\.

Table 3:Ratio of the top two singular valuesσ1/σ2\\sigma\_\{1\}/\\sigma\_\{2\}of the unembedding matrix at five training steps after initialization, for 3B pre\-training experiments\. The post\-initialization value sits near 3 across all runs, sop1≈9​p2p\_\{1\}\\approx 9\\,p\_\{2\}\.Table 4:Ratio of the top two components of the activation distribution,q1/q2q\_\{1\}/q\_\{2\}, at five training steps after initialization, for our 3B pre\-training experiments\. Unlikeσ1/σ2≈3\\sigma\_\{1\}/\\sigma\_\{2\}\\approx 3, this ratio is typically only∼2\\sim 2, indicating thatqqis much more diffuse thanpp\.
### C\.2LAR–generalization relationship through spreads inppandqq

In[Section˜3\.1](https://arxiv.org/html/2605.28975#S3.SS1), we derivedLAR\\mathrm\{LAR\}in terms of the Pearson correlation,corr​\(p,q\)\\text\{corr\}\(p,q\), and the standard deviations,σp\\sigma\_\{p\}andσq\\sigma\_\{q\}ofppandqq, as

LAR=1\+12​logn⁡\(n​σp​σq​corr​\(p,q\)\+1n\),\\mathrm\{LAR\}=1\+\\frac\{1\}\{2\}\\log\_\{n\}\\left\(n\\,\\sigma\_\{p\}\\,\\sigma\_\{q\}\\,\\text\{corr\}\(p,q\)\+\\frac\{1\}\{n\}\\right\),whereσp2=1n​∑i=1npi2−1n2\\sigma\_\{p\}^\{2\}=\\frac\{1\}\{n\}\\sum\\limits\_\{i=1\}^\{n\}p\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\}andσq2=1n​∑i=1nqi2−1n2\\sigma\_\{q\}^\{2\}=\\frac\{1\}\{n\}\\sum\\limits\_\{i=1\}^\{n\}q\_\{i\}^\{2\}\-\\frac\{1\}\{n^\{2\}\}\.

Whenp1≫pip\_\{1\}\\gg p\_\{i\}fori≠1i\\neq 1,σp≈p1n\\sigma\_\{p\}\\approx\\frac\{p\_\{1\}\}\{\\sqrt\{n\}\}\. Similarly, whenq1≫qiq\_\{1\}\\gg q\_\{i\}fori≠1i\\neq 1,σq≈q1n\\sigma\_\{q\}\\approx\\frac\{q\_\{1\}\}\{\\sqrt\{n\}\}\. Ascorr​\(p,q\)\\text\{corr\}\(p,q\)stays approximately constant prior to overfitting,LAR\\mathrm\{LAR\}is thus a proxy forσp​σq≈p1​q1n=n2​L1−3\\sigma\_\{p\}\\sigma\_\{q\}\\approx\\frac\{p\_\{1\}q\_\{1\}\}\{n\}=n^\{2L\_\{1\}\-3\}in language\-modeling pre\-training experiments\. The last equality follows from the definition ofL1L\_\{1\}\.

In[Section˜5](https://arxiv.org/html/2605.28975#S5), we observed thatL1L\_\{1\}contains most of the pre\-overfit signature of overfitting\. The behavior can thus equivalently be described in terms of the product of spreads ofppandqq:σp​σq\\sigma\_\{p\}\\sigma\_\{q\}\. Indeed,σp\\sigma\_\{p\}in[Figure˜12\(c\)](https://arxiv.org/html/2605.28975#A3.F12.sf3)has approximately the same shape asp1p\_\{1\}in[Figure˜8\(a\)](https://arxiv.org/html/2605.28975#S5.F8.sf1)andσq\\sigma\_\{q\}in[Figure˜12\(d\)](https://arxiv.org/html/2605.28975#A3.F12.sf4)has approximately the same shape asq1q\_\{1\}in[Figure˜8\(b\)](https://arxiv.org/html/2605.28975#S5.F8.sf2)\. Furthermore, the pre\-overfit signature of overfitting is visible inσq\\sigma\_\{q\}for Adam and in bothσp\\sigma\_\{p\}andσq\\sigma\_\{q\}for Muon\. This gives a unified interpretation of the LAR–generalization behavior in both pre\-training and grokking experiments in terms of measuring the spread ofppandqq\.

## Appendix DSupplementary Results For Pre\-training Experiments

### D\.1Comparison with other spectral measures

Several spectral diagnostics of weight matrices have been proposed as indicators of generalization\. We compare unembedding LAR against two of them – stable rank and effective rank – below, computing each every10001000training steps for the ten runs of[Section˜5](https://arxiv.org/html/2605.28975#S5)\([Figure˜13](https://arxiv.org/html/2605.28975#A4.F13)\)\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x24.png)Figure 13:Stable rank \(top\) and effective rank \(bottom\) of the unembedding matrix as a function of training step, for the 3B models of[Section˜5](https://arxiv.org/html/2605.28975#S5)\. Vertical dashed lines mark the overfitting onset for each run\.#### Stable rank\.

The stable rank ofWWis defined as

srank​\(W\)=‖W‖F2‖W‖22=1p1\.\\mathrm\{srank\}\(W\)\\;=\\;\\frac\{\\\|W\\\|\_\{F\}^\{2\}\}\{\\\|W\\\|\_\{2\}^\{2\}\}\\;=\\;\\frac\{1\}\{p\_\{1\}\}\.\(10\)Neyshaburet al\.\[[2018](https://arxiv.org/html/2605.28975#bib.bib11)\]\(Theorem 1\) prove a generalization bound for ReLU networks that is monotone in stable rank, predicting that lower stable rank corresponds to better generalization\.Sanyalet al\.\[[2020](https://arxiv.org/html/2605.28975#bib.bib14)\]used this result to introduce stable rank normalization as a regularizer and show empirically that lowering stable rank improves test accuracy across a range of vision classification and GAN tasks\.

For 3B models of[Section˜5](https://arxiv.org/html/2605.28975#S5),srank​\(W\)\\mathrm\{srank\}\(W\)of the unembedding matrix is∼15−20\\sim 15\-20, as top singular values are much larger in magnitude compared to the rest \(cf\.[Figure˜13](https://arxiv.org/html/2605.28975#A4.F13), top row\)\. Hence, it does not distinguish generalization from overfitting regime\.

#### Effective rank\.

The effective rank ofWWis defined inRoy and Vetterli \[[2007](https://arxiv.org/html/2605.28975#bib.bib31)\]as

erank​\(W\)=exp⁡\(−∑i=1rp~i​log⁡p~i\),p~i=si∑jsj\.\\mathrm\{erank\}\(W\)\\;=\\;\\exp\\\!\\left\(\-\\sum\_\{i=1\}^\{r\}\\tilde\{p\}\_\{i\}\\log\\tilde\{p\}\_\{i\}\\right\),\\qquad\\tilde\{p\}\_\{i\}\\;=\\;\\frac\{s\_\{i\}\}\{\\sum\_\{j\}s\_\{j\}\}\.\(11\)Yuniset al\.\[[2024](https://arxiv.org/html/2605.28975#bib.bib15)\]track effective rank across grokking and pre\-training tasks, and conclude that generalization corresponds to smaller values of effective rank\.888They define effective rank as the entropy itself \(the logarithm oferank​\(W\)\\mathrm\{erank\}\(W\)in[Equation11](https://arxiv.org/html/2605.28975#A4.E11)\) and additionally divide by the matrix rankRR\. The qualitative direction is the same\.

[Figure˜13](https://arxiv.org/html/2605.28975#A4.F13), second row, shows that effective rank can distinguish overfitting models from non\-overfitting ones when trained with Muon\. A similar distinction also exists in the case of Adam; however, the effect is quite small\.

### D\.2Unembedding LAR dependence on learning rate schedule

![Refer to caption](https://arxiv.org/html/2605.28975v1/x25.png)\(a\)Effect of warmup\. Without warmup, the LAR peak appears earlier, reflecting the larger learning rates used in the initial optimization steps\.
![Refer to caption](https://arxiv.org/html/2605.28975v1/x26.png)\(b\)Effect of decay\. Replacing cosine decay with a stable schedule \(constant learning rate after warmup\) has no noticeable effect on LAR\.

Figure 14:Effect of the learning rate schedule on unembedding LAR for the 100B dataset size\.To understand the relationship between learning rate schedule and unembedding LAR, we run two additional experiments at the 100B dataset size\.[Figure˜14\(a\)](https://arxiv.org/html/2605.28975#A4.F14.sf1)compares the baseline \(linear warmup over the first 4% of steps\) against a no\-warmup variant\. The LAR peak appears earlier without warmup, which we attribute to the larger learning rates used in the early optimization steps\. In[Figure˜14\(b\)](https://arxiv.org/html/2605.28975#A4.F14.sf2), we keep the 4% warmup but hold the learning rate constant afterward \(a stable schedule\) instead of applying cosine decay\. This has no noticeable effect on unembedding LAR, even though the stable\-schedule model reaches a noticeably worse final loss than the cosine baseline \(1\.80 vs\. 1\.75\)\.

### D\.3Effect of weight tying on unembedding LAR

The models in[Section˜5](https://arxiv.org/html/2605.28975#S5)use untied embeddings, while many modern LLMs use tied embeddings\[Press and Wolf,[2017](https://arxiv.org/html/2605.28975#bib.bib30)\]\. We construct tied\-embedding variants of those models \(∼\\sim2\.5B parameters\) and train them on the 200M and 100B datasets for 2B tokens, matching the setup of[Section˜5](https://arxiv.org/html/2605.28975#S5)otherwise\.[Figure˜15](https://arxiv.org/html/2605.28975#A4.F15)shows training loss, validation loss, and unembedding LAR; to save on evaluation cost, we log validation loss every 1000 steps here instead of every 20\. The overall shape of LAR in the overfitting and non\-overfitting regimes matches our earlier findings, but peak LAR is noticeably smaller \(∼\\sim0\.56 vs\.∼\\sim0\.676–0\.688 in[Section˜5](https://arxiv.org/html/2605.28975#S5)\)\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x27.png)Figure 15:Training loss, validation loss, and unembedding log\-alignment ratio \(LAR\) as functions of training steps for 2\.5B models trained on 100B and 200M dataset sizes for 2B token budgets\. All models used tied embeddings\.
### D\.4Effect of weight decay on unembedding LAR

The models trained with Adam in[Section˜5](https://arxiv.org/html/2605.28975#S5)used no weight decay, while modern LLMs are typically trained with AdamW\[Loshchilov and Hutter,[2019](https://arxiv.org/html/2605.28975#bib.bib22)\]\. To check that this choice does not affect our conclusions, we repeat two experiments of[Section˜5](https://arxiv.org/html/2605.28975#S5)with AdamW and weight decay0\.10\.1\. As shown in[Figure˜16](https://arxiv.org/html/2605.28975#A4.F16), the addition of weight decay has a negligible effect on the unembedding LAR trajectory\. Validation loss is logged every 1000 steps instead of every 20 to reduce evaluation compute\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x28.png)Figure 16:Training loss, validation loss, and unembedding log\-alignment ratio \(LAR\) as functions of training steps for 3B models trained with AdamW on 100B and 200M dataset sizes for 2B token budgets\.
### D\.5LAR for other matrices in the network

LAR is well\-defined for any dense matrix in a neural network\. For the models in[Section˜5](https://arxiv.org/html/2605.28975#S5), we plot LAR for each attention matrix in[Figure˜17](https://arxiv.org/html/2605.28975#A4.F17)\(Adam\) and[Figure˜18](https://arxiv.org/html/2605.28975#A4.F18)\(Muon\), and for each MLP matrix in[Figure˜19](https://arxiv.org/html/2605.28975#A4.F19)\(Adam\) and[Figure˜20](https://arxiv.org/html/2605.28975#A4.F20)\(Muon\)\. We omit the 1\.6B run since its behavior closely matches the 100B run within the 2B\-token training budget\.

We make the following observations:

1. 1\.LAR is approximately12\\tfrac\{1\}\{2\}at initialization for every matrix in the network\.
2. 2\.Adam produces higher LAR than Muon at every internal matrix\. The mean peak LAR under Adam exceeds that under Muon by roughly0\.050\.05at both attention and MLP matrices, in every layer\. We attribute this to Muon’s spectral updates for internal matrices, which make the spectral distribution more uniform compared to Adam\. \(Equivalently,σp\\sigma\_\{p\}is smaller for matrices optimized with Muon as compared to Adam\.\)999Liuet al\.\[[2025](https://arxiv.org/html/2605.28975#bib.bib32)\]report higher SVD entropy of weight matrices under Muon than AdamW, consistent with this observation\.
3. 3\.The overfitting signature appears at a few other matrices in addition to the unembedding layer\. The most notable is the divergence in LAR of the attention output matrices of earlier layers around overfitting epochs\.

We leave a deeper analysis of LAR for other matrices to future work\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x29.png)Figure 17:Log\-alignment ratios of attention matrices in Adam runs\. Each column shows one of the four attention projections \(WqW\_\{q\},WkW\_\{k\},WvW\_\{v\},WoW\_\{o\}\); each row shows one of the eight transformer layers\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x30.png)Figure 18:Log\-alignment ratios of attention matrices in Muon runs\. Each column shows one of the four attention projections \(WqW\_\{q\},WkW\_\{k\},WvW\_\{v\},WoW\_\{o\}\); each row shows one of the eight transformer layers\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x31.png)Figure 19:Log\-alignment ratios of MLP matrices in Adam runs\. Each column shows one of the three MLP projections \(WgateW\_\{\\text\{gate\}\},WupW\_\{\\text\{up\}\},WdownW\_\{\\text\{down\}\}\); each row shows one of the eight transformer layers\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x32.png)Figure 20:Log\-alignment ratios of MLP matrices in Muon runs\. Each column shows one of the three MLP projections \(WgateW\_\{\\text\{gate\}\},WupW\_\{\\text\{up\}\},WdownW\_\{\\text\{down\}\}\); each row shows one of the eight transformer layers\.

## Appendix ESupplementary Results For Grokking Experiments

### E\.1Additional Examples ofppandqq

In[Figure˜2](https://arxiv.org/html/2605.28975#S4.F2), we presented a few examples ofppandqqat different steps during training\. Here, we present additional examples for grokking models \([Figure˜21](https://arxiv.org/html/2605.28975#A5.F21)\) and non\-grokking models \([Figure˜22](https://arxiv.org/html/2605.28975#A5.F22)\)\. The assumptions of property 3 from[Section˜3\.2](https://arxiv.org/html/2605.28975#S3.SS2), i\.e\.pi≈1kp\_\{i\}\\approx\\frac\{1\}\{k\}fori≤ki\\leq kandqi=0q\_\{i\}=0fori\>ki\>k, hold true to varying degree at final checkpoints\.ppandqqare invariably concentrated on the samekkdirections, but the uniformity condition onppis task\-dependent\. As exhibited in[Figure˜3\(a\)](https://arxiv.org/html/2605.28975#S4.F3.sf1),k≈n2​\(1−LAR\)k\\approx n^\{2\(1\-\\mathrm\{LAR\}\)\}approximatesk95k\_\{95\}quite well, nevertheless\. It is also clear thatppandqqare more diffuse at memorization and more concentrated at grokking and final checkpoints in all cases\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x33.png)\(a\)Modular addition, 50% training data
![Refer to caption](https://arxiv.org/html/2605.28975v1/x34.png)\(b\)x/yx/yorx−yx\-y
![Refer to caption](https://arxiv.org/html/2605.28975v1/x35.png)\(c\)Modular subtraction
![Refer to caption](https://arxiv.org/html/2605.28975v1/x36.png)\(d\)x⋅y⋅xx\\cdot y\\cdot xforx,y∈S5x,y\\in S\_\{5\}

Figure 21:Weight distributionppand activation distributionqqat memorization, grokking, and final points during training for tasks that grok\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x37.png)\(a\)Modular addition, 50% training data
![Refer to caption](https://arxiv.org/html/2605.28975v1/x38.png)\(b\)x/yx/yorx−yx\-y

Figure 22:Weight distributionppand activation distributionqqat memorization and final points during training for tasks that do not grok\.
### E\.2Additional Plots of Unembedding LAR and Validation Accuracy

In[Figure˜1](https://arxiv.org/html/2605.28975#S4.F1)of the main text, we showed mean unembedding LAR and mean validation accuracy across different seeds for the two sets of grokking experiments\. Here, we provide these plots for individual seeds in[Figure˜23](https://arxiv.org/html/2605.28975#A5.F23)and[Figure˜24](https://arxiv.org/html/2605.28975#A5.F24)respectively\.

![Refer to caption](https://arxiv.org/html/2605.28975v1/x39.png)\(a\)Seed 1
![Refer to caption](https://arxiv.org/html/2605.28975v1/x40.png)\(b\)Seed 2
![Refer to caption](https://arxiv.org/html/2605.28975v1/x41.png)\(c\)Seed 3

Figure 23:Unembedding LAR and validation accuracy across the 12 binary tasks of[Section˜4](https://arxiv.org/html/2605.28975#S4)for three representative seeds\. The memorization\-to\-grokking pattern seen in the mean across seeds \([Figure˜1\(a\)](https://arxiv.org/html/2605.28975#S4.F1.sf1)\) is also apparent in individual runs\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x42.png)\(a\)Seed 1
![Refer to caption](https://arxiv.org/html/2605.28975v1/x43.png)\(b\)Seed 2
![Refer to caption](https://arxiv.org/html/2605.28975v1/x44.png)\(c\)Seed 3

Figure 24:Unembedding LAR and validation accuracy across the training\-dataset\-fraction\-sweep experiments of[Section˜4](https://arxiv.org/html/2605.28975#S4)for three representative seeds\. The memorization\-to\-grokking pattern seen in the mean across seeds \([Figure˜1\(b\)](https://arxiv.org/html/2605.28975#S4.F1.sf2)\) is also apparent in individual runs\.![Refer to caption](https://arxiv.org/html/2605.28975v1/x45.png)\(a\)12 binary operation tasks\.
![Refer to caption](https://arxiv.org/html/2605.28975v1/x46.png)\(b\)Modular addition with varying training dataset fraction\.

Figure 25:Unembedding LAR \(blue\) and validation accuracy \(orange\) vs\. epochs, with mean and ±1 std shaded bands across 5 seeds\. In contrast to[Figure˜1](https://arxiv.org/html/2605.28975#S4.F1), where x\-axis is clipped, the full range of training epochs is shown here\. The convergence of LAR under prolonged training is hence more apparent\.