Technische Samenvatting: Een Verdediging van het Kwadratische Model

Probleemstelling

Optimalisatietheorie voor Large Language Models (LLM's) vertrouwt vaak op geïdealiseerde modellen om de complexiteit van niet-convexe verlieslandschappen te navigeren. Er bestaat echter een aanzienlijke kloof tussen theoretische behapbaarheid en empirische nauwkeurigheid: modellen die wiskundig eenvoudig zijn, slagen er vaak niet in om de werkelijke dynamiek te beschrijven, terwijl modellen die specifieke probleemstructuren vangen, vaak theoretisch onhandelbaar zijn. De centrale vraag die dit werk adresseert is: In welke mate vangt het simpelste mogbare optimalisatiemodel — het kwadratische model — daadwerkelijk de trainingsdynamiek van LLM's?

Methodologie

De auteurs voeren een rigoureuze "stress test" uit van het kwadratische model met behulp van een 150M-parameter transformer getraind op 3 miljard tokens van FineWeb. De methodologie bestaat uit drie primaire componenten:

1. Lokale Kwadratische Benadering (Sectie 3)

De auteurs evalueren het voorspellende vermogen van lokale Taylor-expansies op intermediaire checkpoints (elke 10% van de training). Ze definiëren twee proxy-verliezen:

• Prox (Lineair): Lineariseert het model $f_\theta$ rond een referentiepunt $\theta_{ref}$ terwijl de oorspronkelijke cross-entropy doelstelling behouden blijft. Dit is equivalent aan multi-class logistische regressie onder het Neural Tangent Kernel (NTK) regime.
• Quad (Kwadratisch): Voert een Taylor-expansie uit op de Prox loss tot de tweede orde, wat resulteert in een convexe kwadratische doelstelling. Cruciaal is dat ze de Gauss-Newton matrix ($G$) gebruiken in plaats van de volledige Hessiaan ($H$) om te garanderen dat de kwadratische doelstelling convex blijft (positief semidefiniet), waardoor onbegrensdheid in richtingen met negatieve kromming wordt vermeden.

Ze trainen deze proxy-modellen voor 10% van het totale tokenbudget uitgaande van elk checkpoint en vergelijken de resulterende optimalisatietrajecten met de werkelijke LLM-trainingsdynamiek.

2. Schatting van het Hessiaans Spectrum (Sectie 4)

Om de structuur van deze lokale kwadratische problemen te analyseren, schatten de auteurs de spectra van de Hessiaan- en Gauss-Newton-matrices (zowel ruw als Adam-gepreconditioneerd) langs het trainingspad.

• Techniek: Ze maken gebruik van Stochastic Lanczos Quadrature (SLQ) met extreem diepe probes ($m=1200$) en rigoureuze foutmargen (Gauss-Radau bounds).
• Schaal: De probes worden uitgevoerd op checkpoints bij 10%, 50% en 100% van de training bij batchgroottes $B \in \{1, 64, 1024\}$.
• Analyse: Ze onderzoeken spectrale verval (machtswetten), de distributie van eigenvectoren over netwerklagen en de uitlijning tussen de gradiënt en eigenvectoren (source conditions).

3. Edge of Stability (EOS) Analyse (Sectie 5)

De auteurs testen empirisch of LLM-training opereert aan de "edge of stability".

• Opzet: Ze bevriezen de Adam-preconditioner en draaien preconditioned SGD op de lokale kwadratische proxy.
• Perturbatie: Ze variëren de leercurve ($\eta$) en batchgrootte ($B$) multipliers op een rooster om de stabiliteitsgrens te bepalen.
• Metriek: Een run wordt als divergerend beschouwd als het verlies een drempelwaarde overschrijdt (100). Ze maken onderscheid tussen Mean Square Stability (gevoelig voor zware staarten) en Lyapunov Stability (typisch gedrag).

Belangrijkste Resultaten

1. Voorspellend Vermogen van Kwadratische Modellen

• Overeenkomstvenster: De lokale kwadratische en lineaire modellen voorspellen de werkelijke optimalisatiedynamiek accuraat voor vensters die tot wel 10% van het totale trainingsbudget duren.
• Temporele Nauwkeurigheid: De kwaliteit van de benadering is aanzienlijk beter in de midden- tot late stadia van de training vergeleken met de vroege stadia.
• Prox vs. Quad: Het lineaire model (Prox) komt over het algemeen beter overeen met de werkelijke loss dan het kwadratische model (Quad) in de vroege fase van de training, terwijl Quad pas na 60% van de training convergeert naar een hoge nauwkeurigheid.

2. Structuur van het Hessiaans Spectrum

• Head-Tail Split: Het spectrum vertoont een scherpe transitie rond index $i \approx 8192$ (de vocabulairegrootte). De "head" (top eigenwaarden) wordt gedomineerd door de unembedding layer, wat consistent is met het "neural collapse" fenomeen.
• Universele Staart: De "tail" van het spectrum (indices $> 10^5$) volgt een machtswet die universeel is: deze is identiek over verschillende batchgroottes en ongevoelig voor de Adam-preconditioner.
• Exponenten:
  • Capaciteitsexponent ($\alpha$): Geschat op $\alpha \approx 0.96$ voor de Gauss-Newton staart en $\alpha \approx 0.65$ voor de Hessiaan staart.
  • Source Exponent ($\beta$): De source condition (gradiënt uitlijning) levert $\beta \in [0.1, 0.2]$ voor Gauss-Newton en $\beta \in [0.3, 0.6]$ voor de Hessiaan. Dit spreekt de gangbare literatuur aan dat $\beta > 1$ tegen, wat suggereert dat het probleem zich gedraagt alsof het einddimensioneel is met een exponentiële bias-afname na een kritieke tijd.
• Negatieve Eigenwaarden: De gepreconditioneerde Hessiaan vertoont grote negatieve eigenwaarden, met name bij grotere batchgroottes, gelokaliseerd in de MLP up-projection layer. De auteurs hypothetiseren dat deze gelinkt zijn aan "dead neurons" en het effect van zeldzame tokens.

3. Edge of Stability Dynamiek

• Stochastisch vs. Deterministisch:
  • Bij kleine tot middelgrote batchgroottes ($B=1, 64$), vindt training plaats aan een stochastische edge of stability. De stabiliteitsgrens volgt een diagonale lijn in de $(\eta, B)$ ruimte, wat aangeeft dat instabiliteit wordt gedreven door stochastiek (ruis).
  • Bij grote batchgroottes ($B=1024$), opereert training aan de deterministische edge of stability, gekenmerkt door een vlakke cutoff in de stabiliteitsregio, consistent met full-batch gradiëntafdaling gedrag.
• Nabijheid van de Grens: In alle settings opereert de training binnen een factor 2 van de stabiliteitsgrens.

Betekenis en Claims

Het artikel betoogt dat het kwadratische model niet slechts een theoretisch speeltje is, maar een geldige, experimenteel geverifieerde proxy voor LLM-pretraining dynamiek.

• Theoretische Behapbaarheid: De bevindingen suggereren dat complexe LLM-optimalisatie lokaal begrepen kan worden als een sequentie van gestructureerde kwadratische problemen. Dit stelt onderzoekers in staat om de rijke gereedschapskist van de kwadratische optimalisatietheorie (bijv. stabiliteitsanalyse, leercurve planning) te gebruiken om over LLM-training te redeneren.
• Herziening van Aannames: De resultaten dagen standaard aannames in de literatuur uit, specifell de sommeerbaarheid van source conditions ($\beta > 1$) en de universaliteit van de Hessiaanse staart onder preconditioning.
• Praktische Implicaties: De identificatie van de "stochastische edge of stability" bij kleine batchgroottes biedt een verfijnd begrip van waarom training divergeert onder specifieke hyperparameter configuraties, waarbij onderscheid wordt gemaakt tussen discretisatiefouten en stochastische ruis.

De auteurs concluderen dat hoewel globale verlieslandschappen niet-convex zijn, de lokale kwadratische benadering een krachtig, hanteerbaar kader biedt voor het analyseren en ontwerpen van trainingspipelines voor grootschalige modellen.