Technische Samenvatting: SPARKLING

Probleemstelling

Progressieve Leering (PL) beoogt de computationele overhead van het trainen van Large Language Models (LLM's) te verminderen door de modelomvang tijdens de training geleidelijk uit te breiden, in plaats van de volledige schaal vanaf het begin te trainen. Terwijl bestaand PL-onderzoek uitgebreid diepte-expansie (bijv. laagstapeling) heeft verkend, blijft breedte-expansie aanzienlijk onderbelicht. Eerdere methoden voor breedte-expansie zijn grotendeels beperkt tot de vroege stadia van de training (bijv. <10–30% van de tokens), wat verwaarloosbare computationele voordelen biedt ten opzichte van het trainen van het doel-breedte model vanaf nul.

Het uitbreiden van de breedte tijdens de middenfase van de pre-training is essentieel voor het maximaliseren van de computationele besparingen, maar brengt formidabele uitdagingen met zich mee. Het artikel identificeert twee kernproblemen die oorzaak zijn van trainingsinstabiliteit en prestatievermindering in dit regime:

1. Falen van Signaalbehoud: Naïeve initialisatie bij het expansiepunt verstoort de statistische distributie van intermediaire activaties (specifiek de root-mean-square of RMS-schaal). Deze mismatch verandert de signaalmagnitudes per laag, wat de optimalisatie destabiliseert, zelfs als de directe loss niet piekt.
2. Symmetrie-vergrendeling (Symmetry Locking): Kopieer-gebaseerde initialisatie, hoewel het de forward-functie behoudt, induceert backward symmetrie. Gedupliceerde kanalen ontvangen identieke gradiënten en, indien de optimizer-states ook symmetrisch zijn, evolueren ze identiek. Dit voorkomt dat de nieuwe capaciteit zich differentieert in unieke features, waardoor de uitgebreide parameters functioneel redundant worden. Dit probleem houdt stand bij verschillende optimizer-families, inclus$\text de element-wise optimizers (AdamW) en spectrale optimizers (Muon).

Methodologie: SPARKLING

De auteurs stellen SPARKLING voor (balans tussen Signal Preservation And symmetrY Keeping voor breedte-progressieve LearnING), een framework ontworpen om een balans te vinden tussen signaalbehoud en symmetrie-doorbreking tijdens mid-stage breedte-expansie. Het framework bestaat uit vier concrete mechanismen:

1. Signaalbehoud via RMS-schaalconsistentie

De auteurs beargumenteren dat echt signaalbehoud vereist dat de statistische distributie van verborgen representaties wordt behouden, specifiek de RMS-schaal, in plaats van alleen functiebehoud (loss-continuïteit).

• Mechanisme: Ze leiden herschalingsfactoren af voor gewichten tijdens expansie om ervoor te zorgen dat de output RMS-schaal invariant blijft.
  • Fan-out expansie: Nieuwe output-kanalen worden geïnitialiseerd door te kopiëren of willekeurig te samplen uit dezelfde distributie als de bestaande gewichten.
  • Fan-in expansie: Nieuwe input-kanalen vereisen specifieke gewichts-herschaling. Als nieuwe coördinaten willekeurig zijn of eenzijdig gekopieerd, worden gewichten geschaald door $\sqrt{d_{in}/d'_{in}}$. Als beide zijden gekopieerd worden (duplicatie van bestaande dimensies), worden gewichten geschaald door $1/\sqrt{1+3c}$ (waarbij $c$ de kopieerratio is) om rekening te houden met de kwadratische bijdrage van covariantie.
  • RMSNorm Expansie: Nieuwe RMSNorm-parameters worden geïnitialiseerd door te kopiëren of willekeurig te samplen zonder extra herschaling, aangezien het aantal parameters verwaarloosbaar is ten opzichte van de lineaire lagen.

2. Symmetrie-doorbreking via Asymmetrische Interventies

Om de "symmetry lock" op te lossen waarbij gekopieerde componenten er niet in slagen te differentiëren, introduceert SPARKLING gerichte interventies die alleen werken op de backward-dynamiek terwijl de forward-functie behouden blijft.

• Asymmetrische Optimizer State Reset: In plaats van de optimizer-states (momentum) te kopiëren of globaal te resetten, behoudt het framework de optimizer-states voor de originele parameters ($W$) maar reset het de states voor de nieuw geïntroduceerde parameters ($\tilde{W}$) naar nul. Dit doorbreekt de gradiëntkoppeling onmiddellijk, waardoor de nieuwe parameters kunnen divergeren.
• Asymmetrische Learning Rate Re-warmup: Bij het expansiepunt vervolgen de originele parameters hun baseline learning rate schedule. De nieuwe parameters ondergaan een re-warmup startend vanaf de huidige learning rate ($\eta_e$) naar een hogere piek ($\hat{\eta}_{max} = \rho \cdot \eta_e$) over een korte window ($\tau_w$), voordat ze weer vervallen naar de gedeelde minimumwaarde. Dit injecteert gecontroleerde optimalisatie-asymmetrie om feature-diversificatie te stimuleren zonder de goed aangepaste pre-trained modellen globaal te verstoren.

Belangrijkste Bijdragen

1. Identificatie van Mid-Stage Uitdagingen: Het artikel stelt vast dat succesvolle mid-stage breedte-expansie berust op twee complementaire principes: Signaalbehoud (stabiliseren van activatie-statistieken via RMS-consistentie) en Symmetrie-doorbreking (oplossen van gradiëntkoppeling in kopieer-gebaseerde initialisatie).
2. Het SPARKLING Framework: Een praktische implementatie die RMS-preserverende schaling, kopieer-gebaseerde initialisatie, asymmetrische optimizer state reset en asymmetrische learning rate re-warmup combineert.
3. Empirische Validatie: Uitgebreide experimenten over dense en Mixture-of-Experts (MoE) architecturen, meerdere breedte-assen (hidden dimension en MoE expert intermediate dimension) en optimizer-families (AdamW en Muon).

Resultaten

• Prestaties: Onder een vastgesteld tokenbudget presteert SPARKLING consistent beter dan het trainen van het volledige model vanaf nul op downstream evaluaties (bijv. ARC, MMLU, HellaSwag) en evenaart of overtreft het de prestaties van naïeve functie-preserverende baselines.
• Efficiëntie: SPARKLING vermindert de trainingskosten aanzienlijk. Bij een 2× breedte-expansie bereikt het tot 35% reductie in training FLOPs en een tot 1.49× wall-clock versnelling vergeleken met het trainen van het uitgebreide model vanaf nul.
• Convergentie: RMS-preserverende schaling verbetert de late-stage convergentie vergeleken met naïeve ongeschaalde initialisatie. Asymmetrische interventies (state reset en re-warmup) zijn cruciaal voor het doorbreken van de symmetrie-vergrendeling, waarbij copy-copy initialisatie gecombineerd met deze interventies de beste uiteindelijke loss behaalt van alle varianten.
• Generaliseerbaarheid: De methode generaliseert naar dense modellen, verschillende optimizers (inclusief Muon) en iteratieve multi-stage expansies (bijv. 256→512→1024).

Betekenis en Claims

Het artikel claimt dat SPARKLING mid-stage breedte-expansie praktisch levensvatbaar maakt, waarmee het de primaire motivatie van Progressieve Leering aanpakt: het verlagen van de trainingskosten terwijl de prestaties van het doel-breedte model gelijk worden gehouden of overtroffen.

De auteurs benadrukken dat hun aanpak de spanning oplost tussen forward continuïteit (het behouden van de pre-trained functie) en backward diversiteit (het toestaan dat nieuwe parameters unieke features leren). Door deze twee aspecten te ontkoppelen—het behouden van signaalstatistieken terwijl de symmetrie in optimalisatiedynamiek wordt doorbroken—maakt SPARKLING efficiënte schaling van LLM's mogelijk. Het werk suggereert dat het framework een gunstiger compute-loss scaling gedrag biedt, waarbij een lagere loss wordt bereikt bij gelijke compute-budgetten vergeleken met training vanaf nul.

Het artikel concludeert door op te merken dat hoewel de resultaten veelbelovend zijn, toekomstig werk nodig is om een verenigd principe voor gelijktijdige breedte- en diepte-expansie vast te stellen en om te onderzoeken of de RMS-preserveringsstrategie voldoet aan de $\mu$P (maximal update parameterization) conditie voor hyperparameter transferbaarheid.