Riepilogo Tecnico: Come Addomesticare il Grokking: La Geometria delle Rappresentazioni come Segnale di Controllo

Definizione del Problema

Il grokking è un fenomeno che si verifica nelle reti neurali sovra-parametrizzate in cui i modelli inizialmente memorizzano i dati di addestramento e mostrano solo una forte generalizzazione dopo una prolungata ottimizzazione. Sebbene studi recenti abbiano collegato il grokking al weight decay, alla compressione delle feature, alla formazione di circuiti e alla dinamica dell'apprendimento delle rappresentazioni, i fattori specifici che governano la tempistica e l'emergere di questa generalizzazione ritardata rimangono non completamente compresi. Un tema comune nella letteratura esistente è che le rappresentazioni apprese si comprimano progressivamente in strutture a dimensionalità inferiore. Tuttavia, il lavoro precedente tratta ampiamente la geometria delle rappresentazioni come una proprietà emergente piuttosto che come un meccanismo di controllo diretto. Questo articolo indaga se la manipolazione esplicita della geometria delle rappresentazioni nascoste possa servire come segnale di controllo per influenzare l'inizio e la velocità del grokking.

Metodologia

Gli autori introducono la Geometric Dimensionality Regularization (GeomDR), un regolarizzatore spettrale a livello di rappresentazione progettato per controllare direttamente la dimensionalità effettiva delle rappresentazioni nascoste durante l'addestramento.

Caratterizzazione Geometrica

Per una matrice di rappresentazione di uno strato nascosto $Z \in \mathbb{R}^{N \times d}$, il metodo prima centra i dati per ottenere $\tilde{Z}$. Successivamente, calcola la matrice di covarianza empirica $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$. La dimensionalità effettiva ($D_{eff}$) viene calcolata utilizzando il rapporto di partecipazione degli autovalori ($\mu_i$) di $C$:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
Inoltre, viene definita una distanza di vicinato locale ($\rho$) come la distanza euclidea media dai $k$ vicini più prossimi delle rappresentazioni normalizzate.

L'Obiettivo GeomDR

GeomDR impone una penalità sulla varianza contenuta nelle direzioni oltre un sottospazio target di dimensione $d^*$. Dati gli autovalori ordinati in modo decrescente ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), la perdita per strato è:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
L'obiettivo di addestramento totale combina la perdita del compito ($L_{task}$) e il regolarizzatore geometrico:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Fondamentalmente, il regolarizzatore non è attivo sin dall'inizio. Viene introdotto dopo una fase iniziale di addestramento non vincolato (fase di attivazione $t_s$) e la sua forza $\lambda(t)$ viene aumentata gradualmente utilizzando una rampa coseno. Questo design imita la naturale traiettoria del grokking, dove la memorizzazione precede la generalizzazione.

Protocollo Sperimentale

Il framework è stato valutato su tre compiti algoritmici noti per esibire il grokking:

1. Addizione Modulare: Predire $(a + b) \mod 97$.
2. Divisione Modulare: Predire $(a \cdot b^{-1}) \mod 97$.
3. Composizione di Permutazioni: Predire la composizione di due permutazioni in $S_5$.

Gli esperimenti hanno utilizzato due architetture: un Multilayer Perceptron (MLP) e un Transformer. Lo studio ha variato sistematicamente gli schemi di intervento (forza $\lambda_{max}$ e durata della rampa $T_{ramp}$), le dimensionalità target ($d^*$) e i tempi di inizio dell'intervento ($t_s$).

Contributi Chiave

1. Osservazione del Collasso Precedente: Gli autori dimostrano che il collasso della dimensionalità precede costantemente la transizione dalla memorizzazione alla generalizzazione (grokking) nelle configurazioni di base. Ciò suggerisce che la compressione geometrica non è solo un correlato, ma un driver meccanicistico della transizione.
2. Introduzione di GeomDR: Propongono GeomDR, un regolarizzatore spettrale che sopprime direttamente lo spettro di covarianza al di fuori di un sottospazio target, controllando così la dimensionalità effettiva delle rappresentazioni nascoste.
3. Controllo Sistematico delle Dinamiche: Il documento fornisce uno studio empirico esaustivo mostrando che gli interventi geometrici possono alterare sostanzialmente le dinamiche di generalizzazione ritardata. A seconda dello schema di intervento e della dimensionalità target, GeomDR può accelerare il grokking, ritardarlo o sopprimerlo interamente.

Risultati

Accelerazione del Grokking

In tutte le configurazioni valutate, GeomDR ha costantemente alterato le dinamiche del grokking:

• Performance MLP: Sull'addizione modulare, GeomDR ha accelerato il grokking fino a 52 volte (riducendo i passi da 362K a ~7K). Miglioramenti simili sono stati osservati per la divisione modulare (39x) e la composizione di permutazioni (~23x). In tutti i casi MLP, GeomDR ha raggiunto un tasso di successo del 100% dove il baseline era instabile o falliva.
• Performance Transformer: Sebbene le dinamiche di grokking baseline per i Transformer fossero già più veloci rispetto agli MLP, GeomDR ha comunque fornito un'accelerazione costante (ad esempio, ~2x sull'addizione modulare), mantenendo un tasso di successo del 100%.

Influenza degli Iperparametri

• Dimensionalità Target ($d^*$): Dimensioni target moderate o grandi ($16 \le d^* \le 64$) hanno prodotto l'accelerazione più forte. Una compressione aggressiva ($d^* \le 8$) ha spesso rallentato o destabilizzato l'apprendimento.
• Tempistica dell'Intervento ($t_s$): È stata identificata una "finestra temporale critica". Gli interventi applicati troppo presto (prima che si formino rappresentazioni utili) o troppo tardi (dopo che la memorizzazione è diventata rigida) erano meno efficaci. Tempi di attivazione intermedi (ad esempio, $t_s = 2000$ per gli MLP) hanno prodotto il grokking più veloce e stabile.
• Programma (Schedule): Le intensità di regolarizzazione e le durate della rampa intermedie hanno generalmente performato meglio; interventi eccessivamente forti potevano destabilizzare l'addestramento.

Correlazioni Geometriche

Lo studio ha confermato che gli interventi che accelerano il grokking hanno accelerato anche l'inizio del collasso della dimensionalità. Nei modelli GeomDR, il divario temporale tra il collasso della dimensionalità e la generalizzazione riuscita è stato significativamente ridotto rispetto ai baseline, suggerendo un accoppiamento più stretto tra la riorganizzazione geometrica e l'emergere della generalizzazione.

Scalabilità ed Efficienza

Gli autori hanno dimostrato che GeomDR rimane efficace anche quando lo spettro di covarianza è stimato da un sottoinsieme casuale di rappresentazioni (ad esempio, il 5%–25% dei dati), offrendo un'accelerazione del tempo di esecuzione di 2,4x–2,7x senza influenzare materialmente l'accelerazione del grokking.

Significato e Rivendicazioni

L'articolo sostiene che la geometria delle rappresentazioni funge da efficace segnale di controllo per il grokking. Modificando direttamente lo spettro di covarianza delle attivazioni nascoste, i ricercatori possono influenzare sistematicamente quando e quanto velocemente una rete transita dalla memorizzazione alla generalizzazione.

Gli autori postulano che questo lavoro fornisca prove del fatto che gli interventi geometrici offrono un approccio pratico per studiare e influenzare la generalizzazione ritardata. Invece di osservare passivamente i cambiamenti rappresentazionali, il metodo permette la manipolazione diretta delle proprietà geometriche per testare ipotesi sulle dinamiche di apprendimento. I risultati suggeriscono che la relazione tra la geometria delle rappresentazioni e il grokking è robusta attraverso diversi compiti e architetture, sebbene l'entità del miglioramento vari (essendo più pronunciata negli MLP che nei Transformer, probabilmente a causa dei loro più forti bias induttivi).

L'articolo conclude che, sebbene non stabilisca la dimensionalità della rappresentazione come l'unico meccanismo causale, dimostra che gli interventi geometrici sono sufficienti per alterare la tempistica del grokking, offrendo una nuova metodologia sperimentale per la ricerca sull'ottimizzazione e la generalizzazione.