Riepilogo Tecnico: Masked Language Flow Models

Problematica

I Modelli Autoregressivi (ARM) dominano la generazione del linguaggio ma soffrono di colli di bottiglia nell'inferenza sequenziale, poiché ogni token dipende da tutti i precedenti. I Masked Diffusion Models (MDM) offrono un'alternativa di decodifica parallela tramite la de-mascheratura iterativa; tuttavia, essi si basano su una transizione inversa fattorizzata che assume l'indipendenza condizionale tra le posizioni mascherate. Questa approssimazione è valida nel limite dei passi infinitesimali, ma degrada significativamente nei regimi di campionamento a pochi passi, dove la generazione parallela è più efficiente, portando a una qualità compromessa a causa delle dipendenze tra i token ignorate.

Al contrario, i Flow Language Models (FLM) operano in spazi di embedding continui, apprendendo un flusso che trasporta il rumore verso sequenze pulite. Sebbene gli FLM supportino una generazione efficiente in un singolo step tramite la distillazione della mappa di flusso, essi tipicamente collassano la generazione in una singola traiettoria. Questa struttura fatica con compiti che richiedono ragionamento multi-step, poiché gli FLM sono costretti a decodificare tutti i token simultaneamente invece di utilizzare stati intermedi parzialmente completati come contesto per i successivi raffinamenti.

Metodologia: Masked Language Flow Models (MLFM)

Gli autori propongono i Masked Language Flow Models (MLFM), un framework che integra il meccanismo di masking degli MDM nella dinamica di flusso continua degli FLM.

1. Framework e Interpolante Stocastico

Gli MLFM definiscono un processo in avanti nello spazio Euclideo utilizzando un ponte di Brownian come interpolante stocastico. Questo ponte collega:

• Sequenze parzialmente mascherate: Rappresentate da embedding puliti per i token non mascherati e da un embedding specifico [MASK] per i token mascherati.
• Sequenze pulite: Rappresentate dai loro embedding di token puliti.

Il vettore $z_{s,t}$ per una sequenza $X$, al tempo di noising $t$ e probabilità di masking $s$, è costruito in modo che:

• Le posizioni non mascherate ($U_s$) rimangano fisse ai loro embedding puliti $E X_\ell$.
• Le posizioni mascherate ($M_s$) evolvano tramite un ponte di Brownian tra l'embedding [MASK] ($m$) e l'embedding pulito ($E X_\ell$).

Questo design permette al modello di apprendere un flusso continuo accoppiato per generare token puliti nelle posizioni mascherate, mantenendo al contempo la struttura di condizionamento esatto per qualsiasi posizione tipica degli MDM.

2. Addestamento e Adattamento

Un contributo chiave è la capacità di warm-start dell'addestramento MLFM utilizzando MDM pre-addestrati.

• Equivalenza Teorica: Al punto finale $t=0$, l'interpolante stocastico MLFM riduce esattamente alla sequenza parzialmente mascherata ed embedded utilizzata negli MDM. Di conseguenza, la perdita (loss) MLFM a $t=0$ coincide con la loss MDM.
• Strategia di Adattamento: Gli autori adattano un transformer bidirezionale pre-addestato (da un MDM) attraverso:
  1. Il mantenimento fisso dello strato di embedding dei token.
  2. Il congelamento dei blocchi transformer e il loro potenziamento con AdaLN (per il condizionamento sul tempo continuo $t$) e adattatori LoRA (per l'adattamento a embedding continui corrotti).
  3. L'adattamento del modulo di output con un adattatore LoRA.
     Questo approccio riduce significativamente i costi computazionali rispetto all'addestramento da zero.

3. Strategia di Campionamento: CCFG con Online Token Promotion

Per affrontare il limite del ragionamento multi-step, gli autori introducono uno schema di campionamento innovativo che combina due meccanismi:

• Context-Corrupted Classifier-Free Guidance (CCFG): Una variante della classifier-free guidance in cui il campo di velocità di riferimento è calcolato corrompendo i token di contesto (applicando la marginale di corruzione in avanti a essi), mentre il campo guidato utilizza il contesto pulito. La differenza isola il contributo del contesto pulito alla direzione prevista.
• Online Token Promotion (OTP): Durante il campionamento, se la moda posteriore per un token mascherato raggiunge una soglia di confidenza ($1-\epsilon$), il token viene immediatamente "promosso". Esso viene fissato al suo embedding pulito e trattato come contesto osservato per tutti i successivi step di denoising. Ciò consente al modello di sfruttare le predizioni ad alta confidenza come contesto immediato, emulando il ragionamento iterativo.

Contributi Chiave

1. Integrazione di Masking e Flow: Gli MLFM colmano il divario tra il masking discreto (MDM) e i flussi continui (FLM) usando un ponte di Brownian, abilitando la generazione condizionale che supporta sia l'efficienza parallela che il raffinamento iterativo.
2. Adattamento Efficiente: La dimostrazione che gli MDM pre-addestrati possono essere convertiti in MLFM tramite un adattamento leggero (LoRA/AdaLN), sfruttando l'equivalenza del punto finale a $t=0$.
3. Nuovo Campionatore: La proposta di CCFG combinata con l'Online Token Promotion, che converte dinamicamente gli stati latenti in contesto pulito, facilitando il ragionamento multi-step.
4. Scalabilità a Task Downstream: Il primo scaling di successo dei modelli di linguaggio basati su flow a compiti di ragionamento (matematica) e follow-up di istruzioni.

Risultati Sperimentali

Gli autori hanno valutato un MLFM da 1.028 miliardi di parametri (adattato da un SMDM pre-addestrato) su GSM8K (ragionamento matematico) e MT-Bench (seguimento di istruzioni).

• MT-Bench: L'MLFM ha ottenuto un punteggio al primo turno di 2.27, superando significativamente il baseline SMDM (1.60) e un baseline AR di dimensioni simili (1.57). Notevolmente, ciò è stato raggiunto con soli 128 step di campionamento (la metà degli SMDM con 256).
• GSM8K: L'MLFM ha raggiunto il 31.24% di accuratezza. Sebbene sia inferiore all'SMDM sottoposto a fine-tuning specifico per il task (58.5%) e a LLaMA-2 (58.6%), gli autori attribuiscono questo dato al protocollo di fine-tuning con miscela di istruzioni più ampio utilizzato per MLFM rispetto al protocollo specifico per il task usato per i baseline.
• Ablazioni: I risultati hanno confermato che sia la CCFG che l'Online Token Promotion aumentano significativamente le prestazioni. Scale di guidance più grandi e più step di campionamento hanno generalmente migliorato i risultati, con l'MLFM che supera i baseline su MT-Bench anche con 16 step di campionamento.

Significato e Rivendicazioni

Il paper sostiene di fornire la prima evidenza che i modelli di linguaggio basati su flow possano scalare per risolvere compiti di ragionamento e follow-up di istruzioni. Combinando il potenziale di generazione parallela dei modelli di flusso con le capacità di ragionamento iterativo dei modelli di masking, gli MLFM offrono un'architettura flessibile che supporta la generazione complessa e multi-step. Il lavoro dimostra che i modelli di flusso non sono limitati alla generazione incondizionale o alla semplice distillazione, ma possono essere adattati per applicazioni condizionali pratiche con un addestramento efficiente tramite l'adattamento da MDM.