Technische Samenvatting: Hidden Decoding op Schaal

Probleemstelling

Het schalen van Large Language Models (LLM's) heeft traditioneel vertrouwd op het vergroten van de Transformer-backbone, wat kostbare rondes van pretraining vereist en de serveerkosten verhoogt. Voor reeds sterke fundamentele modellen is verdere uitbreiding van de backbone vaak onpraktisch. Hoewel recent werk suggereert dat het toewijzen van meer berekening per token (zelfs met vaste parameters) de redenering en nauwkeurigheid kan verbeteren, kampen bestaande benaderingen met aanzienlijke schaalbaarheidsproblemen:

• Depth-Recurrent (Looped) Transformers: Deze hergebruiken dezelfde Transformer-blokken meerdere keren per token om de berekening te verhogen. Ze zijn echter moeilijk te schalen naar frontier-modellen (bijv. 100B+ parameters) omdat de iteratieve aard van de looped berekening de pipeline-parallellisme-aannames doorbreekt die worden gebruikt bij het trainen van de grootste modellen, wat leidt tot GPU-idle tijd.
• Sequence-Length Expansion: Het simpelweg invoeren van langere inputs is compatibel met standaard trainingspipelines, maar vereist een methode om de berekening uit te breiden zonder enkel het contextvenster voor dezelfde token te vergroten.

Het artikel adresseert de uitdaging om een bestaande Transformer-backbone te verbeteren door de berekening per token te verhogen zonder de diepte of breedte van de backbone te wijzigen, terwijl de compatibiliteit met pipeline-parallellisme en grootschalige trainingsinfrastructuur behouden blijft.

Methodologie: Hidden Decoding

De auteurs stellen Hidden Decoding voor, een methode voor sequence-length scaling toegepast tijdens Continued Pretraining (CPT). De kernmechanisme breidt elk input-token uit naar $n$ parallelle stromen die binnen een enkele forward pass worden verwerkt.

1. Multi-Stream Token Expansie

• Embedding Expansie: In plaats van een enkele vocabulaire-embeddingtabel, maakt het model gebruik van $n$ onafhankelijke embeddingtabellen ($E_1, \dots, E_n$). Een input-token $x_i$ wordt uitgebreid naar een sequentie van $n$ representaties: $(E_1(x_i), E_2(x_i), \dots, E_n(x_i))$.
• Sequentie Constructie: Deze representaties worden verweven om een uitgebreide sequentie te vormen van lengte $nL$ (waarbij $L$ de oorspronkelijke sequentielengte is).
• Trainingsdoel: Alleen de laatste stroom ($E_n$) wordt gesuperviseerd. Het model voorspelt de volgende token $x_{i+1}$ op basis van de verborgen toestand van de positie in de laatste stroom. De eerste $n-1$ stromen fungeren als "latente berekeningstoestanden" zonder directe loss, waardoor ze de representatie intern kunnen verfijnen voordat de uiteindelijke voorspelling plaatsvindt.

2. Stream-Factorized Attention

Om te voorkomen dat de attention-kosten kwadratisch ($O(n^2 L^2)$) schalen met de uitgebreide sequentielengte, introduceren de auteurs Stream-Factorized Attention:

• Intra-Stream Lagen: De meeste lagen beperken de attention tot posities binnen dezelfde stroom. Dit reduceert de complexiteit tot ongeveer $O(nL^2)$.
• Cross-Stream Lagen: Een kleine subset van lagen (bijv. full attention of sliding-window lagen uit de basismodel) staat menging toe tussen stromen.
• Efficiëntie: Dit ontwerp houdt de toegevoegde attention-kosten bijna lineair in $n$ in plaats van kwadratisch, wat de methode trainbaar maakt voor modellen met meer dan 100B parameters.

3. Progressieve Expansie en Initialisatie

Om het trainen van grote expansiefactoren te beheren, gebruiken de auteurs Progressieve Expansie (bijv. $n=1 \to 2 \to 4 \to 8$). Ze gebruiken Cyclic Replication Initialization, waarbij nieuwe embeddingtabellen worden geïnitialiseerd door bestaande tabellen te kopiëren (bijv. $E_{n+k} \leftarrow E_k$). Dit stelt het model in staat om aanvankelijk het gedrag van de kleinere configuratie te reproduceren en zich vervolgens soepel aan te passen.

4. Architecturale Optimalisaties (WeLM Specifiek)

Voor de WeLM-backbone maken de auteurs gebruik van een KV-Mirror ontwerp. Omdat tussenliggende stromen geen loss dragen en alleen bijdragen aan de context via Keys en Values (KV), kunnen de gespiegelde lagen (die KV berekenen uit eerdere lagen) de verwerking van tussenliggende stromen volledig overslaan. Dit vermindert de trainingstijd met ongeveer 20% in hun 80B model setup.

Belangrijkste Bijdragen

1. Frontier-Scale Fixed-Backbone Scaling: Het artikel demonstreert de eerste succesvolle toepassing van sequence-length scaling op de 100B+ MoE schaal. Ze trainden WeLM-HD4-80B en WeLM-HD4-617B (met $n=4$), waarmee prestatieverbeteringen werden behaald ten opzichte van hun niet-HD tegenhangers zonder de Transformer-backbone te vergroten.
2. Efficiënte Sequence-Length Expansie: De introductie van Stream-Factorized Attention maakt training mogelijk met uitgebreide sequenties tegen een kostenniveau dat bijna lineair is in $n$ (bijv. een 4x expansie kost ~4.4x–5.1x per batch, ver onder de 16x kosten van dense attention).
3. Expansion-Factor Scaling: Experimenten tonen aan dat het verhogen van de expansiefactor $n$ (van 2 naar 8) consistent de taalmodellerings-loss en downstream nauwkeurigheid verbetert, wat valideert dat $n$ een praktische schaalbare knop is voor vaste backbones.

Experimentele Resultaten

De auteurs evalueerden Hidden Decoding op frontier-scale MoE-modellen en dense baselines:

• Frontier-Scale Prestaties (WeLM-HD4-80B & 617B):

  • WeLM-HD4-80B presteerde beter dan WeLM-80B op alle negen gedeelde harde benchmarks, met significante winsten in SciCode (+4.2), PHYBench (+4.0) en FrontierMath (+3.2).
  • WeLM-HD4-617B verbeterde ten opzichte van WeLM-617B over alle negen benchmarks, inclusief GPQA Diamond (+2.1), HLE (+1.8) en FrontierMath (+2.0).
  • Deze modellen werden getraind met een vroege SFT-only post-training recept om een gecontroleerde vergelijking te garanderen, waarbij het effect van Hidden Decoding werd geïsoleerd.

• Schalen met de Expansiefactor ($n$):

  • Op de dense Qwen3-8B-Base, resulteerde het verhogen van $n$ van 1 naar 8 in monotone verbeteringen in gemiddelde nauwkeurigheid (bijv. MMLU verbeterde van 78.2 naar 81.6).
  • Op de 80B MoE, verhoogde het verhogen van $n$ van 1 naar 8 de MMLU van 85.0 naar 87.5 en verlaagde de Pile-test BPB van 0.386 naar 0.378.

• Ablatie-studies:

  • Supervisie Ontwerp: Alleen de laatste stroom superviseren leverde betere resultaten op dan het toepassen van loss op alle stromen of het sommeren van stroom-outputs.
  • KV Retentie: Het behouden van aparte KV-caches voor elke stroom (in plaats van ze te delen) verbeterde de nauwkeurigheid, wat bevestigt dat tussenliggende stromen een duidelijke, nuttige context bieden.
  • Attention Compositie: Modellen met slechts enkele volledige cross-stream lagen (bijv. 4 van de 27) presteerden bijna net zo goed als modellen met volledige cross-stream mixing in alle lagen, wat de efficiëntie van Stream-Factorized Attention valideert.

• Inference Doorvoersnelheid:

  • Hidden Decoding behoudt een hoge doorvoersnelheid in latency-gevoelige scenario's met kleine batches (het behoudt 83–88% van de baseline doorvoersnelheid bij batch size 1).
  • De doorvoersnelheid neemt af in scenario's met lange context en grote batches (daalt naar ~27% bij batch size 32 met 32k input), wat de toegenomen berekening per token reflecteert.

Betekenis en Claims

Het artikel claimt dat Hidden Decoding een praktische, engineering-realistische weg biedt voor het schalen van frontier LLM's wanneer verdere backbone-expansie te kostbaar is.

• Compatibiliteit: In tegen wordt van looped Transformers, past Hidden Decoding natuurlijk bij pipeline-parallellisme en standaard grote-model trainingsstacks omdat de extra berekening zich manifesteert als een langere sequentie in plaats van iteratieve diepte.
• Effectiviteit: Het is de eerste gedemonstreerde methode die er succesvol in is geslaagd om sequence-length expansion te schalen naar het 100B+ MoE regime, wat meetbare verbeteringen oplevert in redeneer- en kennisbenchmarks.
• Schaalbaarheid: De methode bewijst dat de expansiefactor $n$ een levensvatbare schaalas is, die prestatiewinsten biedt die groeien met $n$ zonder nieuwe modelarchitecturen of massale hertraining van de backbone te vereisen.

De auteurs positioneren dit werk als een oplossing voor de "scaling wall" waar sterke fundamentele modellen voor staan, door een manier te bieden om meer capaciteit uit bestaande parameters te halen via efficiënte latente berekening.