Technische Samenvatting: Kleine Initialisatie is Belangrijk voor Large Language Models

Probleemstelling

Hoewel vooruitgang in Large Language Models (LLM's) doorgaans wordt toegeschreven aan schaal, data en architectuur, is de parameterinitialisatie een kritische ontwerpkeuze die in het tijdperk van grote modellen opnieuw een complex probleem is geworden. Standaard heuristieken (bijv. Xavier, Kaiming) vereisen vaak kostbare trial-and-error afstemming. Een heersende theoretische visie suggereert dat kleine initialisatie netwerken naar een "gecondenseerd" regime drijft waar gewichten uitlijnen langs gedeelde richtingen voordat ze diversifiëren, wat het redeneren en de generalisatie potentieel verbetert. Het blijft echter onduidelijk of deze voordelen standhouden in realistische LLM-pretraining op grote schaal, of dat architecturale componenten de voordelen van kleine initialisatie onbedoeld onderdrukken.

Methodologie

De auteurs onderzoeken systematisch de impact van de initialisatieschaal op de trainingsdynamiek en prestaties van LLM's.

• Initialisatie-parametrisering: Gewichtsmatrices $W$ worden geïnitialiseerd als $W_{ij} \sim \mathcal{N}(0, \sigma^2)$, waarbij $\sigma = d_{in}^{-\gamma}$. De parameter $\gamma$ controleert de schaal: $\gamma = 0.5$ herstelt de standaard Xavier-achtige schaling, terwijl $\gamma > 0.5$ (specifiek $\gamma = 1$) een kleinere initialisatie oplevert.
• Modelarchitecturen: Experimenten werden uitgevoerd op zowel dense decoder-only Transformers als Mixture-of-Experts (MoE) modellen. Architecturale variaties omvatten standaard LayerNorm/RMSNorm en Multi-Head Attention.
• Architecturale Interventies: Om geobserveerde beperkingen aan te pakken, werden twee specifieke aanpassingen getest:
  1. Het reduceren van de RMSNorm-constante ($\epsilon$): Het verlagen van $\epsilon$ van de standaard $10^{-5}$ naar $10^{-12}$ om te voorkomen dat de normalisatiefactor gedomineerd wordt door de constante wanneer de variantie van de verborgen toestand klein is.
  2. Gated Attention: Het introduceren van query-afhankelijke gates naar de outputs van attention heads om de "attention sink" (disproportionele aandacht voor het eerste token) te mitigeren, die wordt verergerd door kleine initialisatie.
• Trainingsopzet: Modellen werden getraind op een hoogwaardige tweetalige (Chinees-Engels) corpus (1 biljoen tokens) met next-token prediction met AdamW. Experimenten besloegen modelgroottes van 0.1B tot 1.5B parameters.
• Analyse-metrieken:
  • Token-niveau: Symmetrische waarschijnlijkheidskloof ($\Delta_{sym}$) en token-moeilijkheidsgraad (gemiddelde loss) om prestatiewinsten te lokaliseren.
  • Mechanistisch: Rij-gewijze cosinus-gelijkenis en Stabiele Rang (Stable Rank $\|W\|_F^2 / \|W\|_2^2$) om de evolutie van de parameterstructuur (condensatie versus expansie) te volgen.

Belangrijkste Resultaten

1. De Schaleringsbeperking en de Oplossing

Onder standaardinstellingen ($\gamma=0.5$ vs. $\gamma=1$) verlaagt kleine initialisatie consequent de validatie-loss, maar het voordeel neemt af naarmate de modelgrootte toeneemt. De auteurs identificeren twee architecturale barrières die dit veroorzaken:

• RMSNorm Verzadiging: Met de standaard $\epsilon=10^{-5}$ wordt de normalisatiefactor ongevoelig voor de verminderde variantie van kleine initialisatie zodra $\sigma^2(h) \lesssim \epsilon$. Het reduceren van $\epsilon$ naar $10^{-12}$ herstelt de gevoeligheid.
• Attention Sink: Kleine initialisatie intensiveert de neiging om aandacht te besteden aan het eerste token. Gated attention mitigeert dit effect effectief, wat de loss aanzienlijk verlaagt wanneer het gecombineerd wordt met kleine initialisatie.
• Resultaat: Wanneer beide barrières worden verwijderd, wordt het prestatieverschil groter. Kleine initialisatie ($\gamma=1$) met deze aanpassingen verbetert de effectieve modelgrootte met ongeveer 44% en levert consistente winsten op over alle modelgroottes.

2. Optimale Initialisatieschaal ($\gamma=1$)

De studie bepaalt dat initialisatie niet willekeurig klein moet zijn.

• Residuele Flow Balans: In een Transformer moet de residuele update vergelijkbaar blijven met de identiteits-pathway. Theoretische analyse laat zien dat de residuele flow schaalt als $d^{1-\gamma}$.
  • Als $\gamma < 1$, domineert de residuele component.
  • Als $\gamma > 1$, verdwijnt de residuele component, waardoor het netwerk nabij een identiteitsmapping blijft.
  • $\gamma = 1$ vertegenwoordigt het kritieke balanspunt waar de residuele update overeenkomt met de embedding-schaal. Empirisch gezien piekt de prestatie bij $\gamma=1$ en verslechtert deze voor $\gamma > 1$.

3. Winsten op Token- en Taakniveau

• Niet-uniforme Winsten: Verbeteringen zijn niet uniform over alle tokens verdeeld. Winsten concentreren zich op "matig tot zeer moeilijke" tokens die contextuele integratie en redeneren vereisen, in plaats van gemakkelijke tokens of extreem ruisige/zeldzame tokens.
• Benchmark Prestaties: Op 1.5B modellen presteert de geoptimaliseerde kleine initialisatie ($\gamma=1$ + aanpassingen) significant beter dan standaard initialisatie ($\gamma=0.5$) op redenerings-intensieve benchmarks (bijv. GSM8K, MATH500, HellaSwag), met absolute winsten van meer dan 4% op verschillende taken.

4. Mechanistisch Traject: Condensatie

Kleine initialisatie induceert een distinct ontwikkelingspad:

• Fase 1 (Condensatie): Vroeg in de training condenseren gewichtsmatrices snel in laag-rang, laag-complexiteit structuren (rijen lijnen zich uit in coherente blokken).
• Fase 2 (Expansie): Naarmate de training vordert, expanderen deze structuren naar rijkere, hogere-complexiteit representaties.
• Contrast: Standaard initialisatie ($\gamma=0.5$) begint met een diffuse, hoog-rang structuur die traag evolueert en de initiële laag-complexiteit fase mist. Dit "laag-naar-hoog" complexiteitstraject wordt verondersteld het model te sturen naar parsimonieuze redeneerregels in plaats van oppervlakkige memorisatie.

Betekenis en Claims

Het artikel betoogt dat parameterinitialisatie een "gen-achtige" determinant is van modelcapaciteit en trainingsdynamiek, en niet louter een detail van de lage-niveau implementatie.

• Ontwerp-as: De auteurs stellen voor om de initialisatieschaal $\gamma$ te behandelen als een expliciete hyperparameter in LLM-ontwerp.
• Praktische Aanbeveling: Ze pleiten voor een $\gamma$-initialisatieregel: behandel de initialisatieschaal als een regelbare knop en hanteer $\gamma=1$ als de standaard. Deze interventie wordt beschreven als "bijna kosteloos" maar in staat om pretraining te verbeteren en redeneren over schalen heen te versterken.
• Theoretisch Inzicht: Het werk levert concreet bewijs voor het idee dat "compressie intelligentie is", door aan te tonen dat een specifiek initialisatieregime modellen door een condensatiefase dwingt die correleert met verbeterde redeneer-capaciteiten.
• Architecturale Interdependentie: De bevindingen benadrukken dat de voordelen van initialisatie afhankelijk zijn van architecturale componenten (specifiek normalisatieconstanten en attention mechanismen), wat suggereert dat deze gezamenlijk afgestemd moeten worden in plaats van onafhankelijk.