Technische Samenvatting: Efficiënt Pre-Training met Token Superpositie

Probleemstelling

Het pre-trainen van Large Language Models (LLM's) wordt steeds onbetaalbaarder en inefficiënter op schaal. Huidige paradigma's vereisen vaak complexe, ingrijpende aanpassingen aan parallelle strategieën, optimalizers, tokenizers of modelarchitecturen om een hoge data-doorvoer te bereiken. Bovendien, terwijl recent onderzoek zich heeft gericht op het verbeteren van de efficiëntie tijdens inferentie of het ontkoppelen van trainings- en inferentie-efficiëntie, ontbreekt het aan methoden die de doorvoer tijdens de training aanzienlijk verhogen zonder de architectuur van het definitieve model of de inferentiedynamiek te wijzigen. Bestaande benaderingen verwarren vaak trainings- en inferentie-efficiëntie of vertrouwen op auxiliaire verliezen die parameters en hyperparameters toevoegen zonder noodzakelijkerwijs de ruwe tokenverwerkingsrate per FLOP te verhogen.

Methodologie: Token Superpositie Training (TST)

De auteurs stellen Token Superpositie Training (TST) voor, een "drop-in" methode ontworpen om de data-doorvoer per FLOP tijdens pre-training te verbeteren zonder de parallelle verwerking, de optimalizer, de tokenizer, de data of de modelarchitectuur te wijzigen. TST werkt in twee distincte fasen:

Fase 1: Superpositiefase

In deze fase verwerkt het model een "zak" van opeenvolgende tokens als één latent eenheid, aangeduid als een "s-token".

• Input Superpositie: Een opeenvolgende reeks van $s$ tokens wordt opgesplitst in niet-overlappende zakken. De embeddings van deze $s$ tokens worden gemiddeld om een enkele gesuperponeerde latente tokenembedding te creëren. Dit vermindert de sequentielengte met een factor $s$ (van $L$ naar $\lfloor L/s \rfloor$).
• Doorvoermechanisme: Om een eerlijke vergelijking met basislijntraining te garanderen (gelijke FLOP's), wordt de data-sequentielengte $L$ tijdens deze fase met een factor $s$ verhoogd. Bijgevolg neemt het model $s$ keer meer datatokens op voor dezelfde rekenkosten per stap.
• Output Superpositie: In plaats van één volgend token te voorspellen, voorspelt het model de volgende "zak" van $s$ tokens. Het standaard one-hot cross-entropy (CE) verlies wordt vervangen door een Multi-Hot Cross-Entropy (MCE) verlies. Het MCE-doel behandelt de doelzak als een multi-hot label, waardoor de waarschijnlijkheid van elke geldige token in de zak naar $1/s$ wordt geduwd. Het verlies wordt geformuleerd als het gemiddelde van de standaard CE-verliezen voor elke token in de zak, min een constante term ($\log |y|$) die geen invloed heeft op de gradiënten.
• Causaliteit: Om causaliteit te behouden, worden standaard labels voor het voorspellen van het volgende token met $s-1$ naar links verschoven voordat ze worden opgesplitst in niet-overlappende zakken.

Fase 2: Herstelfase

Een model dat uitsluitend met TST is getraind, produceert nonsensische output omdat het de tokenvolgorde binnen de zakken verliest en ontbreekt aan standaard autoregressieve bemonsteringsmogelijkheden. Om dit te verhelpen:

• De training keert terug naar het standaard causale voorspellingsdoel voor het volgende token.
• Het model hervat de training vanaf het checkpoint dat aan het einde van de superpositiefase is opgeslagen.
• De TST-code wordt volledig verwijderd om ervoor te zorgen dat de architectuur van het definitieve model en de inferentiedynamiek identiek zijn aan de basislijn.
• De verhouding van stappen die in de superpositiefase worden doorgebracht, wordt aangeduid als $r$ (bijvoorbeeld: $r=0,3$ betekent dat 30% van de stappen TST zijn en 70% standaard).

Belangrijkste Bijdragen

1. Nieuw Trainingsparadigma: Introductie van TST, wat de token-doorvoer $s$-voudig verhoogt tijdens een specifiek trainingsregime zonder de uiteindelijke architectuur of inferentiegedrag van het model te veranderen.
2. Twee-fasen Ontwerp: Een eenvoudig mechanisme dat input superpositie (gemiddelde embeddings) en output superpositie (multi-hot voorspelling) combineert, gevolgd door een herstelfase om het model weer af te stemmen op standaard autoregressief gedrag.
3. Orthogonaliteit met Bestaande Methoden: De auteurs positioneren TST als orthogonaal ten opzichte van methoden met auxiliaire verliezen (zoals Multi-Token Prediction). Waar MTP parameters toevoegt en toekomstige tokens target zonder tokens-per-FLOP te verhogen, verhoogt TST strikt de doorvoer en laat de architectuur onaangetast.
4. Empirische Validatie: Uitgebreide evaluatie over modelgroottes (270M, 600M, 3B en 10B Mixture-of-Experts) die robuustheid aantoont over verschillende hyperparameterinstellingen.

Resultaten

Het artikel rapporteert de volgende bevindingen gebaseerd op experimenten met de DCLM- en FineWeb-Edu-datasets:

• Verliesreductie: TST presteert consistent beter dan basislijnmodellen wat betreft het definitieve trainingsverlies onder gelijke-FLOP-beperkingen.
• Downstream Prestaties: Modellen getraind met TST tonen verbeterde prestaties op downstream benchmarks (HellaSwag, ARC-E/C, MMLU, enz.) in vergelijking met basislijnen die voor hetzelfde aantal stappen of dezelfde rekenkosten zijn getraind.
• Efficiëntiewinsten:
  • Onder gelijk-verliesinstellingen levert TST tot een 2,5x reductie in totale pre-traintijd op bij de schaal van 10B parameters.
  • In het 10B A1B (10B totaal, 1B actief) MoE-experiment zag het TST-model 2T datatokens in vergelijking met de 1,05T tokens van de basislijn, terwijl het een lager definitief verlies bereikte (2,236 versus 2,252) in aanzienlijk minder B200-uren (4768 versus 12311).
• Hypergevoeligheid: De methode is robuust binnen een redelijk bereik van superpositiezakgroottes ($s \in [4, 8]$) en stapverhoudingen ($r \in [0,2, 0,4]$).
• Ablatiestudies:
  • Input versus Output: Zowel input-only als output-only superpositie verbeteren de prestaties ten opzichte van de basislijn, maar de combinatie van beide levert de beste resultaten op, wat suggereert dat het om orthogonale mechanismen gaat.
  • Representatie-uitlijning: Het willekeurig opnieuw initialiseren van de input embeddings en output LM-koppen tussen de superpositie- en herstelfasen elimineert de winsten van TST volledig, wat aangeeft dat het behoud van representatie-uitlijning over de fasen heen kritiek is.

Betekenis en Claims

Het artikel beweert dat TST een aanzienlijke verbetering biedt in pre-trainingsefficiëntie door de doorvoer tijdens de training te maximaliseren, onafhankelijk van de inferentie-architectuur van het model. De auteurs betogen dat deze aanpak het monolithische pre-trainingparadigma uitdaagt door aan te tonen dat een "grof-naar-fijn" granulariteitschema een herbruikbaar ingrediënt is voor efficiënt pre-training.

Belangrijke claims zijn:

• Kostenreductie: Het bereiken van dezelfde modelkwaliteit (verlies) tegen ongeveer de helft van de rekenkosten.
• Eenvoud: De methode vereist geen complexe architecturale wijzigingen, nieuwe optimalizers of extra parameters, waardoor het een praktische "drop-in" oplossing is.
• Ontkoppeling van Efficiëntie: Het koppelt succesvol de efficiëntie tijdens de training los van de efficiëntie tijdens de inferentie, en adresseert zo een gat in de huidige literatuur waar veel methoden de twee verwarren.
• Schaalbaarheid: De methode schaalt effectief van kleine dichte modellen (270M) tot grote Mixture-of-Experts-modellen (10B), wat wijst op brede toepasbaarheid.

De auteurs blijven bescheiden over de onderliggende mechanismen, en merken op dat hoewel ze hypothetiseren dat input superpositie werkt als een vorm van "pre-pre-training" of impliciete regularisatie van de embedding-geometrie, er verder interpretatieonderzoek nodig is. Ze erkennen ook beperkingen, zoals de aanname dat pre-training momenteel reken-gebonden is in plaats van data-gebonden, en het gebrek aan ablaties op grote schaal vanwege rekenbeperkingen.