🤖 AI

xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

Il documento propone xPress, un raffinatore causale leggero che ripristina le dipendenze mancanti nei drafter a diffusione a blocchi attraverso la raffinazione parallela, aumentando significativamente la lunghezza di accettazione e il throughput end-to-end nel decoding speculativo rispetto ai metodi esistenti come dFlash.

Autori originali: Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

Pubblicato 2026-08-04

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

La corsa a parlare più velocemente: perché l'IA ha bisogno di un sistema di bozza migliore

Immaginate di cercare di scrivere una storia, ma con una regola severa: potete scrivere una sola parola alla volta e dovete aspettare che un editor super intelligente controlli la vostra parola prima di poter scrivere la successiva. È così che funzionano oggi la maggior parte dei potenti modelli linguistici IA. Sono incredibilmente accurati, ma sono anche dolorosamente lenti perché devono controllare ogni singola parola una alla volta. Per velocizzare questo processo, gli scienziati hanno inventato un trucco chiamato "decodifica speculativa". Pensatelo come a un assistente veloce e leggermente meno attento che indovina le prossime parole per voi. Se l'editor super intelligente concorda con le intuizioni dell'assistente, potete scrivere tutte quelle parole in una volta sola, saltando il tempo di attesa.

Il problema è che l'assistente è solitamente troppo impaziente. Indovina parole che suonano bene da sole ma che non si incastrano bene in una frase, come uno chef che sceglie ingredienti deliziosi ma dimentica di controllare se stanno bene insieme. Recentemente, è stato creato un nuovo tipo di assistente chiamato "diffusion drafter". Invece di indovinare una parola dopo l'altra, cerca di indovinare un intero blocco di parole tutto in una volta, come se lanciasse un pugno di pezzi di un puzzle sul tavolo. Questo è super veloce, ma poiché li lancia tutti insieme, i pezzi spesso non si collegano correttamente. L'editor super intelligente deve rifiutarne la maggior parte, rallentando tutto di nuovo. La grande domanda che i ricercatori si pongono è: possiamo mantenere la velocità del indovino "tutto in una volta" ma correggere gli errori in modo che l'editor accetti effettivamente le parole?

Entra in scena xPress: il "Raffinatore Parallelo" che risolve il puzzle

Questo articolo introduce una soluzione ingegnosa chiamata xPress. Gli autori, lavorando con modelli come Qwen3-8B, si sono resi conto che, sebbene il rapido "diffusion drafter" sia bravo a indovinare un blocco di parole, trascura la regola cruciale secondo cui le parole dipendono l'una dall'altra (causalità). Per esempio, se la prima parola è "lei", la parola successiva non dovrebbe essere "sono", anche se "sono" è una parola comune di per sé. Il diffusion drafter non lo sa perché indovina tutto simultaneamente.

Per risolvere questo problema, xPress agisce come un raffinatore causale leggero. Immaginate che il diffusion drafter lanci un pugno di pezzi di un puzzle sul tavolo. xPress è una mano veloce e intelligente che guarda l'intero mucchio in una volta sola e sposta i pezzi nella giusta posizione senza smontarli uno per uno. Lo fa utilizzando una tecnica chiamata decodifica Jacobi. Invece di sistemare il puzzle pezzo per pezzo (che è lento), xPress guarda l'intera immagine, effettua un rapido aggiustamento a ogni pezzo simultaneamente e poi guarda di nuovo. Ripete questo ciclo di "osservazione e aggiustamento" solo alcune volte (solitamente circa 4 o 6 volte) finché i pezzi non si incastrano perfettamente.

I risultati sono impressionanti. Utilizzando xPress, il sistema può accettare circa il 30% in più delle parole indovinate in media rispetto all'originale drafter veloce. In alcuni test specifici, come la risoluzione di problemi matematici, l'accelerazione è stata elevata fino al 56%. Misurando la velocità totale del parlato dell'IA, xPress ha reso il processo 1,3 volte più veloce in media, e fino a 1,7 volte più veloce nei casi migliori, rispetto al metodo originale.

L'articolo argomenta esplicitamente contro altri due modi in cui le persone hanno cercato di risolvere questo problema. Un metodo prevede la costruzione di un enorme "albero" di ipotesi, che è complesso e costoso da eseguire. Un altro metodo utilizza una "testa di Markov" che corregge le parole una alla volta in una linea stretta, il che è accurato ma lento perché perde il vantaggio di velocità di indovinare tutto in una volta. Gli autori dimostrano che xPress batte entrambi: è più veloce del correttore lento passo dopo passo e più semplice del complesso metodo ad albero.

Lo studio conferma che xPress funziona iniettando "informazioni causali" (la regola secondo cui le parole dipendono dalle parole precedenti) nel sistema veloce senza rallentarlo. Hanno testato questo approccio su benchmark di matematica, programmazione e chat, scoprendo che xPress ha costantemente superato la concorrenza. L'articolo suggerisce che questo approccio è un miglioramento solido e misurato che consente all'IA di essere sia veloce che accurata, dimostrando che non è necessario sacrificare la qualità per la velocità se si ha il modo giusto per raffinare le proprie intuizioni.

Riepilogo Tecnico: xPress: Raffinamento Parallelo per i Diffusion Drafter nel Decoding Speculativo

1. Definizione del Problema

Il decoding speculativo (SD) accelera la generazione autoregressiva utilizzando un drafter leggero per proporre token da verificare da parte di un modello target più grande. Sebbene i drafter autoregressivi (AR) (ad es. EAGLE) siano efficaci, soffrono di un overhead sequenziale quando generano blocchi lunghi. I drafter a diffusione a blocchi (ad es. dFlash) affrontano questo problema generando un intero blocco di token in un singolo passaggio forward parallelo, riducendo drasticamente la latenza di drafting.

Tuttavia, i drafter a diffusione possiedono un limite fondamentale di accuratezza: predicono i token basandosi su distribuzioni marginali piuttosto che su distribuzioni congiunte. In un processo di diffusione a passaggio singolo, il token $k$ viene campionato senza condizionamento sui valori realizzati dai token da $1 $a$ k-1$. Questa mancanza di causalità intra-blocco porta a sequenze in cui i singoli token sono localmente probabili ma congiuntamente incoerenti (ad es. disaccordo tra soggetto e verbo). Durante la verifica, il modello target rifiuta queste sequenze incoerenti precocemente, limitando severamente la lunghezza di accettazione ( $\tau$ ) e annullando il potenziale aumento di velocità del drafting parallelo.

Le soluzioni esistenti, come il drafting basato su alberi (ad es. PRESTO, DDTree), tentano di mitigare il problema espandendo lo spazio dei candidati in un albero. Tuttavia, questi metodi sono computazionalmente costosi, complessi da servire e spesso cercano di aggirare il problema puntando al richiamo (recall) piuttosto che risolvere l'incoerenza causale sottostante.

2. Metodologia: xPress

Gli autori propongono xPress, un raffinatore causale leggero progettato per ripristinare le dipendenze causali mancanti nei drafter a diffusione senza ricorrere a un lento ciclo autoregressivo token per token.

2.1 Concetto Core

xPress opera come un "raffinatore causale" che prende la proposta iniziale del blocco del drafter a diffusione e la riconcilia attraverso il raffinamento parallelo. Invece di rigenerare i token sequenzialmente, xPress inietta informazioni causali nei logit del drafter a diffusione per correggere le distribuzioni marginali.

2.2 Architettura

Il raffinatore è un modulo leggero che aggiunge un bias di logit appreso ( $\delta_k$ ) ai logit base del drafter ( $s_k$ ) in ogni posizione $k$ . L'architettura consiste in quattro stadi:

Fuse (Fusione): Proietta tre input in uno spazio a basso rango ( $r=256$ $r = 256$ ):
- Lo stato nascosto per posizione del drafter ( $h_k$ ).
- Un riepilogo globale del blocco ( $g$ ) ottenuto tramite mean-pooling.
- L'embedding del token precedente ipotizzato ( $x_{k-1}$ ).
Mix (Miscelazione): Applica un mixer lineare strettamente inferiore-triangolare ( $L$ ) attraverso le posizioni del blocco. Ciò consente alla posizione $k$ di aggregare le informazioni da tutti i predecessori ( $j \le k$ ) in un'unica operazione, imponendo una struttura causale senza computazione sequenziale.
Trans (Trasformazione): Un MLP residuo nello spazio a basso rango.
Read (Lettura): Proietta il risultato nuovamente alla dimensione del vocabolario per produrre il bias $\delta_k$ .

Il logit raffinato finale è $\ell_k = s_k + \delta_k$ .

2.3 Raffinamento Parallelo tramite Jacobi Decoding

Un raffinatore causale diretto richiederebbe un passaggio sequenziale da sinistra a destra, reintroducendo la latenza che xPress mira a evitare. Per risolvere questo, xPress utilizza il decoding di Jacobi:

Inizializzazione: Il blocco viene seminato con le predizioni greedy iniziali del drafter.
Iterazione: Tutte le posizioni vengono aggiornate simultaneamente in parallelo utilizzando lo stato corrente del blocco come prefisso per l'iterazione successiva.
Convergenza: Il processo si ripete per $K$ iterazioni. Teoricamente, per un blocco di dimensione $B$ , il punto fisso sequenziale viene raggiunto in al massimo $B$ iterazioni. Empiricamente, gli autori riscontrano che $K \approx 6$ iterazioni sono sufficienti per eguagliare la lunghezza di accettazione di un decoding sequenziale completo, poiché il seed del drafter è già un tentativo di alta qualità.

2.4 Strategia di Training

xPress è co-addestrato con il drafter a diffusione utilizzando una funzione di perdita composita:

Cross-Entropy (Teacher Forcing): Massimizza la probabilità del token reale dato il prefisso reale.
Total Variation (TV) Loss: Minimizza la distanza tra la distribuzione raffinata e la distribuzione del modello target, ottimizzando direttamente per il tasso di accettazione.
Consistency Loss (Perdita di Consistenza): Una perdita ausiliaria applicata a un secondo passaggio forward dove gli input sono le predizioni dello stesso drafter (auto-condizionamento). Questo allinea la distribuzione di training con il processo di Jacobi durante l'inferenza, evitando che il raffinatore vada in overfitting sui prefissi "gold".
Drafter Anchor: Un termine di perdita applicato all'output del drafter base per garantire che il modello di diffusione sottostante rimanga di alta qualità.

3. Contributi Chiave

Iniezione di Dipendenza Causale: xPress inietta informazioni causali nei drafter a diffusione utilizzando un meccanismo minimo (proiezione a basso rango e un mixer triangolare) che opera sugli stati nascosti del drafer, evitando la necessità di aumenti significativi dei parametri.
Raffinamento Parallelo: Utilizzando il decoding di Jacobi, xPress risolve le dipendenze causali in parallelo. Ottiene l'effetto di $B$ passi autoregressivi dipendenti usando solo $K \ll B$ iterazioni parallele, preservando i vantaggi di velocità del drafting parallelo.
Guadagni End-to-End Coerenti: Il metodo migliora significativamente sia la lunghezza di accettazione che il throughput senza modificare il processo di verifica o richiedere un drafter più grande.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Qwen3-8B utilizzando un drafter dFlash a blocco-16 su sette benchmark (Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP, LiveCodeBench; Chat: MT-Bench).

Lunghezza di Accettazione: xPress aumenta la lunghezza media di accettazione di circa il 30% rispetto al drafter a diffusione dFlash originale.
- Sotto decoding greedy ( $T=0$ ), i guadagni variano dal 18% al 56%.
- Sotto sampling ( $T=1$ ), i guadagni variano dal 19% al 58%.
Speedup del Throughput:
- Greedy ( $T=0$ ): xPress ottiene uno speedup end-to-end medio di 1.3× rispetto al drafter dFlash, con picchi fino a 1.7× (ad es. su GSM8K).
- Sampling ( $T=1$ ): Lo speedup migliora a una media di 1.46×, raggiungendo fino a 1.6×.
Confronto con i Raffinatori Sequenziali: Rispetto al Markov head (un raffinatore causale sequenziale di DSpark), xPress eguaglia o supera la sua lunghezza di accettazione pur essendo significativamente più veloce.
- Con $K=4$ iterazioni di Jacobi, xPress eguaglia l'accuratezza del Markov head ma gira 1.6× più velocemente (379 µs vs 601 µs per blocco).
- Il Markov head richiede un decode seriale fisso di 15 step, mentre xPress converge in meno step paralleli.
Scalabilità: In scenari di serving con grandi batch (vLLM, batch size 16 e 32), xPress mantiene il suo vantaggio. Sebbene lo speedup assoluto rispetto al baseline AR si comprima all'aumentare dei batch (a causa della verifica compute-bound), xPress supera costantemente il semplice drafter dFlash del 1.2–1.4× in throughput, dimostrando che i benefici del raffinatore sopravvivono sotto carico.

5. Significato e Rivendicazioni

L'articolo sostiene che xPress affronta il "disallineamento intrinseco" tra il drafting a diffusione parallelo e la verifica sequenziale, riparando la distribuzione marginale stessa invece di espandere lo spazio di ricerca.

Efficienza: xPress dimostra che l'informazione causale può essere iniettata a un costo quasi nullo rispetto all'overhead di drafting. Il costo di raffinamento è una piccola frazione del tempo totale di drafting, e la natura parallela del raffinamento evita le penalità di latenza del corretto sequenziale.
Ortogonalità: Il metodo è ortogonale ad altri miglioramenti; può essere combinato con drafter più grandi o metodi di verifica basati su alberi, poiché migliora i logit di base su cui questi metodi si affidano.
Generalizzabilità: Gli autori suggeriscono che la filosofia di riconciliare predizioni mutuamente dipendenti tramite pochi iterazioni parallele possa essere applicabile oltre il decoding speculativo, ovunque i modelli emettano output interdipendenti simultaneamente.

L'articolo conclude che xPress riesce a ripristinare la causalità persa nel drafting a diffusione a passaggio singolo, consentendo catene di accettazione più lunghe e un throughput più elevato senza sacrificare l'efficienza parallela che rende attraenti i drafter a diffusione.