Sintesi Tecnica: To Grok Grokking: Provable Grokking in Ridge Regression

Definizione del Problema

Il paper investiga il fenomeno del grokking, definito come l'insorgenza della generalizzazione molto tempo dopo che un modello ha raggiunto una perfetta performance di addestramento (overfitting). Sebbene precedentemente osservato principalmente in reti neurali non lineari complesse e dataset algoritmici, gli autori mirano a stabilire garanzie teoriche rigorose ed end-to-end per questo fenomeno in un contesto classico e lineare. Nello specifico, il paper affronta la lacuna nella letteratura esistente dove i lavori precedenti o mancano di prove formali del "ritardo" tra overfitting e generalizzazione o si affidano a strutture di problemi specifiche e non generalizzabili (ad esempio, l'addizione modulare). La domanda centrale è se il grokking possa essere dimostrato in modo provabile nella regressione ridge sovra-parametrizzata utilizzando il gradiente discendente (GD) con decadimento del peso (weight decay), e se sì, come gli iperparametri controllino quantitativamente il "tempo di grokking".

Metodologia

Gli autori formulano il problema all'interno di un framework teacher-student:

• Setting: Regressione lineare sovra-parametrizzata dove un modello student $N(x; \theta) = \langle \theta, \phi(x) \rangle$ apprende una funzione teacher realizzabile $N^*(x) = \langle \theta^*, \phi(x) \rangle$.
• Ottimizzazione: Lo student viene addestrato utilizzando il Gradiente Discendente (GD) con un passo fisso $\eta$ e un termine di regolarizzazione $\ell_2$ (decadimento del peso) $\lambda$. L'obiettivo è l'errore quadratico medio regolarizzato:
  $$L_n(\theta; \lambda) = \frac{1}{2n}\sum_{i=1}^n (N(x_i; \theta) - N^*(x_i))^2 + \frac{\lambda}{2}\|\theta\|_2^2$$
• Strategia di Analisi: Gli autori decompongono il vettore dei parametri $\theta$ in due componenti ortogonali rispetto allo span dei dati:
  1. $\theta_{\parallel}$: La componente nello spazio riga della matrice delle feature $\Phi$ (subspazio che copre i dati).
  2. $\theta_{\perp}$: La componente nel nucleo (null space) di $\Phi$ (subspazio complementare).
• Approccio Teorico: La prova analizza i tassi di convergenza dell'errore di addestramento (errore empirico) e dell'errore di generalizzazione (errore della popolazione) separatamente. Il meccanismo centrale si basa sulla dimostrazione che $\theta_{\parallel}$ converge rapidamente per adattarsi ai dati di addestramento, mentre $\theta_{\perp}$ decade lentamente a causa del decadimento del peso, causando un prolungato periodo di scarsa generalizzazione.

Contributi Chiave

1. Primo Grokking Provabile End-to-End: Il paper fornisce i primi limiti quantitativi rigorosi che dimostrano le tre fasi distinte del grokking:
   • (i) Rapido overfitting dei dati di addestramento.
   • (ii) Un ritardo persistente in cui l'errore di generalizzazione rimane elevato nonostante l'errore di addestramento sia nullo.
   • (iii) Eventuale convergenza verso un modello con errore di generalizzazione arbitramente piccolo.
2. Limiti Quantitativi sul Tempo di Grokking: Gli autori derivano espliciti limiti inferiori per il "tempo di grokking" ($t_2 - t_1$), dove $t_1$ è il tempo per l'overfitting e $t_2$ è il tempo per la generalizzazione. Questi limiti sono espressi in termini di iperparametri del modello:
   • Decadimento del Peso ($\lambda$): Il tempo di grokking scala inversamente con $\lambda$ ($t_2 \propto 1/\lambda$). Un $\lambda$ più piccolo ritarda significativamente la generalizzazione.
   • Scala di Inizializzazione ($\nu^2$): Aumentare la scala di inizializzazione amplifica il grokking in modo logaritmico ($t_1, t_2 \propto \ln(\nu^2)$).
   • Sovra-parametrizzazione: Il fenomeno mostra di richiedere un setting ad alta dimensionalità ($m \gg n$).
3. Controllo degli Iperparametri: Il paper stabilisce condizioni sufficienti sugli iperparametri (dimensione del campione $n$, dimensione delle feature $m$, decadimento del peso $\lambda$ e scala di inizializzazione $\nu^2$) per indurre o eliminare il grokking.
4. Validazione Empirica su Modelli Non Lineari: Sebbene la teoria sia derivata per la regressione lineare, gli autori dimostrano empiricamente che le dipendenze quantitative sugli iperparametri (ad esempio, l'effetto di $\lambda$ sul ritardo) corrispondono qualitativamente al comportamento osservato nelle reti neurali ReLU a due strati non lineari.

Risultati Chiave

• Tassi di Convergenza: L'errore di addestramento converge a un tasso dominato dal più piccolo autovalore non nullo della covarianza empirica delle feature ($\lambda^+_{\min}(\Phi^\top \Phi)$), che è veloce. Al contrario, l'errore di generalizzazione è limitato inferiormente da un termine che decade al tasso di $(1 - \eta\lambda)^t$, che è arbitrariamente lento per piccoli valori di $\lambda$.
• Il Meccanismo di Grokking: Il ritardo deriva dal fatto che il GD con un piccolo decadimento del peso aggiorna efficacemente solo la proiezione dei pesi sullo span dei dati per minimizzare l'errore di addestramento, mentre la componente ortogonale (che contribuisce all'errore di generalizzazione) decade solo tramite il lento termine di decadimento del peso.
• Teoremi:
  • Teorema 4.1: Stabilisce il grokking per il teacher nullo, fornendo limiti superiori e inferiori corrispondenti sulla convergenza della perdita.
  • Teorema 4.2: Generalizza il risultato a qualsiasi teacher realizzabile, fornendo espliciti limiti per $t_1$ e $t_2$ (Equazioni 6 e 7).
  • Teorema 4.4 & 4.5: Provano la rapida convergenza della perdita di addestramento e la persistenza della scarsa generalizzazione, rispettivamente.
  • Teorema 4.6: Dimostra che il GD converge infine a un minimo globale con buone garanzie di generalizzazione.

Significato e Rivendicazioni

Il paper sostiene di fornire un framework teorico fondamentale per comprendere il grokking. Isolando il fenomeno nel setting semplice e trasparente della regressione lineare, gli autori sostengono che:

• Il grokking non è un fallimento intrinseco delle architetture di deep learning o delle non-linearità, ma è una conseguenza di specifiche condizioni di addestramento (specificamente, l'interazione tra sovra-parametrizzazione e piccolo decadimento del peso).
• Il fenomeno non richiede cambiamenti fondamentali all'architettura del modello o agli algoritmi di apprendimento per essere evitato; può essere controllato tramite la sintonizzazione degli iperparametri.
• I risultati offrono i primi limiti quantitativi rigorosi sul ritardo della generalizzazione, andando oltre le osservazioni empiriche o le spiegazioni euristiche (come le transizioni di regime "lazy-to-rich") trovate nella letteratura precedente.
• Le scoperte suggeriscono che il "plateau dell'errore di test" spesso visto nella classificazione (un artefatto di metriche discrete) esiste anche nella regressione quando valutato tramite metriche di accuratezza basate su soglia, confermando l'universalità del fenomeno attraverso diversi tipi di perdita.

Gli autori mantengono la modestia riguardo all'ambito, notando che sebbene la loro analisi lineare fornisca prove rigorose, estendere questi risultati a setting non realizzabili o provare la transizione "lazy-to-rich" nelle reti neurali profonde rimane una sfida aperta che richiede nuovi strumenti teorici.