Sintesi Tecnica: Induzione in Entrambe le Direzioni

Definizione del Problema

Sebbene i meccanismi interni dei transformer autoregressivi (AR) siano stati ampiamente studiati attraverso la lente dell'interpretabilità meccanicistica, l'operatività dei modelli di linguaggio di diffusione (DLM) rimane in gran parte inesplorata. I DLM generano testo tramite denoising iterativo piuttosto che tramite predizione sequenziale del token successivo, utilizzando l'attenzione bidirezionale e obiettivi di addestramento differenti. Esiste una lacuna critica nel comprendere se i DLM implementino l'induzione — un meccanismo canonico per l'apprendimento in-context in cui un modello identifica un contesto ripetuto e copia il token successivo — nello stesso modo dei modelli AR. Nello specifico, non è chiaro come i DLM gestiscano l'induzione quando il token target è nascosto dietro un token di maschera [M] ed è circondato da un contesto bidirezionale, diversamente dal contesto strettamente causale disponibile per i modelli AR.

Metodologia

Gli autori hanno condotto uno studio meccanicistico controllato confrontando modelli transformer AR attention-only accoppiati con DLM a maschera assorbente (absorbing-mask).

• Architettura del Modello: Lo studio ha utilizzato modelli su piccola scala (profondità $L \in \{1, 2, 3\}$, dimensione nascosta $d_{model}=512$, 16 teste di attenzione) addestrati su sequenze FineWeb-Edu. Entrambe le famiglie di modelli condividevano architetture e ricette di addestramento identiche, differendo solo per le loro maschere di attenzione (causale per AR, bidirezionale per DLM) e gli obiettivi di addestramento (predizione del token successivo vs. diffusione a maschera assorbente).
• Task di Valutazione: Le capacità di induzione sono state misurate utilizzando sequenze sintetiche di lunghezza 512, dove la prima metà era ripetuta nella seconda metà. Il compito consisteva nel predire i token mascherati nella seconda metà basandosi sulla sorgente non mascherata nella prima metà (induzione forward) e viceversa (induzione reverse).
• Tecniche di Analisi:
  • Ablazioni Causali: Sono state eseguite ablazioni medie su singole teste di attenzione per determinare la loro necessità causale per l'induzione.
  • Decomposizione dei Pesi: Lo studio ha impiegato decomposizioni QK (Query-Key) e OV (Output-Value) per analizzare come vengono formati i logit dell'attenzione e come l'informazione fluisce dalla sorgente al target.
  • Linear Probing: È stato addestrato un linear probe sul residual stream per rilevare la presenza di informazioni sulla frazione globale di maschera.
  • Esperimenti di Patching: Tecniche di "patching" sono state utilizzate per intervenire causalmente sulla rappresentazione del tasso di maschera globale per osservare i cambiamenti nell'entropia di predizione.
• Varianti del Modello: Per l'analisi meccanicistica, gli autori hanno utilizzato DLM "folded no-LayerNorm", convertendo modelli addestrati con LayerNorm in mappe lineari per facilitare la decomposizione additiva del residual stream, verificando al contempo che le prestazioni rimanessero comparabili ai modelli originali con LayerNorm.

Contributi Chiave e Risultati

1. Circuito di Induzione Bidirezionale

Lo studio identifica che i DLM apprendono un circuito di induzione bidirezionale che è simmetrico rispetto alla direzione. A differenza dei modelli AR, che si affidano a un singolo percorso causale, i DLoli utilizzano due percorsi distinti:

• Teste del Token Precedente e Successivo: Nel primo strato (Layer 0), teste specifiche prestano attenzione ai vicini immediati a sinistra e a destra del token mascherato, scrivendo queste caratteristiche del contesto locale nel residual stream.
• Teste di Induzione: Negli strati successivi (Layer 1+), le teste di induzione utilizzano i pesi QK per far corrispondere le caratteristiche scritte dalle teste del Layer 0 alla posizione del token mascherato con le corrispondenti caratteristiche alla posizione della risposta sorgente.
• Meccanismo: La decomposizione QK rivela che le teste di induzione utilizzano la Q-composition (query costruita dal contesto locale della maschera) e la K-composition (key costruita dal contesto locale della sorgente) per localizzare la sorgente corrispondente. Il percorso OV copia quindi il token sorgente alla posizione mascherata.

2. Accesso al Contesto vs. Forza del Meccanismo

Gli autori dimostrano che la superiore performance di induzione dei DLM rispetto ai modelli AR non è dovuta a un meccanismo di induzione fondamentalmente più forte, ma piuttosto all'accesso al contesto bidirezionale.

• Quando limitati al solo contesto a sinistra (mimando l'impostazione AR), i DLM performano al livello o al di sotto dei modelli AR.
• I DLM superano i modelli AR solo quando sia il vicino di sinistra che quello di destra del token mascherato sono visibili.
• Il punteggio di induzione nell'impostazione bidirezionale completa approssima la somma dei punteggi delle impostazioni "solo sinistra" e "solo destra", confermando che il vantaggio deriva dall'utilizzo delle informazioni da entrambe le direzioni.

3. Timestep Implicito tramite il Tasso Globale di Maschera

Il documento fornisce prove causali che i DLM calcolano e utilizzano la frazione globale di token mascherati come un timestep implicito, nonostante manchino di embedding espliciti del timestep.

• Rappresentazione: Un linear probe può recuperare la frazione globale di maschera dal residual stream dopo il primo strato di attenzione con alta precisione ($R^2 \approx 0.82-0.91$).
• Utilità Causale: Il patching della "direzione del tasso di maschera" da uno stato di alta corruzione a uno di bassa corruzione aumenta causalmente l'entropia di predizione alla posizione mascherata. Ciò conferma che il modello utilizza il livello di corruzione globale per modulare la propria incertezza, trattando efficacemente il tasso di maschera come un timestep.

4. Robustezza ed Emergenza del Circuito

• Emergenza: L'induzione emerge bruscamente durante l'addestramento (intorno a 8–24k step) e richiede almeno due strati per funzionare, in linea con la necessità di uno strato per scrivere le caratteristiche e di uno strato successivo per recuperarle.
• Finestra Locale: Il meccanismo di induzione si basa fortemente su una finestra locale di circa cinque token (il token mascherato più due vicini per lato). Le prestazioni calano significativamente se questi vicini immediati sono mascherati, anche se il tasso di corruzione globale è basso.

Significato e Rivendicazioni

Questo lavoro rappresenta un passo fondamentale nella comprensione dei meccanismi interni dei DLM, andando oltre l'assunto che essi si limitino a replicare i comportamenti AR. Gli autori affermano che:

1. Divergenza Meccanicistica: I DLM implementano l'induzione attraverso un circuito bidirezionale simmetrico che differisce strutturalmente dai circuiti di induzione causali nei modelli AR, specificamente nel modo in cui le caratteristiche di query e key vengono costruite dalle posizioni mascherate.
2. Vantaggio Contestuale: Il divario di performance tra i DLM e i modelli AR nell'apprendimento in-context è attribuibile alla disponibilità di contesto bidirezionale piuttosto che a una superiorità intrinseca dell'algoritmo di induzione stesso.
3. Gestione dello Stato Implicito: I DLM derivano internamente l'informazione temporale (timestep) dallo stato globale della sequenza (frazione di maschera) senza embedding espliciti, un meccanismo che è causalmente legato all'entropia di predizione.

Il documento conclude che questi risultati forniscono una base per comprendere l'apprendimento in-context nei modelli non autoregressivi, pur riconoscendo i limiti derivanti dall'uso di modelli piccoli, attention-only e task sintetici, notando che ulteriori ricerche sono necessarie per determinare se tali meccanismi scalino su DLM più grandi contenenti moduli MLP e su compiti di linguaggio naturale.