A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

Titolo: Una diagnostica in fase di addestramento per la generalizzazione tramite il Rapporto di Allineamento Logaritmico

Enunciato del Problema
Man mano che i modelli di reti neurali e i dataset aumentano di scala, il costo computazionale delle esecuzioni di addestramento cresce, rendendo lo studio delle dinamiche di addestramento e la progettazione di strumenti diagnostici economici sempre più preziosi. I praticanti necessitano di metriche in grado di segnalare problemi come l'overfitting, l'instabilità o la saturazione utilizzando esclusivamente quantità disponibili durante l'addestramento, senza fare affidamento su dati di validazione tenuti da parte e costosi. Sebbene lavori precedenti abbiano esplorato le proprietà spettrali delle matrici dei pesi (ad esempio, rango effettivo, rango stabile) o le teorie di parametrizzazione (ad esempio, parametrizzazione di aggiornamento massimale, $\mu$P) per comprendere la generalizzazione, le metriche esistenti si concentrano spesso esclusivamente sui pesi o sulle attivazioni, oppure richiedono operazioni computazionalmente costose come la completa Decomposizione ai Valori Singolari (SVD) o l'Analisi delle Componenti Principali (PCA).

Metodologia
Gli autori investigano il Rapporto di Allineamento Logaritmico (LAR), una metrica originariamente introdotta nella teoria della parametrizzazione da Everett et al. [2024] per misurare l'allineamento tra la matrice dei parametri di un livello e le sue attivazioni in ingresso. Il documento riformula teoricamente il LAR per rivelare un secondo significato, precedentemente inesplorato: misura la sovrapposizione tra uno spettro dei pesi ($p$) e uno spettro delle attivazioni ($q$).

• Definizioni:
  • Sia $W$ una matrice con valori singolari $s_i$. Lo spettro dei pesi $p$ è definito come i valori singolari al quadrato normalizzati: $p_i = s_i^2 / \sum s_j^2$.
  • Sia $X$ un batch di input. Lo spettro delle attivazioni $q$ è definito come le proiezioni al quadrato normalizzate degli input sulle direzioni singolari di $W$: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Riformulazione: Gli autori derivano che il LAR può essere espresso come:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Questa formulazione mostra che il LAR è una misura della sovrapposizione logaritmica tra le distribuzioni $p$ e $q$. Può anche essere scritto in termini di covarianza e correlazione di Pearson:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \cdot \text{cov}(p, q) + \frac{1}{n} \right)$$
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \sigma_p \sigma_q \text{corr}(p, q) + \frac{1}{n} \right)$$
  dove $\sigma_p$ e $\sigma_q$ rappresentano la "dispersione" o concentrazione delle distribuzioni.
• Efficienza Computazionale: Il LAR è calcolabile a partire da quantità disponibili durante il passaggio in avanti (norme delle matrici) con un sovraccarico trascurabile ($O(|B|m + |B|n + mn)$), non richiedendo dati tenuti da parte o una SVD completa.

Principali Contributi

1. Riformulazione Distribuzionale: Il documento dimostra che il LAR non è semplicemente una misura di allineamento, ma un quantificatore della sovrapposizione tra gli spettri dei pesi e delle attivazioni. Fornisce forme statistiche equivalenti (covarianza, correlazione) che rendono il comportamento della metrica interpretabile in termini di quantità elementari.
2. Predizione della Dimensione Effettiva: Gli autori mostrano che quando una matrice dei pesi concentra l'energia su $k$ direzioni singolari e le attivazioni si allineano ad esse, il LAR è legato alla dimensione effettiva $k$ tramite $k \approx n^2(1 - \text{LAR})$. Questa formula è verificata empiricamente negli esperimenti di "grokking".
3. Diagnostica dell'Overfitting su Larga Scala: Nell'addestramento pre su larga scala (modelli da 3B parametri), il documento stabilisce che il LAR di deembedding si stabilizza durante la generalizzazione ma declina bruscamente all'avvicinarsi dell'overfitting. La deviazione del LAR da una baseline non affetta traccia il divario di generalizzazione.
4. Proxy per la Dispersione Spettrale: Lo studio rileva che l'ottimizzazione porta $p$ e $q$ ad essere altamente correlati. Di conseguenza, il LAR agisce come un proxy per la varianza (dispersione) di queste distribuzioni. Nell'addestramento pre, la direzione singolare dominante è quella principale, rendendo il contributo della componente principale ($L_1$) un indicatore primario della firma pre-overfit.

Risultati Sperimentali
Gli autori validano queste scoperte in due contesti distinti:

1. Esperimenti di Grokking (Piccoli Compiti Algoritmici):

   • Addestramento di Transformer a 2 livelli su 12 compiti di operazioni binarie.
   • Osservazione: Durante la fase di memorizzazione, il LAR è più basso (indicando $p$ e $q$ diffuse). Man mano che il modello transita verso il "grokking" (generalizzazione), il LAR aumenta mentre $p$ e $q$ si concentrano su un insieme condiviso di direzioni.
   • Dimensionalità: La dimensione predetta dal LAR, $k_{pred} = n^2(1 - \text{LAR})$, corrisponde strettamente a $k_{95}$ (il numero di componenti principali che spiegano il 95% della varianza) ai checkpoint finali, con un basso Errore Quadratico Medio (RMSE).
   • Differenziazione: I modelli che effettuano il grokking raggiungono valori finali di LAR più elevati ($>0.65$) rispetto ai modelli che non effettuano il grokking ($\le 0.65$).

2. Addestramento Pre su Larga Scala (Modelli Linguistici da 3B):

   • Addestramento di modelli basati su Gemma-2 (3B parametri) utilizzando gli ottimizzatori Adam e Muon su dataset che vanno da 200M a 100B token.
   • Generalizzazione vs Overfitting: In regimi non affetti (ad esempio, 100B token), il LAR sale a un picco e poi declina lentamente, stabilizzandosi. In regimi affetti (dataset più piccoli), il LAR declina bruscamente all'avvicinarsi dello step di overfitting.
   • Divario di Generalizzazione: La differenza tra il LAR di una specifica esecuzione e la baseline da 100B (differenziale LAR) rispecchia da vicino il divario di generalizzazione (perdita di validazione meno perdita di addestramento).
   • Meccanismo: Il declino del LAR precedente all'overfitting è guidato principalmente dalla componente singolare principale ($L_1$). Con Adam, ciò è guidato dalla diffusione delle attivazioni ($q_1$); con Muon, coinvolge sia i pesi ($p_1$) che le attivazioni.

Significato e Affermazioni
Il documento afferma che il LAR offre una diagnostica pratica, computazionalmente economica e priva di dati per la generalizzazione.

• Efficienza: A differenza del rango effettivo o di altre misure spettrali che richiedono SVD, il LAR può essere calcolato con un sovraccarico trascurabile durante il passaggio in avanti.
• Potere Predittivo: Traccia con successo la transizione dalla memorizzazione alla generalizzazione nel grokking e funge da segnale di allarme precoce per l'overfitting nell'addestramento pre su larga scala, potenzialmente abilitando metodi di addestramento privi di valutazione.
• Interpretabilità: Collegando il LAR alla dispersione degli spettri dei pesi e delle attivazioni, la metrica fornisce una visione unificata di come le dinamiche di ottimizzazione (concentrazione vs diffusione degli spettri) si relazionino alla generalizzazione del modello.

Gli autori riconoscono le limitazioni, notando che le loro evidenze sono correlazionali, si concentrano principalmente sul livello di deembedding e non hanno ancora testato se regolarizzare direttamente la dispersione spettrale prevenga l'overfitting. Suggeriscono che lavori futuri dovrebbero esplorare il LAR per altri tipi di matrici e il suo comportamento al variare delle dimensioni del batch.