Fast Generalization after Interpolation via Critically Damped Momentum Optimization
2026-06-01 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem that machine learning models often fit training data perfectly but don't perform well on new data, especially when there are many possible solutions. They propose a two-step method called GROKtimizer that first quickly finds solutions that fit the training data and then refines these solutions to be simpler and potentially better at generalizing. Their method uses a special momentum technique to speed up this refinement and is proven to be faster than standard approaches. They test GROKtimizer on both simulated and real datasets and connect their results to existing ideas about why some solutions work better than others.
generalizationinterpolation thresholdoptimization dynamicsCritically Damped Momentumnorm minimizationgradient descentflat minima hypothesismachine learning optimizationhigh-dimensional data
Authors
Luca Muscarnera, Silas Ruhrberg Estévez, Yuanzhang Xiao, Mihaela Van der Schaar
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.