MARS: Enabling Autoregressive Models Multi-Token Generation

Probleemstelling

Autoregressieve (AR) taalmodellen genereren tekst token voor token, waarbij elke forward pass precies één token produceert. Dit resulteert in een uniforme rekencost, ongeacht hoe voorspelbaar de volgende token is. Bijvoorbeeld, het genereren van een voor de hand liggende zin als "het antwoord is" kost evenveel rekentijd als het voorspellen van een onzekere, creatieve woordkeuze.

Bestaande oplossingen voor multi-token generatie hebben aanzienlijke nadelen:

• Speculative Decoding: Vereist een apart 'draft'-model naast het doelmodel, wat het geheugengebruik verdubbelt en de orkestratie complex maakt.
• Multi-head benaderingen (bijv. Medusa, EAGLE): Voegen extra voorspellingshoofden toe aan de architectuur, wat extra parameters vereist en specifieke training nodig heeft.
• Block-masked Diffusion: Bestaande methoden die proberen meerdere tokens per stap te voorspellen, leiden vaak tot een aanzienlijke kwaliteitsdaling, vooral bij redenering en codering, vanwege een mismatch met de oorspronkelijke AR-training.

De kernvraag is: Kan een AR-model leren om meerdere tokens per forward pass te genereren zonder de architectuur te wijzigen, extra parameters toe te voegen of de oorspronkelijke kwaliteit te verliezen?

Methodologie: MARS (Mask AutoRegreSsion)

MARS is een lichtgewicht fine-tuning methode die een instructie-geoptimaliseerd AR-model leert om optioneel meerdere tokens per stap te voorspellen, terwijl het volledig compatibel blijft met de standaard AR-inferentie.

1. Analyse van de "Gaps" (Kloven)
De auteurs identificeren vier verschillen tussen traditionele AR-modellen en block-masked predictie (zoals in Block Diffusion):

1. Token Masking: Toekomstige tokens worden vervangen door [MASK] (onvermijdelijk voor multi-token).
2. Attention Pattern: Sommige methoden gebruiken bidirectionele attention binnen blokken. MARS behoudt strikt causale (links-naar-rechts) attention.
3. Logits Alignment: AR-modellen voorspellen $x_{t+1}$ op positie $t$ (rechts-geschoven logits). MARS behoudt deze conventie.
4. Generatievolgorde: Diffusiemodellen accepteren tokens soms in willekeurige volgorde. MARS accepteert tokens strikt links-naar-rechts.

De auteurs concluderen dat alleen gap (1) inherent is aan multi-token predictie. Gaps (2), (3) en (4) zijn onnodige ontwerpkeuzes in eerdere werken die de AR-compatibiliteit breken. MARS sluit deze drie gaten.

2. Training
MARS start vanuit een bestaand AR SFT-checkpoint. Het trainingsschema gebruikt twee parallelle streams binnen één forward pass:

• Clean Stream: De originele tokens worden intact gelaten en getraind met standaard next-token predictie.
• Noisy Stream: Blokken van $B$ tokens worden vervangen door [MASK] en het model moet deze voorspellen op basis van de clean prefix van eerdere blokken.
• Verliesfunctie: Het totale verlies is een som van het masker-verlies ($L_{mask}$) en het standaard AR-verlies ($L_{AR}$) op de clean stream.
  • Belangrijkste innovatie: De toevoeging van $L_{AR}$ is cruciaal. Zonder deze zou het AR-signaal verdwijnen naarmate de blok-grootte ($B$) toeneemt (signal decay), wat leidt tot kwaliteitsverlies. De combinatie houdt het AR-signaal boven de 50%, ongeacht de blok-grootte.

3. Inferentie
Tijdens inferentie worden $B$ [MASK]-tokens toegevoegd aan de huidige prefix. Het model voert één forward pass uit. Tokens worden sequentieel (links-naar-rechts) geaccepteerd zolang de waarschijnlijkheid boven een drempelwaarde ($\tau$) ligt.

• Als $\tau \to 1.0$, accepteert het model maximaal 1 token per stap (identiek aan standaard AR).
• Bij lagere $\tau$ accepteert het model meerdere tokens per stap, wat de doorvoer verhoogt.
• De drempelwaarde kan live worden aangepast tijdens het serveren zonder het model opnieuw te laden.

4. Block-level KV Caching
Om de doorvoer in batch-inferentie te maximaliseren, wordt een blok-niveau KV-cache strategie ontwikkeld. In plaats van de volledige sequentie opnieuw te berekenen, wordt de cache per blok geüpdatet. Dit elimineert de $O(T^2)$ herberekening en zorgt voor lineaire schaalbaarheid.

Belangrijkste Resultaten

De resultaten zijn getest op Qwen2.5-0.5B en Qwen2.5-7B modellen op zes benchmarks (IFEval, BBH, MMLU-Pro, GPQA, GSM8K, HumanEval).

1. Kwaliteitbehoud (One-Token Mode):

   • Wanneer MARS slechts één token per stap genereert, presteert het gelijk aan of beter dan de originele AR SFT-baseline.
   • Bij Qwen2.5-7B steeg de gemiddelde score van 56.6 (AR SFT) naar 58.1 (MARS), met significante verbeteringen op GSM8K (+4.5) en HumanEval (+3.0).
   • Dit bewijst dat multi-token training geen kwaliteitskost met zich meebrengt; het fungeert zelfs als data-augmentatie.

2. Multi-Token Doorvoer:

   • Met multi-token generatie bereikt MARS een doorvoer van 1.5x tot 1.7x tokens per forward pass.
   • De nauwkeurigheidsdaling is minimaal en voorspelbaar. Bij Qwen2.5-7B daalt de gemiddelde score slechts met 1.3 punten (van 58.1 naar 56.8) terwijl de doorvoer met 1.68x toeneemt.
   • De trade-off tussen snelheid en kwaliteit is glad (geen "cliff" waar kwaliteit plotseling instort).

3. Wall-Clock Snelheid:

   • Met de block-level KV-cache strategie bereikt MARS een 1.71x snelheidswinst in wall-clock tijd ten opzichte van standaard AR met KV-cache op Qwen2.5-7B (bij batch size 4).
   • Zonder deze cache zou MARS trager zijn dan AR bij grotere batch sizes door herberekening.

4. Ablatie Studies:

   • Zonder de extra SFT-loss ($L_{AR}$) stort de kwaliteit bij grotere blokken (bijv. $B=16$) in, vooral bij redeneringstaken. De SFT-loss is dus essentieel voor stabiliteit.
   • Vergelijking met Block Diffusion (die bidirectionele attention gebruikt) toont aan dat deze methoden falen op redeneringstaken, wat de noodzaak van causale attention bevestigt.

Bijdragen en Significantie

De belangrijkste bijdragen van dit werk zijn:

1. Architectuurloze Multi-Token Generatie: MARS vereist geen extra parameters, geen extra hoofden en geen aparte modellen. Het is een puur fine-tuning methode.
2. Strikte Superset: Het resulterende model is een "strict superset" van het origineel; het kan exact hetzelfde worden gebruikt als een standaard AR-model (zonder kwaliteitsverlies) of worden omgeschakeld naar multi-token modus voor snelheid.
3. Oplossing voor Signal Decay: De auteurs identificeren en oplossen het probleem van verdwijnend AR-signaal bij grotere blokken door een gecombineerde loss-functie te gebruiken.
4. Dynamische Snelheidsregeling: Het systeem biedt een praktische "knop" voor deployment: de snelheid kan live worden aangepast aan de vraag (bijv. meer doorvoer onder hoge druk) zonder modelwissel of herstart.

Conclusie:
MARS demonstreert dat het mogelijk is om autoregressieve modellen te versnellen door meerdere tokens per stap te genereren, zonder de complexe infrastructuur van speculative decoding of de architecturale ingrepen van multi-head methoden. Het biedt een praktische, schaalbare oplossing voor het verminderen van inferentie-latentie in productieomgevingen, met behoud van de hoge kwaliteit van het oorspronkelijke model.