重思扩散Transformer中的跨层信息路由 Rethinking Cross-Layer Information Routing in Diffusion Transformers
用可学习、时步自适应的softmax聚合替代固定残差，缓解DiT的PreNorm稀释。
前置知识
扩散Transformer (Diffusion Transformer, DiT)
用Transformer替代U-Net做扩散/流匹配去噪的架构。每个block含self-attention+MLP,通过adaLN-Zero把扩散时步t与类别条件通过LayerNorm的仿射参数注入。本文实验对象SiT-XL/2即28-block的DiT类模型,约675M参数。
DAR是面向DiT的残差替代方案,只有理解DiT由同质sublayer堆叠、且每个sublayer都通过adaLN接受t调制,才能理解为什么DAR要把时步t显式注入跨层路由权重。
PreNorm 与 PreNorm 稀释
PreNorm指在子层前做RMSNorm的Transformer设计: $h_{l+1}=h_l+f_l(\text{RMSNorm}(h_l))$。PreNorm Dilution指随深度增加残差流被新加项稀释,表现为隐藏态RMS膨胀、梯度逐层衰减。
本文核心诊断就是PreNorm Dilution在DiT上的具体表现:forward RMS从block 1的约15.5增长到block 28的约1576(100×膨胀),backward梯度在前5个block后锐降一个数量级。
自适应层归一化Zero (adaLN-Zero)
DiT中用MLP把t与类别c映射为每层LayerNorm的缩放$\gamma$、偏移$\beta$与残差门控$\alpha$,且$\alpha$初始化为0使初始输出等于输入。
DAR的dynamic query $q_l(t)=W_q^{(l)} v_{l-1}$ 之所以获得时步信息,正是因为$v_{l-1}$已经经过adaLN-Zero调制,Fig. 5的linear probe验证了$h_l$中$t$的$R^2\approx 1$。
SiT 与 REPA
SiT (Scalable Interpolant Transformers)统一扩散与rectified flow训练目标,用velocity prediction loss训DiT,是DAR的基线。REPA把DiT隐藏态与DINOv2特征对齐以加速收敛。
DAR以SiT为基线并把REPA作为正交加速器,理解二者才能体会DAR是在'残差路径'这一未被改造的轴上做文章,与REPA这种'目标函数'轴上的改造天然正交。
Softmax跨层注意力与Attention Residuals
AttnRes用softmax注意力替代固定残差: $h_l=\sum_i\alpha_{i\to l}v_i$,$\alpha_{i\to l}=\text{softmax}(q_l^\top k_i/\sqrt{d})$,$k_i=\text{RMSNorm}(v_i)$。源自LLM研究。
DAR正是AttnRes在DiT上的重新实例化并加入时步注入与chunked aggregation。理解AttnRes才能识别DAR的'DiT特异性':把t这一DiT独有维度显式注入路由权重。
研究动机
Diffusion Transformers (DiTs) 在 tokenization、attention、条件化、训练目标、latent autoencoder 等几乎所有设计轴上都被反复重访,但控制信息如何跨层累积的残差流却直接照搬自原始 Transformer。本文对 DiT 的跨层信息流做系统性诊断,识别出 PreNorm 残差在 DiT 上与深度和去噪时步联合作用下出现的三个具体症状:(1)前向幅值膨胀——以 SiT-XL/2 为例,隐藏态 RMS 从 block 1 的约 15.5 单调增长到 block 28 的约 1576,接近 100× 膨胀;(2)反向梯度衰减——反向传播梯度 RMS 在前 5 个 block 之后锐降超过一个数量级;(3)块间冗余——相邻 block 的逐 token 余弦相似度在深层稳定高于 0.9。这三个症状在去噪时步方向上也呈系统性变化,但标准残差流对此完全无能为力,因为它对所有历史子层输出施加单位系数、做时间无关的固定聚合。
本文的目标是本文的具体目标是提出 Diffusion-Adaptive Routing (DAR),一种可直接替换 DiT 标准残差的即插即用模块,把跨层信息路由从"继承来的惯例"提升为一个显式的设计轴,使路由权重同时满足三个性质:可学习(timestep-adaptive,能随噪声水平切换路由模式)、时步自适应(time-varying,响应去噪过程的时间动态)、非增量式(non-incremental,不再依赖逐层累加)。作者希望 DAR 能同时带来收敛加速与最终 FID 改善,且保持 Transformer stack 的各向同性,使现有增强方法(REPA 等)可以无侵入地叠加上去。
与已有工作不同的是，已有对 DiT 深度结构的改造主要走两条路:一是 U-Net 风格的长跳连(U-ViT、U-DiT),把浅深层配对后用固定权重融合,这种手工配对破坏了 Transformer 的均质性,且权重是时步不变的;二是把残差流整体替换为多流 recurrence(Hyper-Connections、mHC),但仍属于固定拓扑的循环。另一类 LLM 侧的 Attention Residuals 把残差换成 softmax 跨层注意力,但它是为 LLM 设计的、query 不感知时步。本文识别出 DiT 与 LLM 的一个本质差异——去噪时步 $t$ 是 DiT 的第一类控制信号,adaLN-Zero 已把它注入到每一层的隐藏态中——并据此在 AttnRes 框架内为 DiT 定制 timestep-aware query,再通过 chunked aggregation 控制激活显存,使之成为首个面向 DiT 的可学习、时步自适应、非增量式残差替代方案。
核心方法
DAR 的整体思路非常直觉:既然标准残差 $h_l = h_0 + \sum_{i<l} v_i$ 把所有历史子层输出 $v_i=f_i(h_i;t)$ 等权累加,不妨把等权求和换成按权重加权求和,权重 $\alpha_{i\to l}(t)$ 由跨层 softmax 注意力动态生成。技术路线分四步:(1)每个 sublayer 输出 $v_i$ 经 RMSNorm 得到 key $k_i=\text{RMSNorm}(v_i)$;(2)为每个 sublayer 配 query $q_l(t)$(static / explicit t-injection / dynamic 三种);(3)用 $\alpha_{i\to l}(t)=\text{softmax}(q_l^\top k_i/\sqrt{d})$ 做跨源注意力;(4)加权求和的 $h_l$ 喂给 sublayer 变换 $f_l(h_l;t)$。为控制显存,引入 chunked aggregation,激活从 $O(Ld)$ 降到 $O((S+N)d)$。
DAR 与已有方法的本质区别在于把跨层信息路由从一个固定惯例提升为第一个针对 DiT 显式设计的学习维度,在三个意义上都是新的:(1)相对于 LLM 侧的 AttnRes,DAR 的 query 显式依赖去噪时步 $t$,通过 static($w_l$)/explicit($w_l+e(t)$)/dynamic($W_q^{(l)} v_{l-1}$)三种注入方式,使路由权重在不同噪声水平下切换强调的源——Fig. 3 显示 SiT baseline 反事实重要性图本身就随 $t$ 系统性变化,作者据此论证 timestep-aware 路由是 DiT 残差路径的内在需求;(2)相对于 U-ViT/U-DiT 等 U-Net 风格跳连,DAR 不再手工指定浅深配对、不再破坏 isotropic Transformer stack;(3)相对于 Hyper-Connections/mHC 的多流 recurrence,DAR 是非增量式聚合,绕开 running stream 的幅值膨胀。
方法步骤详情
训练与推理分五步。第一步:对 $v_i=f_i(h_i;t)$ 做 RMSNorm 得 key $k_i$。第二步 query 构造:static 用 $q_l=w_l$;explicit 用 $q_l(t)=w_l+e(t)$($e(t)$ 复用 DiT t-embedder,zero-init 使训练初退化为 static);dynamic 用 $q_l(t)=W_q^{(l)} v_{l-1}$,经 adaLN-Zero 隐式继承 $t$。第三步 softmax 路由:$\alpha_{i\to l}=\exp(q_l^\top k_i/\sqrt{d})/\sum_j\exp(q_l^\top k_j/\sqrt{d})$,再 $h_l=\sum_{i<l}\alpha_{i\to l} v_i$。第四步 chunked aggregation:把 $L$ 个 sublayer 划成长度 $S$ 的 chunk,每 chunk 末尾 $c_n=v_{nS}$ 作前序汇总。第五步:$v_l=f_l(h_l;t)$ 用标准 attention+MLP+adaLN-Zero。
技术新颖性
技术新颖性可归纳为四点。第一,首次在 DiT 上系统诊断 PreNorm dilution 的三个症状(forward RMS 膨胀 100×、backward 梯度降一个数量级、cos 相似度 > 0.9),且联合沿"深度 × 去噪时步"二维展开。第二,把 AttnRes 为 DiT 重新实例化并加入 timestep 注入,用 Fig. 5 linear probe 证明 dynamic query 输入 $v_{l-1}$ 在前 5 个 block 后对 $t$ 解码 $R^2$ > 0.95、深栈接近 1.0。第三,显式给出 chunked aggregation 的率失真分解 $L(S)=\log(L/S)+S+\alpha\log S$ 并严格证明 U 形曲线有唯一极小值 $S^\star=\sqrt{(1-\alpha)/(1+\alpha)}\cdot\sqrt{L}$,把"为什么是 $S=4$"从经验观察升级为可计算最优值。第四,与 REPA 正交:DAR+REPA 在 100K iter 已达 REPA 单跑 200K 的水平,即 2× 早期加速。
实验结果
实验在 ImageNet $256\times 256$ 上以 SiT-XL/2 为基线展开,核心发现分四层。第一层最终质量:Static c4 ODE 下 FID 7.56(基线 9.67,降 2.11)、SDE 下 6.92(基线 8.61);Dynamic c4 ODE+CFG 下最佳 FID 2.05(基线 2.15)。第二层收敛加速:Static c4 50K iter 即达 9.35 FID,匹配基线 1.75M iter 的 8.61 仅需 200K iter,即 8.75× 训练加速。第三层诊断-性能关联:Static c4 同时把 forward RMS 爆炸式增长压平、backward 梯度衰减抬高、相邻 cos 相似度降低,三个症状缓解与 FID 改善同步发生。第四层正交性:DAR+REPA@100K 已等于 REPA@200K 即 2× 早期加速;chunk size 扫描 $S=4$ 得 FID 8.39 优于 $S=1$ 的 10.41 与 $S=8$ 的 11.14,证实 U 形曲线。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256x256 无条件生成 (ODE, w/o CFG) | FID-50K (越低越好) | Static c4: 7.56 (600K iter, 675M) | SiT-XL/2: 9.67 (1.75M iter, 675M) | FID 下降 2.11(-21.8%),同时训练迭代数减少 65.7% |
| ImageNet 256x256 无条件生成 (SDE, w/o CFG) | FID-50K (越低越好) | Static c4: 6.92 (600K iter, 675M) | SiT-XL/2: 8.61 (1.75M iter, 675M) | FID 下降 1.69(-19.6%),且 DAR 训练仅用基线 1/3 迭代即超基线 |
| ImageNet 256x256 无条件生成 (ODE, w/ CFG w=1.5) | FID-50K (越低越好) | Dynamic c4: 2.05 (500K iter, 751M) | SiT-XL/2: 2.15 (1.75M iter, 675M) | FID 下降 0.10,同时参数仅 +11%、迭代数 -71.4% |
| ImageNet 256x256 与 U-Net 风格跳连对比 (SDE, w/ CFG) | FID-50K (越低越好) | Static c4: 2.23 (675M) | U-DiT-L: 3.00 (810M), U-ViT-H/2: 2.29 (585M) | Static c4 比 U-DiT-L FID 低 0.77 且参数仅 0.83×,比 U-ViT-H/2 在 ODE 下低 0.24 |
| ImageNet 256x256 与 REPA 兼容性 (100K iter) | FID-50K (越低越好) | DAR + REPA: 7.09 | SiT + REPA: 9.89 | FID 下降 2.80(-28.3%),DAR+REPA@100K 即可匹配 REPA@200K 的 6.89,即 2× 早期加速 |
| Chunk size 扫描 (300K iter, no CFG) | FID-50K (越低越好) | S=4: 8.39 (IS 121.7) | S=1: 10.41, S=8: 11.14 | U 形曲线在 S=4 取极小值,与率失真命题1预测的 S*≈4 一致 |
局限与改进
作者在正文没有专门 Limitations 小节,附录 F 简短讨论,且承认 DAR 主要在 class-conditional ImageNet 上验证,T2I 上的扩展以 DMD 视觉效果示意为主未给完整指标。从技术细节可观察到以下几点限制:(1)Chunk size $S=4$ 的最优值依赖于率失真式中的 $\alpha\in(0,1)$,而 $\alpha$ 没有显式可计算形式,跨任务(视频 DiT、3D DiT)迁移时需重新 sweep;(2)实验主要集中在 ImageNet $256\times 256$,T2I 蒸馏实验只展示定性视觉对比,没有提供 FID/IS 等定量指标;(3)与 REPA 之外的表示对齐方法(RAE、DMD2)的兼容性未验证,正交性结论只覆盖了 REPA 一家;(4)Tab. 1 显示 Dynamic c4 参数多 11.4%,但收益上 Static c4 的 SDE FID 反而更优(6.92 vs 7.39),作者未分析原因;(5)推理时 DAR 的额外 softmax 计算与显存开销相比基线没有显式报告。
独立分析的弱点
独立审视后有以下弱点值得改进。第一,chunked aggregation 把前序 chunk 硬选为末尾 $c_n=v_{nS}$,可能丢失 chunk 内早期 sublayer 的重要特征,改进方向是允许 summary 由小型 attention pool 学习得到。第二,timestep-aware query 的三种变体在 Tab. 2 早期差距很大(100K iter: 22.36 vs 17.39 vs 13.95),但 400K iter 接近(8.10 vs 7.97),作者未分析为什么 explicit 在长训练下追平 dynamic,也未给出选择判据。第三,论文没有给出推理时延与显存的实测数字,forward pass 中每次 sublayer 多一次 softmax 的代价没量化。第四,DAR+REPA 的 2× 早期加速在 100K-200K 区间成立,300K 时差距从 2.80 收窄到 0.61,作者未讨论长期边际收益曲线。第五,Fig. 3 只在 block 3 与 block 13 上可视化源权重,缺乏对所有 28 个 block 的系统分析。
未来方向
作者在附录 F 提到把 DAR 扩展到更长的 horizon 与更多模态。基于此可延伸出五条值得研究的方向:(1)Video DiT 与长 horizon 扩展:把 chunked aggregation 推广到"空间 chunk + 时间 chunk"的二维分块;(2)与更多表示对齐目标兼容:验证 DAR 是否与 RAE、REPA-E、Score Distillation 等目标兼容;(3)可解释路由:对所有 28 个 block 的路由权重画热图并按 $t$ 切片,可视化"高噪声强调哪些 block、低噪声强调哪些 block",把 DAR 升级为 DiT 内部信息流的可解释工具;(4)非线性 query/key 与门控:目前 $q_l^\top k_i$ 是点积,可加入 FFN 非线性、或对 key 加 sigmoid 门控;(5)动态 chunk size:对不同深度使用不同 $S$(浅层 $S$ 小、深层 $S$ 大),可能进一步压低率失真曲线。
复现评估
复现评估分四方面。开源情况:论文未声明代码与 checkpoint 开源,仅给项目主页 https://dar-route.github.io,但 SiT/REPA 官方代码完全开源,DAR 模块本身仅约 50 行 PyTorch(用 F.softmax 在 source set 做 attention),自行实现门槛不高。数据:ImageNet-1K 需自行下载,4096 张诊断子集可直接从 ImageNet val 采样,REPA 的 DINOv2 模型从 huggingface 下载。算力:SiT-XL/2 baseline 训 1.75M iter 需 64-128 张 A100 跑 7-14 天,DAR Static c4 仅需 600K iter、Dynamic c4 需 500K iter,折算单次完整复现约需 30-60 张 A100 跑 7-10 天。复现难度:技术栈主要是 Transformer+CFG+评估 pipeline(FID-50K、sFID、IS),都是社区成熟工具,DAR 改动局部可独立单元测试。中小团队可在 4-8 张 A100 完成小规模重现。
论文图表
三栏并列展示 SiT-XL/2 baseline(蓝)与 DAR Static b4(红)在 28 个 Transformer block 上的诊断量:左栏 forward RMS 从约 15.5 单调增长到约 1576,DAR 显著压平;中栏 backward gradient RMS 在前 5 个 block 后从约 5e-7 降到约 1e-7 以下,DAR 抬高早期信号;右栏相邻 block 的 cos 相似度全程 > 0.9,DAR 把深层冗余明显降低。
这是论文诊断侧的奠基图,把'PreNorm dilution 在 DiT 上具体长什么样'第一次直观展示出来,后续 DAR 的所有收益都需要回到这张图的三个症状来解释。
对 SiT baseline(左两列)与 DAR(右两列),分别在浅层 block 3 与深层 block 13 上画 4 张子图,横轴是去噪时步 t∈{0.01, 0.5, 0.99},纵轴是 source index。Baseline 的反事实重要性图已显示时步依赖(高噪声与低噪声的偏好源不同),DAR 的 softmax 权重则把这种偏好明确化:权重集中在少数历史源,且该选择随 t 平滑迁移。
这是论证'timestep-aware 路由是 DiT 残差路径的内在需求而非外部归纳偏置'的最核心证据,直接把 baseline 的反事实模式与 DAR 的实际路由模式并列。