Source-linked AI summary
Is Grokking a Loss of Normal Hyperbolicity of the Interpolation Manifold?
Suvinava Basak
TL;DR
The paper asks whether grokking reflects a loss of normal hyperbolicity in the interpolation manifold or smooth drift along a uniformly attracting manifold. Using the smallest nonzero residual-Jacobian singular value as a diagnostic, it finds no collapse during modular-addition grokking, favoring the smooth-drift picture while leaving bifurcations unproved.
Problem
Existing work does not determine whether grokking’s sharp generalization transition reflects lost normal hyperbolicity of the interpolation manifold or smooth slow drift.
Method
The paper uses the smallest nonzero singular value of the residual Jacobian as the squared-loss diagnostic for the slowest normal restoring rate.
Results
The interpolation manifold remains normally hyperbolic through modular-addition grokking, with no global or low-spectrum collapse across seeds.
Takeaways & Limitations
The findings constrain explanations of grokking’s sharp transition by favoring smooth contraction over a bifurcation-like loss of normal hyperbolicity.
Takeaways & Limitations
The AdamW, moderately gradual-transition setting and global spectral diagnostic do not exclude bifurcations in sharper regimes or narrow test-relevant subspaces.
Abstract
from arXiv · showhide
A recent line of work recasts the post-memorization phase of grokking as constrained optimization: once a network interpolates the training set, weight decay drives a slow drift along the zero-loss manifold toward lower norm. In the language of dynamical systems, this is a fast-slow system in which the interpolation manifold plays the role of a slow manifold. We ask a question that this framing makes natural but the existing literature does not address: is the sharp generalization transition a loss of normal hyperbolicity of that manifold: a fold- or bifurcation-like event in which a normal restoring direction goes flat? Or does the manifold stay uniformly attracting while generalization happens by smooth drift? We propose a simple, optimizer-agnostic diagnostic: the smallest nonzero singular value $σ_{\min}^{+}(\mathbf J)$ of the residual Jacobian, which, for the squared loss, equals the slowest normal restoring rate of the manifold. On a two-layer ReLU network trained to grok modular addition under squared loss, $σ_{\min}^{+}(\mathbf J)$ does not collapse at the transition; it is near zero only before memorization and attains its largest values during the transition. The result holds across five seeds, and the six smallest singular values behave identically; there is no subspace-local collapse either. This is preliminary evidence against the bifurcation hypothesis and in favor of the smooth-contraction picture. We are explicit that a single-setting, gradual-transition experiment under Adam optimizer does not prove the absence of a bifurcation; it constrains where one could hide.
1 Introduction
The paper frames post-memorization grokking as fast-slow constrained optimization on an interpolation manifold and asks whether generalization reflects loss of normal hyperbolicity or smooth slow drift.
- Background: Grokking describes perfect training accuracy with near-chance test accuracy followed by generalization only after prolonged additional training.It is related to, but distinct from, double descent.
- Related work: Recent accounts describe post-memorization training as minimizing parameter norm subject to remaining on the zero-loss manifold.These accounts include constrained optimization and Riemannian norm flow formulations.
- Motivation: The shared structure is a fast process approaching interpolating solutions followed by weight-decay-driven slow motion along the interpolation manifold.The paper identifies this as a singularly perturbed, or fast-slow, dynamical system.
- Hypotheses: The paper contrasts H1, in which a normal restoring direction softens toward zero and the fast/slow separation breaks, with H0, in which the manifold remains uniformly normally hyperbolic.Under H1, the transition could resemble a fold or pitchfork; under H0, it is a feature of slow drift.
2 The diagnostic
The diagnostic uses the smallest nonzero singular value of the residual Jacobian to measure the slowest normal restoring rate of the interpolation manifold. A transition-time dip toward zero would support loss of normal hyperbolicity, whereas a value bounded away from zero would support smooth drift.
- The diagnostic: Squared loss makes the interpolation manifold an exact object with well-defined tangent and normal structure, while decoupled weight decay induces slow drift along it.For small λ, the gradient term rapidly pulls toward the manifold and −λθ supplies the slow drift.
- The diagnostic: The residual Jacobian’s row space is normal to the interpolation manifold, its kernel is tangent, and the nonzero eigenvalues of J⊤J are squared singular values.Thus, the smallest nonzero singular value determines the slowest normal restoring rate.
- The diagnostic: Normal hyperbolicity requires σ_min^+(J)^2 to remain uniformly bounded away from zero along the slow drift; loss of this condition is prerequisite for the fold or pitchfork hypothesis.Under normal hyperbolicity, Fenichel theory guarantees persistence of the perturbed slow manifold and a reduced norm-flow description.
- The diagnostic: A dip of σ_min^+(J(θ(t))) toward zero at the transition supports H1, whereas a value bounded away from zero supports H0 when read against test accuracy.The diagnostic directly tests whether normal contraction weakens during generalization.
- The diagnostic: Because σ_min^+(J) depends only on the model and data at θ, the diagnostic is defined for any optimizer.It is not restricted to the gradient-flow setting in which the slow-drift theory is cleanest.
3 Experimental setup
The experiment uses a two-layer ReLU network on modular addition with a 70%/30% split, squared loss, and full-batch optimization. Snapshots measure σ_min^+(J) separately via float32 SVD, across five initialization seeds.
- Network and training: The study trains a width-96 two-layer ReLU network without second-layer bias on modular addition mod p=11 using one-hot inputs and targets.It uses a random 70%/30% train/test split, squared loss, full-batch optimization, and Kaiming ×3.5 initialization.
- Jacobian diagnostic: Every 500 steps, the experiment computes σ_min^+(J) through a separate float32 SVD pass, avoiding training slowdown.The Jacobian diagnostic is evaluated on parameter snapshots rather than during optimization.
- Trajectory and replication: Although ReLU activation switches can change J, the resulting σ_min^+(J) trajectory is smooth at snapshot resolution and is repeated across five initialization seeds.The smoothness reflects the observed snapshot-level trajectory, not a claim that J is globally unchanged.
4 Results
The network memorizes by about step 4,000, then generalizes between roughly steps 7,000–17,000 while the smallest nonzero singular value of the residual Jacobian rises rather than collapses. Across the six smallest singular values and five seeds, the transition shows no normal-restoring-rate dip, though AdamW-specific norm behavior and timescale limitations remain caveats.
- Grokking dynamics: Training accuracy reaches 1.0 by about step 4,000, while test accuracy rises from 0.0 to ≈0.95 over roughly steps 7,000–17,000 before settling near 0.86.Test accuracy remains below the ≈0.09 chance level through about step 6,000.
- Normal restoring rate: σ_min^+(J) does not dip during the transition; it is small only before memorization and largest while the model generalizes.This result uses a two-layer ReLU network with squared loss, AdamW, and modular addition p=11.
- Normal restoring rate: ≈1 × 10−2 is the global minimum of σ_min^+(J), occurring at step 1,000 before memorization; the rate then rises monotonically and peaks mid-transition.The diagnostic registers a genuinely near-degenerate Jacobian before memorization, showing sensitivity to such collapse.
- Robustness checks: The six smallest singular values form a tight cluster bounded away from zero throughout the transition, so no low-dimensional subspace collapses either.The smallest singular value dips only pre-memorization.
- Robustness checks: 0.18–0.23 is the range of σ_min^+(J) during each seed’s transition window, with no seed showing a dip and no dip in the seed-averaged curve.Transition timing varies across seeds: test accuracy first exceeds 0.5 between steps ≈10,000 and 16,000.
- Caveats: The parameter norm stays near 10 and drifts slightly upward through the transition, while the σ_min^+(J)^2/λ timescale ratio is not reported for AdamW.AdamW’s decoupled weight decay and coordinate-wise step rescaling make the raw ratio unsuitable; a clean ratio belongs to gradient flow.
5 Discussion and limitations
The experiment favors H0: the interpolation manifold remains normally hyperbolic through grokking, with no softening of a normal restoring direction. This evidence is preliminary because AdamW, a moderately gradual transition, global spectral measures, and discrete trajectories leave sharper or localized failures untested.
- Discussion: The interpolation manifold remains robustly normally hyperbolic across grokking, and the transition’s suddenness is not mirrored by softening of a normal restoring direction.This favors H0 over H1 and supports the smooth-contraction picture.
- Discussion: The negative result also removes the motivation for a delayed-loss-of-stability route that applies only if a bifurcation exists.The cited delay law is included for context and was not independently reproduced.
- Limitations: AdamW, moderately gradual transitions, global σ_min^+(J) measurements, and discrete trajectories do not establish that sharper or localized loss of normal hyperbolicity is absent.The untested regime includes gradient descent with large initialization and small weight decay; a narrow test-relevant subspace could evade the global spectrum.
- Future experiments: The decisive follow-ups are gradient-descent tests in sharp-transition regimes, subspace-resolved curvature diagnostics, and identifying a geometric quantity that moves monotonically if no dip appears.Under AdamW, that quantity might be ℓ∞- or sparsity-based rather than ℓ2 norm.
6 Conclusion
The study asks whether grokking reflects a loss of normal hyperbolicity and uses the smallest nonzero residual-Jacobian singular value as a measurable diagnostic. In modular-addition experiments, the evidence appears to argue against that hypothesis robustly across seeds and the low end of the spectrum.
- 6 Conclusion: The smallest nonzero singular value of the residual Jacobian converts the normal-hyperbolicity question into a single measurable curve.For squared loss, it is justified as the slowest normal restoring rate.
- 6 Conclusion: The modular-addition results appear inconsistent with a loss of normal hyperbolicity, robustly across seeds and across the low end of the spectrum.The conclusion is framed as preliminary evidence rather than a definitive exclusion of the hypothesis.
- Reproducibility: Entire training and diagnostic code, Fig. 1 and 2 data, and necessary scripts are available at GitHub.The authors provide these materials for reproducibility.