Deciphering Two Training Clocks in Grokking via Deep Linear Network Theory with Conditional ReLU Reduction

Abstract: Grokking suggests that fitting the training data and learning a simple underlying rule may occur on different time scales. We formalize this phenomenon by separating the fast decay of the classification loss from the slower simplification of the learned representation, and we call the resulting pair of stopping times two training clocks. For deep linear networks, we show that a post-margin gap-growth or one-step tail-contraction condition reduces the cross-entropy loss to level epsilon on a logarithmic time scale. In contrast, when layerwise weight decay is present, the induced regularization on the end-to-end map can be expressed as a Schatten-type penalty; under a sharp late-time Kurdyka-Lojasiewicz tail, this structural energy closes on a polynomial time scale. The two clocks, therefore, separate fitting from representation simplification. We then explain how the same mechanism can appear in ReLU MLPs. In regions where the activation patterns on the training set remain fixed, the network reduces to a linear model in the active coordinates. In a two-layer ReLU embedding model, chain-rule estimates further show that the classifier head can receive larger effective gradients than the embedding block under controlled downstream norms. This supports a two-stage mechanism in which the classifier fits first, while the representation continues to simplify later. We use modular addition as the main experimental setting. The deep linear theory provides the rigorous core of the analysis. But the ReLU results are formulated as conditional reductions that account for empirical behavior without claiming a global proof for nonlinear training dynamics.

• The paper demonstrates that the classifier clock converges logarithmically, quantifying rapid reduction in training loss.
• It introduces a representation clock that converges polynomially, linking low-rank spectral simplification with delayed generalization.
• Conditional ReLU reduction extends these deep linear insights to nonlinear MLPs, explaining the grokking plateau observed in practice.

Deciphering the Two-Clock Mechanism Underlying Grokking: Deep Linear Theory and Conditional ReLU Reduction

Introduction and Motivation

This paper provides a rigorous mathematical framework for separating two temporally distinct processes underlying grokking in overparameterized neural networks: rapid classifier fitting and slow representation simplification. Empirically, in tasks such as modular addition, it is well documented that a model can achieve near-zero training error long before its test generalization improves, with the test accuracy only rising after a protracted plateau. The authors formalize this by defining two training clocks: (1) a classifier clock measuring the time required to interpolate the training set (drive the cross-entropy tail below a threshold), and (2) a representation clock that quantifies the time required for the internal map to simplify, for instance, via spectral regularization and effective low rank. The analysis is anchored in deep linear networks, where the dynamics are tractable, and is then extended via conditional reduction to stabilized ReLU MLPs. The main theoretical claims are formal convergence-time separations: the classifier clock is logarithmic in $1/\varepsilon$ (target loss), whereas the representation clock is polynomial in $1/\eta$ (target structural error). These results systematically account for the empirical delays characteristic of grokking.

Empirical Evidence for Time-Scale Separation

The authors demonstrate the time-scale-separation phenomenon using ReLU MLPs trained on modular addition modulo a prime. The model architecture utilizes trainable token embeddings for each integer, a symmetric combination layer, and a ReLU-activated MLP with a softmax output. A representative experimental run is illustrated below.

Figure 1: Loss curves for modular addition (mod 113) in a finite-width ReLU MLP. Training loss drops rapidly, whereas test loss and representation measures improve only later, in accordance with the two-clock description.

Training loss (on the left axis) decays quickly for all weight decay choices, but test loss (orange) and corresponding stable rank statistics (Figure 2) improve only after continued training. This directly motivates the study of time-separated training dynamics and the subsequent theoretical analysis.

Formalization: Two Training Clocks

Let $A_t$ denote the effective map at time $t$. The classifier clock is defined as

$$T_{\mathrm{cls}}(\varepsilon) = \inf \{ t : \mathcal{L}_{\mathrm{CE}}(A_t) \leq \varepsilon \}$$

and the representation clock as

$$T_{\mathrm{rep}}(\eta) = \inf \{ t : \mathcal{S}(A_t) \leq \eta \}$$

where $\mathcal{S}(A_t)$ is a structural gap—for example, the difference in a Schatten-$p$ quasi-norm induced by weight decay or a stable rank proxy. The authors' main claim is that these clocks are decoupled, with the classifier clock corresponding to rapid fitting and the representation clock characterizing slow simplification.

Deep Linear Surrogate: Fast Classifier Clock

The first core theoretical result concerns the classifier clock. In deep linear networks trained via cross-entropy and layerwise weight decay, once the effective classifier has reached a 'post-margin' regime where all incorrect-class logits are separated by a gap that grows linearly in training time, the tail of the cross-entropy loss contracts exponentially. This is formalized as follows: under a (dynamical) post-margin gap-growth condition, the time to reach a loss $\varepsilon$ is logarithmic in $1/\varepsilon$.

The paper emphasizes that this cannot be deduced from a fixed margin alone; a persistent dynamical margin growth or direct tail-contraction property is needed. This distinction is substantiated by precise value and gradient estimates for the softmax and its derivatives.

Slow Representation Clock: Spectral Dynamics and Low-Rank Emergence

After the classifier clock has effectively saturated, the residual training dynamics are governed by the representation clock. In the deep linear surrogate, layerwise weight decay induces a nonconvex Schatten-$1/\eta$0 penalty on the end-to-end map $1/\eta$1 with $1/\eta$2. On smooth spectral strata, the singular value dynamics are governed by

$1/\eta$3

where the second term disproportionately suppresses small singular values, generating a low-rank bias.

The convergence toward a low-rank solution is controlled by a quantitative Kurdyka–Łojasiewicz (KL) tail condition for the limiting energy gap. If the KL exponent exceeds $1/\eta$4, the rate at which the structural gap decreases is polynomial in time: $1/\eta$5 yielding the main claim: the representation clock is polynomial in $1/\eta$6, in sharp contrast to the classifier clock.

Empirical Analysis of Cyclic and Low-Rank Structure

The emergence of low-rank, rule-aligned structure is illustrated by visualizing both the learned token embeddings and the output weights after suitable label reordering, consistent with the underlying modular addition group.

Figure 3: Learned cyclic structure in token embeddings (A) and output weights (B) in a modular addition task, reflecting the gradual simplification tracked by the representation clock.

The late decrease in stable rank during training, measured in Figure 2, further corroborates the slow formation of simplified, low-dimensional representations after perfect fitting.

Figure 2: Evolution of test loss and stable rank for the modular addition task, illustrating the delay between classifier fitting and the completion of spectral simplification.

Conditional ReLU Reductions

To extend these conclusions to nonlinear settings, the authors prove conditional reduction results for ReLU networks. Under gates-stability—i.e., on intervals where the binary activation patterns (gates) for all training samples remain fixed—the network reduces to an affine (linear) subsystem, and the deep linear theory applies exactly. Additionally, chain rule analyses reveal that gradient magnitudes for head parameters are strictly larger (up to operator norm factors) than those for embedding parameters, particularly in two-layer architectures. This justifies a model in which head fitting precedes embedding and feature simplification—mirroring the observed two-stage training.

However, these results are conditional on the network reaching a stable gate regime; the theory does not guarantee that entry into such a regime occurs in every run for arbitrary widths and initializations.

Theoretical and Practical Implications

Generalization Bounds via Robust Low-Rank Representations

Building on robustness-based generalization theory, the authors formalize how effective low rank in the learned map reduces the covering number of the induced feature space, thus improving generalization guarantees for the same empirical risk. The representation clock can thus be coupled to provable generalization improvement, not merely descriptive geometric simplicity.

Implications for Training Protocols and Model Selection

The decoupled clocks framework explains why extended training—even after empirical loss reaches zero—can yield further generalization gains, especially in small-data or synthetic regimes. It provides a formal explanation of the "grokking" plateau in terms of polynomial time scales for representation simplification vis-à-vis the fast, exponential convergence of the classifier.

Limitations and Directions for Future Work

While the deep linear analysis is theorem-level and the ReLU reduction is precise under conditional stability, the theory does not encompass global nonlinear training dynamics. Additionally, low effective rank is one possible, but not necessary, route to grokking: invariance formation, neural collapse, or group-structured representations may govern other tasks. Deriving sufficient conditions for the post-margin regime and late KL geometry in finite-width, truly nonlinear networks is posed as a challenge for further theory.

Conclusion

This paper establishes a mathematically explicit mechanism for the temporal separation of classifier fitting and representation simplification during neural network training, explaining empirical grokking as a natural consequence of distinct polynomial and logarithmic convergence clocks. The analysis is compelling in the deep linear setting, and, under suitable dynamical conditions, extends to nonlinear models. This time-scale perspective reframes delayed generalization not as a dynamical paradox, but as an expected outcome of the structural regularization acting after fitting, and highlights future directions for mechanistic interpretability and training acceleration in deep learning.

---

Reference:

Tan, H., Gai, K., Zhang, S. "Deciphering Two Training Clocks in Grokking via Deep Linear Network Theory with Conditional ReLU Reduction" (2606.05863)