DBLAST: Dependent Block Drafting for Stochastic Speculative Decoding

Riepilogo Tecnico: DBlast – Dependent Block Drafting per lo Stochastic Speculative Decoding

Problematica

Lo speculative decoding accelera l'inferenza dei Large Language Model (LLM) utilizzando un drafter leggero che propone molteplici token futuri, i quali vengono poi verificati in parallelo da un modello target. Sebbene i recenti drafter a blocchi (ad esempio, DFlash) possano predire diverse posizioni in un singolo passaggio, essi si affidano tipicamente al campionamento di blocchi indipendenti. Questo approccio assume che i token all'interno di un blocco di bozza siano condizionalmente indipendenti dato il prefisso.

Questa assunzione regge bene per il decoding greedy o a bassa entropia, ma diventa fragile nei regimi di decoding stocastico (non-greedy), che sono essenziali per compiti aperti come il dialogo, la scrittura creativa e l'esplorazione del reinforcement learning. Nei contesti stocastici, la distribuzione target è deliberatamente diversificata e molteplici continuazioni sono plausibili. Il documento identifica un disallineamento fondamentale: mentre il drafter campiona i token indipendentemente per corrispondere alle distribuzioni marginali del target, il modello target li verifica sequenzialmente, condizionando ogni accettazione ai token precedentemente accettati. Di conseguenza, un blocco con marginali per posizione accurate può fallire nel formare una traiettoria condizionale coerente, portando a una significativa degradazione della lunghezza del draft accettato all'aumentare dell'entropia della distribuzione di campionamento del target.

Metodologia

Gli autori propongono DBlast (Dependent Block Drafting), un framework che introduce dipendenze nei drafter a diffusione di blocchi senza sacrificare l'efficienza della predizione parallela. Il metodo consiste in due componenti principali:

1. Campionamento di Blocchi Dipendenti tramite Mixture Latente:
   Invece di fattorizzare la distribuzione di proposta attraverso le posizioni ($q(x_{1:b}|y) = \prod q_i(x_i|y)$), DBlast introduce una variabile latente categorica $z$ (un "ramo" o "modalità") per indurre correlazioni all'interno del blocco. La distribuzione congiunta è modellata come una mistura canonica di tipo CP (canonical polyadic) a basso rango:
   $$q_{dep}(x_{1:b} | y) = \sum_{z=1}^{K} q(z | y) \prod_{i=1}^{b} q_i(x_i | y, z)$$

   • Architettura: Il metodo modifica un drafter in stile DFlash aggiungendo un leggero "espansore di offset nascosti per categoria" che mappa gli stati nascosti del transformer in $K$ offset specifici per categoria, e una "testa di priorità della categoria" per predire $q(z|y)$.
   • Inferenza: Durante l'inferenza, viene campionato una categoria latente (opzionalmente con una temperatura $Z_T$), e il relativo ramo viene decodificato in modo greedy in un singolo passaggio parallelo. Ciò preserva l'efficienza del passaggio singolo del block drafting pur permettendo al modello di rappresentare continuazioni a livello di blocco coerenti e alternative.

2. Obiettivo di Training Orientato all'Accettazione:
   Il training standard minimizza la Negative Log-Likelihood (NLL) dei blocchi target, ottimizzando la verosimiglianza della proposta ma ignorando la natura sequenziale della verifica speculativa. DBlast introduce una funzione surrogata basata sulla Lunghezza Accettata Attesa (Expected Accepted Length - AL).

   • L'obiettivo è derivato dall'aspettativa della lunghezza del prefisso accettato, decomposta in eventi di accettazione del prefisso.
   • Per gestire l'instabilità del weighting di importanza quando il drafter e le distribuzioni target hanno una bassa sovrapposizione (nella fase iniziale del training), gli autori propongono un surrogato nel dominio logaritmico con troncamento a soglia. Questa perdita somma il rapporto logaritmico delle probabilità e il logaritmo della lunghezza condizionale accettata sui prefissi mantenuti dove il rapporto di importanza supera una soglia $\tau$.
   • Questo approccio mira direttamente al comportamento di accettazione sequenziale del verificatore, complementando la modellazione dipendente quando la sola verosimiglianza del blocco non è sufficiente.

Contributi Chiave

1. Diagnosi del Disallineamento: Gli autori dimostrano empiricamente che il campionamento di blocchi indipendente fallisce sempre di più man mano che il decoding del target diventa meno deterministico. Mostrano che la lunghezza accettata degrada significativamente nei regimi ad alta entropia perché le proposte indipendenti mancano della coerenza necessaria per la verifica sequenziale.
2. Campionamento di Blocchi Dipendenti: Introducono la prima applicazione di proposte a mistura latente di tipo CP ai drafter a diffusione di sottoparole (subword block-diffusion drafters) per il campionamento speculativo non-greedy. Questo metodo induce correlazioni tra le posizioni del blocco mantenendo la predizione parallela dei token.
3. Loss Orientata all'Accettazione: Propongono una nuova funzione surrogata di training motivata dalla lunghezza accettata attesa. Questo obiettivo, combinato con il troncamento a soglia, fornisce un segnale di training stabile che allinea il drafter con il processo di verifica sequenziale.

Risultati Sperimentali

Gli esperimenti sono stati condotti utilizzando modelli target Qwen3-4B e Qwen3-8B su quattro benchmark: GSM8K (matematica), MT-Bench (chat), HumanEval (coding) e un benchmark di scrittura creativa. L'evaluzione ha coperto tre regimi di campionamento target che vanno dall'alta determinazione (Temperatura 0.7) alla bassa determinazione (Temperatura 1.5).

• Guadagni di Performance: DBlast ha costantemente superato il campionamento di blocchi indipendente (DFlash) e i baseline addestrati con NLL.
  • Nel setting ad alta entropia (Target Temp 1.5), DBlast ha ottenuto un guadagno medio macro di 12,1% nella lunghezza accettata per Qwen3-8B rispetto al baseline indipendente.
  • I guadagni sono stati più pronunciati nelle regioni a bassa determinazione, confermando che la modellazione della dipendenza è più preziosa quando molteplici continuazioni sono plausibili.
• Studi di Ablazione:
  • Categorie Latenti ($K$): L'aumento di $K$ ha migliorato le performance, con rendimenti decrescenti osservati dopo $K=4$.
  • Obiettivo di Training: La loss orientata all'accettazione (AL) ha costantemente superato il training NLL, anche per i drafter indipendenti ($K=1$), sebbene la combinazione di dipendenza e AL abbia prodotto i risultati migliori.
  • Temperatura di Inferenza: Calibrare la temperatura di campionamento della categoria ($Z_T$) ha permesso al drafter di corrispondere meglio alla stocasticità della distribuzione target, migliorando ulteriormente la lunghezza accettata.

Significatività e Rivendicazioni

Il paper sostiene che DBlast affronti una lacuna critica nel decoding speculativo: l'incapacità degli attuali drafter a blocchi di gestire efficientemente la verifica stocastica. Dimostrando che la modellazione della dipendenza intra-blocco e il training allineato all'accettazione sono complementari, gli autori mostrano che è possibile ottenere alti tassi di accettazione in scenari di decoding diversificati e non-greedy senza sacrificare l'efficienza parallela che rende attraente lo speculative decoding.

Gli autori notano con modestia che, sebbene la loro funzione surrogata non sia un estimatore non distorto o un limite inferiore formale, essa funge da segnale di training pratico ed efficace. Riconoscono inoltre che la relazione tra la mistura latente "soft" utilizzata durante il training e il decoding a rami greedy utilizzato durante l'inferenza rimane un'area aperta per una caratterizzazione formale, sebbene i loro risultati empirici dimostrino un trasferimento efficace. Il lavoro suggerisce che i futuri sistemi di speculative decoding per la generazione aperta debbano andare oltre la predizione di token indipendenti per catturare la coerenza strutturale delle sequenze future.