🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

Questo articolo introduce le Rappresentazioni di Relè Apprese (Relay), un metodo che consente ai Modelli di Diffusione Mascherati di propagare informazioni latenti attraverso i passaggi di denoising tramite un canale differenziabile per token, migliorando così le prestazioni nei compiti di pianificazione e codifica e riducendo significativamente la latenza di inferenza.

Autori originali: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Pubblicato 2026-05-25

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Il Problema: L'Amnesia del "Reset Rigido"

Immagina di dover risolvere un puzzle complesso, come un Sudoku o scrivere un pezzo di codice. Stai usando un assistente molto intelligente (un modello di IA) per aiutarti.

Attualmente, questo assistente funziona come una persona affetta da amnesia che dimentica tutto nel momento in cui finisce una singola frase.

L'assistente guarda una pagina vuota piena di punti interrogativi.
Pensa intensamente, esegue calcoli mentali complessi e decide di inserire una o poche lettere.
Il Reset Rigido: Non appena scrive quelle lettere, scarta immediatamente tutti i pensieri complessi, i calcoli e la "sensazione" che aveva riguardo al resto della pagina.
Per il passo successivo, ricomincia da zero, guardando solo le lettere che ha appena scritto, dimenticando completamente il ricco lavoro interno che ha appena svolto.

Questo è chiamato il problema del "Reset Rigido". Il documento sostiene che ciò è inefficiente. È come uno chef che assaggia una zuppa, scrive "serve sale" e poi dimentica immediatamente il sapore del brodo prima di aggiungere il sale. Deve riasaggiare l'intera zuppa da capo ogni volta che aggiunge un ingrediente.

La Soluzione: Il "Bastone" della Staffetta

Gli autori propongono un nuovo metodo chiamato Relay (Staffetta).

Immagina una gara di staffetta. Invece di fermarsi, dimenticare la gara e ricominciare, il corridore passa un bastone al prossimo corridore. Questo bastone trasporta l'impulso, la strategia e la fatica del corridore precedente, così che il successivo possa riprendere esattamente da dove si era interrotto.

Nel mondo dell'IA, il "bastone" è un flusso continuo di informazioni (uno stato nascosto) che il modello porta avanti.

Prima del Relay: Il modello calcola, scrive un token e cancella il calcolo.
Con il Relay: Il modello calcola, scrive un token, e passa un "promemoria" a se stesso per il passo successivo. Questo promemoria dice: "Ehi, stavo pensando a questo specifico schema e sono sicuro al 80% su questa prossima parte".

Come Funziona: Imparare a Passare il Bastone

Il documento introduce alcuni trucchi chiave per far funzionare questo sistema:

Addestramento "Forward-Thinking" (Pensiero Avanti):
Di solito, i modelli di IA sono addestrati a ottenere giusto solo la risposta corrente. Gli autori hanno addestrato il modello a essere "forward-thinking". Hanno insegnato al modello: "Non scrivere solo la lettera giusta oggi; scrivi un promemoria (il relay) che ti aiuterà a scrivere la lettera giusta domani."
Backpropagation Troncata (BPTT):
Questo è un termine matematico sofisticato per "guardare avanti per imparare". Immagina di praticare una routine di danza. Invece di praticare solo un movimento, ne pratichi una sequenza di tre. Se inciampi sul terzo movimento, ti rendi conto: "Oh, ho sbagliato il primo movimento perché non mi sono preparato per il terzo".
Il documento utilizza un metodo chiamato BPTT troncata per permettere al modello di guardare avanti di qualche passo durante l'addestramento. Impara a passare informazioni in avanti specificamente per rendere i passi futuri più facili.
Il Canale "Soft" (Morbido):
Le informazioni passate in avanti non sono solo una parola o una lettera; sono un numero "soft" (un valore continuo). Pensaci come a un dimmer per le luci piuttosto che a una luce on/off. Permette al modello di trasportare indizi sottili e probabilità che non sono ancora stati trasformati in una decisione finale.

I Risultati: Più Veloce e Più Intelligente

Gli autori hanno testato questo su due cose:

Puzzle Sudoku: Hanno trattato il Sudoku come un compito di pianificazione. Il modello "Relay" ha risolto i puzzle con meno tentativi (meno "forward passes") e ha commesso meno errori rispetto al modello standard. Era migliore nel mantenere a mente il "quadro generale" mentre riempiva la griglia.
Coding (Scrittura di Software): Hanno preso un'IA all'avanguardia per la scrittura di codice (Fast-dLLM v2) e hanno aggiunto il sistema Relay.
- Migliore Accuratezza: Ha scritto codice migliore.
- Velocità Maggiore: Ha dovuto eseguire il suo "motore di pensiero" 32% in meno per ottenere lo stesso risultato.

Perché Questo È Importante

Il documento afferma che, interrompendo il "Reset Rigido" e permettendo al modello di portare avanti i propri "pensieri" come un bastone da staffetta, possiamo rendere i modelli di IA:

Più Intelligenti: Possono ragionare meglio su sequenze lunghe (come codice complesso o puzzle logici).
Più Veloci: Non devono ricalcolare tutto da capo ogni volta, risparmiando tempo ed energia.

In breve, il documento insegna ai modelli di IA a smettere di dimenticare il proprio lavoro e a iniziare a costruirvi sopra, passo dopo passo, proprio come fa un essere umano quando risolve un problema difficile.

Riepilogo Tecnico: Rappresentazioni di Relay Apprese per Modelli di Diffusione Discreta "Forward-Thinking"

1. Enunciato del Problema: Il "Reset Rigido" nei Modelli di Diffusione Mascherata

I Modelli di Diffusione Mascherata (MDM) generano sequenze discrete attraverso un processo iterativo di denoising. Partendo da una sequenza completamente mascherata, il modello sblocca iterativamente una frazione di posizioni. Tuttavia, nelle inferenze standard degli MDM esiste un'inefficienza critica: dopo ogni passaggio in avanti, il modello scarta i ricchi calcoli interni continui (stati nascosti) generati per le posizioni mascherate. Il passo successivo inizia esclusivamente dai nuovi token discreti impegnati.

Gli autori definiscono questo il problema del reset rigido. Sebbene gli MDM eseguano più passaggi in avanti per generazione, l'impossibilità di riutilizzare il calcolo continuo intermedio impedisce al modello di ammortizzare efficacemente il ragionamento tra i passi. Ciò contrasta con il calcolo ricorrente, dove lo srotolamento di un modello a parametri fissi attraverso i passi ha dimostrato di espandere la classe di funzioni che il modello può approssimare, migliorando le prestazioni su compiti di ragionamento complessi. La domanda centrale affrontata è come abilitare gli MDM a supportare un calcolo ricorrente che trasporti informazioni più ricche attraverso i passi sequenziali di sblocco, senza scartare gli stati latenti intermedi.

2. Metodologia: Rappresentazioni di Relay Apprese (Relay)

La soluzione proposta, Relay, introduce uno stato latente continuo e differenziabile che viene trasmesso in avanti attraverso i passi di decodifica, consentendo al modello di essere "forward-thinking" (prospettico).

Meccanismo Principale

Traiettoria dello Stato Augmentato: Ad ogni passo di inferenza $k$ , il modello mantiene uno stato augmentato $s_k = (x_{t_k}, h_k)$ , dove $x_{t_k}$ è la sequenza parzialmente mascherata corrente e $h_k$ è uno stato di relay appreso (vettore latente continuo).
Architettura: Il modello è composto da un backbone $f_\theta$ $f_{θ}$ , embedding dei token $Emb_\theta$ $E m b_{θ}$ , un capo di unembedding $UnEmb_\theta$ $U n E m b_{θ}$ e un specifico modulo relay $R_\theta$ $R_{θ}$ .
- Il prossimo stato nascosto è calcolato come: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ .
- Lo stato di relay $h_k$ viene proiettato e aggiunto ai token incorporati prima di entrare nel backbone.
- I logit di output $\ell_k$ sono derivati da $h_{k+1}$ tramite $UnEmb_\theta$ .
Addestramento tramite BPTT Troncato: A differenza degli MDM standard addestrati con passi indipendenti, Relay viene addestrato utilizzando la Backpropagation Through Time (BPTT) Troncata.
- Il modello esegue un rollout di $K$ passi (ad esempio, $K=2$ ) partendo da una sequenza completamente mascherata.
- Durante il rollout, le posizioni vengono sbloccate in base a una politica (ad esempio, soglie di confidenza) e i token impegnati vengono "teacher-forced" ai valori di verità fondamentale per prevenire l'accumulo di errori durante il rollout.
- La perdita è la somma delle perdite di cross-entropy sui $K$ passi.
- I gradienti vengono retropropagati attraverso il percorso del relay ( $h_k \to h_{k+1}$ ), permettendo al modello di imparare non solo a prevedere il passo corrente, ma a codificare informazioni in $h_k$ che siano massimamente utili per i passi futuri.

Inferenza

Al momento dell'inferenza, il processo rimane sostanzialmente invariato in termini di programma di sblocco o politica di campionamento. L'unica aggiunta è che il modello trasporta lo stato di relay $h_k$ in avanti insieme ai token impegnati. Lo stato di relay funge da canale di memoria differenziabile per le informazioni non ancora impegnate in token discreti.

3. Contributi Chiave

Framework Relay: Un metodo generale per incorporare il calcolo ricorrente negli MDM addestrando il modello a trasmettere un relay latente appreso in avanti attraverso i passi di decodifica tramite BPTT troncata. È agnostico rispetto all'architettura e compatibile con le tecniche esistenti di accelerazione degli MDM come la diffusione a blocchi e la cache KV.
Scalabilità e Adattamento: Validazione di Relay alla scala dei Large Language Model (LLM) attraverso l'adattamento a parametri completi di Fast-dLLM v2 (1.5B). Il metodo adatta con successo un Modello di Linguaggio Discreto (DLM) all'avanguardia con un numero minimo di FLOP aggiuntivi di addestramento.
Miglioramenti Prestazioni-Latenza: Dimostrazione empirica che Relay avanza il fronte di Pareto prestazioni-latenza, superando il fine-tuning supervisionato standard (SFT) sui compiti di codifica riducendo significativamente la latenza di inferenza.
Analisi dello Spazio di Progettazione: Estesi studi di ablazione su un compito di pianificazione basato su Sudoku, che isolano i contributi della procedura di rollout, del canale relay e dell'obiettivo di addestramento BPTT.

4. Risultati Sperimentali

Compito di Pianificazione Sudoku

Utilizzando il dataset Sudoku-Extreme, gli autori hanno confrontato quattro obiettivi di addestramento:

MLM: Diffusione mascherata uniforme standard (singolo passo, nessun relay).
Rollout: Rollout multi-passo con teacher forcing ma senza canale relay ( $R_\theta \equiv 0$ ).
Relay (sg): Include il percorso relay ma utilizza stop-gradients (nessuna BPTT attraverso il relay).
Relay: Metodo completo con $K=2$ BPTT attraverso il relay.

Risultati:

Frontiera Accuratezza-NFE: Relay ha raggiunto il miglior compromesso tra accuratezza e Numero di Valutazioni in Avanti (NFE).
Impatto delle Componenti: Il passaggio da MLM a Rollout ha migliorato le prestazioni. L'aggiunta del canale relay (Relay sg) ha fornito un salto significativo. Infine, l'abilitazione della BPTT (Relay) ha prodotto i migliori risultati, dimostrando che il modello impara a mantenere la coerenza interna della scacchiera parziale sotto sblocchi più aggressivi.
Efficienza: A una soglia di confidenza fissa, Relay ha impegnato più celle per passaggio in avanti mantenendo la legalità (nessuna violazione di riga/colonna/box), raggiungendo la stessa accuratezza in un numero inferiore di passaggi in avanti totali rispetto ai baseline.

Adattamento del Modello Pre-addestrato (Fast-dLLM v2)

Gli autori hanno adattato Fast-dLLM v2 (1.5B) utilizzando Relay su un mix di dataset di codice e matematica.

Prestazioni di Codifica: Sui benchmark HumanEval e MBPP, Relay ha superato lo SFT standard.
- HumanEval: Relay ha raggiunto un'accuratezza del 42,1% (vs 38,4% per SFT) con 88,3 NFE (vs 130,7 per SFT).
- MBPP: Relay ha raggiunto un'accuratezza del 46,6% (vs 43,9% per SFT) con 78,8 NFE (vs 84,8 per SFT).
Latenza: Relay ha ridotto la latenza di inferenza fino al 32% rispetto allo SFT standard, migliorando al contempo l'accuratezza.
Sovraccarico di Memoria: Il profiling su una GPU A100 ha mostrato che, sebbene Relay aumenti la memoria di attivazione in vita durante il secondo passaggio in avanti, l'uso massimo di memoria rimane quasi identico allo SFT vanilla (circa 20,1 GiB vs 21,2 GiB) perché il collo di bottiglia è il passaggio all'indietro della proiezione del vocabolario, non l'archiviazione dello stato di relay.

5. Significato e Affermazioni

Il documento afferma che Relay cambia fondamentalmente la natura della generazione iterativa non autoregressiva consentendo agli MDM di accumulare calcoli continui intermedi.

Forward-Thinking: Il metodo permette ai modelli di ottimizzare esplicitamente le rappresentazioni intermedie per le decisioni future di sblocco, invece di trattare ogni passo come un nuovo problema di previsione.
Avanzamento del Fronte di Pareto: Riducendo il numero di passaggi in avanti necessari per raggiungere un'accuratezza target, Relay avanza il fronte prestazioni-latenza, rendendo i modelli di diffusione discreta più efficienti per compiti complessi di ragionamento e generazione.
Compatibilità: L'approccio è presentato come un potenziamento "drop-in" che funziona con i DLM all'avanguardia esistenti e le tecniche di accelerazione (ad esempio, diffusione a blocchi, cache KV) senza richiedere rivoluzioni architetturali.

Gli autori concludono che Relay rappresenta un passo significativo verso una generazione non greed e forward-thinking, offrendo un substrato continuo per il ragionamento latente che può essere ulteriormente esplorato per l'interpretabilità e l'applicazione in altre modalità (ad esempio, immagini, molecole).