Sintesi Tecnica: BlockGen: Modellazione di Sequenze a Blocchi Flessibile con Campionatori Ibridi

Problema

I Transformer autoregressivi (AR) dominano la modellazione di sequenze ma soffrono di limiti nella generazione sequenziale (throughput/latenza) e di vincoli di attenzione causale che ostacolano i compiti di ragionamento che richiedono un contesto bidirezionale. I modelli di diffusione discreta offrono un'alternativa, raffinando iterativamente le sequenze rumorose, consentendo aggiornamenti paralleli. Tuttavia, recenti confronti tra i Modelli di Diffusione Mascherata (MDM) e i Modelli di Diffusione a Stato Uniforme (USDM) hanno prodotto intuizioni contrastanti a seconda del regime di campionamento:

1. Campionamento Ancestrale: Gli USDM spesso superano gli MDM nei task a valle nonostante una perplexity peggiore, grazie alla capacità degli USDM di correggere i propri errori (i token possono transire tra qualsiasi stato), mentre gli MDM fissano i token una volta de-mascherati.
2. Generazione a Blocchi: Il lavoro precedente confronta principalmente la diffusione a sequenza intera. Resta da chiarire se il vantaggio degli USDM si mantenga quando i token vengono generati blocco per blocco (semi-autoregressivo), un setting che abilita il KV caching e una generazione più veloce.
3. Correttori Informati: I confronti esistenti utilizzano spesso campionatori predictor-corrector (PC) "non informati" che reinjectano rumore in posizioni casuali. Non è noto come i correttori "informati" (che mirano agli errori probabili) influenzino la classifica MDM vs USDM, in particolare senza modelli di verifica ausiliari.

Metodologia: BlockGen

Gli autori introducono BlockGen, un modello di sequenza a blocchi addestrato su una miscela di dimensioni di blocco, istanziato con diffusione mascherata o uniforme all'interno dei blocchi.

1. Formulazione della Miscela e Addestramento

BlockGen definisce la densità della sequenza come una miscela su $M$ dimensioni di blocco $S = \{s_1, \dots, s_M\}$:
$$p_{\theta}^{\text{BlockGen}}(x) = \sum_{i=1}^M \gamma_i p_{\theta}^{(s_i)}(x)$$
dove $p_{\theta}^{(s_i)}$ è una densità valida fattorizzata in blocchi di dimensione $s_i$.

• Denoiser Condiviso: Invece di addestrare modelli separati, un singolo denoiser condiviso $\hat{x}_\theta$ riceve in input la dimensione del blocco $L'$, regolando i pattern di attenzione di conseguenza.
• Limiti di Verosimiglianza: Il paper deriva due limiti per la densità di miscela. Un limite log-sum-exp più stretto viene utilizzato per la valutazione, mentre un limite di verosimiglianza di miscela più lasso (un'aspettativa sulle dimensioni dei blocchi) viene usato per l'addestramento per mantenere l'efficienza computazionale.
• Campionamento Stratificato: Per ridurre la varianza durante l'addestramento multi-GPU, le dimensioni dei blocchi sono selezionate tramite campionamento stratificato anziché campionamento i.i.d.

2. Predictor-Corrector AR-Informato (ARPC)

Un contributo chiave è il campionatore ARPC, che sfrutta la miscela di BlockGen per eseguire correzioni informate senza un verificatore ausiliario.

• Meccanismo: Durante la generazione, dopo gli standard step di diffusione ancestrale, il campionatore esegue un passaggio di "proposta pulita". Esegue quindi un secondo passaggio forward in modalità AR ($L'=1$) utilizzando lo stesso denoiser per valutare i token proposti.
• Scoring: I token con bassa log-verosimiglianza AR vengono identificati come "improbabili" e riportati (corretti) a un livello di rumore precedente.
• Vantaggio: Ciò consente al modello di usare le proprie predizioni AR per guidare le correzioni di diffusione, intersecando efficacemente i paradigmi AR e di diffusione all'interno di un singolo modello.

3. Attenzione a Livello di Blocco

BlockGen utilizza l'attenzione block-causal: i token all'interno di un blocco sono soggetti ad attenzione bidirezionale, mentre il blocco è soggetto ad attenzione causale rispetto a tutti i blocchi precedenti. Questo abilita il KV caching attraverso i blocchi decodificati, preservando i benefici di efficienza della generazione semi-autoregressiva.

Risultati Chiave

1. Ragionamento Matematico (GSM8K)

Valutato su GSM8K utilizzando modelli addestrati su TinyGSM:

• Campionamento Ancestrale: Nel setting blocco per blocco, gli USDM superano gli MDM, specialmente nel regime a pochi step (basso NFE). Questo conferma che il vantaggio della diffusione a sequenza intera degli USDM si trasferisce alla generazione a blocchi.
• Campionamento ARPC: Il divario di performance tra USDM e MDM si restringe e si inverte ad alto NFE. Sotto ARPC con blocchi di 16 token, gli MDM superano leggermente gli USDM sull'accuratezza di GSM8K.
• ARPC vs Baselines: ARPC supera costantemente sia il campionamento ancestrale che l'Entropy-Informed PC (EIPC) nella maggior parte dei budget NFE. ARPC con MDM si avvicina alle prestazioni dell'AR con decoding greedy.

2. Language Modeling (OpenWebText & LM1B)

• Likelihood: L'addestramento su una miscela di dimensioni di blocco migliora significativamente la perplexity rispetto ai modelli a blocco fisso. BlockGen con diffusione mascherata raggiunge una Val. PPL di 17.5 su OpenWebText (rispetto a 16.7 per AR e 21.6 per il miglior BDM a blocco fisso), riducendo il divario tra modelli di diffusione e autoregressivi.
• Qualità del Campionamento:
  • Sotto campionamento ancestrale, gli USDM ottengono una Generative Perplexity (Gen. PPL) inferiore rispetto agli MDM.
  • Sotto ARPC, la tendenza si inverte ad alto NFE (es. 64 NFE/blocco), dove il masked-ARPC raggiunge una frontiera di Gen. PPL migliore rispetto al uniform-ARPC.
  • ARPC riduce il divario con l'AR ma non lo supera in termini di Gen. PPL nel regime pratico a bassa temperatura.

Significato e Claim

Il paper sostiene di risolvere due questioni aperte riguardanti la diffusione discreta:

1. Validità a Blocchi: Il vantaggio della diffusione a stato uniforme rispetto alla diffusione mascherata osservato nella generazione a sequenza intera si trasferisce alla generazione blocco per blocco sotto campionamento ancestrale.
2. Correzione Informata: La classifica tra MDM e USDM è dipendente dal regime. Mentre gli USDM eccellono con campionatori ancestrali semplici o non informati, i correttori informati (specificamente ARPC) possono colmare questo divario e persino invertire la classifica, favorendo la diffusione mascherata ad alti budget computazionali (NFE).

Gli autori sottolineano che BlockGen fornisce un framework flessibile che interpola tra AR e pura diffusione, permettendo a un singolo modello di agire sia come campionatore a blocchi che come verificatore AR. Notano che, sebbene le loro affermazioni empiriche siano robuste entro i budget e i campionatori testati, non stabiliscono se questi trade-off persistano alla scala dei grandi modelli linguistici (LLM), poiché i loro esperimenti sono limitati a backbone da 170 milioni di parametri. Il lavoro suggerisce che la scelta del paradigma di diffusione (mascherata vs uniforme) non è universalmente superiore ma dipende fortemente dalla strategia di campionamento e dal budget computazionale.