Unlocking Feature Learning in Gated Delta Networks at Scale

Technische Samenvatting: Het ontsluiten van feature learning in Gated Delta Networks op schaal

Probleemstelling
De schaling van Large Language Models (LLMs) staat voor twee primaire knelpunten: de steile computationele kosten van het trainen van grotere modellen en de kwadratische complexiteit van standaard Transformer self-attention. Lineaire recurrente architecturen, zoals Gated Delta Networks (GDN), zijn naar voren gekomen als efficiënte alternatieven die lineaire-tijd inferentie behouden terwijl ze langetermijnafhankelijkheden vastleggen. Er bestaat echter een kritieke kloof in het effectief schalen van deze modellen. Hoewel Maximal Update Parametrization (µP) succesvol heeft gezorgd voor zero-shot hyperparameter transfer (specifiek leer snelheden) voor standaard Transformers en sommige lineaire modellen, blijft de toepassing ervan op Gated Delta Networks nog onverkend. De kernuitdaging ligt in de recurrente staatdynamiek van de architectuur, die wordt bijgewerkt via data-afhankelijke gating-mechanismen en full-matrix delta-regels, in plaats van de scalaire recurrences of standaard feedforward-structuren die door bestaande µP-afleidingen worden verondersteld. Eerder werk aan Structured State Space Models (SSMs) zoals Mamba heeft aangetoond dat vanilla µP faalt voor diagonale SSM's, maar het gebruik van full-matrix state updates door GDN maakt die specifieke correcties onbruikbaar.

Methodologie
De auteurs adresseren deze kloof door de µP-formulering voor Gated Delta Networks rigoureus af te leiden. De methodologie omvat het propageren van coördinaat-grootte schattingen door de volledige forward pass, inclusief de unieke gating-mechanismen en recurrente staatsovergangen, om de noodzakelijke schalingsregels voor initialisatie, forward multipliers en leer snelheden te bepalen.

Belangrijke analytische stappen zijn:

1. Forward Pass Analyse: De auteurs leiden de coördinaat-groottes af van geprojecteerde features, de latente staat $S_t$, en de readout. Ze identificeren dat de L2-normalisatie van queries en keys, gecombineerd met de rank-one write updates, resulteert in specifieke variantie-schalingen. Cruciaal is dat zij een $\sqrt{d}$ multiplier voorstellen vóór de RMSNorm-laag om ervoor te zorgen dat de input aan de normalisatielaag een $\Theta(1)$ coördinaat-grootte behoudt, wat de verstoring van de gradiënt-schaling die nodig is voor feature learning voorkomt.
2. Gating Mechanisme Analyse: Het artikel analyseert de niet-lineaire gating parameters $\alpha_t$ (decay) en $\beta_t$ (write strength). In tegen tegenstelling tot standaard gewichten, zijn deze afgeleid van scalaire pre-activaties die trainbare gewichtsvectoren ($W_\alpha, W_\beta$) en scalaire parameters ($a_{log}, b$) bevatten. De analyse onthult dat deze gating-componenten niet de standaard µP schalingswetten volgen vanwege hun specifieke rollen in de state update vergelijking.
3. Backward Pass en Leer Snelheid Afleiding:
   • AdamW: Vanwege de normalisatie van gradiënten door AdamW door hun tweede moment, is de effectieve update-grootte consistent over verschillende gewichtsklassen, wat de eisen voor de leer snelheid vereenvoudigt.
   • SGD: Onder plain SGD variëren de ruwe gradiënt-magnitudes aanzienlijk tussen verschillende gewichtsklassen. De auteurs leiden af dat de gating gewichtsmatrices ($W_\alpha, W_\beta$) een leer snelheid schaling van $\Theta(1/\sqrt{d})$ vereisen, en de scalaire gating parameters ($a_{log}, b$) een schaling van $\Theta(\sqrt{d})$ vereisen. Deze schalingen wijken af van de standaard µP setting waar leer snelheden doorgaans uniform zijn of een eenvoudige $1/n$ regel volgen.

Kernbijdragen

1. Theoretische Afleiding: Het artikel biedt de eerste volledige µP-formulering voor Gated Delta Networks. Het leidt coördinaat-grootte schattingen af door de volledige forward pass en stelt principiële initialisatie-varianties en forward multipliers vast voor alle gewichtsklassen.
2. Nieuwe Schalingsregels: De auteurs ontdekken dat de gating-componenten niet-standaard leer snelheid schalingen vereisen onder SGD:
   • Gating gewichtsmatrices ($W_\alpha, W_\beta$) vereisen $\Theta(1/\sqrt{d})$ schaling.
   • Scalaire gating parameters ($a_{log}, b$) vereisen $\Theta(\sqrt{d})$ schaling.
3. Empirische Validatie: De auteurs pretrainen GDN taalmodellen over meerdere breedtes (variërend van 256 tot 1536 hidden dimensies) met behulp van zowel AdamW als SGD optimizers.
   • Resultaten: Onder de voorgestelde µP-configuratie transfeert de optimale leer snelheid zero-shot over alle model breedtes voor beide optimizers. In tegenstelling hiermee falen Standard Parametrization (SP) en de originele µP-configuratie bij transfer, waarbij de optimale leer snelheden significant verschuiven naarmate de model breedte toeneemt.
   • Stabiliteit: De voorgestelde configuraties maken stabiele leer-snelheid transfer mogelijk, wat de theoretische afleidingen bevestigt.

Betekenis en Claims
Het artikel claimt dat de analyse het openstaande probleem van het parametriseren van Gated Delta Networks voor feature learning op schaal oplost. Door aan te tonen dat de afgeleide µP-formulering zero-shot hyperparameter transfer mogelijk maakt, valideert het werk de juistheid van de theoretische afleiding en de praktische bruikbaarheid ervan. De auteurs benadrukken dat dit het eerste werk is dat een µP-consistente parametrisering specifiek voor Gated Delta Networks afleidt en valideert, wat het onderscheidt van eerder werk op diagonale SSM's. De bevindingen suggereren dat een correcte parametrisering essentieel is voor het schalen van lineaire recurrente modellen en dat de specifieke gating mechanismen in GDN unieke schalingsregels vereisen die verschillen van zowel standaard Transformers als andere lineaire architecturen. De auteurs hopen dat deze afleidingen toekomstig onderzoek naar de scaling laws van andere lineaire of hybride architecturen zullen begeleiden.