Diffusion Transformer 中空间关系生成的电路机制

01 动机Motivation

文生图模型在生成单一物体属性（颜色、形状）方面已取得显著进步，但在生成"A 在 B 的左边"这类组合空间关系时仍然频繁出错。已有研究从两个角度归因：一是 cross-attention 机制不足，二是文本编码器的局限性。本文提出：文本编码器的设计从根本上塑造了空间关系生成所对应的内部计算电路，从而将两种视角统一起来。 Text-to-image models have made marked progress on single-object attributes such as color and shape, yet they still fail frequently on compositional spatial relations such as "A is to the left of B". Prior work attributes this to two causes: an inadequate cross-attention mechanism, and limitations of the text encoder. This paper argues that the design of the text encoder fundamentally shapes the internal computational circuit used for spatial relation generation, thereby unifying the two views.

"We analyze how Diffusion Transformers (DiTs) generate correct spatial relations … and discover that the underlying mechanisms differ significantly based on text encoder choice."

teaser figure — **图1：论文核心发现概览。** 左列 (RTE-DiT)：使用随机 token 嵌入训练的 DiT，通过两个专用注意力头构成模块化电路；右列 (T5-DiT)：使用预训练 T5 编码器时，空间信息通过 contextual embedding 的向量运算传递。两种模型均能达到约 80–84% 的空间关系准确率，但内部机制截然不同。 **Figure 1: overview of the paper's core findings.** Left column (RTE-DiT): a DiT trained with random token embeddings builds a modular circuit out of two dedicated attention heads. Right column (T5-DiT): with a pretrained T5 encoder, spatial information is carried by vector arithmetic over contextual embeddings. Both models reach roughly 80-84% spatial relation accuracy, but their internal mechanisms are entirely different.

84.3%RTE-DiT 空间关系准确率 (DiT-B)RTE-DiT spatial relation accuracy (DiT-B)

80.8%T5-DiT 空间关系准确率 (DiT-B)T5-DiT spatial relation accuracy (DiT-B)

2 阶段2 stagesRTE 电路的注意力头数量Number of attention heads in the RTE circuit

5%DiT-nano 空间关系准确率（失败基线）DiT-nano spatial relation accuracy (failure baseline)

核心问题：空间关系是如何在 DiT 内部被编码的？Core question: how are spatial relations encoded inside a DiT?

空间关系（"上方"、"左侧"等）在语义上是非交换的（"A 在 B 左边" ≠ "B 在 A 左边"），且需要将文本语义和图像空间位置绑定在一起——这对模型来说是一个根本性挑战。以往工作只关注现象而未揭示机制，本文通过在受控数据集上从头训练 DiT，系统地追踪了空间关系是如何在内部被计算的。 Spatial relations ("above", "left of", etc.) are semantically non-commutative ("A is left of B" ≠ "B is left of A"), and they require binding textual semantics to image-space position - a fundamental challenge for the model. Earlier work only documented the phenomenon without revealing the mechanism; here the authors train DiTs from scratch on a controlled dataset and systematically trace how spatial relations are computed internally.

02 方法Method

研究采用"受控数据集 + 从头训练"的范式，配合提出的 Attention Synopsis 方法，系统地在扩散时间步和空间 token 两个维度上定位关键注意力头，再通过消融和因果干预验证其功能。 The study adopts a "controlled dataset + training from scratch" paradigm, together with the proposed Attention Synopsis method, to systematically localize the key attention heads along both the diffusion timestep and the spatial token axes, and then verifies their function through ablation and causal intervention.

受控数据集构建Building the controlled dataset

dataset schematic — **图2：数据集设计。** 图像包含两个有颜色的几何形状（3 种形状：circle / triangle / square；2 种颜色：red / blue），背景为灰色，物体位置避免碰撞。 Prompt 格式为 `[color A] [shape A] [relation] [color B] [shape B]`，共 8 种空间关系（left / right / above / below / upper-left / upper-right / lower-left / lower-right）。评估集包含 96 条 prompts（8 关系 × 12 物体对）。 **Figure 2: dataset design.** Each image contains two colored geometric shapes (3 shapes: circle / triangle / square; 2 colors: red / blue) on a gray background, with object positions kept collision-free. Prompts follow the format `[color A] [shape A] [relation] [color B] [shape B]`, covering 8 spatial relations (left / right / above / below / upper-left / upper-right / lower-left / lower-right). The evaluation set contains 96 prompts (8 relations × 12 object pairs).

PixArt 风格模型架构PixArt-style model architecture

模型采用 PixArt-style 架构，由三部分组成： The model follows a PixArt-style architecture with three components:

文本编码器：分别使用 T5-XXL（预训练）、Random Token Embedding (RTE)、或 RTE 无位置编码三种配置Text encoder: three configurations - T5-XXL (pretrained), Random Token Embedding (RTE), or RTE without positional encoding
VAE：来自 Stable Diffusion 的预训练 VAEVAE: the pretrained VAE from Stable Diffusion
Diffusion Transformer 主干：测试 DiT-B (12L, 12H, 768d)、mini (6L, 6H, 384d)、micro (6L, 3H, 192d)、nano (3L, 3H, 192d) 四种规模Diffusion Transformer backbone: four scales are tested - DiT-B (12L, 12H, 768d), mini (6L, 6H, 384d), micro (6L, 3H, 192d), nano (3L, 3H, 192d)

Attention Synopsis 方法The Attention Synopsis method

本文提出 Attention Synopsis：对所有扩散时间步的注意力图取汇总统计，并在所有 prompt 条件下聚合，从而识别"对特定语义变量（如 relation）有选择性响应"的注意力头。这避免了只看单一时间步带来的偶然性，能系统筛查模型中的稀疏功能电路。 The paper proposes Attention Synopsis: summary statistics of the attention maps are taken across all diffusion timesteps and aggregated over all prompt conditions, so as to identify attention heads that respond selectively to a particular semantic variable such as the relation. This avoids the arbitrariness of inspecting a single timestep and makes it possible to systematically screen a model for sparse functional circuits.

circuit diagram — **图3：RTE-DiT 中发现的两阶段空间电路示意图。** 第一阶段（Stage 1）：Layer 2 Head 8（L2H8）为"空间关系头"，读取 relation token 的 key，将正弦位置编码映射为空间梯度（如"above"对应竖向梯度）；第二阶段（Stage 2）：Layer 4 Head 3（L4H3）为"物体生成头"，接收来自 L2H8 的位置标记，将其与物体 shape token 匹配并在对应位置生成物体。 **Figure 3: schematic of the two-stage spatial circuit found in RTE-DiT.** Stage 1: Layer 2 Head 8 (L2H8) is the "spatial relation head"; it reads the key of the relation token and maps the sinusoidal positional encoding into a spatial gradient (for example "above" maps to a vertical gradient). Stage 2: Layer 4 Head 3 (L4H3) is the "object generation head"; it receives the positional marker from L2H8, matches it against the object shape token, and renders the object at the corresponding location.

Stage 1：空间关系头 (L2H8)Stage 1: spatial relation head (L2H8)

通过 QK circuit 读取 relation text tokenReads the relation text token through the QK circuit
将正弦位置编码变换到 query 空间，relation embedding 变换到 key 空间Maps the sinusoidal positional encoding into query space and the relation embedding into key space
输出平滑的空间梯度（如"above"→竖向梯度）Outputs a smooth spatial gradient (for example "above" → a vertical gradient)
在采样步骤 step 0 即开始激活Already active from sampling step 0

Stage 2：物体生成头 (L4H3)Stage 2: object generation head (L4H3)

接收来自关系头的位置标记Receives the positional marker from the relation head
将标记位置与物体 shape token 匹配Matches the marked position against the object shape token
在对应空间位置生成物体Generates the object at the corresponding spatial position
在采样步骤 4–8 后期激活Activates late, after sampling steps 4-8

T5-DiT 中的不同机制A different mechanism inside T5-DiT

T5 word vector manipulation — **图4：T5-DiT 中的向量运算干预实验。** 通过从第二个物体的 T5 contextual embedding 中减去原 relation 向量并加入目标 relation 向量，可以因果性地改变生成物体的空间位置，同时保持形状和颜色不变——证明空间信息融合在 contextual token 中。 **Figure 4: vector-arithmetic intervention experiment in T5-DiT.** Subtracting the original relation vector from the T5 contextual embedding of the second object and adding a target relation vector causally changes the spatial position of the generated object while leaving its shape and color unchanged - evidence that the spatial information is fused into the contextual token.

在 T5-DiT 中，空间信息并非由独立的 attention head 处理，而是通过 T5 自注意力将 relation 信息融合进第二个物体 token 的 contextual embedding 中。方差分解（Variance Partitioning）显示： DiT MLP 映射前，shape2 解释约 37.5% 的方差，relation 仅占 12%；经 DiT MLP 投影后，relation 的贡献上升至 21.3%。 In T5-DiT the spatial information is not handled by a dedicated attention head; instead T5 self-attention fuses the relation information into the contextual embedding of the second object token. Variance partitioning shows that before the DiT MLP mapping, shape2 explains about 37.5% of the variance while the relation accounts for only 12%; after the DiT MLP projection the contribution of the relation rises to 21.3%.

03 实验结果Experimental results

在受控数据集上从头训练，使用 4 项二值指标衡量生成质量： 颜色准确率、形状准确率、绑定准确率（正确的形状-颜色对应）、空间关系准确率。评估在 96 条 prompt 上进行（8 关系 × 12 物体对）。 Models are trained from scratch on the controlled dataset and generation quality is measured with 4 binary metrics: color accuracy, shape accuracy, binding accuracy (correct shape-color correspondence), and spatial relation accuracy. Evaluation runs over 96 prompts (8 relations × 12 object pairs).

主要性能对比（DiT-B）Main performance comparison (DiT-B)

模型配置Model configuration	形状 ShapeShape	颜色 ColorColor	绑定 BindingBinding	空间关系 SpatialSpatial	严格空间 StrictStrict spatial
RTE-DiT	0.877	0.928	0.855	0.843	0.758
T5-DiT	0.857	0.892	0.820	0.808	0.749
CLIP-DiT-B	0.806	0.900	0.772	0.759	—
RTE（无位置编码）RTE (no positional encoding)	—	—	0.415	0.207	—
DiT-nano（RTE）DiT-nano (RTE)	—	—	—	0.050	—

注：RTE 无位置编码时空间准确率崩溃至 20.7%，证明位置信息对空间关系生成是必要条件。 DiT-nano（3 层）空间准确率仅 5%，说明模型容量达到阈值以下时电路无法形成。 Note: without positional encoding, RTE spatial accuracy collapses to 20.7%, showing that positional information is a necessary condition for spatial relation generation. DiT-nano (3 layers) reaches only 5% spatial accuracy, indicating that below a certain model-capacity threshold the circuit cannot form.

消融实验：电路头的因果验证（RTE-DiT）Ablation: causal validation of the circuit heads (RTE-DiT)

消融条件Ablation condition	受影响指标Affected metric	消融前Before	消融后After	下降幅度Drop
移除 L2H8（空间关系头）Remove L2H8 (spatial relation head)	空间关系准确率Spatial relation accuracy	67%	33%	−34pp
移除 L4H3（物体生成头）Remove L4H3 (object generation head)	形状生成准确率Shape generation accuracy	90%	76%	−14pp

训练动态Training dynamics

鲁棒性与扰动分析Robustness and perturbation analysis

relation head visualization — **图6：L2H8 空间关系头的注意力可视化。** 对于不同的 relation 词（left / right / above / below 等），该头输出不同方向的平滑空间梯度，在采样 step 0 即产生，为后续物体生成阶段提供空间位置先验。 **Figure 6: attention visualization of the L2H8 spatial relation head.** For different relation words (left / right / above / below and so on) the head outputs a smooth spatial gradient in a different direction, already present at sampling step 0, providing a spatial position prior for the later object generation stage.

RTE-DiT 鲁棒性RTE-DiT robustness

对填充词（"the"）鲁棒：添加 "the" 不影响性能Robust to filler words such as "the": adding "the" does not affect performance
对 relation 词删除敏感：移除 relation 词后准确率显著下降Sensitive to deletion of the relation word: accuracy drops markedly once the relation word is removed
电路依赖 relation token 的 key 向量，因此对 relation 词本身敏感The circuit relies on the key vector of the relation token, hence its sensitivity to the relation word itself

T5-DiT 鲁棒性T5-DiT robustness

对填充词敏感：在物体前添加 "the" 后，空间准确率从 0.808 → 0.498Sensitive to filler words: adding "the" before the objects moves spatial accuracy from 0.808 → 0.498
原因：T5 contextual embedding 对分布偏移脆弱，单一 token 的编码受上下文影响大Reason: T5 contextual embeddings are fragile under distribution shift, and the encoding of a single token depends heavily on context
向量运算干预实验证实信息存储在 object2 的 contextual embedding 中The vector-arithmetic intervention confirms that the information is stored in the contextual embedding of object2

迁移到预训练模型 PixArt-SigmaTransfer to the pretrained PixArt-Sigma model

将电路分析工具应用于公开预训练的 PixArt-Sigma 模型：在所测试的物体对中，约 27% 展现出非平凡的空间关系生成能力，同样能定位到稀疏的空间关系电路，与在受控模型上发现的机制一致。 Applying the circuit analysis toolkit to the publicly released pretrained PixArt-Sigma model: about 27% of the tested object pairs exhibit non-trivial spatial relation generation, and a sparse spatial relation circuit can likewise be localized, consistent with the mechanism found in the controlled models.

04 局限性Limitations

注：本文没有设置专门的 Limitations 章节。以下各点为基于论文设计和文中讨论推断（inferred）的局限性，少数为论文结果中明确提及的范围限制（标注"stated"）。 Note: the paper has no dedicated Limitations section. The points below are inferred from the paper's design and its discussion; a few are scope limits explicitly stated in the results (marked "stated").

受控数据集过于简化，与真实世界场景存在较大差距（inferred）The controlled dataset is heavily simplified and far from real-world scenes (inferred)

实验仅使用 3 种形状、2 种颜色、8 种空间关系的极简数据集，背景为纯灰色、物体无纹理。虽然这种受控设置便于机制发现，但其结论能否直接推广到真实场景的文生图模型（如复杂背景、多物体、连续颜色）尚不明确。 The experiments use a minimal dataset of only 3 shapes, 2 colors and 8 spatial relations, with a plain gray background and untextured objects. This controlled setting makes mechanism discovery tractable, but whether the conclusions transfer directly to real-world text-to-image models (complex backgrounds, many objects, continuous colors) remains unclear.

预训练大模型中电路分析成功率有限（stated）Circuit analysis succeeds only to a limited extent in large pretrained models (stated)

在预训练 PixArt-Sigma 上，仅约 27% 的物体对展示出非平凡的空间生成能力，稀疏度更高、机制更难识别。在通用大模型中进行完整的电路逆向工程仍面临巨大挑战。 On the pretrained PixArt-Sigma, only about 27% of object pairs show non-trivial spatial generation ability; the circuits are sparser and harder to identify. Full circuit-level reverse engineering of general large models remains a formidable challenge.

T5-DiT 的鲁棒性明显弱于 RTE-DiT（stated）T5-DiT is clearly less robust than RTE-DiT (stated)

T5 contextual embedding 的空间信息高度依赖 prompt 的词语组合，仅添加填充词 "the" 就能使空间准确率从 80.8% 跌至 49.8%（下降 31pp）。这表明基于预训练语言模型的 T5 电路存在分布偏移脆弱性，在面向更多样化 prompt 时泛化能力受限。 The spatial information in T5 contextual embeddings depends heavily on the exact wording of the prompt: merely adding the filler word "the" drops spatial accuracy from 80.8% to 49.8% (a fall of 31pp). This indicates that the T5 circuit, built on a pretrained language model, is fragile under distribution shift and generalizes poorly to more diverse prompts.

仅研究了空间关系，未涵盖其他组合推理类型（inferred）Only spatial relations are studied; other kinds of compositional reasoning are not covered (inferred)

本文聚焦于"左/右/上/下"等 8 种二维空间关系，未涉及数量关系、属性绑定以外的组合推理（如大小比较、遮挡等）。是否存在类似的模块化电路机制需要进一步研究。 The paper focuses on 8 two-dimensional spatial relations such as left / right / above / below, and does not address counting or compositional reasoning beyond attribute binding (such as size comparison or occlusion). Whether analogous modular circuit mechanisms exist there requires further study.

DiT-nano 以下规模模型无法形成有效电路（stated）Models below the DiT-nano scale cannot form a working circuit (stated)

DiT-nano（3 层，3 头，192 维）的空间关系准确率仅为 5%，说明存在某个模型容量阈值。当模型规模低于该阈值时，两阶段空间电路无法形成，空间关系学习完全失败。具体的阈值条件尚未被系统研究。 DiT-nano (3 layers, 3 heads, 192 dimensions) reaches only 5% spatial relation accuracy, indicating the existence of a model-capacity threshold. Below that threshold the two-stage spatial circuit fails to form and spatial relation learning breaks down entirely. The exact threshold condition has not yet been studied systematically.