Riepilogo Tecnico: Domino

Enunciato del Problema

Lo speculative decoding accelera l'inferenza dei Large Language Model (LLM) utilizzando un modello bozza leggero per proporre più token futuri, che vengono poi verificati in parallelo dal modello target. L'efficacia di questo approccio è governata da un compromesso fondamentale tra qualità della bozza e costo di generazione:

1. Generazione Autoregressiva (es. serie EAGLE): Questi metodi generano i token bozza in sequenza, modellando esplicitamente le dipendenze causali all'interno del blocco bozza. Ciò produce bozze di alta qualità e lunghezze di accettazione elevate. Tuttavia, generare $k$ token richiede $k$ passaggi in avanti sequenziali e $k$ proiezioni LM-head su vocabolario completo, comportando un sovraccarico sequenziale significativo che limita l'accelerazione.
2. Generazione Parallela (es. DFlash): Questi metodi generano un intero blocco bozza in un singolo passaggio in avanti, eliminando il sovraccarico sequenziale e massimizzando l'utilizzo della GPU. Tuttavia, rimuovendo le dipendenze sequenziali, indeboliscono la modellazione causale intra-blocco, risultando spesso in lunghezze di accettazione inferiori rispetto ai metodi autoregressivi.

Gli approcci esistenti faticano a bilanciare questi fattori. Ridurre i costi di proiezione del vocabolario (es. FR-Spec) o utilizzare la verifica basata su alberi non risolve completamente la tensione tra il costo sequenziale della modellazione causale e la perdita di qualità della generazione parallela. La domanda centrale affrontata è: Possiamo mantenere i benefici di qualità della bozza derivanti dalla modellazione delle dipendenze causali preservando al contempo l'efficienza a basso costo della generazione parallela a blocchi?

Metodologia: Domino

Gli autori propongono Domino, un framework di speculative decoding che disaccoppia la modellazione delle dipendenze causali dall'esecuzione autoregressiva costosa della bozza. L'architettura è composta da due componenti principali:

1. Backbone Bozza Parallelo

Domino utilizza un backbone bozza parallelo (istanziato tramite l'architettura DFlash) per generare distribuzioni bozza preliminari per l'intero blocco in un singolo passaggio in avanti non autoregressivo.

• Input: Il backbone riceve le feature del contesto del modello target dal prefisso verificato e un blocco bozza mascherato (contenente l'ultimo token verificato seguito da token [MASK]).
• Output: Produce stati nascosti per l'intero blocco, dai quali vengono calcolati i logit di base ($L^{base}_i$) utilizzando una LM-head del target congelata.

2. Testa Domino (Correzione Causale Leggera)

Per reintrodotre informazioni causali senza esecuzione sequenziale, Domino applica una testa di correzione leggera sopra i logit di base paralleli. Questa testa opera in sequenza ma è computazionalmente economica:

• Codificatore Causale: Per ogni posizione bozza $i$, un leggero GRU riassume gli embedding dei token bozzati precedenti (o token ground-truth durante l'addestramento) in uno stato causale $S_{i-1}$. Questo stato cattura informazioni dipendenti dal prefisso senza invocare un modello autoregressivo completo.
• Testa di Correzione a Bassa Rango: La testa calcola una correzione residua nello spazio dei logit utilizzando un collo di bottiglia a bassa rango. Concatena lo stato nascosto di base $H_i$ e lo stato causale $S_{i-1}$, li proietta in uno spazio a bassa rango (dimensione $r=256$) e li mappa nuovamente nello spazio del vocabolario per produrre logit di correzione $\Delta L_i$.
• Logit Finali: I logit bozza finali sono calcolati come $L_i = L^{base}_i + \Delta L_i$. I token vengono campionati da questa distribuzione raffinata.

Crucialmente, la correzione viene applicata nello spazio dei logit piuttosto che nello spazio nascosto. Questo evita la necessità di rieseguire l'intera LM-head dopo ogni aggiornamento causale, mantenendo il calcolo di base parallelo e limitando il ramo sequenziale a un aggiornamento residuo a bassa rango.

Strategia di Addestramento

Per stabilizzare l'addestramento e impedire al ramo di correzione di "scorciatoiare" il backbone, gli autori introducono due strategie chiave:

1. Codifica Causale Teacher-Forced: A differenza dei metodi che utilizzano prefissi auto-generati (Testing durante l'Addestramento), Domino fornisce al codificatore causale gli embedding dei token ground-truth. Questo allinea la distribuzione di addestramento con il meccanismo di accettazione dello speculative decoding, dove le correzioni contano solo quando i token precedenti sono verificati come corretti.
2. Curriculum Ancorato alla Base: Per impedire al backbone parallelo di collassare (cioè diventare debole perché il ramo di correzione lo compensa), la funzione di perdita di addestramento è una somma ponderata della perdita sui logit di base e della perdita sui logit finali. Il peso $\lambda_t$ viene annealato linearmente da 1 a 0, costringendo inizialmente il backbone a imparare una forte distribuzione di base prima di spostare gradualmente l'ottimizzazione verso la distribuzione finale corretta.

Contributi Chiave

• Disaccoppiamento della Modellazione Causale dall'Esecuzione: Domino separa la modellazione delle dipendenze causali dall'esecuzione costosa dei passaggi autoregressivi, permettendo un affinamento causale senza latenza sequenziale.
• Correzione Causale Leggera: Utilizzando un codificatore basato su GRU e un collo di bottiglia a bassa rango nello spazio dei logit, Domino inietta informazioni dipendenti dal prefisso con un sovraccarico di parametri minimo (aggiungendo solo ~5.3% di parametri al backbone) e un aumento di latenza trascurabile.
• Curriculum di Addestramento Stabilizzato: L'introduzione del teacher forcing combinato con una funzione di perdita curriculum ancorata alla base previene il collasso del backbone e garantisce un apprendimento robusto sia della distribuzione di base parallela che delle correzioni causali.

Risultati Sperimentali

Gli esperimenti sono stati condotti sui modelli Qwen3-4B e Qwen3-8B su benchmark di ragionamento matematico (GSM8K, MATH, AIME25), generazione di codice (HumanEval, MBPP, LiveCodeBench) e dialogo (MT-Bench, Alpaca).

• Accelerazione End-to-End: Con il backend Transformers, Domino raggiunge un'accelerazione fino a 5.49× su Qwen3-8B (decodifica greedy), superando DFlash (4.66×) ed EAGLE-3 (2.26×). Su GSM8K specificamente, raggiunge un'accelerazione di 7.92×.
• Lunghezza di Accettazione: Domino migliora significativamente la lunghezza media di accettazione rispetto alle baseline parallele (ad esempio, aumentando da 6.59 a 10.03 su Qwen3-8B per GSM8K), riducendo il divario con i metodi autoregressivi mantenendo al contempo l'efficienza parallela.
• Throughput ad Alta Concorrenza: Utilizzando SGLang, Domino dimostra un throughput superiore a vari livelli di concorrenza (da 2 a 32 richieste), raggiungendo un'accelerazione del throughput fino a 5.8× rispetto alla decodifica autoregressiva.
• Studi di Ablazione:
  • Rimuovere la testa Domino riduce l'accelerazione media da 3.31× a 2.84×, confermando l'efficacia della correzione causale.
  • Utilizzare teacher forcing invece del testing durante l'addestramento migliora la lunghezza di accettazione (da 3.80 a 3.96).
  • Il curriculum ancorato alla base aumenta ulteriormente la lunghezza di accettazione (da 3.96 a 4.19) e previene il collasso della perdita del backbone.

Significato

Il paper afferma che Domino dimostra che le informazioni causali possono essere efficacemente reintrodotte nella generazione parallela per ottenere un'inferenza LLM più rapida. Disaccoppiando la modellazione delle dipendenze dal costo di esecuzione, Domino raggiunge un equilibrio superiore tra qualità della bozza e sovraccarico di generazione. I risultati suggeriscono che una correzione causale leggera è una via percorribile per superare i limiti delle strategie di speculative decoding puramente autoregressive e puramente parallele, offrendo una soluzione pratica per accelerare il servizio LLM senza compromettere la fedeltà della distribuzione di output.