Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

Probleemstelling

Hoewel Diffusion Transformers (DiT) aanzienlijke vooruitgang hebben geboekt in tekst-naar-beeld generatie, worstelen deze modellen nog steeds met het correct genereren van ruimtelijke relaties tussen meerdere objecten zoals beschreven in een tekstprompt (bijv. "een rode vierkant boven een blauwe cirkel"). Bestaande modellen kunnen vaak de attributen van individuele objecten (kleur, vorm) goed genereren, maar falen bij het correct positioneren van deze objecten ten opzichte van elkaar. Bestaande oplossingen, zoals layout-conditioning of het manipuleren van attention-maps tijdens inferentie, zijn vaak ad-hoc en bieden geen inzicht in de onderliggende neurale circuits die verantwoordelijk zijn voor deze relaties.

Methodologie

De auteurs hanteren een mechanistische interpretabiliteitsbenadering om te onderzoeken hoe DiT-modellen ruimtelijke relaties leren en implementeren.

1. Dataset en Training:

   • Ze hebben een synthetische dataset opgezet met prompts van de vorm [object A] [relatie] [object B] (bijv. "rode vierkant links van blauwe cirkel").
   • Ze trainen DiT-modellen van verschillende groottes (van nano tot B) vanaf nul (from scratch) om deze taak te leren.
   • Vergelijking van Tekstencoders: Ze vergelijken drie soorten tekstconditioning:
     • RTE (Random Token Embeddings): Random geinitialiseerde vectoren met sinusoidale positiële encoding. Dit elimineert semantische structuur en test of het model relaties puur via de architectuur kan leren.
     • T5: Een voorgeïntegreerde, semantisch rijke tekstencoder.
     • CLIP: Een andere veelgebruikte encoder (in de appendix).

2. Analyse-technieken:

   • Attention Synopsis: Een schaalbare methode om cross-attention maps te aggregeren over tijdstappen en token-categorieën (bijv. "relatie-woorden" vs. "object-tokens") om stabiele communicatiepatronen te identificeren.
   • Weight-Space Screening: Een methode om heads te vinden die een lineaire relatie leggen tussen positiële embedding van afbeeldingstokens en tekst-features (via QK-interacties), zonder dat er beelden gegenereerd hoeven te worden.
   • Causale Manipulatie: Ablatiestudies (het verwijderen van specifieke attention heads) en vector-arithmetic (het manipuleren van embeddingvectoren) om de causaliteit van de gevonden circuits te verifiëren.

Kernbevindingen en Resultaten

De studie onthult dat de onderliggende circuits drastisch verschillen afhankelijk van de gekozen tekstencoder, hoewel beide benaderingen hoge nauwkeurigheid kunnen bereiken.

1. Circuit in RTE-DiT (Random Embeddings): Een Twee-Staps Schakeling

Wanneer het model trainen met random embeddings (RTE), ontwikkelt het een modulair, ontvleesd circuit bestaande uit twee gespecialiseerde cross-attention heads:

• Stap 1: Ruimtelijke Relatie Head (bijv. L2H8): Deze head leest de relatie-woorden (bijv. "boven", "links") en transformeert deze via QK-interacties met de positiële embedding van de afbeeldingstokens. Dit creëert ruimtelijke gradiënten (positional tags) op het canvas die aangeven waar het eerste object moet verschijnen.
• Stap 2: Object Generatie Head (bijv. L4H3): Deze head leest de vorm- en kleur-woorden (bijv. "vierkant", "rood") en plaatst het object op de locaties die zijn gemarkeerd door de eerste stap.
• Conclusie: Het proces is strikt gescheiden: eerst wordt de locatie bepaald, daarna wordt het object geplaatst. Dit circuit is zeer robuust tegen kleine veranderingen in de prompt (zoals het toevoegen van vulwoorden zoals "de").

2. Circuit in T5-DiT: Informatiefusie in Tokens

Bij gebruik van de T5-encoder (die semantische context biedt) werkt het model op een fundamenteel andere manier:

• Geen gescheiden relatie-heads: De T5-encoder integreert de volledige zinscontext in de token-embeddings. De auteurs vinden dat ruimtelijke informatie niet via een specifieke relatie-woord-token wordt gelezen, maar is gefuseerd in de embedding van het tweede object-token (bijv. het woord voor het tweede object bevat informatie over de relatie ten opzichte van het eerste).
• Mechanisme: De DiT leest de ruimtelijke relatie en de objectattributen tegelijkertijd uit één enkele token-embedding (de contextuele embedding van het tweede object).
• Fragiliteit: Hoewel dit efficiënt is, is het circuit kwetsbaar. Kleine perturbaties in de prompt (zoals het toevoegen van "the" of het veranderen van woordvolgorde) verstoren de contextuele embedding van het object-token. Dit leidt tot een drastische daling in de nauwkeurigheid van ruimtelijke relaties, terwijl de attributen van individuele objecten intact blijven.

3. Robuustheid en Generalisatie

• RTE-modellen zijn robuust tegen lexische variaties omdat ze expliciete relatie-tokens direct koppelen aan ruimtelijke gradiënten.
• T5-modellen zijn gevoelig voor perturbaties omdat ze afhankelijk zijn van de precieze contextuele samenstelling van de token-embeddings. Een kleine wijziging in de tekst verschuift de embeddingvector in een richting die de DiT interpreteert als een andere ruimtelijke relatie.

Bijdragen

1. Mechanistisch Inzicht: Het paper biedt het eerste concrete bewijs van de neurale circuits die verantwoordelijk zijn voor ruimtelijke redenering in Diffusion Transformers. Het toont aan dat deze circuits niet universeel zijn, maar afhankelijk van de input-embeddings.
2. Nieuwe Analyse-methoden: De auteurs introduceren de "Attention Synopsis" en "Weight-Space Screening" als systematische tools om attention-patronen in iteratieve generatieve modellen te analyseren.
3. Rol van Embeddings: Het paper toont aan dat de keuze van de tekstencoder (en de daaruit voortvloeiende inductieve bias) bepalend is voor hoe het model relaties leert. Dit suggereert dat verbeteringen in embedding-modellen misschien belangrijker zijn dan het aanpassen van de DiT-architectuur zelf voor ruimtelijke taken.
4. Verklaring van Falen: Het biedt een mechanistische verklaring waarom T2I-modellen falen bij ruimtelijke relaties in real-world scenario's: de afhankelijkheid van contextuele embeddings (zoals bij T5) maakt het systeem kwetsbaar voor de variabiliteit van natuurlijke taal.

Significantie

Deze studie is cruciaal voor het veld van generatieve AI omdat het de "black box" van ruimtelijke redenering openlegt. Het laat zien dat hoge nauwkeurigheid niet altijd betekent dat het model robuust is. Voor toekomstige ontwerpen van tekst-naar-beeld modellen suggereert het dat een ontvleesde (disentangled) representatie van relaties (zoals in RTE) wenselijker kan zijn voor controle en stabiliteit, terwijl semantisch rijke encoders (zoals T5) voordelen bieden voor taalgeneralisatie maar ten koste gaan van ruimtelijke precisie bij kleine tekstvariaties. Dit inzicht kan leiden tot betere trainingstechnieken of hybride architecturen die de sterktes van beide benaderingen combineren.