Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

Sintesi Tecnica: Regimi di Decadimento dei Pesi nei Transformer Grokking

Enunciato del Problema
Grokking si riferisce al fenomeno per cui i transformer addestrati su compiti algoritmici (specificamente l'aritmetica modulare) mostrano una transizione netta dalla memorizzazione alla generalizzazione molto dopo che la perdita di addestramento ha raggiunto la saturazione. Sebbene studi meccanicistici abbiano identificato la formazione di circuiti interni (ad esempio, rappresentazioni di Fourier) e le "induction heads" come motori di questa transizione, la letteratura esistente manca di una "diagnostica online economica" per tracciare queste dinamiche in tempo reale. Inoltre, il ruolo del decadimento dei pesi ($\lambda$) come parametro di controllo per questi regimi non è stato mappato sistematicamente attraverso le scale dei modelli, né è stata isolata la specifica contribuzione del coordinamento delle testate di attenzione dagli effetti della magnitudine dei pesi.

Metodologia
Gli autori addestrano transformer decoder-only su compiti di aritmetica modulare (mod+, mod-, mod$\times$, mod$\div$) utilizzando l'ottimizzatore AdamW. Lo studio impiega una scansione densa di valori di decadimento dei pesi e varia le scale dei modelli (da 0,82M a 85M parametri) e le architetture (Transformer, MLP, LSTM, Mamba).

I componenti metodologici chiave includono:

1. Parametri d'Ordine Online: Due diagnostici scalari sono calcolati ad ogni passo di addestramento utilizzando esclusivamente le attivazioni di attenzione:
   • Media della Similarità Cosine tra le Testate di Attenzione ($\bar{s}(t)$): Misura il coordinamento (sincronizzazione) tra le testate di attenzione.
   • Deviazione Standard dell'Entropia ($\sigma_H(t)$): Misura la dispersione dell'entropia di Shannon tra le testate, tracciando la differenziazione.
2. Scalatura di Dimensione Finita: Gli autori adattano una funzione logistica alla probabilità di grokking ($P(\text{grok})$) in funzione del decadimento dei pesi per stimare una soglia critica $\lambda_c$. Adattano inoltre una legge di potenza al "tempo per grokkare" ($t_{\text{grok}}$) per determinare l'esponente critico $\nu$.
3. Interventi Causali: Per distinguere tra la struttura del pattern delle testate e la magnitudine dei pesi, gli autori eseguono esperimenti accoppiati in cui le testate di attenzione vengono re-inizializzate rispetto a casi in cui le norme dei pesi vengono tagliate al picco della fase di differenziazione.
4. Verifica Formale: Le identità diagnostiche (limiti e proprietà algebriche) sono verificate meccanicamente in Lean 4 per garantire la correttezza matematica.

Contributi Chiave

1. Soglia Critica Quantitativa: Lo studio identifica una soglia critica netta per il decadimento dei pesi per il transformer canonico a 4 livelli e 8 testate su compiti mod+: $\lambda_c = 0,0158$ (IC 95% [0,0109, 0,0200]).
2. Diagramma dei Regimi: Una mappa a due assi ($\lambda$, $N$) rivela tre regimi distinti:
   • Memorizzazione: $\lambda < \lambda_c$ (grokking vicino allo zero).
   • Grokking Sviluppo: $\lambda \ge \lambda_c$ (raggiungimento di un tasso di grokking di ~100%, con il tempo per grokkare che diminuisce all'aumentare di $\lambda$).
   • Collasso: $\lambda = 10$ (tutte le testate convergono verso pattern identici, $\bar{s}=1,0$).
3. Dinamiche a Due Fasi: La traiettoria di addestramento è caratterizzata da:
   • Fase 1 (Sincronizzazione): Le testate convergono ($\bar{s}$ sale da ~0,93 a 0,995) coincidendo con l'inizio del grokking.
   • Fase 2 (Differenziazione): Le testate divergono ($\bar{s}$ scende a ~0,88, $\sigma_H$ aumenta) mentre l'accuratezza sul test rimane alta.
4. Dimensione per Testata come Modulatore di Ampiezza: L'ampiezza della differenziazione ($\sigma_H$) è empiricamente dominata dalla dimensione per testata ($d/H$), saturando monotonicamente all'aumentare di $d/H$.
5. Portata Cross-Architettura: La transizione controllata dal decadimento dei pesi è osservata in architetture non basate su attenzione (MLP, LSTM, Mamba), sebbene la soglia critica $\lambda_c$ si sposti significativamente (coprendo circa un ordine di grandezza) a seconda dell'architettura.

Risultati

• Esponente di Scalatura: Il tempo per grokkare segue una legge di potenza $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ con un esponente empirico $\nu = 0,757$ (IC 95% [0,725, 0,799]). Gli autori notano esplicitamente che questo valore non corrisponde agli esponenti di riferimento standard (ad esempio, $\nu=1/2$ o Ising 3D $\nu \approx 0,63$) sotto la loro attuale densità di griglia, rinviando l'identificazione della classe di universalità a lavori futuri.
• Fragilità a Lungo Orizzonte: A orizzonti estesi (20.000 epoche), in alcune traiettorie si osserva un collasso "anti-grokking" dipendente dal seme, caratterizzato da un pattern a cinque stadi che termina con un decadimento dell'accuratezza. Questo è identificato come una fragilità piuttosto che un ciclo universale.
• Isolamento Causale: La re-inizializzazione delle testate di attenzione riduce significativamente l'ampiezza della Fase 2 ($\sigma_H$) rispetto ai controlli, mentre il taglio delle norme dei pesi non lo fa. Questo isola l'effetto alla struttura del pattern delle testate piuttosto che alla magnitudine dei pesi.
• Replicazione Multi-Task: Il pattern di controllo tramite decadimento dei pesi vale su quattro operazioni modulari e due scale di modelli, confermando la robustezza della mappa dei regimi oltre l'addizione.

Significato e Affermazioni
Il paper afferma di fornire un framework diagnostico economico e calcolabile online per tracciare le dinamiche di grokking utilizzando esclusivamente le attivazioni di attenzione, complementando diagnostici più costosi del paesaggio di perdita come il coefficiente di apprendimento locale raffinato (rLLC). Stabilisce il decadimento dei pesi come un parametro di controllo empirico primario che separa i regimi di memorizzazione, grokking sviluppo e collasso.

Gli autori sono modesti nel loro ambito:

• Affermano nessuna classe di universalità per l'esponente osservato $\nu \approx 0,76$, notando che gli esponenti di riferimento testati cadono al di fuori dei loro intervalli di confidenza.
• Non affermano che il decadimento dei pesi sia necessario per la generalizzazione in tutte le architetture, ma solo che agisce come parametro di controllo all'interno delle architetture di attenzione standard.
• Rinviando esplicitamente le affermazioni sui meccanismi causali riguardanti il perché avviene la differenziazione a lavori successivi, concentrandosi invece su quando e come si manifesta.
• I risultati sono limitati all'aritmetica modulare in piccoli modelli di attenzione transformer; le affermazioni riguardanti i grandi modelli linguistici o una universalità più ampia sono fuori ambito.

Il lavoro si posiziona come un contributo fenomenologico e diagnostico, fornendo una "mappa dei regimi" e "parametri d'ordine" che completano il lavoro teorico concorrente sulla scalatura di dimensione finita, la contrazione di Lyapunov e il rilassamento vetroso.