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

扩散Transformer（DiT）极大地推动了文本到图像生成，但模型仍难以按文本提示生成对象之间正确的空间关系。在本研究中，我们采用机制可解释性方法，研究DiT如何在对象之间生成正确的空间关系。我们从零开始训练不同规模、使用不同文本编码器的DiT，让它们学习生成包含两个对象的图像，对象的属性与空间关系由文本提示指定。我们发现，尽管所有模型都能以近乎完美的准确率学会该任务，其底层机制却因文本编码器的选择而有巨大差异。当使用随机文本嵌入时，我们发现空间关系信息经由一个两阶段电路传递给图像token，该电路涉及两个交叉注意力头，分别读取文本提示中的空间关系和单对象属性。当使用预训练文本编码器（T5）时，我们发现DiT使用另一条电路，利用文本token中的信息融合，从单个文本token中一并读取空间关系与单对象信息。我们进一步表明，尽管两种设置在域内性能相似，但它们对域外扰动的鲁棒性不同，这可能暗示在真实场景中生成正确空间关系的难度。