跳转至

DifFRACT: 基于时序感知转码器与归因图的多模态 Diffusion Transformer 机制可解释性与回路追踪

会议: NeurIPS 2026
arXiv: 2606.15796
代码: https://github.com/Artalmaz31/DifFRACT
领域: 可解释性
关键词: 机制可解释性, Diffusion Transformer, 回路追踪, 转码器, 特征归因

一句话总结

DifFRACT 将机制可解释性中的转码器回路追踪拓展至多模态 Diffusion Transformer（FLUX.1），通过构建带去噪步调制的时间感知转码器与局部替换线性化模型，实现了精确的跨层特征级归因与精细干预。

研究背景与动机

近年来，基于 Transformer 架构的多模态扩散模型（MM-DiT，如 FLUX.1）逐渐取代传统 U-Net 成为高质量文生图的核心基石。然而，文生图扩散模型的内部计算过程依然是一个难以透视的黑盒。我们既不清楚高层语义信息如何在连续去噪时间步之间逐步固化并沉降到像素空间，也缺乏工具精确刻画双流（Image-Text）架构内部注意力交互与跨模态特征转化的深层回路机制。现有的注意力图只能粗浅地反映 token 间的交互权重，而稀疏自编码器（SAE）虽然能在特定层分离出单特征，却仅重建某单一节点的激活，无法揭示特征穿过非线性 MLP 子层时如何进行非线性组合与特征演变。

这一困境的核心矛盾在于：SAE 提取的特征本质上是神经元的密集线性组合，在面对由多层残差连接与非线性 MLP 子层堆叠而成的多步去噪网络时，无法实现跨层、端到端的因果归因追踪。大语言模型领域虽然探索了使用转码器（Transcoder）拟合 MLP 子层的输入输出映射以构建特征回路，但在多模态扩散模型中，多去噪步导致激活分布剧烈偏移，双流（文本流与图像流）在不同去噪阶段的作用存在动态重构，直接套用单向语言模型回路追踪方法会遭遇结构性失效。

本文的切入角度是：针对双流 MM-DiT 架构的 MLP 子层，显式引入时间步 FiLM 调制来训练时序感知转码器，并构建冻结分母与概率的局部替换模型（LRM），从而在保持稀疏性-保真度帕累托前沿的同时将非线性网络完全线性化。核心 idea：通过时间感知转码器将 FLUX 双流块的 MLP 映射解耦为稀疏特征，并构建局部替换线性模型导出跨去噪步的特征到特征因果归因图，实现从被动单特征提取到全局回路追踪与精准因果干预的跨越。

方法详解

整体框架

DifFRACT 针对 FLUX.1[schnell] 前 16 层双流块（双流结构包含独立处理图像与文本的权重，仅通过联合注意力交互）展开机制分析。整体流程分为三阶段：首先针对每个双流块及双流（文本流与图像流），分别训练引入去噪时间步 FiLM 调制的时间感知转码器，精确拟合 MLP 子层的输入到输出映射；其次，在选定 prompt 和时间步下构建局部替换模型（LRM），通过冻结 LayerNorm 归一化分母、冻结联合注意力概率并引入误差残差补齐，使整个前向计算转化为关于源特征激活的仿射变换；最后，从选定目标特征的反向传播梯度出发，迭代构建紧凑、稀疏的特征级因果回路图，并实施精细概念引导干预。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：Prompt 文本嵌入 + 扩散隐变量噪声"] --> B["时间感知转码器<br/>双流独立拟合 + 去噪步 FiLM 调制"]
    B --> C["局部替换模型<br/>冻结注意力概率与归一化分母"]
    C --> D["精确特征归因与守恒校验<br/>单步反向梯度推导虚拟权重"]
    D --> E["预算贪心回路扩展与两阶段剪枝<br/>位置聚合与按流影响力筛选"]
    E --> F["输出：紧凑因果回路图与精准引导干预"]

关键设计

1. 时间感知转码器：解决多时间步激活分布漂移与双流重构拟合难题

为了使转码器能适应 4 步去噪过程中激活均方根跨越十余倍、极值相差两个数量级的动态变化，论文在每一层 \(\ell \in \{0, \dots, 15\}\) 为图像流和文本流分别构建独立的转码器 \(TC_\ell^s\)（其中 \(s \in \{\mathrm{img}, \mathrm{txt}\}\)）。转码器不采用静态编码器，而是将时间步标量 \(t\) 经由正弦位置编码和两层 SiLU 激活的 MLP 映射为时间嵌入 \(e_t \in \mathbb{R}^{d_t}\)，再通过全零初始化的仿射层生成 FiLM 调制系数对 MLP 输入向量 \(x\) 进行特征缩放与平移： $\(x_{\mathrm{mod}} = x \odot (1 + \mathrm{scale}(e_t)) + \mathrm{shift}(e_t)\)$ 随后通过宽维度隐藏层（\(d_{\mathrm{feat}} = 16 \times d_{\mathrm{model}} = 49{,}152\)）提取稀疏激活 \(z(x, t) = \mathrm{ReLU}(W_{\mathrm{enc}} x_{\mathrm{mod}} + b_{\mathrm{enc}})\)，并由单位球约束的解码器投影 \(TC_\ell^s(x, t) = W_{\mathrm{dec}} z(x, t) + b_{\mathrm{dec}}\) 预测原始 MLP 输出。在训练目标中，引入了以激活方差为分母的归一化保真度损失，确保不同深度与不同去噪步具有一致的优化刻度： $\(\mathcal{L}_\ell^s = \frac{\mathbb{E}_{x, t} \| \mathrm{MLP}_\ell^s(x) - TC_\ell^s(x, t) \|_2^2}{\sum_{j=1}^{d_{\mathrm{model}}} \mathrm{Var}_{x, t} (\mathrm{MLP}_\ell^s(x)_j) + \varepsilon} + \lambda^s \mathbb{E}_{x, t} \| z(x, t) \|_1\)$ 由于图像流与文本流内部计算复杂度差异巨大，论文对两流设定不同的稀疏惩罚系数（\(\lambda^{\mathrm{img}} = 3 \times 10^{-4}, \lambda^{\mathrm{txt}} = 5 \times 10^{-5}\)），避免图像端退化为密集激活或文本端重建坍塌。

2. 局部替换模型与精确因果归因：将非线性交互仿射化并消除中间计算阻断

为了在多层网络中追踪特征因果路径，论文为选定的 Prompt 和时间步 \(t\) 构造局部替换模型（LRM）。LRM 保持原始模型残差流结构，但做三项关键线性化替换：将所有 LayerNorm 的方差分母以缓存值常数冻结（均值中心化保留为输入线性变换）；将联合注意力机制的注意力概率矩阵 \(P^\ell\) 冻结为常数矩阵，仅让值向量 \(V^\ell\) 保留为前驱状态的线性投影，并附加注意力重建残差 \(\varepsilon_{\mathrm{attn}}\) 补偿浮点截断；将 MLP 子层替换为转码器加上预先缓存的拟合残差 \(\varepsilon_{\mathrm{mlp}}\)。在此局部线性系统下，目标特征的预激活值 \(h^*\) 对所有上游源特征构成了严格的加性分解： $\(h^* - b_{\mathrm{eff}}^* = \sum_{\mathrm{src}} A_{\mathrm{src} \to f^*}\)$ 其中 \(b_{\mathrm{eff}}^*\) 为目标特征的有效偏置（包含编码器偏置、FiLM 平移、下游模块常数贡献累加）。对于任意浅层激活特征 \((\ell, s, p, i)\)，只需单次反向传播提取当前块入口的残差流梯度 \(g^{\ell+1, s}(p)\)，即可直接闭式计算该特征对目标特征的精确归因值： $\(A_{(\ell, s, p, i) \to f^*} = z^{(\ell, s, i)}(p) \cdot \bigl(g^{\ell+1, s}(p) \odot \mathrm{gate}_{\mathrm{mlp}}^{\ell, s}\bigr)^\top f_{\mathrm{dec}}^{(\ell, s, i)}\)$ 由于冻结联合注意力层天然跨双流混洗（Token 拼接），梯度能够直接在文本流与图像流之间反向贯通，自然导出跨模态归因边。

3. 位置聚合、预算回路扩展与双流解耦剪枝：控制组合爆炸并保留核心拓扑

因图像流空间包含 \(1024\) 个 patch token，单特征在多位置激活会引发节点数爆炸（原始图包含超过 \(10^4\) 个节点与 \(10^6\) 条边）。DifFRACT 采用无损保持总归因值的位置聚合规则，将同一特征在不同空间位置的激活贡献代数求和 \(\bar{A}_{(\ell, s, i) \to f^*} = \sum_p A_{(\ell, s, p, i) \to f^*}\)，并将未展开的边界特征吸收进截断误差节点。在拓扑构建上，采用有限预算贪心扩展（默认 \(N_{\max}=1000\) 次 VJP，每轮扩展前 50 个间接影响力最高的候选特征）。最后实施两阶段独立剪枝：首先依据间接影响力得分保留覆盖各模态流 80% 影响力的核心特征节点；随后在保留节点构成的子图中，按照贡献得分保留各流 98% 累积权重的关键因果边。对两流分别设置剪枝阈值，避免了全局剪枝时深度图像流节点完全吞没浅层关键文本控制节点的问题。

一个完整示例

以目标提示词为“A cat sitting on a red couch”时层 12 的图像流“猫眼”特征 \(f_{\mathrm{cat\_eyes}}^{(\mathrm{img}, 12)}\) 为例。在 \(t=0\) 步，DifFRACT 首先运行 LRM 前向传播并缓存各层激活。以 \(f_{\mathrm{cat\_eyes}}^{(\mathrm{img}, 12)}\) 的预激活值触发单次 VJP 反向传播，算法在一轮内捕获上游直接前驱边；评分模块评估未展开特征的累积间接影响，挑出文本流中对“cat”和“couch”产生高响应的特征 \(f_{\mathrm{cat}}^{(\mathrm{txt}, 7)}\) 以及浅层低级边缘特征 \(f_{\mathrm{edge}}^{(\mathrm{img}, 4)}\)。通过 50 步批次迭代扩展，系统抓取了跨文本流与图像流的 337 个关键节点。位置聚合将分散在眼部对应数十个 patch 的激活合并为单一概念节点；最终得到的回路明确指示出语义如何在 \(t=0\) 依靠文本流的 \(f_{\mathrm{cat}}^{(\mathrm{txt}, 7)}\) 通过联合注意力注入图像流，再逐层由低阶斑块驱动高阶动物特征成型。

实验关键数据

主实验

论文首先对比了时序感知转码器与稀疏自编码器（SAE）在 FLUX.1[schnell] 早期（\(\ell=6\)）、中期（\(\ell=12\)）、晚期（\(\ell=18\)）双流中的稀疏度（\(L_0\)）与归一化重建均方误差（nMSE）；随后在大规模随机测试提示词下验证了全转码器替换模型与 LRM 局部替换模型的端到端保真度。

表 1：双流 MM-DiT 模块转码器拟合与端到端潜在空间保真度评测

评测对象 / 配置 评估指标 本文转码器 / LRM 对比基线 (SAE / 原模型) 性能特征与说明
双流块 \(\ell=6, 12, 18\) 稀疏-保真度前沿 均等或优于 SAE 相同参数 SAE 转码器可完整拟合跨 MLP 输入输出非线性映射
文本流转码器 (全 16 层) 训练收敛 nMSE 0.001 - 0.011 - 文本流特征变化平缓，拟合保真度极高
图像流转码器 (全 16 层) 训练收敛 nMSE 0.040 - 0.300 - 空间特征变化剧烈，\(L_0\) 占用 \(0.05\% \sim 1.2\%\)
全转码器端到端替换 (512 测试句) Latent Cosine Sim \(\uparrow\) 0.7839 (中位数 0.7960) 1.0000 (原模型基准) 32 个转码器完全替换原网络 MLP，全局布局精准一致
全转码器端到端替换 (512 测试句) Latent MSE \(\downarrow\) 0.4786 (中位数 0.4313) 0.0000 纹理细节轻微漂移，宏观物体位置结构完全保留
LRM 替换模型 (\(t=0\) 步单步替换) Latent Cosine Sim \(\uparrow\) 0.9854 (MSE \(2.5 \times 10^{-2}\)) 1.0000 仅产生浮点舍入级别的极微弱数值漂移
LRM 替换模型 (\(t=3\) 步单步替换) Latent Cosine Sim \(\uparrow\) 1.0000 (MSE \(7.9 \times 10^{-5}\)) 1.0000 与原始模型输出在浮点精度上严格完全一致

表 2：回路图单源消融在原始模型中的因果忠实度与回路规模统计

指标维度 / 目标类型 原始图规模 (Raw) 聚合图规模 (Aggregated) 剪枝回路 (Pruned) 预测-真实因果相关性 (Spearman \(\rho\) / Pearson \(r\))
图像目标节点 (n=51) 节点: 10,640 / 边: 1,814,092 节点: 888 / 边: 274,041 节点: 337 / 边: 20,713 Spearman: 0.676 (中位数 0.693) / Pearson: 0.769
文本目标节点 (n=35) 节点: 5,744 / 边: 1,466,397 节点: 1,007 / 边: 414,681 节点: 473 / 边: 55,381 Spearman: 0.545 (中位数 0.563) / Pearson: 0.744

消融实验

表 3：反向传播预算 \(N_{\max}\) 对回路质量与规模的消融影响

VJP 预算 \(N_{\max}\) 未剪枝守恒相对误差 \(\delta\) (%) \(\downarrow\) 因果预测 Spearman \(\rho \uparrow\) 剪枝后回路节点数 说明
500 9.3% 0.658 317 搜索计算轻量，但丢失部分长程次要前驱
1000 (默认基准) 7.8% 0.680 412 质量与开销的最佳平衡平衡点
1500 6.7% 0.691 652 节点数膨胀超过一倍，但因果忠实度提升趋于饱和

表 4：去噪时间步阶段与语义干预有效性消融

干预时间窗口 文本语义超节点 (Cat / Watercolor) 抑制效果 图像流内部感知演化贡献比例 结构演化核心机制
早期去噪步 (\(t \in \{0, 1\}\)) 完全消除对应主体与画风 文本流贡献率: 89.9% \(\to\) 跨模态交互主导 语义概念确立与空间绑定阶段
晚期去噪步 (\(t \in \{2, 3\}\)) 对画面主体与风格无任何可见影响 图像流贡献率上升至 94.6% (跨模态仅 2.0%) 局限于图像流内部的高频纹理自感知细化

关键发现

  • 两阶段去噪范式：回路追踪揭示了扩散过程内在机制的本质转变。文本流特征贡献在 \(t=0\) 步高达 89.9%，到 \(t=3\) 步单调衰减至 5.4%，跨模态交互边占比从 14.9% 骤降至 2.0%。在晚期时间步修改文本流语义无法影响图像生成，干预必须施加于早期。
  • 主动抑制特征（Suppressor Features）的存在：回路揭示了扩散模型不仅使用正向激励，还普遍依赖负向阻断。例如文本流存在“狗抑制猫”特征 \(f_{\mathrm{dog-suppressor}}^{(\mathrm{txt}, 7)}\)，在输入包含狗时强烈抑制图像端猫特征的生成；若要将生成的狗可靠转变为猫，必须同时抑制狗特征并阻断该抑制特征。
  • 概念与上下文解耦控制：传统单节点 SAE 无法分离核心词与周围语义关联。利用回路能够提取上下文特征并进行联合抑制，使棒球棒直接退化为无特征木棍而非被替换为棒球，有效解决了属性泄漏与强先验偏见。

亮点与洞察

  • 多模态扩散网络全仿射线性化：利用双流 MM-DiT 结构特性，通过冻结注意力概率与归一化分母构建 LRM，使深层非线性 Transformer 网络转化为严格保真的局部仿射系统，推导出了跨层特征间闭式虚拟权重。
  • 从孤立特征观测迈向因果回路溯源：改变了过去仅能靠 SAE 寻找离散特征激活的局限，首次在扩散模型中重现了特征级因果图的自动导出，使模型幻觉、颜色泄漏（Color Leakage）以及属性绑定失败的机制能够在单条边上追根溯源。
  • 回路引导的协同干预范式：发现单纯压制异常特征往往被上游强先验冲刷失效，通过回路顺藤摸瓜同时压制其关联的先验源节点，能够在反事实属性生成中取得远超常规单特征 Steering 的成功率。

局限与展望

  • 单流模块尚未覆盖：当前分析仅针对 FLUX.1[schnell] 前 16 层双流块，后 38 层单流块权重完全共享且结构更密集，尚未纳入转码器回路分解中。
  • 注意力 \(Q\)-\(K\) 路径被视为静态外生变量：模型当前仅在线性化 \(O\)-\(V\) 路径上传播梯度与归因，注意力权重分配本身的非线性因果机理仍未被特征化解耦。
  • 多步时序回路尚未直接跨步贯通：目前单张图仅对应单步去噪的层间因果，多时间步之间的跨步累积影响依然依靠经验阶段采样进行，未形成全时间轴的统一动态时序有向无环图。

相关工作与启发

  • vs LLM 转码器与回路追踪（Dunefsky et al. / Ameisen et al.）：LLM 仅处理单一离散自回归序列，而本文必须克服扩散模型的多步去噪激活漂移与双流联合注意力带来的跨模态交互，提出了时间感知 FiLM 调制和双流平衡独立剪枝策略。
  • vs 扩散模型稀疏自编码器（Cywiński et al. / Huang et al.）：传统扩散 SAE 仅重构同一层的单点激活，特征间的跨层影响被非线性 MLP 完全阻断；本文转码器直接建模 MLP 映射，建立起端到端的特征因果边。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将转码器因果回路追踪引入现代双流多模态 Diffusion Transformer，填补了生成模型机制可解释性的一大空白。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 86 组目标因果验证、端到端潜空间保真度评测、两阶段时序消融以及颜色泄漏/概念抑制等丰富实证。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导严谨自洽，公式定义清晰，实验逻辑环环相扣。
  • 价值: ⭐⭐⭐⭐⭐ 为解决文生图中的属性绑定失败、偏见固化与模型可控编辑提供了坚实的底层分析工具。

相关论文