Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

TL;DR

This work investigates how diffusion transformers generate spatial relations between multiple objects by applying mechanistic interpretability to models trained from scratch on a minimal two-object task. It reveals two distinct circuits governed by the text encoder: random embeddings rely on a two-stage modular circuit with a dedicated spatial relation head and an object-generation head, while a pretrained T5 encoder uses information fusion within text tokens so relations are decoded from a single token. The findings show similar end-task performance but different robustness to prompt perturbations, highlighting how embedding design shapes inductive bias and interpretability. These insights inform how to design more robust, controllable, and explainable T2I systems where spatial composition is crucial.

Abstract

Diffusion Transformers (DiTs) have greatly advanced text-to-image generation, but models still struggle to generate the correct spatial relations between objects as specified in the text prompt. In this study, we adopt a mechanistic interpretability approach to investigate how a DiT can generate correct spatial relations between objects. We train, from scratch, DiTs of different sizes with different text encoders to learn to generate images containing two objects whose attributes and spatial relations are specified in the text prompt. We find that, although all the models can learn this task to near-perfect accuracy, the underlying mechanisms differ drastically depending on the choice of text encoder. When using random text embeddings, we find that the spatial-relation information is passed to image tokens through a two-stage circuit, involving two cross-attention heads that separately read the spatial relation and single-object attributes in the text prompt. When using a pretrained text encoder (T5), we find that the DiT uses a different circuit that leverages information fusion in the text tokens, reading spatial-relation and single-object information together from a single text token. We further show that, although the in-domain performance is similar for the two settings, their robustness to out-of-domain perturbations differs, potentially suggesting the difficulty of generating correct relations in real-world scenarios.

Figures (23)

• Figure 1: Schematics of the model and task. Our T2I model architecture adopted the design of PixArt chen2023PixArtAlpha. There are three main components: the text encoder that processes tokenized natural language prompts into text embeddings, the VAE that processes image inputs into image tokens, and the Diffusion Transformer (DiT) which is the backbone of the denoising diffusion process. The text information routes through the cross attention mechanism in each DiT block and influence the denoising of image tokens. The task is to generate two objects with a specified spatial relation.
• Figure 2: Training dynamics of the T2I models (DiT-B). A) and B) Both models trained with random token embedding and T5 can achieve good accuracy on the task. Solid lines shows the result of model using exponential moving averaged (ema) weights, while dashed line shows the non-averaged weights. C) The task is learned in distinct stages. In both models, they first learn to generate objects but with wrong attributes binding, then they the correct binding of single-object attributes (e.g. red square), finally they learn the correct spatial relation. (see App.\ref{['app:eval_learn_dynamics']} for other models).
• Figure 3: Illustration of our method to find relevant heads, which we name as "attention synopsis". The giant attention tensor is first reduced to those only between two interested groups of tokens (e.g. the relation token regardless specific words, or an object token regardless where or what it is). Then the reduced attention tensor is averaged over diffusion time steps, resulting in a layer $\times$ head map which we use to pinpoint relevant heads.
• Figure 4: The spatial relation heads in random-embedding-based T2I. A) We find specialized cross attention heads that contributes to the object image tokens (top: the object1 in the text; bottom: the object2 in the text) attending to the relation text tokens. B) We show the activation of this head across images tokens and sampling steps. The map for the composite relation “below and right” decomposes cleanly as the sum of the maps for “below” and “right”, C) The observed attention patterns can be induced by positional embedding.
• Figure 5: The object generation heads in random-embedding-based T2I. A) We find specialized heads in the synopses of cross-attention, computed from image tokens of each object to its own shape tokens. B) We show the activation of this head across images tokens and sampling steps for the prompt “red square is below and to the right of the blue circle”: tokens at the eventual square location attend to “square,” while the other object attends to “circle”; selectivity sharpens from Step 0→12. C) Injecting the VO output of the relation head (L2H8) into positional embeddings is sufficient to elicit selective attention from tagged locations to the “square” token (left); without the tag the pattern is weak (right). This indicates the object generation head reads the relational tag generated by the spatial relation head.