Sintesi Tecnica: Modelli di Linguaggio a Diffusione Multi-Blocco (MBD-LM)

1. Definizione del Problema

I Modelli di Linguaggio a Diffusione (DLM) offrono una decodifica parallela nativa come alternativa alla generazione autoregressiva, ma storicamente faticano a causa della mancanza di un caching KV e del supporto per la generazione a lunghezza dinamica. Recenti Modelli di Linguaggio a Diffusione a Blocchi (BD-LM) hanno affrontato questo problema introducendo la generazione a causalità di blocco, abilitando la Diffusione a Blocco Singolo (SingleBD). In SingleBD, il modello decodifica un blocco rumoroso condizionato da un prefisso pulito e memorizzato in cache, consentendo il caching KV e il parallelismo intra-blocco. Tuttavia, i blocchi vengono ancora elaborati sequenzialmente, creando "bolle di archiviazione" (storing bubbles) in cui il sistema attende che un blocco completato venga memorizzato prima di decodificare il successivo, bloccando il parallelismo inter-blocco.

Sebbene la Discrete Diffusion Forcing (D2F) abbia introdotto la visibilità tra molteplici blocchi rumorosi per abilitare la Diffusione Multi-Blocco (MultiBD), rimane un significativo disallineamento tra addestramento e inferenza (train–inference mismatch). La D2F addestra su suffissi rumorosi lunghi e monotoni, mentre la pratica MultiBD in fase di inferenza opera su un insieme corrente limitato (bounded running-set) di blocchi consecutivi con pattern di rumore slot-wise eterogenei. Gli esistenti BD-LM addestrati con Teacher Forcing (TF) o D2F non si allineano a questi specifici stati di inferenza, portando a prestazioni subottimali quando si tenta di decodificare più blocchi simultaneamente. Inoltre, le implementazioni naive di MultiBD spesso si affidano a forme di input dinamiche, che sono incompatibili con i meccanismi di esecuzione a forma statica, come la cattura di CUDA Graph, ostacolando i guadagni pratici di throughput.

2. Metodologia

Il documento propone i Modelli di Linguaggio a Diffusione Multi-Blocco (MBD-LM), un framework unificato e una strategia di addestramento per colmare il divario tra gli esistenti BD-LM e l'efficiente inferenza MultiBD.

2.1 Formulazione Unificata: MBD-LM

Gli autori definiscono i MBD-LM come una generalizzazione dei BD-LM in cui un insieme corrente ($R_s$) di blocchi consecutivi viene decodificato simultaneamente. A differenza di SingleBD (dove $|R_s|=1$) o della D2F a sequenza completa, i pratici MBD-LM operano in un regime intermedio limitato dove $|R_s|$ è piccolo (ad esempio, 2–4 blocchi) per bilanciare il parallelismo con l'efficienza computazionale.

2.2 Multi-block Teacher Forcing (MultiTF)

Per allineare gli stati di addestramento con il regime di inferenza a insieme corrente limitato, gli autori introducono il Multi-block Teacher Forcing (MultiTF), un metodo di post-addestramento per gli esistenti BD-LM.

• Costruzione dei Gruppi di Rumore: Invece di corrompere un singolo blocco (TF) o un lungo suffisso monotono (D2F), il MultiTF costruisce gruppi di rumore: gruppi consecutivi e limitati di blocchi rumorosi condizionati da un prefisso pulito.
• Layout dei Gruppi: I dati di addestramento sono partizionati utilizzando sia layout sistematici (enumerando tutti gli shift e le dimensioni dei gruppi fino a un valore massimo $G_{max}$) sia layout casuali per garantire una copertura diversificata delle dimensioni dell'insieme corrente e delle posizioni relative ai blocchi.
• Scheduler del Rumore Chain-Uniform: All'interno di ogni gruppo di rumore, il MultiTF impiega uno scheduler del rumore chain-uniform. A differenza dello schema monotono della D2F, questo metodo campiona un livello base (floor) a livello di gruppo e poi campiona il rapporto di rumore di ogni blocco successivo condizionato dal rapporto del blocco precedente. Ciò crea gap di rumore slot-wise eterogenei che corrispondono meglio ai pattern di rumore empirici osservati durante l'inferenza MultiBD.
• Maschera Dual-Stream Consapevole del Gruppo: Una maschera di attenzione specializzata permette ai blocchi rumorosi all'interno di un gruppo di prestare attenzione tra loro, mantenendo al contempo la visibilità causale di blocco verso il prefisso pulito, prevenendo la fuga di informazioni dai token rumorosi futuri verso il prefio pulito.

2.3 Motore di Inferenza Ottimizzato: Block Buffer

Per abilitare un'inferenza MultiBD pratica e ad alto throughput, gli autori progettano una pipeline ottimizzata basata sul meccanismo Block Buffer:

• Esecuzione a Forma Statica: Inveve di appendere dinamicamente i blocchi all'input (il che interrompe la compatibilità con CUDA Graph), il motore mantiene un buffer fisico di dimensione fissa ($N_{buf}$).
• Transizioni di Stato: I blocchi all'interno del buffer transitano attraverso quattro stati: dummy (inattivo), active (in fase di decodifica), to-cache (completato ma non ancora memorizzato) e in-cache (impegnato nel cache KV).
• Meccanismo: I futuri blocchi entrano nel buffer attivando gli slot dummy esistenti anziché estendere la sequenza di input. I blocchi completati vengono impegnati nel cache KV e sostituiti da nuovi slot dummy alla fine della coda. Questo preserva la forma statica dell'input necessaria per l'esecuzione efficiente su GPU, mantenendo al contempo l'insieme corrente logico per il parallelismo inter-blocco.
• Preservazione del Prefix Cache: Questo design separa rigorosamente il prefisso immutabile memorizzato in cache dal buffer attivo, garantendo la compatibilità con le semantiche di prefix caching standard dei BD-LM, che la D2F nativa spesso viola.

3. Contributi Chiave

1. Formulazione Unificata MBD-LM: Il documento formula i MBD-LM come un framework unificato parametrizzato da un insieme corrente limitato, identificando la pratica MultiBD come un regime intermedio tra SingleBD e la D2F a sequenza completa.
2. Post-Training MultiTF: Una nuova strategia di addestramento che allinea i BD-LM con l'inferenza MultiBD addestrandoli su gruppi di rumore limitati con pattern di rumore eterogenei, riducendo significativamente il disallineamento tra addestramento e inferenza.
3. Motore di Inferenza Ottimizzato: Una pipeline basata su Block Buffer che consente l'esecuzione a forma statica per i CUDA Graph, preserva il caching KV del prefisso e sovrappone la decodifica con l'archiviazione del cache per ottenere accelerazioni nel tempo di esecuzione (wall-clock).

4. Risultati Sperimentali

Gli autori hanno valutato i MBD-LM su benchmark di ragionamento matematico (GSM8K, MATH500) e generazione di codice (MBPP+, HumanEval+) utilizzando le famiglie di modelli LLaDA2 e SDAR.

• Parallelismo e Accuratezza: MBD-LLaDA2-Mini ha aumentato i Token Per passaggio Forward (TPF) medi da 3.47 (SingleBD) a 6.19 (+78.4%) migliorando al contempo l'accuratezza media dal 79.95% all'81.03%.
• Combinazione con DMax: Combinato con la strategia di decodifica DMax, MBD-LLaDA2-Mini-DMax ha raggiunto un TPF medio di 9.34 (un miglioramento del +47.1% rispetto a SingleBD DMax) con una perdita di accuratezza di solo l'1.02%.
• Throughput (TPS): Il motore di inferenza ottimizzato ha tradotto i guadagni di TPF in throughput di tempo reale. MBD-LLaDA2-Mini-DMax ha raggiunto 926.67 TPS in media, rispetto ai 779.49 TPS del baseline SingleBD DMax.
• Studi di Ablazione:
  • Allineamento Addestramento-Inferenza: La MultiBD senza addestramento (applicando l'inferenza MultiBD a modelli non modificati) ha migliorato il TPF ma ha degradato l'accuratezza. Il post-addestramento MultiTF ha recuperato e migliorato l'accuratezza, confermando la necessità dell'allineamento degli stati.
  • Design dello Scheduler: Il proposto scheduler chain-uniform ha superato gli scheduler monotoni di tipo D2F e altre varianti uniformi, evidenziando l'importanza dei gap di rumore eterogenei.
  • Diversità dei Layout: La combinazione di layout di gruppo sistematici e casuali ha prodotto il miglior compromesso tra TPF e accuratezza.

5. Significato

Il documento sostiene che una diffusione Multi-Blocco affidabile ed efficiente richiede un approccio olistico che combini l'allineamento degli stati al tempo di addestramento e il supporto di sistema al tempo di inferenza.

• Dimostra che applicare semplicemente l'inferenza MultiBD a modelli esistenti è insufficiente a causa dei disallineamenti distribuzionali nei pattern di rumore e nelle dimensioni dell'insieme corrente.
• Prova che i MBD-LM possono raggiungere il parallelismo inter-blocco senza sacrificare i benefici del caching KV e del riutilizzo del prefisso che rendono efficienti i BD-LM.
• Il lavoro stabilisce che l'"insieme corrente limitato" è il punto di equilibrio pratico per la generazione basata su diffusione, offrendo una via percorribile per accelerare i DLM oltre i colli di bottiglia sequenziali di SingleBD, evitando al contempo la complessità e l'incompatibilità del cache della diffusione a sequenza completa.