Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Technische Samenvatting: Snelle Generalisatie na Interpolatie via Kritisch Gedempte Momentum-optimalisatie

Probleemstelling
In hoogdimensionele, laag-steekproef regimes (veelvoorkomend in genomica, geneeskunde en financiën), bereiken machine learning-modellen vaak bijna perfecte trainingsprestaties (interpolatie), terwijl ze falen in het generaliseren naar ongeziene data. Deze kloof ontstaat omdat het optimalisatielandschap veel parameterconfiguraties bevat die identieke trainingsverliezen opleveren, maar significant verschillen in complexiteit en testprestaties. Hoewel regularisatiemethoden (bijv. weight decay) traditioneel worden gebruikt om de optimalisatie te sturen naar eenvoudigere oplossingen, onthult het fenomeen van "grokking" een temporele disconnectie: modellen memoriseren vaak eerst de data en gaan pas veel later over naar generalisatie. Deze vertraagde transitie is computationeel inefficiënt en vereist orders van grootte meer epochs voor generalisatie dan voor het fitten. Het kernprobleem is hoe men direct na het bereiken van de interpolatiedrempel een oplossing met lage complexiteit efficiënt kan selecteren, in plaats van te vertrouwen op trage, impliciete dynamiek.

Methodologie
De auteurs stellen GROKtimizer voor, een bifasische optimalisatiestrategie die ontworpen is om de fit-fase expliciet te scheiden van de generalisatiefase.

1. Theoretisch Kader:

   • Post-Interpolatie Dynamiek: Het artikel modelleert de post-interpolatie regime als een lokaal kwadratisch dynamisch systeem. Onder de aanname dat het verlieslandschap nabij een interpolerende oplossing benaderd kan worden door een kwadratische vorm bestaande uit de Hessiaan-matrix ($H$) en een $L_2$-regularisatieterm ($\lambda$), lijken de dynamiek langs de vlakke richtingen (nulruimte van $H$) op een gedempte oscillator.
   • Kritische Demping: De auteurs demonstreren dat Kritisch Gedempte Momentum (CDM) de optimale traject biedt voor het minimaliseren van de parameternorm binnen deze vlakke basin. Door de momentumparameter $\beta$ af te stemmen op $\beta = 1 - 2\sqrt{\lambda\eta}$ (waarbij $\eta$ de leersnelheid is), bereikt de optimizer kritische demping.
   • Convergentie Versnelling: Theoretisch biedt CDM een kwadratische versnelling ten opzichte van klassieke Gradient Descent (GD) in dit regime. Terwijl GD convergeert met een tijdschaal proportioneel aan $\lambda^{-1}$, convergeert CDM met een tijdschaal proportioneel aan $\lambda^{-1/2}$. Dit resulteert in een bewezen kwadratische reductie in het aantal benodigde epochs om de minimum-norm oplossing te bereiken.
   • Schatting van de Leersnelheid: Om CDM praktisch te implementeren zonder de volledige Hessiaan te berekenen, stellen de auteurs een Aitken-gecorrigeerde Power Iteration-methode voor om de grootste eigenwaarde ($\lambda_{max}$) van de Hessiaan te schatten. Dit maakt de online berekening mogelijk van de optimale leersnelheid $\eta_{max}$ aan de rand van stabiliteit.

2. Algoritmische Implementatie (GROKtimizer):

   • Fase 1 (Interpolatie): De optimizer draait met nul weight decay om snelle convergentie naar de interpolatiedrempel te prioriteren. Deze fase maakt gebruik van standaard momentum om de trainingsdata snel te fitten.
   • Fase 2 (Generalisatie): Zodra interpolatie is bereikt, schakelt de optimizer over naar Kritisch Gedempte Momentum met weight decay. De momentumparameter wordt afgestemd op de conditie van kritische demping afgeleid uit de theorie, en weight decay wordt toegepast om de parameters naar de minimum-norm oplossing binnen het interpolerende manifold te drijven.

Belangrijkste Bijdragen

1. Karakterisering van het Post-Interpolatie Regime: De auteurs karakteriseren de vertraagde generalisatiefase formeel als een gedempte oscillator-probleem langs de vlakke richtingen van het verlieslandschap, waarbij zij het fenomeen van grokking koppelen aan de geometrie van de nulruimte van de Hessiaan.
2. Theoretische Versnelling: Zij bewijzen dat kritisch gedempte momentum de convergentie naar low-norm interpolerende oplossingen kwadratisch sneller versnelt dan standaard gradient descent, waarmee zij optimaliteit vaststellen onder eerste-orde methoden onder het lokale kwadratische model.
3. GROKtimizer: Zij introduceren een praktisch, bifasisch optimalisatieschema dat dynamisch schakelt van ongeregulariseerde snelle fitting naar kritisch gedempte norm-minimalisatie.
4. Empirische Validatie: De methode wordt gevalideerd over synthetische grokking-benchmarks en real-world hoogdimensionele datasets, waarbij consistent verbeteringen worden getoond ten opzichte van standaard optimizers (Adam, AdamW, SGD, Muon).

Experimentele Resultaten
Het artikel evalueert GROKtimizer op verschillende benchmarks:

• Synthetische Benchmarks: Op taken zoals Modular Addition, Binary Addition en Sparse Parity presteert GROKtimizer significant beter dan de baselines. Zo verminderde het op de Gaussian-taak de validatieverlies met meerdere grootheden vergeleken met de sterkste baseline.
• Real-World Hoogdimensionele Data: In biomedische classificatietaken (Leukemia, TCGA) en moleculaire eigenschapsprognose (QM9) behaalde GROKtimizer de hoogste nauwkeurigheid en het laagste validatieverlies. De winsten waren het meest uitgesproken in data-arme settings waar het aantal kenmerken de hoeveelheid samples ver overstijgt.
• Steekproef-efficiëntie: In low-data MNIST-experimenten (10–50 samples per klasse) toonde GROKtimizer het grootste voordeel, wat bevestigt dat expliciete norm-minimalisatie de steekproef-efficiëntie verbetert wanneer de oplossingsruimte onderbepaald is.
• Taalmodel Pretraining: Experimenten op WikiText2 en BabyLM modellen lieten zien dat hoewel de directe correlatie tussen parameternorm en validatieverlies minder duidelijk is in transformers (door schaal-symmetrieën en adaptieve updates), het bifasische patroon nog steeds standhoudt. Het schakelen naar de norm-minimalisatiefase na interpolatie produceerde verdere reducties in validatieverlies.

Betekenis en Claims
Het artikel claimt dat GROKtimizer een principieel oplossing biedt voor het "grokking" probleem door post-interpolatie training te behanden als een expliciet optimalisatieprobleem in plaats van een passieve wachttijd. Door gebruik te maken van de geometrie van het verlieslandschap, versnelt de methode de selectie van oplossingen met lage complexiteit.

De auteurs verzoenen hun bevindingen met de flat-minima hypothese, waarbij zij argumenteren dat de dimensionaliteit van de nulruimte van de Hessiaan dient als een robuuste maatstaf voor vlakheid en complexiteit. Zij suggereren dat het minimaliseren van de norm binnen deze nulruimte equivalent is aan het navigeren naar vlakkere, meer generaliseerbare minima.

Beperkingen
De auteurs erkennen enkele beperkingen:

• De aanpak is primair ontworend voor data-beperkte regimes en kan minder kritisch zijn in grootschalige visie- of taalpretraining waar data-overvloed de prestaties domineert.
• De methode gaat ervan uit dat het model voldoende capaciteit heeft om interpolatie te bereiken; het is geen oplossing voor onder-geparameteriseerde modellen.
• De theoretische analyse steunt op een lokale kwadratische benadering van het verlieslandschap, hoewel empirische resultaten suggereren dat de methode robuust is tegen afwijkingen van deze aanname.
• Parameternorm wordt gebruikt als een proxy voor complexiteit, wat mogelijk niet alle aspecten van modelgeneralisatie in elke architectuur volledig vangt.