DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

Sintesi Tecnica: DSL-LLaDA: Scalare il Denoising Continuo a 8B Masked Diffusion LMs

Problematica
I Modelli di Linguaggio a Diffusione Discreta (DLM), come LLaDA, generano testo tramite una decodifica parallela iterativa. Sebbene questi modelli siano stati scalati fino a miliardi di parametri, la qualità dell'inferenza degrada drasticamente sotto budget di step bassi (generazione a pochi step). Il problema centrale è un compromesso tra la lunghezza dell'output e la qualità: con un numero fisso di passaggi in avanti, l'unmasking iterativo standard produce o output brevi e di alta qualità, oppure testi lunghi, ripetitivi e incoerenti. Questo fallimento deriva dalla fattorizzazione della transizione inversa, dove ogni posizione di token viene decodificata indipendentemente dalla propria distribuzione marginale condizionata dal contesto, trascurando le correlazioni inter-token. Quando i budget di step sono bassi, ogni step deve impegnare molte posizioni simultaneamente, amplificando questi errori e portando a una terminazione prematura o a loop di ripetizione.

Il denoising continuo nello spazio degli embedding offre una soluzione teorica evolvendo tutte le posizioni congiuntamente per evitare la fattorizzazione nello spazio dei token, ma costruire un tale modello da zero su larga scala rimane un problema di design aperto a causa delle complessità nella geometria degli embedding, nello scheduling del rumore e negli obiettivi di addestramento.

Metodologia
Gli autori propongono DSL-LLaDA, un metodo per adattare leggermente un esistente DLM pre-addestrato (specificamente LLaDA-8B-Instruct) per supportare il denoising continuo nello spazio degli embedding, invece di addestrare un modello continuo da zero.

1. Adattamento Discrete Stochastic Localization (DSL):

   • Il metodo sostituisce l'input binario "pulito-o-mascherato" dei classici DLM con stati continui rumorosi. Per un token pulito $x_i$, viene campionato un Rapporto Segnale-Rumore (SNR) $\gamma_i$, e uno stato rumoroso $z_i$ viene costruito come una miscela dell'embedding del token e del rumore Gaussiano.
   • Un convertitore softmax mappa questi stati DSL continui nuovamente nello spazio degli embedding di input del backbone. Questo convertitore produce una "soft mask" dove un basso SNR corrisponde a uno stato simile a una maschera, un SNR medio a token casuali inaffidabili e un alto SNR a embedding di token chiari.
   • Il backbone viene sottoposto a continue-pretraining per soli 1.000 step su FineWeb-Edu. L'obiettivo di addestramento rimane la classica cross-entropy dei token, ma gli stati di input sono ora miscele continue piuttosto che maschere binarie. Fondamentalmente, il convertitore permette al modello di elaborare questi input senza cambiamenti architettonici, poiché l'invarianza rispetto all'SNR del denoiser ottimale elimina la necessità di input espliciti del timestep.

2. Inferenza Continua (Campionatore SDE):

   • Durante l'inferenza, il modello mantiene uno stato continuo per tutte le posizioni. Utilizza un campionatore basato su Equazioni Differenziali Stocastiche (SDE) (step Heun) per far evolvere lo stato da un basso SNR (alta incertezza) a un alto SNR.
   • Il modello predice una distribuzione di probabilità sui token, che viene convertita in una media pesata degli embedding dei token (utilizzando i top-$M$ token più probabili). Questa stima dell'embedding aggiorna lo stato continuo.
   • L'impegno del token "hard" è differito fino all'ultimo step. Ciò consente al modello di raffinare congiuntamente tutte le posizioni, preservando le correlazioni inter-token durante l'intero processo di generazione.

3. Generazione Condizionale:

   • Per compiti come la summarization, i token di contesto sono bloccati ai loro embedding puliti, mentre la regione di risposta subisce il denoising continuo. Il backbone attende bidirezionalmente, condizionando il denoising della risposta sul contesto pulito.

Contributi Chiave

1. Conversione Leggera: Ha dimostrato che un DLM mascherato pre-addestrato da 8B può essere convertito per supportare il denoising nello spazio continuo tramite un minimo adattamento di 1.000 step usando DSL, evitando di dover addestrare un modello continuo da zero.
2. Evitare i Modi di Fallimento a Basso NFE: Ha mostrato che l'inferenza SDE continua evita efficacemente la terminazione prematura e i loop di ripetizione caratteristici dell'unmasking iterativo a bassi budget di step (NFE $\le$ 16).
3. Robustezza Selettiva dello Stato Rumoroso: Ha identificato che l'adattamento induce una robustezza in cui il modello può correggere i token corrotti preservando quelli puliti (correzione selettiva), un comportamento assente nelle baseline discrete con pari computazione.

Risultati

• Zero-Shot Summarization: Su quattro benchmark (XSum, CNN/DailyMail, arXiv, PubMed), DSL-LLaDA-SDE ha ottenuto i migliori punteggi ROUGE-1 a bassi budget di step ($\le$ 16 passaggi in avanti) senza euristiche di controllo della lunghezza. Al contrario, le baseline discrete hanno troncato gli output precocemente o hanno sofferto di alti tassi di ripetizione.
• Ripetizione e Lunghezza: Nella generazione aperta, DSL-LLaDA-SDE ha mantenuto i tassi di ripetizione sotto il 9% in tutti i budget di step, mentre le baseline discrete hanno superato il 60% di ripetizione a basso NFE o si sono troncate a lunghezze molto brevi.
• Correzione dell'Errore: Il modello adattato ha dimostrato una "correzione selettiva", recuperando i token corrotti preservando oltre il 98% dei token puliti. In confronto, le baseline discrete con corruzione di token casuali (XDLM) hanno danneggiato una parte significativa dei token puliti, e il masked diffusion standard (MDM-CPT) ha mostrato una capacità di correzione trascurabile.
• Esperimenti di Controllo: Esperimenti di controllo utilizzando lo standard masked diffusion training con la stessa computazione (MDM-CPT) e corruzione di token casuali (XDLM) non sono riusciti a replicare i benefici del denoising continuo o la robustezza selettiva, confermando che i miglioramenti derivano dall'esposizione al rumore continuo piuttosto che dalla sola computazione aggiuntiva.

Significatività e Rivendicazioni
Il paper sostiene che DSL-LLaDA fornisce una procedura di conversione general-purpose per i DLM mascherati, consentendo loro di sfruttare i benefici del denoising continuo (evoluzione congiunta dei token, impegno ritardato) con un overhead di addestramento minimo. I risultati suggeriscono che l'adattamento al rumore continuo può risolvere il compromesso tra lunghezza e qualità inerente all'unmasking iterativo a bassi budget di step.

Gli autori rimangono modesti riguardo alle implicazioni più ampie:

• La valutazione è limitata principalmente alla generazione ancorata (summarization e instruction following). La scrittura creativa aperta mostra ancora tassi più elevati di degenerazione (24% vs $\le$4% sulla summarization).
• Il metodo è attualmente validato solo sulla scala LLaDA-8B; le prestazioni su altre architetture o scale maggiori non sono state testate.
• Sebbene l'SDE superi i metodi discreti a basso NFE, il LLaDA vanilla eguaglia o supera le prestazioni SDE in compiti con riferimento lungo (high NFE, dove la terminazione prematura non è un problema critico).
• Gli autori non rivendicano la parità con i modelli autoregressivi (AR) della stessa scala, notando che i modelli AR sono ancora superiori nei compiti con riferimento lungo ad alti step budget.
• La valutazione umana è identificata come lavoro futuro.