Mixture of Training: Recombining Small-Scale Scaffolded Pretraining Runs into a Larger Language Model

Technische Samenvatting: Mixture of Training (MoT)

Probleemstelling
De pre-training van Large Language Models (LLM's) is momenteel georganiseerd als een enkele, monolithische end-to-end optimalisatierun. Deze architectuur creëert een schaalbaarheidsbottleneck waarbij alle lagen gelijktijdig getraind moeten worden, wat het proces kwetsbaar maakt voor fouten (waarbij een enkele fout de gehele run stopt) en inflexibel voor iteratief onderzoek (waarbij verbeteringen een herstart of voortzetting van een massief gekoppeld systeem vereisen). Het artikel vraagt zich af of pre-training kan worden gedecomposeerd in kleinere, onafhankelijk trainbare jobs die later kunnen worden gerecomposeerd tot een coherent, groter model zonder kwaliteitsverlies.

Methodologie: Mixture of Training (MoT)
De auteurs introduceren Mixture of Training (MoT), een gescaffolded modulaire pre-training procedure ontworpen om een doel-Transformer te partitioneren in aaneengesloten laagblokken. De methode werkt in drie stadia:

1. Stadium 0 (Voorbereiding): Een vooraf getrainde aligner ($A$) wordt geselecteerd of voorbereid. De doel-Transformer ($F$) en de aligner worden beide gepartitioneerd in $K$ aaneengesloten laagblokken. De aligner is vorm-compatibel met het doelmodel, waarbij ze dezelfde globale breedte, aandachtshoofd-dimensionaliteit, feed-forward breedte, token-embeddings en output-heads delen.
2. Stadium 1 (Scaffolded Training): Het doelmodel wordt gedecomposeerd in $K$ submodellen ($f_1, \dots, f_K$). Voor elk submodel $f_i$ wordt een scaffolded netwerk $S_i$ geconstrueerd door $f_i$ in de bevroren aligner-structuur in te voegen:
   $$S_i = a_K \circ \dots \circ a_{i+1} \circ f_i \circ a_{i-1} \circ \dots \circ a_1$$
   In deze opstelling is alleen het doelblok $f_i$ trainbaar; alle aligner-segmenten zijn bevroren. Elk scaffold wordt onafhankelijk geoptimaliseerd met de standaard next-token prediction loss. Cruciaal is dat er tijdens deze fase geen gradiënten worden uitgewisseld tussen de doelblokken, maar ze leren allemaal binnen de stabiele representationele context die door de gedeelde aligner wordt geboden.
3. Stadium 2 (Recompositie en Adaptatie): De aligner wordt weggegooid. De getrainde doelblokken worden gerecomposeerd tot het volledige model $\hat{F} = f_K \circ \dots \circ f_1$. Dit "cold-composed" model kan direct worden geëvalueerd, of het kan een korte end-to-end adaptatie-pass ondergaan om resterende representationele mismatches op te lossen.

Belangrijkste Bijdragen

1. Scaffolded Modulaire Training: Een procedure die doel-diepteslices onafhankelijk traint terwijl een gedeelde representationele interface behouden blijft via een bevroren aligner, wat recompositie mogelijk maakt zonder achteraf lagen aan elkaar te naaien.
2. Bewijs van Mechanisme: Een kleinschalige studie die aantoont dat onafhankelijk getrainde diepte-slices succesvol kunnen worden gerecomposeerd tot een bruikbaar 1,3B-parameter taalmodel.
3. Rekentechnische en Planning-verantwoording: Een gedetailleerde analyse van de trade-offs tussen rekenkracht (FLOPs), token-expositie en het geïdealiseerde laag-equivalent kritieke pad. Dit omvat expliciete behandeling van de kosten van de aligner, waarbij onderscheid wordt gemaakt tussen "directe" trainingskosten en "volledig doorbelaste" kosten (waarbij de volledige aligner-kost aan een enkele run wordt toegewezen) versus geamortiseerde kosten over meerdere runs.
4. Ablatie-studies: Empirisch bewijs dat de noodzaak van de aligner aantoont, de impact van disjuncte datastromen, en de kwaliteit-efficiëntie trade-off geassocieerd met het verhogen van het aantal splits ($K$).

Experimentele Resultaten
De auteurs evalueerden MoT op een 12-laags, 1,3B-parameter Gemma-stijl model getraind op de C4-dataset.

• Kwaliteitspariteit: Een "quality-parity" schema, dat de training van submodellen verlengde en een 30k-step adaptatie-pass bevatte, bereikte dezelfde perplexiteit (15,0) als de monolithische baseline.
• Rekentechnische Trade-offs:
  • Het quality-parity schema verwerkte meer geaggregeerde tokens (47,1B vs. 33,6B) en had een korter geïdealiseerd kritiek pad (1,7x versnellingspotentieel), maar vereiste meer totale FLOPs (285,0 EFLOPs vs. 268,4 EFLOPs) wanneer de aligner-kosten volledig werden doorbelast aan een enkele run.
  • Echter, als de aligner wordt hergebruikt over $R$ onafhankelijke runs, daalt de effectieve kosten per run. Het artikel merkt op dat voor $R \ge 3$, de effectieve kosten onder de monolithische baseline liggen.
  • Een lager-rekenkracht schema (MoT + 15k adaptatie) bereikte een perplexiteit van 15,9 bij 189,4 EFLOPs (volledig doorbelast), wat onder het baseline-budget ligt, hoewel er geen gelijkwaardige-rekenkracht monolithische baseline werd geleverd voor directe vergelijking bij dit specifieke budget.
• Ablaties:
  • Noodzaak van de Aligner: Zonder de aligner stort de kwaliteit van de cold-compositie in (PPL 38,9), wat bevestigt dat de aligner de interface-mismatch voorkomt.
  • Datastromen: Disjuncte datastromen voor submodellen verbeterden de kwaliteit van de cold-compositie (PPL 19,3) vergeleken met gedeelde data (PPL 20,3) wanneer een aligner aanwezig was.
  • Split Count: Het verhogen van splits van $K=2$ naar $K=4$ verminderde de rekenkracht maar degradeerde de kwaliteit van de cold-compositie (PPL 24,8), wat de trade-off tussen parallellisme en interface-complexiteit benadrukt.

Betekenis en Claims
Het artikel positioneert MoT niet als een algemene vervanging voor monolithische pre-training, maar als een kleinschalig framework voor het bestuderen van de vraag of gescaffolded sub-runs kunnen fungeren als herbruikbare trainingsunits.

De auteurs claimen drie primaire conclusies:

1. Onafhankelijk getrainde submodellen kunnen worden gerecomposeerd tot een coherent taalmodel als ze binnen een gedeelde aligner-scaffold worden getraind; zonder deze scaffold faalt cold-compositie.
2. De mismatch tussen onafhankelijk getrainde blokken is grotendeels herstelbaar; een korte adaptatie-pass overbrugt het grootste deel van de kloof, en een langer schema kan de pariteit met monolithische training bereiken.
3. MoT verandert de engineering-vorm van pre-training door kleinere, onafhankelijk programmeerbare, herstartbare en ablateerbare jobs te creëren. Dit is bijzonder geschikt voor kleinschalig trainingsonderzoek, wat het potentieel goedkoper, reproduceerbaarder en gemakkelijker itereerbaar maakt.

De auteurs stellen expliciet dat het huidige bewijs modulaire composabiliteit aantoont en een geamortiseerd-hergebruik-planningsregime identificeert, in plaats van een uniforme rekenkracht of gelijkwaardig-hardware voordeel vast te stellen. De studie is opzettelijk beperkt in omvang (één modelfamilie, één dataset, perplexiteit-alleen metrieken) om te dienen als een bewijs van mechanisme voor een concrete designruimte van herbruikbare scaffolded sub-runs.