Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Technical Summary: Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Problem Statement
In high-dimensional, low-sample regimes (common in genomics, medicine, and finance), machine learning models often achieve near-perfect training performance (interpolation) while failing to generalize to unseen data. This gap arises because the optimization landscape contains many parameter configurations that yield identical training loss but differ significantly in complexity and test performance. While regularization methods (e.g., weight decay) are traditionally used to bias optimization toward simpler solutions, the phenomenon of "grokking" reveals a temporal disconnect: models often memorize data first and only transition to generalization much later. This delayed transition is computationally inefficient, requiring orders of magnitude more epochs for generalization than for fitting. The core problem is how to efficiently select a low-complexity interpolating solution immediately after the model reaches the interpolation threshold, rather than relying on slow, implicit dynamics.

Methodology
The authors propose GROKtimizer, a biphasic optimization strategy designed to explicitly separate the fitting phase from the generalization phase.

1. Theoretical Framework:

   • Post-Interpolation Dynamics: The paper models the post-interpolation regime as a local quadratic dynamical system. Under the assumption that the loss landscape near an interpolating solution can be approximated by a quadratic form involving the Hessian matrix ($H$) and an $L_2$ regularization term ($\lambda$), the dynamics along the flat directions (null space of $H$) resemble a damped oscillator.
   • Critical Damping: The authors demonstrate that Critically Damped Momentum (CDM) provides the optimal trajectory for minimizing the parameter norm within this flat basin. By tuning the momentum parameter $\beta$ to $\beta = 1 - 2\sqrt{\lambda\eta}$ (where $\eta$ is the learning rate), the optimizer achieves critical damping.
   • Convergence Speedup: Theoretically, CDM offers a quadratic speedup over classical Gradient Descent (GD) in this regime. While GD converges with a timescale proportional to $\lambda^{-1}$, CDM converges with a timescale proportional to $\lambda^{-1/2}$. This results in a provable quadratic reduction in the number of epochs required to reach the minimum-norm solution.
   • Learning Rate Estimation: To implement CDM practically without computing the full Hessian, the authors propose using an Aitken-corrected Power Iteration method to estimate the largest eigenvalue ($\lambda_{max}$) of the Hessian. This allows for the online calculation of the optimal learning rate $\eta_{max}$ at the edge of stability.

2. Algorithmic Implementation (GROKtimizer):

   • Phase 1 (Interpolation): The optimizer runs with zero weight decay to prioritize rapid convergence to the interpolation threshold. This phase utilizes standard momentum to fit the training data quickly.
   • Phase 2 (Generalization): Once interpolation is reached, the optimizer switches to Critically Damped Momentum with weight decay. The momentum parameter is tuned to the critical damping condition derived in the theory, and weight decay is applied to drive the parameters toward the minimum-norm solution within the interpolating manifold.

Key Contributions

1. Characterization of Post-Interpolation Regime: The authors formally characterize the delayed generalization phase as a damped oscillator problem along the flat directions of the loss landscape, linking the phenomenon of grokking to the geometry of the Hessian's null space.
2. Theoretical Speedup: They prove that critically damped momentum accelerates convergence to low-norm interpolating solutions quadratically faster than standard gradient descent, establishing optimality among first-order methods under the local quadratic model.
3. GROKtimizer: They introduce a practical, biphasic optimization schedule that dynamically switches from unregularized rapid fitting to critically damped norm minimization.
4. Empirical Validation: The method is validated across synthetic grokking benchmarks and real-world high-dimensional datasets, demonstrating consistent improvements over standard optimizers (Adam, AdamW, SGD, Muon).

Experimental Results
The paper evaluates GROKtimizer on several benchmarks:

• Synthetic Benchmarks: On tasks like Modular Addition, Binary Addition, and Sparse Parity, GROKtimizer significantly outperforms baselines. For instance, on the Gaussian task, it reduced validation loss by several orders of magnitude compared to the strongest baseline.
• Real-World High-Dimensional Data: In biomedical classification tasks (Leukemia, TCGA) and molecular property prediction (QM9), GROKtimizer achieved the highest accuracy and lowest validation loss. The gains were most pronounced in data-scarce settings where the number of features far exceeds the number of samples.
• Sample Efficiency: In low-data MNIST experiments (10–50 samples per class), GROKtimizer showed the largest advantage, confirming that explicit norm minimization improves sample efficiency when the solution space is underdetermined.
• Language Model Pretraining: Experiments on WikiText2 and BabyLM models showed that while the direct correlation between parameter norm and validation loss is less clear in transformers (due to scale symmetries and adaptive updates), the biphasic pattern still holds. Switching to the norm-minimization phase after interpolation produced further reductions in validation loss.

Significance and Claims
The paper claims that GROKtimizer offers a principled solution to the "grokking" problem by treating post-interpolation training as an explicit optimization problem rather than a passive waiting period. By leveraging the geometry of the loss landscape, the method accelerates the selection of low-complexity solutions.

The authors reconcile their findings with the flat-minima hypothesis, arguing that the dimensionality of the Hessian's null space serves as a robust measure of flatness and complexity. They suggest that minimizing the norm within this null space is equivalent to navigating toward flatter, more generalizable minima.

Limitations
The authors acknowledge several limitations:

• The approach is primarily designed for data-constrained regimes and may be less critical in large-scale vision or language pretraining where data abundance dominates performance.
• The method assumes the model has sufficient capacity to reach interpolation; it is not a solution for under-parameterized models.
• The theoretical analysis relies on a local quadratic approximation of the loss landscape, though empirical results suggest the method is robust to deviations from this assumption.
• Parameter norm is used as a proxy for complexity, which may not capture all aspects of model generalization in every architecture.