To Grok Grokking: Provable Grokking in Ridge Regression

SAI replication review · Referee report

Summary

The paper takes on a genuinely appealing target: give the first end-to-end mathematical proof of grokking — early overfitting, a long plateau of poor generalization, then eventual small generalization error — in a classical, transparent setting. The chosen setting is over-parameterized ridge regression with a fixed feature map, a realizable teacher $N^*(x) = \langle \theta^*, \phi(x)\rangle$ , Gaussian random initialization, and vanilla GD with weight decay. Within that scope the argument is clean: decompose $\theta = \theta_\parallel + \theta_\perp$ along and orthogonal to the row space of $\Phi$ ; observe that $\theta_\parallel$ contracts at rate $(1 - \eta\lambda_{\min}^+(\Phi^\top\Phi)/n - \eta\lambda)^t$ while $\theta_\perp$ contracts only at rate $(1-\eta\lambda)^t$ ; combine with a Rademacher-based uniform-convergence bound for the eventual generalization. The three-stage result is packaged as Theorems 4.4/4.5/4.6 and assembled into Theorem 4.2, and a Gaussian-features specialization (Eq. 8) turns the rates into explicit closed forms in $(n,m,\lambda,\nu^2)$ . Independent re-implementation of the experiments (no author code was released) matches the paper's headline numbers at the reported hyperparameters: on the ridge Figure 1 setup, training loss crosses the threshold near step $\approx 200$ while test loss does not until $\approx 2.3 \times 10^5$ ; the Eq. (8) upper and lower bounds hold at all 23 sweep points in Figure 2, with $t_2$ following $1/\lambda$ at a log-log slope of $-1.0005$ ; the Theorem 4.1 envelopes hold at every recorded step across all 50 seeds.

The main conceptual worry is scope: the mechanism proved here — slow weight-decay-driven shrinkage of an over-parameterized null-space component — is essentially a well-known behaviour of ridge regression, and it is not obvious that this is the same phenomenon as the delayed emergence of generalization on modular arithmetic or sparse-parity tasks, where feature learning, circuit formation, and lazy-to-rich transitions are the standard explanations. Fully realizable, noiseless labels leave the theorems silent on precisely the setting where grokking looks most striking; the paper acknowledges this openly, but the current framing ("provable grokking") is broader than what the proofs support. A related question is whether the result distinguishes grokking from ordinary slow convergence: making $\lambda \to 0$ inflates both $t_1$ and $t_2$ , and no lower bound on the ratio $t_2/t_1$ is stated. The step from the ridge story to non-linear networks in Section 5 is presented as suggestive rather than derived, and the replication surfaces a directional inversion in one panel and a silent reproducibility failure in Figure 5 that the paper would need to address.

Strengths

Clean, transparent setting. Ridge regression with a fixed feature map is possibly the simplest possible arena for a rigorous grokking analysis, and the paper makes good use of that simplicity — the mechanism (rate separation of $\theta_\parallel$ vs $\theta_\perp$ ) is legible and lends itself to explicit hyperparameter dependence.
End-to-end structure. Unlike prior work that proves parts of the picture (fast overfitting, or late-stage convergence, or rate transitions), the paper stitches all three stages into one theorem with matching high-probability guarantees.
Explicit quantitative bounds. Equations (6)–(8) turn the informal statement "grokking time can be made large" into concrete, closed-form dependencies on $\lambda$ , $n$ , $m$ , $\nu^2$ , which is genuinely useful for anyone trying to design controlled grokking experiments.
Careful hyperparameter empirics. Figure 2 goes beyond qualitative shape and directly overlays the theoretical bound from Eq. (8), and the linear-ridge experiments reproduce the predicted scaling of $t_2 \propto 1/\lambda$ and $t_1, t_2 \propto \ln\nu^2$ — independent re-implementation gives a $t_2$ -vs- $\lambda$ log-log slope of $-1.0005$ and near-perfect linear-in- $\ln\nu^2$ fits ( $R^2 \ge 0.9986$ ).
Theorem 4.1 envelopes hold quantitatively. The training-loss upper bound and population-loss lower bound of Theorem 4.1 hold at every recorded step across 50 independent seeds in the reproduction, with the empirical training-loss decay exponent about $950\times$ larger than $\eta\lambda$ and the empirical population-loss exponent equal to $\eta\lambda$ to five significant figures — a strong quantitative confirmation of the rate-separation mechanism.
Honest scope reporting. The paper explicitly names its open problems (non-realizable ridge, agnostic random features, lazy-to-rich transitions) and identifies where its techniques break, which is more useful for follow-up work than a broader claim.

Weaknesses

Realizable, noiseless setting is a strong scope restriction. The proofs only speak to fully realizable teachers. Classical grokking observations (Power et al., 2022) sit alongside noisy or agnostic settings, and Figure 6 (label noise) and Figure 3 (non-realizable random features) both show grokking-like behaviour that the theorems do not cover. The formal contribution should be positioned more precisely against this gap.
Ambiguity about whether the mechanism is "grokking". The proved mechanism is a rate separation intrinsic to over-parameterized ridge regression. It is unclear whether this is the same phenomenon studied in modular arithmetic / sparse parity, where feature learning and lazy-to-rich transitions are the operative story. The paper does not really argue that the two are the same.
No lower bound on the ratio $t_2/t_1$ . The construction that makes $t_2 - t_1$ arbitrarily large also makes both $t_1$ and $t_2$ large in absolute terms. Grokking is usually understood as a disproportionate delay; without an explicit bound on the ratio, the theorems technically allow the null hypothesis "regularization is weak so learning is slow uniformly".
Marchenko–Pastur invocation is non-standard. The limit $\lambda_{\min}^+(\Phi^\top\Phi) \to \gamma^{-1}(1-\sqrt\gamma)^2$ in Remark 4.3 does not obviously match the standard MP edges for the declared normalisation ( $\phi$ entries variance $1/m$ , $\gamma = m/n > 1$ ). Either the exact scaling should be derived, or the formula should be replaced by the standard $(\sqrt\gamma - 1)^2$ form with clearly stated normalizations.
Step-size condition uses a loose smoothness surrogate. The stability bound $\eta \le 2/(L + 2\lambda)$ uses $L = \frac{1}{n}\|\Phi\|_F^2$ rather than the top eigenvalue of $\frac{1}{n}\Phi^\top\Phi$ . In the over-parameterized regime this can be substantially conservative, and the experiments (which use $\eta = 1$ ) likely violate the stated condition even though GD is empirically stable.
Overclaims in framing. "Fully and quantitatively controlled" is stronger than the one-sided lower bounds actually established, and the abstract's "amplified or eliminated" claim is only demonstrated on the amplification side — no panel shows the plateau disappearing, and re-running the Figure 2 $\lambda$ sweep confirms that even at the largest swept $\lambda = 10^{-3}$ the gap $t_2/t_1$ is still $\approx 12$ , more than a decade. Similarly, the "misspecified to well-specified prior" narrative in the discussion is asserted rather than derived from the analysis.
Section 5.2 handwaves the residual approximation error. Framing the random-features experiment as "arbitrarily close to realizable" is not the same as realizable — the residual training loss at any finite $m$ is positive, and Theorem 4.6 gives $L \le 2 L_n + \epsilon$ , which the tight generalization narrative depends on being small. The apparent contradiction with the Section 6 statement that "arbitrarily small weight decay in the agnostic setting may impede generalization entirely" is not reconciled.
The Xu et al. (2024) and Boursier et al. (2025) contrasts are load-bearing but under-argued. Both comparisons are one sentence each, and it is not clear whether the gap being closed is a proof-technique gap or a behavioural gap. Making these more specific would strengthen the novelty claim.
Local expositional slips. $\theta_\perp$ is described as "remaining close to its initialization" while the eventual-generalization guarantee requires it to decay; the loss-rate $(1-\eta\lambda)^{2t}$ and parameter-rate $(1-\eta\lambda)^t$ are used side by side without signposting; the identification of the convergence target with $\theta^* = \mathbf{0}$ implicitly assumes $\Sigma$ is full rank. A duplicate "Left Lower" appears in the Figure 3 caption.

Reproducibility & code

The manuscript ships no code, so the observations below come from an independent re-implementation reconstructed from the paper text and run at the stated hyperparameters. The picture is broadly positive — the theoretical experiments are amenable to reconstruction and the headline claims reproduce — but a few underspecified choices caused silent failures.

Headline linear-ridge results reproduce cleanly. Figures 1 and 2 come out as claimed at the reported settings: $t_1 \approx 197$ vs $t_2 \approx 2.3 \times 10^5$ over 50 seeds on the Figure 1 Left setup; the Eq. (8) upper and lower bounds hold at all 23 sweep points; the $t_2$ -vs- $\lambda$ log-log slope is $-1.0005$ ; the Theorem 4.1 envelopes hold at every recorded step across all 50 seeds; and the $\nu^2$ dependence is near-perfectly linear in $\ln\nu^2$ ( $R^2 \ge 0.9986$ ).
Figure 5's teacher is missing from the specification. Appendix B gives $m, d, n, \eta, \lambda, \epsilon$ but not the teacher. Under the natural reading — inheriting the non-realizable single-ReLU-neuron teacher of Section 5.2 — the surrogate-accuracy curve fails silently to leave the low plateau, staying in $[0.088, 0.208]$ for the full $10^7$ steps. Only after switching to a realizable linear teacher over the same feature map (implied by Appendix B's phrase "linear regression") does the figure's phenomenon appear. This should be a one-line clarification in the caption.
Figure 3 population loss floors above the paper's $c = 0.01$ threshold. Because the Section 5.2 teacher is non-realizable, the test loss never reaches $c$ in any of the 16 replicated sweep configurations, so the paper's own definition of $t_2$ is unresolved throughout. The qualitative delayed-descent shape is still visible, but Figure 3 does not exhibit the sharp late collapse of the realizable ridge picture.
Figure 4 has a directional discrepancy on the $\nu^2$ panel. Three of four Figure 4 dependencies match Figure 2 in the reproduction, but increasing $\nu^2$ decreases $t_1$ in Figure 4 (2501 / 1914 / 1196 for $\nu^2 = 0.25/1/4$ ) where Figure 2 raises it (123 to 352 across $\nu^2 = 0.25$ to $8$ ). The trend is monotone with roughly non-overlapping run-to-run spread, so it is not noise; the likely cause is the hidden-layer init $w_j \sim \mathcal{N}(\mathbf{0}, \nu^2/d \cdot I_d)$ , which makes $\nu^2$ scale feature magnitudes as well as output-layer scale. The paper's "qualitative match" claim needs to be narrowed to $\lambda, n, m$ or the mismatch should be discussed.
Elimination of grokking (abstract's second-half claim) is not evidenced. The Figure 2 $\lambda$ sweep stops at $10^{-3}$ , at which $t_2/t_1 \approx 12$ (still a full decade of visible plateau), and the codebase's own diagnostic already flags $\lambda = 10^{-3}$ as outside the required rate-separation regime $\lambda \ll \lambda_{\min}^+(\Phi^\top\Phi)/n$ . No no-grokking setting is exercised in either theory or experiment.
Fixed $\eta = 1$ silently constrains the admissible $\nu^2$ range. For the random ReLU features setting, $\nu^2 = 16$ at the reported $\eta = 1$ violates the Theorem 4.1 step-size condition and losses diverge to $\sim 10^{305}$ within a few hundred steps. Section 5.2 states no $\nu^2$ range and no stability condition, so extending the sweep past $\nu^2 = 4$ is a trap for reproducers.
Label-noise "insensitivity" holds only for low std. At the paper's std $\in \{0, 0.01, 0.1\}$ the median $t_2$ values are near-identical ( $\approx 133$ k, $134$ k, $170$ k); extending to std $\in \{0.5, 1\}$ (values the paper does not claim) makes the test loss rise above $c$ and $t_2$ becomes unresolved — a boundary the caption should acknowledge.
Setup-paragraph specification is incomplete. Figure 1's caption lists $n, m, \lambda, \nu^2, \eta, d$ but not the input/feature distribution (Section 5.1's $\mathcal{N}(\mathbf{0}, I_m/m)$ has to be assumed); the number of runs and dispersion-band definitions are stated for Figure 1 (50 runs) but not for Figures 2, 3, 4; the test-loss estimation protocol (analytical via $\Sigma$ vs held-out sample) is unstated for the non-linear experiments; the Figure 6 caption does not state the pre-noise teacher.

Recommended Changes

Essential

Clarify scope of "grokking" claim. Add a sentence to the introduction and abstract distinguishing the mechanism the paper proves (over-parameterized ridge null-space slow-mode) from the feature-learning / lazy-to-rich stories in the modular-arithmetic grokking literature — as raised under "Ambiguity about whether the mechanism is 'grokking'" and "Realizable, noiseless setting is a strong scope restriction".
State an explicit bound on $t_2/t_1$ (or its absence). As raised under "No lower bound on the ratio $t_2/t_1$ ", either prove that the ratio grows as $\lambda \to 0$ (arguing this is genuine grokking rather than uniform slow convergence) or state clearly that the current bounds do not separate the two.
State the Figure 5 teacher explicitly. As raised under "Figure 5's teacher is missing from the specification", add a one-line statement that Figure 5 uses a realizable linear teacher over the specified feature map, and note that realizability is what makes the surrogate-accuracy plateau appear.
Reduce framing overclaims and add an elimination panel (or drop the claim). As raised under "Overclaims in framing" and "Elimination of grokking is not evidenced", either add an experimental panel showing $t_2/t_1$ collapsing under sufficiently large $\lambda$ (with a matching theoretical prediction), or drop "eliminated" from the abstract and downgrade "fully and quantitatively controlled" in Section 1.
Reconcile the small- $\lambda$ agnostic claim in Section 6 with the empirics. As raised under "Section 5.2 handwaves the residual approximation error" and "Section 6's negative claim about the agnostic case is too vague", give the mechanistic reason (e.g., $\lambda \to 0$ recovers the ridgeless interpolator that overfits noise) so the apparent conflict with Figures 3 and 6 is resolved.
Fix or derive the Marchenko–Pastur limit in Remark 4.3. As raised under "Marchenko–Pastur invocation is non-standard", either replace with the standard $(\sqrt\gamma-1)^2$ form under a clearly stated normalisation for $\Phi^\top\Phi$ , or add the two-line derivation that yields the currently-stated $\gamma^{-1}(1-\sqrt\gamma)^2$ .
Discuss the Figure 4 $\nu^2$ directional discrepancy or narrow the qualitative-match claim. As raised under "Figure 4 has a directional discrepancy on the $\nu^2$ panel", either explain why the two-layer ReLU network moves $t_1$ opposite to the ridge case under increasing $\nu^2$ , or narrow the Section 5.3 claim to the $\lambda, n, m$ panels.
Release code with the exact Eq. (8) plotting formulas, sweep grids, seeds, and $t_1/t_2$ extraction rules. As raised throughout Reproducibility & code — reconstruction from the manuscript alone works but is fragile at the specifications above.

Suggested

Sharpen the Xu et al. (2024) and Boursier et al. (2025) contrasts. Add one sentence each specifying which part of their result is missing (proof-technique gap vs behavioural gap), as raised under "The Xu et al. (2024) and Boursier et al. (2025) contrasts are load-bearing but under-argued".
State input/feature distributions, run counts, dispersion bands, and test-loss protocol in every figure caption. In-place statements for Figures 1, 3, 4, 5, 6, 7 covering $\mathcal{D}_x$ , teacher distribution where applicable, number of runs, dispersion-band definition, and the analytical-vs-held-out test-loss estimator (as raised throughout Reproducibility & code).
Note the step-size gap between theory and experiment. Comment that the trace-based bound $\eta \le 2/(L+2\lambda)$ is conservative relative to the operator norm, and that experiments with $\eta = 1$ still lie in the intended regime (as raised under "Step-size condition uses a loose smoothness surrogate").
Add a stability caveat for the Section 5.2 $\nu^2$ sweep. State the admissible $\nu^2$ range at the fixed $\eta = 1$ so re-implementers do not step into the divergent $\nu^2 = 16$ regime (as raised under "Fixed $\eta = 1$ silently constrains the admissible $\nu^2$ range").
Signpost loss-rate vs parameter-rate. One sentence clarifying that $(1-\eta\lambda)^{2t}$ is a loss rate while $(1-\eta\lambda)^t$ is a parameter rate would prevent readers from hunting for a nonexistent inconsistency (as raised under "Local expositional slips").
Qualify " $\theta_\perp$ remains close to its initialization" and the identification of the convergence target with $\theta^* = \mathbf{0}$ . These are correct with mild qualifiers ("for many steps", "under $\lambda_{\min}(\Sigma) > 0$ ") that would prevent misreading (as raised under "Local expositional slips").
Soften "misspecified to well-specified prior" to interpretive language. Replace "Hence" with "This can be viewed as", since the correspondence is a narrative rather than a formal consequence (as raised under "Overclaims in framing").
Fix the duplicate "Left Lower" in the Figure 3 caption. Trivial correction, as raised under "Local expositional slips".
Qualify the label-noise "insensitivity" claim to the low-noise range. Note that the claim holds within the range shown and that at higher std the threshold-based $t_2$ definition ceases to resolve (as raised under "Label-noise insensitivity holds only for low std").
Cite non-asymptotic bounds on $\lambda_{\min}^+(\Phi^\top\Phi)$ where they apply. Bai–Yin / Silverstein and Rudelson–Vershynin bounds give finite- $n$ information the paper currently disclaims (as raised under " $\lambda_{\min}^+(\Phi^\top\Phi)$ is treated as unknown while known non-asymptotic bounds apply").
Note the " $L \le 2 L_n + \epsilon$ " leading factor. Point out that item (iii)'s "arbitrarily small generalization error" claim depends on realizability making $L_n(\theta_\lambda^*) \to 0$ , and that the factor of 2 would matter outside the realizable case (as raised under "Uniform-convergence bound gives $L \le 2 L_n + \epsilon$ ").