Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

AI summary

Overview

Research area: mechanistic interpretability of text-to-image diffusion transformers (DiTs), applied to how spatial relations between objects are generated.

Technical level: Advanced — the paper assumes familiarity with diffusion models, cross-attention, transformers, and interpretability terminology, though the core circuit story is stated in relatively concrete terms.

Scope: A controlled, from-scratch training study that reverse-engineers the attention circuits a PixArt-style diffusion transformer uses to place two objects in a requested spatial relation, comparing models conditioned on random token embeddings versus a pretrained T5 encoder.

What This Paper Is About

Text-to-image diffusion transformers can now render individual objects well, but they still frequently fail to honor the spatial relations between multiple objects as described in a prompt (for example, "red square above and to the left of blue circle"). Existing remedies manipulate layouts or fine-tune models, but few works ask what internal computation actually produces a correct relation. This paper trains minimal DiT models from scratch on a synthetic two-object relational task and then traces the cross-attention circuitry the models use to solve it, showing that the circuit depends heavily on which text encoder supplied the prompt embedding.

Key Contributions

1. A mechanistic circuit account of non-commutative spatial relation generation in a diffusion transformer: for random token embedding (RTE) models, a two-head, two-stage circuit in which one cross-attention head ("spatial relation head," L2H8) reads relation words and writes positional tags onto image tokens, and another ("object generation head," L4H3) reads those tags to place the correct shape.
2. An alternative, contrasting circuit for T5-conditioned models: relation information is fused by T5 into object token embeddings (especially the second shape token), and the DiT decodes both relation and object identity from essentially one or two tokens, with head L3H7 producing spatial gradients from the relation factor inside the shape2 embedding.
3. Two general analysis tools for DiTs: "Attention Synopsis," which aggregates category-to-category cross-attention into a layer x head summary over time steps, and a weight-space head screening method that scores QK interactions between image positional features and text features against ideal spatial gradients.
4. A causal comparison of the two regimes: head ablation and VO-output injection for the RTE circuit, and vector arithmetic on the T5 shape2 embedding for the T5 circuit, plus evidence that the two circuits differ sharply in robustness to out-of-domain prompt perturbations.

Main Findings

• A minimal task reproduces the problem in a controlled setting. Each sample pairs a prompt of the form "[descriptor A] [object A] [relation] [descriptor B] [object B]" (e.g., "red square above and to the left of blue circle") with an image of two objects on a gray background, using three shapes (circle, triangle, square), two colors (red, blue), and eight spatial relations (left, right, above, below, upper left, upper right, lower left, lower right). Shape and color of the two objects are always distinct, positions avoid collisions, color descriptors are randomly dropped, and relations are paraphrased.

• Text encoder choice does not prevent high accuracy; positional encoding does matter. Both RTE and T5 models at DiT-B reach strong unique_binding and spatial_relation accuracy, showing pretrained semantic structure is not strictly required to learn object relations. RTE without positional encoding is much worse on these metrics, because without positional cues the output is invariant to permutation of text embeddings and "red A on top of blue B" collapses with "blue B on top of red A." Accuracy rises with model size up to DiT-mini, with only marginal gains to DiT-B.

• Training proceeds in stages. Color accuracy converges first, then shape, then unique_binding, with spatial_relation learned slowest. T5-DiT models converge faster on all four metrics with tighter gaps between them, an early hint that the two families use different internal mechanisms.

• RTE-DiT uses a two-stage, modularized circuit. The spatial relation head (L2H8) dominates for the first object in the text, and a small number of heads play that role for the second object. These heads activate at step 0, before image information exists, and their QK circuit transforms sinusoidal positional embeddings of image tokens into queries and MLP-projected relation-word embeddings into keys, producing smooth spatial gradients whose orientation encodes the relation (e.g., "above" gives a vertical gradient with highest energy on top). For the composite relation "below and right," the attention map decomposes cleanly as the sum of the "below" and "right" maps. The object generation head (L4H3) then links each object's image tokens to its shape word, active later in sampling (steps 4–8), invariant to the object's position in the sentence and to the relation.

• Ablation and causal injection confirm the RTE circuit. Ablating relation attention specifically in L2H8 drops spatial relation accuracy from 67% to 33%, with other heads showing negligible effects. Ablating cross-attention to object text in L4H3 drops shape generation accuracy from 90% to 76%. Injecting the VO output of L2H8 together with positional embeddings as the "image tokens" of L4H3 is sufficient to elicit selective attention from tagged locations to the "square" token; without the injected VO inputs, no obvious attention pattern appears — evidence that the object head reads the relational tag from the relation head.

• The RTE circuit generalizes across sizes, except the smallest. Spatial relation heads emerge consistently in DiT-B, mini, and micro. DiT-nano shows no such head and performs poorly on spatial relations (accuracy 5%).

• T5-DiT uses a different, object-centric circuit. Attention Synopsis reveals no clear spatial relation head pattern in T5-DiT. Token masking shows relation words, filler, and color words have little effect on generation; ablating <end_of_sentence> disrupts denoising and collapses all metrics; masking shape1 reduces relation accuracy by 15%; masking shape2 lowers shape, binding, and relation accuracies by 50%.

• Relation information is linearly present in the T5 shape2 embedding. Dimension-reduction visualizations (UMAP, tSNE, PCA) of the T5 contextual embedding of the second shape token show separation by spatial relation and, within each relation, by object2's color. Variance partitioning fits the embedding as a sum of four factor vectors (V_shape2 + V_color2 + V_shape1 + V_rel). In the T5 embedding, shape2 explains most variance (about 37.5% partial R²) with relation contributing about 12%; after DiT MLP projection, relation becomes dominant (about 21%) while shape2 falls to about 14.9%, suggesting DiTs amplify relation information in the shape2 token.

• Weight-space screening finds a T5 relation head. Screening QK interactions between image positional encodings and learned relation factor vectors against ideal spatial gradients across the eight relations identified one or two cross-attention heads producing strong spatial gradients matching the intended relation (L3H7). The same screening applied to RTE, T5, and CLIP DiTs found similar relation heads.

• Vector arithmetic on the T5 embedding changes generated relations causally. In the 4096-dimensional prompt embedding space, subtracting the factor vector for the original relation from the shape2 embedding and adding a scaled vector for an alternative relation (e.g., 3 * V_lower right or 3 * V_upper right) consistently shifts generated object positions to reflect the new relation while preserving shape and color.

• Robustness differs sharply despite similar in-domain accuracy. Small prompt changes such as inserting the filler word "the" before the object, reversing sentence order, or using a color synonym cause large drops in T5-DiT relation accuracy, while RTE-DiT stays much more stable. In Table 2's evaluation over 264 prompts (50 samples each, checkpoint at 4000 epochs with EMA), T5 DiT-B scores 0.808 spatial relation on the "O1 is Rel O2" template but 0.498 on "O1 is Rel the O2" and 0.400 on "O1 Rel the O2", with large coordinate biases such as Dx = -35.1 and Dy = -37.0. RTE DiT-B stays in the 0.82–0.86 range across templates. The authors explain this as a consequence of decoding relations from a contextual embedding that nuisance words perturb.

• The tools transfer to a pretrained model. PixArt-Sigma has only weak spatial relation ability: among 30 object pairs, only 8 show nontrivial object and relation accuracy. Token ablation indicates spatial information is present in both relation and object words, consistent with the synthetic-dataset findings, and head screening identifies a small number of heads strongly associated with spatial generation.

Methodology in Plain English

The authors build a deliberately small, controllable version of the text-to-image problem. They generate synthetic prompts and images with two objects, defined by shape and color, arranged according to one of eight spatial relations, then train PixArt-style diffusion transformers from scratch on this data. They train several sizes (DiT-B at 12 layers, 12 heads, 768 dimensions; mini at 6 layers, 6 heads, 384 dimensions; micro at 6 layers, 3 heads, 192 dimensions; nano at 3 layers, 3 heads, 192 dimensions, all with patch size 2) and vary only the text encoder: T5-XXL, a random token encoder with sinusoidal positional encoding, and the same random encoder without positional encoding. Images go through a VAE pretrained from Stable Diffusion.

Evaluation uses 96 prompts spanning the eight relations and 12 object pairs during training, with the DPM-Solver++ sampler at 14 steps and classifier-free guidance of 4.5. Four binary metrics are scored per sample using classic cv2 segmentation and classification: color correctness, shape correctness, unique attribute binding, and spatial relation correctness. Because the token categories are known on both sides of the cross-attention, the authors can aggregate attention maps into category-to-category summaries, average over time steps, and reduce a huge tensor (over 10 million attention maps) to a layer-by-head synopsis that reveals which heads carry which information. To confirm causality, they ablate individual heads and inject the output of one head into the input of another. For the T5 model, they decompose the shape2 embedding into factor vectors, then add and subtract those vectors to shift the generated relation.

Why This Matters

Impact on research: the paper connects two previously separate explanations of spatial failures — that cross-attention controls placement and that text encoders fail to preserve relations — by showing that the text encoder determines which attention circuit is even available. It gives a concrete mechanistic example of how a diffusion transformer implements a non-commutative relation, and offers reusable tooling (Attention Synopsis, weight-space head screening) for analyzing DiTs despite the complication of iterative sampling.

Real-world applications:

• Prompt-adherence debugging for image generation products, where failures on multi-object layouts are a common user complaint.
• Text-encoder selection and design for text-to-image systems, since the paper suggests the embedding model may be the bottleneck rather than the DiT.
• Layout-conditioned or attention-guidance editing tools that could target the identified relation and object heads directly.
• Robustness testing of generative models against trivial prompt variation such as added filler words or synonyms.

Industry relevance: teams building or fine-tuning diffusion transformers can use the finding that a compact, disentangled relation-to-location pathway arises with non-semantic embeddings while pretrained encoders fuse relation into object tokens at the cost of brittleness. That trade-off between interpretability, prompt robustness, and reliance on explicit relation tokens is directly relevant to architecture and encoder choices in production systems.

Future Directions

• Determine whether the circuit differences observed on a synthetic two-object task carry over to real-world datasets and natural prompts at scale, given that PixArt-Sigma showed weak relation ability with only 8 of 30 object pairs nontrivial.
• Explain why the T5 circuit is so sensitive to nuisance words, and whether training or embedding interventions can restore prompt robustness without losing the compact object-centric representation.
• Use the weight-space screening method more broadly across encoders and pretrained models to test whether sparse, interpretable spatial relation heads are a general property of DiTs.
• Investigate whether the identified relation and object heads can be edited or steered to improve spatial control in pretrained image generators.

Target Audience

Mechanistic interpretability researchers studying diffusion and multimodal models; text-to-image researchers working on compositional generation, attribute binding, and spatial control; engineers choosing or designing text encoders for generative systems; and graduate students or advanced practitioners comfortable with cross-attention, diffusion sampling, and ablation methodology.

Authors’ 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.