Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

Titolo: Meccanismi Circuituali per la Generazione di Relazioni Spaziali nei Diffusion Transformers

1. Il Problema

I modelli di generazione di immagini basati su Diffusion Transformers (DiT) hanno raggiunto risultati eccezionali nella sintesi di immagini ad alta fedeltà da prompt testuali. Tuttavia, persiste una significativa difficoltà nel generare correttamente le relazioni spaziali tra più oggetti specificati nel prompt (es. "un quadrato rosso sopra un cerchio blu").
Mentre la generazione di attributi singoli (colore, forma) è migliorata notevolmente, la composizione spaziale rimane un collo di bottiglia. Le ricerche precedenti hanno attribuito questo fallimento a diverse cause: mancanza di guida nell'attenzione incrociata (cross-attention) o limitazioni negli encoder testuali (che non preservano sufficientemente le relazioni spaziali). Questo studio si propone di indagare il problema attraverso la interpretabilità meccanicistica, cercando di comprendere i circuiti interni (neuroni, teste di attenzione, pesi) che i DiT utilizzano per implementare queste relazioni e perché falliscono in determinate configurazioni.

2. Metodologia

Gli autori hanno adottato un approccio controllato e sistematico:

• Dataset Sintetico: Hanno creato un dataset minimale contenente immagini con due oggetti (forme e colori distinti) su sfondo grigio, descritti da prompt testuali che specificano una delle otto relazioni spaziali possibili (es. sinistra, destra, sopra, sotto, diagonali).
• Architettura e Training: Hanno addestrato da zero DiT di diverse dimensioni (da Nano a B) seguendo l'architettura PixArt.
• Variabile Chiave (Encoder Testuale): Per isolare il meccanismo di ragionamento, hanno confrontato tre tipi di encoder testuali:
  1. T5-XXL: Un encoder pre-addestrato semantico che produce embedding contestuali (le parole sono mescolate semanticamente).
  2. RTE (Random Token Embedding): Un encoder che mappa i token a vettori casuali fissi, privi di struttura semantica, ma con codifica posizionale sinusoidale.
  3. RTE senza posizione: Una variante senza codifica posizionale (usata come controllo negativo).
• Strumenti di Analisi:
  • Attention Synopsis: Un metodo scalabile per aggregare e riassumere le mappe di attenzione incrociata tra categorie di token (es. token di relazione vs. token di oggetto) su tutti i passaggi di denoising, riducendo la complessità per identificare i circuiti stabili.
  • Screening nello Spazio dei Pesi: Una tecnica per identificare le teste di attenzione responsabili delle relazioni spaziali analizzando l'interazione QK (Query-Key) tra le posizioni spaziali dell'immagine e i vettori di feature testuali, senza bisogno di generare immagini.
  • Manipolazione Causale: Ablazione di specifiche teste di attenzione e manipolazione vettoriale degli embedding (aritmetica vettoriale) per testare la causalità.

3. Risultati Chiave e Meccanismi Scoperti

Lo studio rivela che, sebbene entrambi i modelli (T5 e RTE) raggiungano un'alta accuratezza nel task, utilizzano circuiti interni radicalmente diversi:

A. Modelli con Random Token Embedding (RTE-DiT): Un Circuito Modulare a Due Stadi
In assenza di struttura semantica preesistente, il modello impara un circuito altamente disaccoppiato:

1. Testa di Relazione Spaziale (Spatial Relation Head): Una testa di attenzione specifica (es. Layer 2, Head 8) legge i token di relazione (es. "sopra") e interagisce con le codifiche posizionali sinusoidali dei token immagine. Questa testa genera gradienti spaziali (mappe di attenzione) che agiscono come "tag" posizionali, indicando dove l'oggetto deve apparire.
2. Testa di Generazione dell'Oggetto (Object Generation Head): Una seconda testa (es. Layer 4, Head 3) legge i token di forma/colore e, basandosi sui "tag" posizionali generati dalla prima testa, assegna gli attributi corretti alle posizioni indicate.

• Conclusione: Il processo è sequenziale e modulare: prima si definisce la posizione, poi si assegna l'identità.

B. Modelli con Encoder T5 (T5-DiT): Fusione Contestuale
Con un encoder semantico come T5, il meccanismo è diverso e più compatto:

1. Fusione nell'Embedding: L'attenzione self-attention di T5 fonde le informazioni dell'intera frase. Di conseguenza, l'informazione spaziale non risiede più solo nel token della relazione, ma viene incorporata nell'embedding contestuale dei token degli oggetti (in particolare il token della seconda forma, shape2).
2. Decodifica da Singolo Token: Il DiT legge le informazioni di relazione e di oggetto direttamente dall'embedding del token shape2. Le teste di attenzione non leggono più il token "sopra" separatamente, ma decodificano la relazione geometrica direttamente dal vettore contestuale dell'oggetto.
3. Meccanismo: Le teste di attenzione (es. Layer 3, Head 7) interagiscono con i fattori di relazione estratti dall'embedding di shape2 per generare gradienti spaziali.

4. Robustezza e Generalizzazione

Una scoperta cruciale riguarda la robustezza ai disturbi del prompt (out-of-domain perturbations):

• RTE-DiT: È molto robusto. Poiché il circuito è modulare e basato su token specifici, l'aggiunta di parole riempitive (es. "the") o variazioni sintattiche non rompe la logica spaziale.
• T5-DiT: È estremamente fragile. L'aggiunta di una semplice parola riempitiva o una variazione lessicale altera l'embedding contestuale di shape2. Poiché il DiT decodifica la relazione da questo embedding unico, la perturbazione sposta il vettore in una direzione che il modello interpreta come una relazione spaziale diversa, causando un crollo dell'accuratezza (fino al 40-50% di perdita).

5. Contributi e Significato

1. Mappatura dei Circuiti di Relazione: Il paper fornisce la prima evidenza meccanicistica dettagliata di come i DiT implementino relazioni non commutative tra oggetti, identificando circuiti specifici (gradienti spaziali generati da interazioni QK).
2. Ruolo Critico dell'Encoder: Dimostra che la scelta dell'encoder testuale non influenza solo le prestazioni, ma definisce l'architettura interna del ragionamento. Gli encoder semantici (T5) creano un collo di bottiglia di robustezza perché fondono troppe informazioni in un singolo vettore, rendendo il sistema sensibile a piccole perturbazioni linguistiche.
3. Metodologia Generale: Introduce l'uso combinato di "Attention Synopsis" e "Weight-space screening" come strumenti potenti per analizzare i modelli di diffusione, superando la complessità temporale del campionamento iterativo.
4. Implicazioni per il Design: Suggerisce che per migliorare la generazione spaziale in scenari reali, potrebbe essere più critico migliorare o progettare meglio gli encoder di embedding (per renderli più robusti o disaccoppiati) piuttosto che modificare l'architettura del DiT stesso.

In sintesi, lo studio rivela che la capacità di generare relazioni spaziali non è un comportamento monolitico dei modelli, ma emerge da strategie circuituali distinte che dipendono fortemente da come il testo viene codificato in input, con un compromesso (trade-off) tra efficienza della rappresentazione (T5) e robustezza/interpretabilità (RTE).