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

01 动机

文生图模型在生成单一物体属性（颜色、形状）方面已取得显著进步，但在生成"A 在 B 的左边"这类组合空间关系时仍然频繁出错。已有研究从两个角度归因：一是 cross-attention 机制不足，二是文本编码器的局限性。本文提出：文本编码器的设计从根本上塑造了空间关系生成所对应的内部计算电路，从而将两种视角统一起来。

"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% 的空间关系准确率，但内部机制截然不同。

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

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

2 阶段RTE 电路的注意力头数量

5%DiT-nano 空间关系准确率（失败基线）

核心问题：空间关系是如何在 DiT 内部被编码的？

空间关系（"上方"、"左侧"等）在语义上是非交换的（"A 在 B 左边" ≠ "B 在 A 左边"），且需要将文本语义和图像空间位置绑定在一起——这对模型来说是一个根本性挑战。以往工作只关注现象而未揭示机制，本文通过在受控数据集上从头训练 DiT，系统地追踪了空间关系是如何在内部被计算的。

02 方法

研究采用"受控数据集 + 从头训练"的范式，配合提出的 Attention Synopsis 方法，系统地在扩散时间步和空间 token 两个维度上定位关键注意力头，再通过消融和因果干预验证其功能。

受控数据集构建

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 物体对）。

PixArt 风格模型架构

模型采用 PixArt-style 架构，由三部分组成：

文本编码器：分别使用 T5-XXL（预训练）、Random Token Embedding (RTE)、或 RTE 无位置编码三种配置
VAE：来自 Stable Diffusion 的预训练 VAE
Diffusion Transformer 主干：测试 DiT-B (12L, 12H, 768d)、mini (6L, 6H, 384d)、micro (6L, 3H, 192d)、nano (3L, 3H, 192d) 四种规模

Attention Synopsis 方法

本文提出 Attention Synopsis：对所有扩散时间步的注意力图取汇总统计，并在所有 prompt 条件下聚合，从而识别"对特定语义变量（如 relation）有选择性响应"的注意力头。这避免了只看单一时间步带来的偶然性，能系统筛查模型中的稀疏功能电路。

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 匹配并在对应位置生成物体。

Stage 1：空间关系头 (L2H8)

通过 QK circuit 读取 relation text token
将正弦位置编码变换到 query 空间，relation embedding 变换到 key 空间
输出平滑的空间梯度（如"above"→竖向梯度）
在采样步骤 step 0 即开始激活

Stage 2：物体生成头 (L4H3)

接收来自关系头的位置标记
将标记位置与物体 shape token 匹配
在对应空间位置生成物体
在采样步骤 4–8 后期激活

T5-DiT 中的不同机制

T5 word vector manipulation — **图4：T5-DiT 中的向量运算干预实验。** 通过从第二个物体的 T5 contextual embedding 中减去原 relation 向量并加入目标 relation 向量，可以因果性地改变生成物体的空间位置，同时保持形状和颜色不变——证明空间信息融合在 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%。

03 实验结果

在受控数据集上从头训练，使用 4 项二值指标衡量生成质量： 颜色准确率、形状准确率、绑定准确率（正确的形状-颜色对应）、空间关系准确率。评估在 96 条 prompt 上进行（8 关系 × 12 物体对）。

主要性能对比（DiT-B）

模型配置	形状 Shape	颜色 Color	绑定 Binding	空间关系 Spatial	严格空间 Strict
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（无位置编码）	—	—	0.415	0.207	—
DiT-nano（RTE）	—	—	—	0.050	—

注：RTE 无位置编码时空间准确率崩溃至 20.7%，证明位置信息对空间关系生成是必要条件。 DiT-nano（3 层）空间准确率仅 5%，说明模型容量达到阈值以下时电路无法形成。

消融实验：电路头的因果验证（RTE-DiT）

消融条件	受影响指标	消融前	消融后	下降幅度
移除 L2H8（空间关系头）	空间关系准确率	67%	33%	−34pp
移除 L4H3（物体生成头）	形状生成准确率	90%	76%	−14pp

训练动态

training dynamics — **图5：四项指标的训练动态。** 颜色准确率最先收敛，其次是形状，再次是绑定，空间关系准确率收敛最慢。 T5-DiT 模型（橙色曲线）整体收敛更快，指标之间差距更小，暗示两种文本编码器下的内部机制存在本质差异。

鲁棒性与扰动分析

relation head visualization — **图6：L2H8 空间关系头的注意力可视化。** 对于不同的 relation 词（left / right / above / below 等），该头输出不同方向的平滑空间梯度，在采样 step 0 即产生，为后续物体生成阶段提供空间位置先验。

RTE-DiT 鲁棒性

对填充词（"the"）鲁棒：添加 "the" 不影响性能
对 relation 词删除敏感：移除 relation 词后准确率显著下降
电路依赖 relation token 的 key 向量，因此对 relation 词本身敏感

T5-DiT 鲁棒性

对填充词敏感：在物体前添加 "the" 后，空间准确率从 0.808 → 0.498
原因：T5 contextual embedding 对分布偏移脆弱，单一 token 的编码受上下文影响大
向量运算干预实验证实信息存储在 object2 的 contextual embedding 中

迁移到预训练模型 PixArt-Sigma

将电路分析工具应用于公开预训练的 PixArt-Sigma 模型：在所测试的物体对中，约 27% 展现出非平凡的空间关系生成能力，同样能定位到稀疏的空间关系电路，与在受控模型上发现的机制一致。

04 局限性

注：本文没有设置专门的 Limitations 章节。以下各点为基于论文设计和文中讨论推断（inferred）的局限性，少数为论文结果中明确提及的范围限制（标注"stated"）。

受控数据集过于简化，与真实世界场景存在较大差距（inferred）

实验仅使用 3 种形状、2 种颜色、8 种空间关系的极简数据集，背景为纯灰色、物体无纹理。虽然这种受控设置便于机制发现，但其结论能否直接推广到真实场景的文生图模型（如复杂背景、多物体、连续颜色）尚不明确。

预训练大模型中电路分析成功率有限（stated）

在预训练 PixArt-Sigma 上，仅约 27% 的物体对展示出非平凡的空间生成能力，稀疏度更高、机制更难识别。在通用大模型中进行完整的电路逆向工程仍面临巨大挑战。

T5-DiT 的鲁棒性明显弱于 RTE-DiT（stated）

T5 contextual embedding 的空间信息高度依赖 prompt 的词语组合，仅添加填充词 "the" 就能使空间准确率从 80.8% 跌至 49.8%（下降 31pp）。这表明基于预训练语言模型的 T5 电路存在分布偏移脆弱性，在面向更多样化 prompt 时泛化能力受限。

仅研究了空间关系，未涵盖其他组合推理类型（inferred）

本文聚焦于"左/右/上/下"等 8 种二维空间关系，未涉及数量关系、属性绑定以外的组合推理（如大小比较、遮挡等）。是否存在类似的模块化电路机制需要进一步研究。

DiT-nano 以下规模模型无法形成有效电路（stated）

DiT-nano（3 层，3 头，192 维）的空间关系准确率仅为 5%，说明存在某个模型容量阈值。当模型规模低于该阈值时，两阶段空间电路无法形成，空间关系学习完全失败。具体的阈值条件尚未被系统研究。