MARS: Enabling Autoregressive Models Multi-Token Generation

1. Il Problema

I modelli linguistici autoregressivi (AR) attuali generano testo un token alla volta, eseguendo un passaggio in avanti (forward pass) per ogni singolo token, indipendentemente da quanto quel token sia prevedibile dal contesto precedente. Questo approccio comporta un costo computazionale uniforme e spesso inefficiente, specialmente per sequenze altamente prevedibili.

Le soluzioni esistenti per la generazione multi-token presentano significativi svantaggi:

• Speculative Decoding: Richiede un modello "bozza" separato, raddoppiando l'uso di memoria e aggiungendo complessità all'orchestrazione.
• Approcci Multi-Head (es. Medusa, EAGLE): Aggiungono teste di predizione extra all'architettura, richiedendo parametri aggiuntivi e addestramento specifico, complicando le pipeline di produzione.
• Diffusione a Blocchi (Block Diffusion): Le conversioni esistenti da modelli AR a modelli di diffusione a blocchi spesso causano un degrado significativo della qualità, specialmente in compiti di ragionamento e coding, a causa di incompatibilità architetturali.

L'obiettivo è creare un metodo che permetta la generazione multi-token senza modifiche architetturali, senza parametri aggiuntivi e mantenendo il modello come un modello AR standard funzionante.

2. Metodologia: MARS (Mask AutoRegreSsion)

MARS è un metodo di fine-tuning leggero che insegna a un modello AR già istruito (instruction-tuned) a prevedere più token per passaggio, mantenendo la capacità di generare un token alla volta senza perdita di qualità.

Analisi dei "Gap" (Divari)

Gli autori identificano quattro divari tra la predizione AR standard e la predizione mascherata a blocchi (block-masked prediction). Tre di questi sono scelte di design eliminabili che causano il degrado della qualità:

1. Pattern di Attenzione: MARS mantiene l'attenzione causale (sinistra-destra) rigorosa, a differenza di alcuni approcci di diffusione che usano attenzione bidirezionale all'interno dei blocchi.
2. Allineamento dei Logit: MARS preserva la convenzione AR di prevedere il token $t+1$ dalla posizione $t$ (logit spostati a destra).
3. Ordine di Generazione: MARS accetta i token strettamente da sinistra a destra, evitando l'ordine basato sulla confidenza che rompe la causalità.
4. Mascheramento dei Token: Questo è l'unico divario intrinseco necessario per la predizione parallela (sostituire i token futuri con [MASK]).

Addestramento

• Input Concatenato: Per ogni sequenza di risposta $x$, il modello elabora in parallelo due stream: uno "pulito" (token originali) e uno "rumoroso" (blocchi di token sostituiti da [MASK]).
• Maschera di Attenzione Strutturata: Viene utilizzata una maschera di attenzione che permette allo stream rumoroso di vedere solo i token puliti dei blocchi precedenti e i token [MASK] all'interno del proprio blocco (attenzione causale intra-blocco).
• Funzione di Perdita Ibrida: Per evitare che la capacità AR si degradi man mano che la dimensione del blocco ($B$) aumenta, MARS combina due termini di perdita:
  • $L_{mask}$: Cross-entropy sui token mascherati.
  • $L_{AR}$: Perdita standard di predizione del prossimo token sullo stream pulito.
  • Importanza: La componente $L_{AR}$ è cruciale. Senza di essa, il segnale di addestramento AR si diluisce ($1/B$) con l'aumentare della dimensione del blocco, portando a un crollo delle prestazioni su compiti di ragionamento. Con la perdita combinata, il segnale AR rimane sopra il 50% indipendentemente da $B$.

Inferenza

• Finestra Scorrevole: Durante l'inferenza, vengono aggiunti $B$ token [MASK] dopo il prefisso corrente. Il modello esegue un singolo passaggio in avanti.
• Soglia di Confidenza ($\tau$): I token vengono accettati consecutivamente da sinistra a destra finché la probabilità massima $p(x_t)$ supera una soglia $\tau$.
• Adattabilità: Se il modello non è sicuro, accetta solo un token (comportamento AR standard). Se è sicuro, ne accetta più di uno. La soglia $\tau$ può essere regolata dinamicamente al momento della richiesta (on-the-fly) per bilanciare velocità e qualità senza ricaricare il modello.

3. Contributi Chiave

1. Analisi dei Gap: Dimostrano che il degrado della qualità nelle predizioni a blocchi non è intrinseco, ma causato da scelte di design (attenzione bidirezionale, allineamento logit errato, ordine non causale) che MARS risolve.
2. Metodo Senza Modifiche: MARS non richiede nuove teste di predizione, nuovi parametri o modifiche architetturali. È un singolo checkpoint che funge da "supersetto" del modello AR originale.
3. Ruolo della Perdita SFT: Identificano la perdita di addestramento supervisionato (SFT loss) sullo stream pulito come ingrediente fondamentale per mantenere le capacità AR a blocchi grandi.
4. Caching KV a Livello di Blocco: Sviluppano una strategia di caching KV specifica per l'inferenza in batch che sincronizza i blocchi, ottenendo significativi guadagni di velocità nel tempo reale.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Qwen2.5-0.5B e Qwen2.5-7B.

• Qualità (Modalità 1 Token): Quando configurato per generare un token alla volta ($\tau=1.0$), MARS eguaglia o supera il baseline AR SFT su 6 benchmark standard (IFEval, BBH, MMLU-Pro, GPQA, GSM8K, HumanEval). Ad esempio, su Qwen2.5-7B, MARS ottiene un punteggio medio di 58.1 contro 56.6 dell'AR SFT, con miglioramenti specifici su GSM8K (+4.5) e HumanEval (+3.0).
• Velocità vs Qualità (Modalità Multi-Token): Abilitando la generazione multi-token, MARS raggiunge un throughput di 1.5–1.7× token per passaggio in avanti, mantenendo una precisione vicina al baseline.
  • Su Qwen2.5-7B, a $\tau=0.95$, il modello genera 1.68 token/passaggio con una perdita di accuratezza media minima (-1.3 punti), superando comunque il baseline AR puro.
• Velocità Reale (Wall-Clock Speedup): Utilizzando la strategia di caching KV a blocchi, MARS ottiene un speedup di 1.71× rispetto all'AR con KV cache su Qwen2.5-7B (batch size 4), riducendo il tempo totale di inferenza da 276.2s a 161.2s.
• Robustezza: Senza la perdita SFT aggiuntiva, le prestazioni crollano drasticamente all'aumentare della dimensione del blocco (da 30.4 a 22.2 in media su 0.5B passando da B=4 a B=16). Con la perdita SFT, le prestazioni rimangono stabili.

5. Significato e Impatto

MARS rappresenta un cambiamento di paradigma nella generazione di testo efficiente:

• Semplificazione Operativa: Elimina la necessità di modelli multipli (draft/target) o architetture complesse, permettendo di distribuire un singolo modello che può adattarsi dinamicamente al carico di lavoro.
• Controllo Granulare: Offre un "knob" (manopola) pratico per il deployment: i sistemi di servizio possono aumentare il throughput in tempo reale durante picchi di richieste abbassando la soglia di confidenza, senza riavviare o cambiare modello.
• Efficienza Computazionale: Dimostra che la generazione parallela non è incompatibile con l'architettura AR pura, purché vengano rispettati i vincoli di causalità e allineamento.

In sintesi, MARS trasforma i modelli AR esistenti in sistemi ibridi capaci di accelerazione multi-token senza compromettere la qualità o richiedere risorse infrastrutturali aggiuntive.