Riepilogo Tecnico: Diffusion a Blocchi Tripletta RWKV (B3D-RWKV)

Enunciato del Problema

I modelli linguistici Transformer causali affrontano due limitazioni strutturali primarie: la decodifica strettamente sequenziale, che impedisce la parallelizzazione, e i costi quadratici di attenzione per passo, che rendono l'inferenza su contesti lunghi computazionalmente costosa. Sebbene i modelli causali a tempo lineare (ad esempio RWKV) e i modelli di diffusione discreti affrontino ciascuna di queste debolezze in modo indipendente, la loro integrazione è stata intrinsecamente incoerente. I modelli di diffusione richiedono attenzione bidirezionale per denoisare blocchi di token in parallelo, mentre i modelli causali sono strettamente unidirezionali. I tentativi esistenti di combinarli spesso richiedono modifiche architetturali che compromettono l'efficienza del backbone a tempo lineare o necessitano di un addestramento da zero.

Metodologia

Gli autori propongono B3D-RWKV, una variante RWKV in stile diffusione che unisce l'efficienza di inferenza $O(L)$ con la diffusione discreta bidirezionale e parallela, senza alterare il backbone causale sottostante. L'innovazione centrale è il Layout a Blocchi Tripletta, una strategia di addestramento e inferenza progettata per simulare l'accesso bidirezionale all'interno di un'architettura strettamente causale.

1. Layout a Blocchi Tripletta (Addestramento)

Per abilitare la diffusione su un modello causale, ogni blocco logico di generazione di dimensione $B$ viene espanso in tre blocchi fisici contigui all'interno del campione di addestramento:

• $b_1$ (Copia Mascherata): Una versione mascherata del blocco logico.
• $b_2$ (Copia Mascherata con Perdita): Una copia mascherata identica su cui viene calcolata la perdita di denoisaggio.
• $b_3$ (Copia Pura di Verità Fondamentale): Una versione non mascherata che aggiorna lo stato ricorrente per il blocco successivo.

Accesso Pseudo-Bidirezionale:
Poiché il backbone causale elabora i token strettamente da sinistra a destra, lo stato nascosto che raggiunge qualsiasi posizione mascherata in $b_2$ ha già elaborato l'intera $b_1$. Poiché $b_1$ e $b_2$ condividono lo stesso pattern di mascheramento, i token non mascherati in $b_1$ vengono efficacemente "assorbiti" nello stato nascosto prima che il modello tenti di prevedere i token mascherati in $b_2$. Questo conferisce al modello accesso pseudo-bidirezionale al proprio contesto non mascherato all'interno del blocco logico, soddisfacendo i requisiti bidirezionali della diffusione mentre mantiene la natura unidirezionale del backbone RNN/SSM.

Obiettivo di Addestramento:
Il modello è addestrato utilizzando un obiettivo di previsione mascherata. La perdita è composta da:

• Perdita Cross-Entropy ($L_{CE}$): Calcolata sulle posizioni supervisionate dove i token sono mascherati.
• Perdita Parallela Consapevole della Confidenza ($L_{CAP}$): Seguendo LLaDA-2.0, questo termine affina la distribuzione di probabilità sulle posizioni già correttamente previste, incoraggiando il campionatore di inferenza a impegnare più token per passo.

2. Inferenza: Denoisaggio Iterativo a Blocchi

Durante l'inferenza, il modello genera un blocco logico alla volta. Il processo comporta:

1. Inizializzazione di un blocco con tutti i token mascherati.
2. Esecuzione di passaggi iterativi di denoisaggio in cui il modello prevede le probabilità top-1 per le posizioni mascherate.
3. Impegno dei token la cui probabilità top-1 supera una soglia $\tau$.
4. Un meccanismo di fallback garantisce il progresso impegnando i $k_{min}$ token più sicuri se meno di $k_{min}$ superano la soglia.
5. Una volta impegnate tutte le posizioni, il blocco pulito viene aggiunto al prefisso e inizia il prossimo blocco logico.

Contributi Chiave

1. B3D-RWKV-7.2B: Il primo modello linguistico RNN a tempo lineare in stile diffusione addestrato alla scala 7B con un obiettivo di previsione mascherata. Integra la selezione parallela dei token nel backbone RWKV-7 senza modificare i suoi parametri originali.
2. Framework a Blocchi Tripletta: Una ricetta di addestramento universale che converte backbone strettamente causali (che soddisfano la causalità rigorosa e le proprietà di propagazione in avanti dello stato) in modelli di diffusione a blocchi. Questo approccio evita la complessità architetturale di modificare il backbone in un blocco bidirezionale.
3. Efficienza e Prestazioni: Il metodo ottiene significativi miglioramenti nel throughput rispetto alla decodifica causale standard, mantenendo capacità di ragionamento competitive.

Risultati Sperimentali

Il modello è stato valutato su una suite di 8 compiti, inclusi conoscenza generale (MMLU, ARC, PIQA, RACE) e compiti di ragionamento/matematica (GSM8K, MATH, GPQA).

• Accuratezza: B3D-RWKV-7.2B raggiunge un'accuratezza paragonabile alla baseline RWKV-7 e ad altri modelli di diffusione di scala simile (ad esempio LLaDA-8B, Dream-7B). Supera notevolmente le baseline in compiti di ragionamento come ARC-Challenge e RACE, attribuito alla percezione pseudo-bidirezionale che potenzia il ragionamento.
• Throughput: Il modello dimostra un accelerazione media di 1.6× nel throughput di decodifica rispetto alla baseline RWKV-7 su lunghezze di contesto che vanno da 1K a 512K token. Con parametri di campionamento aggiustati, è raggiungibile un'accelerazione di 2.02× con un leggero compromesso sulla qualità.
• Ragionamento Matematico: Il documento nota un leggero calo delle prestazioni in compiti matematici complessi (ad esempio MATH) rispetto alla baseline causale. Gli autori attribuiscono ciò alla natura della decodifica parallela, che fatica con la verifica altamente strutturata e passo-passo richiesta per i problemi di matematica simbolica, dove non viene assegnato credito parziale.

Significato e Affermazioni

Il documento afferma di offrire una via efficiente per trasformare modelli linguistici causali pre-addestrati in modelli linguistici di diffusione senza cambiamenti architetturali. Sfruttando il layout a blocchi tripletta, B3D-RWKV unisce con successo l'efficienza di inferenza $O(L)$ dei modelli a tempo lineare con le capacità di decodifica parallela della diffusione.

Gli autori posizionano il lavoro come una soluzione strutturale all'incompatibilità tra i paradigmi causale e di diffusione. Riconoscono le limitazioni, inclusa la convalida attuale su un singolo backbone (RWKV-7), il costo della lunghezza della sequenza fisica 3× durante l'addestramento e l'uso di un dataset SFT relativamente piccolo senza allineamento RL, che potrebbero contribuire a regressioni di accuratezza su compiti specifici. Tuttavia, il contributo centrale rimane la dimostrazione che la generazione parallela di diffusione ad alto throughput è realizzabile su backbone strettamente causali e a tempo lineare attraverso il layout dei dati piuttosto che la riprogettazione architetturale.