Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Sintesi Tecnica: Generalizzazione Rapida dopo l'Interpolazione tramite Ottimizzazione con Momentum Criticamente Smorzato

Problema
Nei regimi ad alta dimensionalità e basso numero di campioni (comuni nella genomica, in medicina e nella finanza), i modelli di machine learning spesso raggiungono prestazioni di addestramento quasi perfette (interpolazione) pur fallendo nel generalizzare su dati non visti. Questo divario deriva dal fatto che il panorama dell'ottimizzazione contiene molte configurazioni di parametri che producono la stessa perdita di addestramento (loss), ma differiscono significativamente per complessità e prestazioni di test. Mentre i metodi di regolarizzazione (ad esempio, il weight decay) vengono tradizionalmente utilizzati per orientare l'ottimizzazione verso soluzioni più semplici, il fenomeno del "grokking" rivela un distacco temporale: i modelli spesso memorizzano i dati prima e solo successivamente transitano verso la generalizzazione molto più tardi. Questa transizione ritardata è computazionalmente inefficiente, richiedendo ordini di grandezza di epoch in più rispetto all'interpolazione. Il problema centrale è come selezionare efficientemente una soluzione interpolante a bassa complessità immediatamente dopo che il modello ha raggiunto la soglia di interpolazione, piuttosto che fare affidamento su dinamiche implicite e lente.

Metodologia
Gli autori propongono GROKtimizer, una strategia di ottimizzazione bifasica progettata per separare esplicitamente la fase di fitting dalla fase di generalizzazione.

1. Quadro Teorico:

   • Dinamiche Post-Interpolazione: Il documento modella il regime post-interpolazione come un sistema dinamico quadratico locale. Sotto l'assunto che il panorama della perdita vicino a una soluzione interpolante possa essere approssimato da una forma quadratica che coinvolge la matrice Hessiana ($H$) e un termine di regolarizzazione $L_2$ ($\lambda$), la dinamica lungo le direzioni piatte (spazio nullo di $H$) assomiglia a un oscillatore smorzato.
   • Smorzamento Critico: Gli autori dimostrano che il Momentum Criticamente Smorzato (CDM) fornisce la traiettoria ottimale per minimizzare la norma dei parametri all'interno di questo bacino piatto. Regolando il parametro di momentum $\beta$ in $\beta = 1 - 2\sqrt{\lambda\eta}$ (dove $\eta$ è il tasso di apprendimento), l'ottimizzatore raggiunge lo smorzamento critico.
   • Accelerazione della Convergenza: Teoricamente, il CDM offre un'accelerazione quadratica rispetto al classico Gradient Descent (GD) in questo regime. Mentre il GD converge con una scala temporale proporzionale a $\lambda^{-1}$, il CDM converge con una scala temporale proporzionale a $\lambda^{-1/2}$. Ciò risulta in una riduzione provabile della quantità di epoch necessarie per raggiungere la soluzione a norma minima.
   • Stima del Tasso di Apprendimento: Per implementare il CDM in modo pratico senza calcolare l'intera matrice Hessiana, gli autori propongono di utilizzare un metodo di Power Iteration corretto con Aitken per stimare l'autovalore massimo ($\lambda_{max}$) della matrice Hessiana. Ciò consente il calcolo online del tasso di apprendimento ottimale $\eta_{max}$ al limite della stabilità.

2. Implementazione Algoritmica (GROKtimizer):

   • Fase 1 (Interpolazione): L'ottimizzatore opera con zero weight decay per dare priorità a una rapida convergenza alla soglia di interpolazione. Questa fase utilizza il momentum standard per adattare rapidamente i dati di addestramento.
   • Fase 2 (Generalizzazione): Una volta raggiunta l'interpolazione, l'ottimizzatore passa al Momentum Criticamente Smorzato con weight decay. Il parametro di momentum è tarato sulla condizione di smorzamento critico derivata dalla teoria, e il weight decay viene applicato per guidare i parametri verso la soluzione a norma minima all'interno del manifold interpolante.

Contributi Chiave

1. Caratterizzazione del Regime Post-Interpolazione: Gli autori caratterizzano formalmente la fase di generalizzazione ritardata come un problema di oscillatore smorzato lungo le direzioni piatte del panorama della perdita, collegando il fenomeno del grokking alla geometria dello spazio nullo della matrice Hessiana.
2. Accelerazione Teorica: Dimostrano che il momentum criticamente smorzato accelera la convergenza verso soluzioni interpolanti a norma ridotta in modo quadraticamente più veloce rispetto al normale gradient descent, stabilendo l'ottimalità tra i metodi del primo ordine sotto il modello quadratico locale.
3. GROKtimizer: Introducono uno schema di ottimizzazione bifasico che passa dinamicamente da un rapido fitting non regolarizzato al minimo della norma con smorzamento critico.
4. Validazione Empirica: Il metodo è validato attraverso benchmark di grokking sintetici e dataset reali ad alta dimensionalità, dimostrando miglioramenti costanti rispetto agli ottimizzatori standard (Adam, AdamW, SGD, Muon).

Risultati Sperimentali
Il documento valuta GROKtimizer su diversi benchmark:

• Benchmark Sintetici: Su compiti come l'Addizione Modulare, l'Addizione Binaria e la Parità Sparsa, GROKtimizer supera significativamente i baseline. Ad esempio, sul compito Gaussiano, ha ridotto la perdita di validazione di diversi ordini di grandezza rispetto al baseline più forte.
• Dati Reali ad Alta Dimensionalità: Nei compiti di classificazione biomedica (Leukemia, TCGA) e nella predizione delle proprietà molecolari (QM9), GROKtimizer ha raggiunto l'accuratezza più elevata e la perdita di validazione più bassa. I guadagni sono stati più pronunciati nei contesti con scarsità di dati, dove il numero di feature supera di gran lunga il numero di campioni.
• Efficienza dei Campioni: Negli esperimenti MNIST a bassi dati (10–50 campioni per classe), GROKtimizer ha mostrato il vantaggio maggiore, confermando che la minimizzazione esplicita della norma migliora l'efficienza dei campioni quando lo spazio delle soluzioni è sottodeterminato.
• Pre-training di Modelli Linguistici: Gli esperimenti su WikiText2 e modelli BabyLM hanno mostrato che, sebbene la correlazione diretta tra norma dei parametri e perdita di validazione sia meno chiara nei transformer (a causa delle simmetrie di scala e degli aggiornamenti adattivi), il pattern bifasico regge comunque. Il passaggio alla fase di minimizzazione della norma dopo l'interpolazione ha prodotto ulteriori riduzioni della perdita di validazione.

Significato e Rivendicazioni
Il documento sostiene che GROKtimizer offre una soluzione principata al problema del "grokking", trattando l'addestramento post-interpolazione come un problema di ottimizzazione esplicito piuttosto che come un periodo di attesa passiva. Sfruttando la geometria del panorama della perdita, il metodo accelera la selezione di soluzioni a bassa complessità.

Gli autori riconciliano le loro scoperte con l'ipotesi dei minimi piatti (flat-minima hypothesis), sostenendo che la dimensionalità dello spazio nullo della matrice Hessiana funge da misura robusta di piattezza e complessità. Suggeriscono che minimizzare la norma all'interno di questo spazio nullo sia equivalente a navigare verso minimi più piatti e generalizzabili.

Limitazioni
Gli autori riconoscono alcune limitazioni:

• L'approccio è progettato principalmente per regimi con vincoli di dati e potrebbe essere meno critico nella pre-formazione su larga scala di visione o linguaggio dove l'abbondanza di dati domina le prestazioni.
• Il metodo assume che il modello abbia capacità sufficiente per raggiungere l'interpolazione; non è una soluzione per modelli sotto-parametrizzati.
• L'analisi teorica si basa su un'approssimazione quadratica locale del panorama della perdita, sebbene i risultati empirici suggeriscano che il metodo sia robusto alle deviazioni da tale assunto.
• La norma dei parametri è usata come proxy per la complessità, il che potrebbe non catturare tutti gli aspetti della generalizzazione del modello in ogni architettura.