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

Titel: Een diagnose tijdens het trainen voor generalisatie via de Log-Alignement Ratio

Probleemstelling
Naarmate neurale netwerkmodellen en datasets groeien, neemt de rekenkosten van trainingsruns toe, waardoor het bestuderen van trainingsdynamiek en het ontwerpen van goedkope diagnostische hulpmiddelen steeds waardevoller wordt. Beoefenaars hebben behoefte aan metrieken die problemen zoals overfitting, instabiliteit of verzadiging kunnen signaleren met uitsluitend grootheden uit het trainingsproces, zonder afhankelijk te zijn van dure, vastgehouden validatiegegevens. Hoewel eerder werk spectrale eigenschappen van gewichtsmatrices (bijv. effectieve rang, stabiele rang) of parametrisatietheorieën (bijv. maximale update-parametrisatie, $\mu$P) heeft onderzocht om generalisatie te begrijpen, richten bestaande metrieken zich vaak uitsluitend op óf gewichten óf activeringen, of vereisen ze rekenintensieve bewerkingen zoals volledige Singular Value Decomposition (SVD) of Principal Component Analysis (PCA).

Methodologie
De auteurs onderzoeken de Log-Alignement Ratio (LAR), een metriek die oorspronkelijk in de parametrisatietheorie door Everett et al. [2024] werd geïntroduceerd om de alignering tussen de parametrisatiematrix van een laag en de invoeractiveringen te meten. Het artikel formuleert LAR theoretisch opnieuw om een tweede, tot nu toe onontdekte betekenis bloot te leggen: het meet de overlap tussen een gewichtsspectrum ($p$) en een activeringsspectrum ($q$).

• Definities:
  • Laat $W$ een matrix zijn met singuliere waarden $s_i$. Het gewichtsspectrum $p$ wordt gedefinieerd als de genormaliseerde kwadratische singuliere waarden: $p_i = s_i^2 / \sum s_j^2$.
  • Laat $X$ een batch van invoer zijn. Het activeringsspectrum $q$ wordt gedefinieerd als de genormaliseerde kwadratische projecties van invoer op de singuliere richtingen van $W$: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Herformulering: De auteurs leiden af dat LAR kan worden uitgedrukt als:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Deze formulering toont aan dat LAR een maat is voor de log-overlap tussen verdelingen $p$ en $q$. Het kan ook worden geschreven in termen van covariantie en Pearson-correlatie:
  $$\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)$$
  waarbij $\sigma_p$ en $\sigma_q$ de "spreiding" of concentratie van de verdelingen vertegenwoordigen.
• Rekenefficiëntie: LAR is berekenbaar uit grootheden die beschikbaar zijn tijdens de forward pass (matrixnormen) met verwaarloosbare overhead ($O(|B|m + |B|n + mn)$), zonder dat vastgehouden data of volledige SVD nodig is.

Belangrijkste Bijdragen

1. Verdelingsherformulering: Het artikel toont aan dat LAR niet slechts een maat voor alignering is, maar een kwantificering van de overlap tussen gewichts- en activeringsspectra. Het biedt equivalente statistische vormen (covariantie, correlatie) die het gedrag van de metriek interpreteerbaar maken in termen van elementaire grootheden.
2. Voorspelling van Effectieve Dimensie: De auteurs tonen aan dat wanneer een gewichtsmatrix energie concentreert op $k$ singuliere richtingen en activeringen hiermee aligneren, LAR gerelateerd is aan de effectieve dimensie $k$ via $k \approx n^2(1 - \text{LAR})$. Deze formule is empirisch geverifieerd in grokking-experimenten.
3. Diagnose voor Overfitting op Schaal: Bij grootschalig pre-training (modellen met 3B parameters) stelt het artikel vast dat de unembedding-LAR stabiliseert tijdens generalisatie, maar scherp daalt naarmate overfitting nadert. De afwijking van LAR van een niet-overfitting baseline volgt de generalisatiekloof.
4. Proxy voor Spectrale Spreiding: De studie vindt dat optimalisatie $p$ en $q$ ertoe brengt sterk gecorreleerd te zijn. Bijgevolg fungeert LAR als een proxy voor de variantie (spreiding) van deze verdelingen. Bij pre-training domineert de top singuliere richting, waardoor de bijdrage van het topcomponent ($L_1$) een primaire indicator wordt van het pre-overfit-signatuur.

Experimentele Resultaten
De auteurs valideren deze bevindingen in twee verschillende settings:

1. Grokking-experimenten (Kleine Algorithmische Taken):

   • Getrainde 2-laags Transformers op 12 binaire operationele taken.
   • Observatie: Tijdens het memorisatiefase is LAR lager (wat wijst op diffuse $p$ en $q$). Naarmate het model overgaat naar "grokking" (generalisatie), neemt LAR toe naarmate $p$ en $q$ zich concentreren op een gedeelde set richtingen.
   • Dimensionaliteit: De door LAR voorspelde dimensie $k_{pred} = n^2(1 - \text{LAR})$ komt nauw overeen met $k_{95}$ (het aantal hoofdcomponenten dat 95% van de variantie verklaart) bij de laatste checkpoints, met een lage Root Mean Squared Error (RMSE).
   • Differentiatie: Grokking-modellen bereiken hogere uiteindelijke LAR-waarden ($>0.65$) in vergelijking met niet-grokking-modellen ($\le 0.65$).

2. Grootschalig Pre-training (3B Taalmodellen):

   • Getrainde Gemma-2-gebaseerde modellen (3B parameters) met Adam- en Muon-optimizers op datasets variërend van 200M tot 100B tokens.
   • Generalisatie versus Overfitting: In niet-overfitting regimes (bijv. 100B tokens) stijgt LAR naar een piek en daalt vervolgens langzaam, om te stabiliseren. In overfitting regimes (kleinere datasets) daalt LAR scherp naarmate de overfitting-stap nadert.
   • Generalisatiekloof: Het verschil tussen de LAR van een specifieke run en de baseline 100B-run (LAR-differentie) weerspiegelt nauw de generalisatiekloof (validatieverlies minus trainingsverlies).
   • Mechanisme: De daling in LAR voorafgaand aan overfitting wordt voornamelijk gedreven door het top singuliere component ($L_1$). Bij Adam wordt dit gedreven door de diffusie van activeringen ($q_1$); bij Muon zijn zowel gewichten ($p_1$) als activeringen betrokken.

Betekenis en Beweringen
Het artikel beweert dat LAR een praktische, rekenkundig goedkope en datavrije diagnose voor generalisatie biedt.

• Efficiëntie: In tegenstelling tot effectieve rang of andere spectrale maatstaven die SVD vereisen, kan LAR worden berekend met verwaarloosbare overhead tijdens de forward pass.
• Voorspellende Kracht: Het volgt succesvol de overgang van memorisatie naar generalisatie bij grokking en dient als een vroeg waarschuwingssignaal voor overfitting bij grootschalig pre-training, wat mogelijk evaluatievrije trainingsmethoden mogelijk maakt.
• Interpreteerbaarheid: Door LAR te koppelen aan de spreiding van gewichts- en activeringsspectra, biedt de metriek een verenigd beeld van hoe optimalisatiedynamiek (concentratie versus diffusie van spectra) relateert aan modelgeneralisatie.

De auteurs erkennen beperkingen, waarbij zij opmerken dat hun bewijs correlatief is, dat zij zich voornamelijk richten op de unembedding-laag, en dat zij nog niet hebben getest of het direct regulariseren van de spectrale spreiding overfitting zou voorkomen. Zij suggereren dat toekomstig werk LAR voor andere matrixtypes en het gedrag onder variërende batchgroottes moet onderzoeken.