Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

会议: CVPR 2026
arXiv: 2601.06338
代码: 无
领域: Image Generation / Mechanistic Interpretability
关键词: 扩散Transformer, 空间关系生成, 机械可解释性, 注意力电路, 文本编码器

一句话总结

通过机械可解释性方法揭示了扩散Transformer（DiT）生成空间关系的内部电路机制：随机嵌入模型使用两阶段模块化电路（关系头+物体生成头），T5编码器模型则将关系信息融合到物体token中通过单token解码，两种机制的鲁棒性差异显著。

研究背景与动机

文本到图像（T2I）扩散模型在生成高质量图像方面取得了巨大进展，但在组合多个物体的空间关系时经常失败（例如"红色方块在蓝色圆形的左上方"）。虽然单物体属性生成的准确率在快速提升，但空间关系的生成改进却相对缓慢。

现有工作提出了多种补救措施（布局条件、交叉注意力引导、课程学习等），但很少有工作从模型内部机制的角度理解为什么空间关系生成会失败。本文的动机是：

1. 神经网络如何编码和使用物体间的非交换关系（如"A在B上方"≠"B在A上方"）？
2. 扩散模型的迭代采样本质使注意力图分析复杂化，如何系统化地总结和定位关键头？
3. 空间关系生成的瓶颈在交叉注意力还是文本编码？需要一个整体性的研究视角

方法详解

整体框架

作者构建了一个最小化的文本-图像数据集，训练多种规模的DiT模型从零开始学习生成两个物体（含组合形状和颜色属性）按指定空间关系排列的图像。使用3种形状×2种颜色×8种空间关系的组合。模型架构采用PixArt风格DiT，对比三种文本编码器：T5-XXL、随机token嵌入（RTE）、无位置编码的RTE。

关键设计

1. 注意力概要（Attention Synopsis）：从上千万张注意力图里定位关键头

扩散模型的迭代采样让注意力分析爆炸——层×头×时间步×条件/无条件×token 数，动辄上千万张图，没法逐张看。作者提出一种可扩展的聚合范式：先把 token 按类别分组（图像 token 按物体分割分组、文本 token 按语义属性分组），在类别粒度上聚合注意力得到可解释的类别间交互，再对时间步取平均，最终把整个注意力张量压成一张 [层数, 头数] 的概要图。这样就能从超过 1000 万张注意力图中快速定位出真正负责空间关系的关键头。

2. RTE-DiT 的两阶段模块化电路：先标位置、再放物体

在随机 token 嵌入（RTE）模型里，作者发现空间关系是分两步生成的。空间关系头（L2H8）在采样最初一步（step 0）就激活，通过 QK 电路让图像 token 的正弦位置编码与关系词嵌入交互——"above" 产生垂直梯度、"left" 产生水平梯度，相当于先在画布上刷出一层位置标签标记物体该放哪（这一机制惊人地类似胚胎发育中分子梯度引导细胞分化）。物体生成头（L4H3）则在采样后期（step 4-8）才激活，读取这些位置标签，把带匹配标签的区域连到对应形状 token，从而在正确位置长出正确物体；物体头只管传递形状身份，与空间位置和关系本身无关。两个头各司其职，这种模块化正是 RTE 电路鲁棒的来源。

3. T5-DiT 的融合单 token 解码电路：关系信息藏在物体词里

换成 T5 编码器，机制完全不同。T5 的自注意力把整句上下文融进每个 token，于是 DiT 反而从非关系词 token（尤其第二个形状词 shape2）里解码空间关系。方差分解能看到这一点：T5 嵌入里 shape2 解释 37.5% 方差、关系只贡献 12.1%，但经 DiT MLP 投影后关系信息被放大到 21.3%。作者进一步用向量算术做因果验证——在 T5 嵌入上减去原关系向量、加上新关系向量，生成物体的空间位置随之改变，坐实了关系信息确实编码在融合的物体 token 里。

4. 权重空间头筛选：不生成样本也能找到关系头

逐样本跑生成来找头太贵，作者给了个纯权重空间的快速筛选法：直接计算图像位置特征与文本关系特征的 QK 交互，看产生的空间图是否与参考关系梯度对齐，对齐的就是候选关系头，全程不需要采任何样本。

损失函数 / 训练策略

• 标准扩散训练，使用DPM-Solver++ (14步) 采样，CFG=4.5
• 训练多种模型尺寸：DiT-B (12层12头768维)、mini (6层6头384维)、micro (6层3头192维)、nano (3层3头192维)
• EMA权重用于评估
• 四维评估指标：颜色、形状、唯一绑定、空间关系

实验关键数据

主实验

模型	文本编码器	颜色↑	形状↑	唯一绑定↑	空间关系↑
DiT-B	T5	99%	97%	93%	89%
DiT-B	RTE	99%	96%	90%	86%
DiT-B	RTE w/o pos	99%	96%	41%	15%
DiT-nano	RTE	-	-	-	5%

消融实验

配置	关键指标	说明
消融L2H8的关系注意力	空间关系准确率 67%→33%	关系头对空间布局至关重要
消融L4H3的物体注意力	形状准确率 90%→76%	物体头对形状生成有因果作用
T5-DiT消融关系词	几乎无影响	T5将关系信息融合到其他token
T5-DiT消融shape2	关系准确率降低50%	关系信息主要编码在shape2
T5-DiT插入filler词"the"	关系准确率大幅下降	T5电路对微小词汇变化敏感
RTE-DiT插入filler词	保持稳定	模块化电路对扰动更鲁棒

关键发现

1. 电路机制取决于文本编码器：RTE使用模块化两阶段电路（关系→位置标签→物体），T5使用融合单token解码电路
2. 位置编码是必要的：无位置编码的RTE空间关系准确率仅15%，因为无法区分"A在B上方"和"B在A上方"
3. 学习动态呈阶段性：颜色→形状→属性绑定→空间关系，关系学习最慢
4. 鲁棒性差异显著：RTE-DiT对关系词消融敏感但对filler词鲁棒；T5-DiT相反
5. 可迁移到预训练模型：在PixArt-Sigma上也能识别出稀疏的空间电路

亮点与洞察

• 机械可解释性方法论：Attention Synopsis 和权重空间头筛选为理解大规模DiT提供了可扩展的分析工具
• 生物学类比：空间关系头的梯度机制与胚胎发育中的分子梯度有惊人相似
• 统一视角：首次将"交叉注意力是瓶颈"和"文本编码器是瓶颈"两种观点统一，展示它们在不同配置下各自成立
• 设计启示：模块化（RTE）vs 融合（T5）的权衡——模块化更鲁棒、更可解释，融合更紧凑但更脆弱
• 实用意义：改善空间关系生成可能需要优先改进嵌入模型而非DiT本身

局限性 / 可改进方向

1. 实验在极简数据集（3形状×2颜色×8关系）上进行，真实世界场景的复杂度远高于此
2. 仅研究了两个物体的空间关系，多物体（3+）组合的电路机制待探索
3. RTE和T5的对比可能受训练数据量和训练充分度的影响
4. 未探索如何利用发现的电路机制来改进空间关系生成（如通过注意力干预）
5. 对预训练模型的分析（PixArt-Sigma）较浅，空间关系性能本身就很弱

相关工作与启发

• 与Transformer电路分析（Elhage et al., 2021）的方法论一致，但首次应用于扩散模型
• Attend-and-Excite等方法通过操控交叉注意力改善组合性，本文的发现为这些方法提供了机制解释
• 文本编码器的选择对模型行为的影响被低估——CLIP、T5、随机嵌入导致根本不同的内部计算
• 为未来设计更鲁棒的T2I架构提供了指导：模块化电路可能优于融合电路

评分

• 新颖性: ⭐⭐⭐⭐⭐ — 首次揭示DiT空间关系生成的具体电路机制
• 实验充分度: ⭐⭐⭐⭐ — 因果操控和消融设计严谨，但受限于简化设置
• 写作质量: ⭐⭐⭐⭐⭐ — 叙事清晰，图示精美，逻辑严密
• 价值: ⭐⭐⭐⭐ — 为理解和改进T2I模型的组合生成提供重要基础

相关论文

• [CVPR 2026] Just-in-Time: Training-Free Spatial Acceleration for Diffusion Transformers
• [CVPR 2026] PixelDiT: Pixel Diffusion Transformers for Image Generation
• [CVPR 2026] MakeAnything: Harnessing Diffusion Transformers for Multi-Domain Procedural Sequence Generation
• [CVPR 2026] SPREAD: Spatial-Physical REasoning via geometry Aware Diffusion
• [CVPR 2026] Region-Adaptive Sampling for Diffusion Transformers