Fast Generalization after Interpolation via Critically Damped Momentum Optimization

2026-05-31

AI Summary

This work addresses the poor generalization of interpolating solutions in high-dimensional, small-sample settings by proposing GROKtimizer, a two-stage optimization method. In the first stage, the optimizer rapidly converges to an interpolating solution; in the second stage, it employs critically damped momentum (CDM) to minimize the norm of the solution, thereby enhancing generalization. Notably, this is the first application of CDM for complexity control after interpolation, achieving a quadratic acceleration over standard gradient descent under a locally quadratic model while retaining the theoretical optimality of first-order methods. Empirical evaluations demonstrate that GROKtimizer significantly improves generalization on both synthetic grokking tasks and multiple real-world datasets, providing strong support for the flat minima hypothesis.

Abstract

A central problem in machine learning is that models can achieve near-perfect training performance while generalizing substantially less well to unseen examples. This gap is especially acute in high-dimensional, low-sample regimes, where many interpolating solutions exist and optimization must implicitly select among minima with different generalization properties. Following recent theoretical advances on optimization dynamics near the interpolation threshold, we note that the two-regime structure of risk minimization, with loss minimization followed by complexity minimization, motivates a biphasic optimization schedule. We thus theoretically demonstrate that GROKtimizer, a biphasic strategy that combines rapid convergence to interpolation with Critically Damped Momentum (CDM)-based post-interpolation norm minimization, offers a natural solution for selecting low-norm interpolating solutions. Under a local quadratic model of the post-interpolation basin, GROKtimizer provides a quadratic speedup over classical gradient descent, with provable optimality among first-order optimizers. To showcase the applicability of our method, we evaluate GROKtimizer on several synthetic benchmarks common in the classical grokking literature and on various real-world datasets. Finally, we reconcile our findings with the flat-minima hypothesis, highlighting the importance of post-interpolation dynamics in the construction of high-quality, generalizing models.

Problem

• generalization
• interpolation
• optimization
• high-dimensional
• flat minima

Innovation

• Critically Damped Momentum
• biphasic optimization
• interpolation
• norm minimization
• generalization