Sintesi Tecnica: LESS: Mutual-Stability Sampling per i Diffusion Language Models

Definizione del Problema

I Diffusion Large Language Models (dLLM) offrono un'alternativa promettente al decoding autoregressivo (AR), consentendo aggiornamenti paralleli dei token e condizionamento bidirezionale attraverso la denoising iterativa di sequenze mascherate. Tuttavia, la loro efficienza pratica è attualmente limitata da procedure di sampling che si affidano a programmi di reverse denoising fissi. I sampler standard eseguono un numero predeterminato di step, spesso sprecando computazione su posizioni che hanno già stabilizzato le proprie predizioni, rischiando al contempo un impegno prematuro verso token instabili. Una volta che un token viene de-mascherato in un processo di diffusione, viene tipicamente trattato come fisso; se un token errato viene impegnato precocemente a causa di un programma rigido, i successivi step inversi non possono revisionarlo, portando a un degrado della qualità. Gli esistenti sampler adattivi training-free affrontano questo problema solo parzialmente utilizzando criteri basati sulla confidenza o sul margine, ma spesso non riescono a catturare l'intera stabilità temporale e distributiva della traiettoria predittiva di un token.

Metodologia

Gli autori propongono LESS (Less is More), un sampler adattivo training-free e agnostico rispetto al modello che inquadra l'impegno del token come un problema di stopping online per singola posizione. LESS implementa il mutual-stability sampling, dove una posizione mascherata è idonea alla de-mascherazione solo quando tre specifici segnali di stabilità indicano congiuntamente che la predizione si è assestata.

La Regola della Stabilità Congiunta

Per ogni posizione mascherata $i$ allo step inverso $t$, LESS valuta tre segnali derivati dalla distribuzione predittiva del modello $p_{t,i}$:

1. Confidenza Predittiva Top-1: La probabilità assegnata all'attuale token top-1 ($w_{t,i}$). Una posizione deve superare una soglia di confidenza $c$.
2. Persistenza del Token Top-1: L'attuale token top-1 deve rimanere invariato durante una finestra recente di $P$ step inversi. Ciò evita di impegnarsi in un token che ha appena subito un cambio, anche se la sua attuale confidenza è elevata.
3. Divergenza di Jensen–Shannon (JSD) Inter-Step Top-K: Questo è il contributo innovativo di LESS. Misura la stabilità della distribuzione predittiva stessa tra passi inversi consecutivi.
   • Invece di utilizzare la divergenza di Kullback-Leibler (KL) asimmetrica, LESS utilizza la JSD simmetrica e limitata.
   • Per garantire l'efficienza computazionale e concentrarsi sui candidati rilevanti, la JSD è calcolata su un vocabolario approssimato costituito dall'unione dei top-$K$ token degli step corrente e precedente, con la massa rimanente collassata in un simbolo residuo.
   • La JSD deve essere inferiore a una soglia $d$.

Una posizione è accettata per la de-mascherazione se e solo se:
$$\text{accept}_{t,i} = (\text{conf}_{t,i} \ge c) \land \text{persist}_{t,i} \land (\text{JSD}_{t,i} \le d)$$

Politica di De-mascherazione

LESS impiega una politica frontier-first per i modelli di diffusione a sequenza completa. Tenta di de-mascherare la posizione mascherata più a sinistra se è idonea. In caso contrario, scansiona una finestra di look-ahead per trovare la prima posizione idonea. Se nessuna posizione è idonea, lo step viene saltato. Per garantire la terminazione, un fallback di de-mascherazione forzata viene attivato dopo $\Pi$ step saltati consecutivi, sebbene questo fallback non sia conteggiato come un impegno "stabile".

Contributi Chiave

1. Formulazione dell'Impegno del Token: Il paper inquadra la de-mascherazione adattiva nei dLLM come un problema di stopping online per singola posizione indotto dall'irreversibilità della de-mascherazione.
2. JSD Inter-Step Top-K: Gli autori introducono la divergenza di Jensen–Shannon inter-step top-K come un segnale di stabilità simmetrico e limitato. Ciò affronta i limiti dei lavori precedenti in cui la sola confidenza non può rilevare il movimento distributivo in corso e la divergenza KL è direzionale e illimitata.
3. Sampler LESS: Un sampler training-free che combina confidenza, persistenza del token e JSD inter-step in una regola di stabilità congiunta. Non richiede il riaddestramento del modello ed è compatibile con gli esistenti checkpoint di dLLM.
4. Validazione Empirica: Valutazione completa su tre famiglie di dLLM (Dream-7B, LLaDA-8B, LLaDA-1.5-8B) e sette benchmark che coprono matematica, codice e conoscenza generale.

Risultati

La valutazione dimostra che LESS migliora significativamente il compromesso accuratezza-efficienza rispetto sia al decoding a budget fisso che ad altri forti sampler training-free (Prophet e KLASS).

• Efficienza: Attraverso tre famiglie di modelli, LESS utilizza in media il 72,1% in meno di step inversi rispetto al decoding a budget fisso (Base).
• Accuratezza: LESS ottiene un guadagno medio di accuratezza di +0,91 punti rispetto al baseline a budget fisso pur utilizzando meno step. Supera inoltre Prophet e KLASS in termini di accuratezza media in tutte le famiglie di modelli.
• Performance nei Task: I guadagni sono più pronunciati nei task di generazione matematica e di codice. Ad esempio, su Dream-7B, LESS ottiene una accuratezza media del 65,18% (rispetto al 63,30% di Base) con soli 64,3 step (rispetto ai 256 di Base).
• Risparmio Computazionale: Poiché ogni step inverso richiede un forward pass del Transformer, la riduzione degli step si traduce direttamente in una minore computazione di inferenza. Su Dream-7B GSM8K, LESS riduce la latenza di wall-clock da 19,45s a 5,16s e riduce il numero di valutazioni forward di circa il 75% con budget di step elevati.
• Studi di Ablazione:
  • Complementarità dei Segnali: Nessun singolo segnale (confidenza, persistenza o JSD) performa bene quanto la regola congiunta. La confidenza è il driver principale, ma la persistenza e la JSD forniscono salvaguardie essenziali contro l'impegno prematuro e il drift distributivo.
  • Robustezza: Il metodo performa in modo coerente attraverso diversi modelli e regimi di decoding (sequenza completa vs. blockwise semi-autoregressivo) utilizzando una singola configurazione globale $(c=0,75, d=0,040)$.

Significato e Claim

Il paper sostiene che LESS fornisca un meccanismo pratico e training-free per accelerare il decoding dei modelli di linguaggio a diffusione, impegnando i token solo quando le loro predizioni sono fiduciose, temporalmente persistenti e distribuzionalmente stabili.

Gli autori sottolineano che la stabilità è più della semplice confidenza. Una predizione ad alta confidenza può comunque essere inaffidabile se il token top-1 è cambiato di recente o se la massa di probabilità si sta ancora spostando tra le alternative. Trattando la stabilità come una proprietà congiunta della predizione attuale e della sua recente traiettoria, LESS evita le insidie dell'impegno prematuro.

Il lavoro posiziona LESS come ortogonale e complementare ad altri metodi di accelerazione:

• Complementa i metodi in fase di training (es. distillazione) che accorciano il processo inverso.
• Complementa i metodi di sistema (es. caching, speculative decoding) che riducono la latenza per singolo step.
• Agisce come una policy di inferenza "drop-in" che può essere combinata con questi metodi senza modificare il modello sottostante.

Il paper conclude che la stabilità reciproca è un criterio valido per un robusto impegno precoce dei token, permettendo ai dLLM di raggiungere una maggiore efficienza senza sacrificare, e spesso migliorando, la qualità della generazione. Gli autori riconoscono che il metodo espone un compromesso qualità-efficienza controllato dalle soglie di stabilità, che potrebbero richiedere una calibrazione per modelli o task specifici, sebbene una configurazione globale fissa si sia dimostrata efficace nella loro diversificata suite di valutazione.