一句话总结

为应对扩散 Transformer 中显著的反向梯度衰减与明显的块级冗余问题，作者提出了扩散自适应路由（Diffusion-Adaptive Routing, DAR）。该方案作为一种即插即用的残差替代模块，能够对各子层输出进行可学习、时间步自适应且非增量的聚合。在 ImageNet 256×256 数据集上，该方法使 SiT-XL/2 的 FID 指标提升了 2.11（7.56 对比 9.67），同时仅用 8.75 倍的更少训练迭代次数便达到了与基线相当的生成质量。

核心贡献

• 本研究系统性地实证分析了扩散 Transformer 中跨层信息流在深度与去噪时间步维度上的变化规律，诊断出标准残差加法存在的三种失效模式：正向幅度单调膨胀、反向梯度急剧衰减以及显著的块级冗余。
• 本文引入了扩散自适应路由（DAR），这是一种即插即用的残差替代机制。它针对历史子层输出执行可学习、时间步自适应且非增量的聚合操作，同时保持与现代 Transformer 增强技术（如 REPA）的兼容性。
• 在 ImageNet 256x256 数据集上的评估表明，DAR 使 SiT-XL/2 基线的 FID 指标提升了 2.11 分，并以 8.75 倍的更少训练迭代次数实现了等效的收敛质量。

引言

扩散 Transformer 通过以基于 token 的去噪器取代卷积 U-Net，已迅速成为视觉生成的标准架构。然而，有效捕捉去噪过程中随时间变化的动态特性仍是一项关键挑战。大多数现有架构从语言模型中继承了固定且与时间无关的残差流。随着网络深度增加，这会导致隐藏状态膨胀、反向梯度衰减，并使相邻层产生冗余。这种静态路由机制无法适应去噪时间线。在去噪过程中，最优特征提取需要从高噪声阶段的粗略结构逐渐转移至低噪声阶段的精细细节。作者通过首次系统性地分析跨网络深度与去噪时间步的跨层信息流，针对性地解决了上述局限。本文提出了 DAR，这是一种即插即用的残差替代方案，用可学习的时间步自适应路由取代了均匀聚合。该方法显著加快了训练收敛速度，提升了生成质量，并能与现代表征对齐目标无缝集成。

方法

作者提出了扩散自适应路由（DAR），这是一种新型残差替换机制，旨在解决扩散 Transformer（DiT）中标准残差累积的局限性。其核心思想是用可学习且时间步自适应的聚合机制，取代对子层输出进行固定增量求和的做法。该机制会根据历史子层输出的贡献动态分配权重。该方法将跨层信息路由视为一项基础设计维度，突破了从原始 Transformer 继承的固定残差模式。

DiT 中的标准残差路由由递推公式 $h_{l+1} = h_l + f_l(h_l; t)$hl+1​=hl​+fl​(hl​;t) 定义，其中 $h_l$hl​ 表示第 $l$l 层的隐藏状态，$f_l$fl​ 为变换函数（注意力机制或多层感知机），$t$t 为去噪时间步。展开该递推式可知，输出 $h_l$hl​ 仅是初始输入与所有前面子层输出的简单求和：$h_l = h_0 + \sum_{i=0}^{l-1} f_i(h_i; t)$hl​=h0​+∑i=0l−1​fi​(hi​;t)。该方法存在固有局限，因为它对所有历史表征均应用单位权重，导致模型无法根据当前网络深度或去噪阶段选择性提取或抑制特定特征。这种固定路由模式也与类 U-Net 的跳跃连接不同。后者通过手动配对深层与浅层进行特征融合，但缺乏纯 Transformer 架构的扩展性与同构性。

DAR 从根本上重新审视了这一累积过程。它使用基于 softmax 的加权聚合取代无权重求和，作用于子层输出的历史序列。令 $v_i = f_i(h_i; t)$vi​=fi​(hi​;t) 表示第 $i$i 个子层的输出，且 $v_0 = h_0$v0​=h0​。第 $l$l 层的聚合隐藏状态 $h_l$hl​ 计算如下：

$$h_l = \sum_{i=0}^{l-1} \alpha_{i \to l}(t) v_i \quad \text{with} \quad \alpha_{i \to l}(t) = \frac{\exp\bigl(q_l(t)^\top k_i / \sqrt{d}\bigr)}{\sum_{j=0}^{l-1} \exp\bigl(q_l(t)^\top k_j / \sqrt{d}\bigr)} \,,$$hl​=i=0∑l−1​αi→l​(t)vi​withαi→l​(t)=∑j=0l−1​exp(ql​(t)⊤kj​/d​)exp(ql​(t)⊤ki​/d​)​,

其中 $k_i = \text{RMSNorm}(v_i)$ki​=RMSNorm(vi​) 是源 $v_i$vi​ 对应的键向量，softmax 函数作用于集合 $\mathcal{S}_l = \{v_0, v_1, \ldots, v_{l-1}\}$Sl​={v0​,v1​,…,vl−1​}。该机制使模型能够学习哪些历史表征对当前子层最为相关，从而实现非增量且自适应的信息检索。

用于确定路由权重的查询向量 $q_l(t)$ql​(t) 采用三种参数化方式以融入时间步感知能力。纯静态变体使用一个可学习且与时间无关的向量 $w_l$wl​。动态变体根据最近的子层输出计算查询向量：$q_l(t) = W_q^{(l)} v_{l-1}$ql​(t)=Wq(l)​vl−1​，这隐式注入了 $v_{l-1}$vl−1​ 中包含的时间步信息。显式变体在静态查询基础上增加时间步嵌入 $e(t)$e(t)：$q_l(t) = w_l + e(t)$ql​(t)=wl​+e(t)，提供直接且显式的信号。实验发现，具备去噪时间步感知能力的动态与显式变体均显著优于静态变体，凸显了时间步感知对高效路由的重要性。

为管理保留全部 $L$L 个子层输出所带来的计算开销，DAR 实现了分块聚合策略。网络被划分为 $N$N 个大小为 $S = L/N$S=L/N 的块。每个块 $n$n 由其最终子层的输出 $c_n = v_{nS}$cn​=vnS​ 进行摘要表示。对于位于块 $n$n 中的子层 $l$l，聚合的源集合变为：

$$\mathcal{S}_l = \{c_0, c_1, \ldots, c_{n-1}\} \cup \{v_{(n-1)S+1}, \cdots, v_{l-1}\} \,,$$Sl​={c0​,c1​,…,cn−1​}∪{v(n−1)S+1​,⋯,vl−1​},

该集合由之前块的摘要与当前块的原始输出组成。这将源数据的内存占用从 $O(Ld)$O(Ld) 降低至 $O((S + N)d)$O((S+N)d)。理论推导得出的最优块大小为 $S^* = \sqrt{L \cdot \frac{1 - \alpha}{1 + \alpha}}$S∗=L⋅1+α1−α​​，其中 $\alpha$α 是控制路由熵项的超参数。实证分析证实，对于 SiT-XL/2 模型，块大小 $S=4$S=4 为最优配置。

预测层的最终聚合过程得到进一步优化。系统采用专用最终聚合器，除访问所有先前块的摘要外，还能获取最后一个块的全部原始子层输出。这一设计基于一个直观认知：最近的输出包含最具任务特异性的信号。保留其原始形式可使最终层恢复在单一摘要中会丢失的细粒度信息。在将 DAR 与 REPA 结合时，最终聚合复用最后一个块中 MLP 聚合器的参数。这一基于实证的设计进一步提升了性能。

实验

在 ImageNet 上使用标准扩散训练协议进行的评估验证了：DiT 中的传统残差路由存在特征幅度膨胀、梯度不平衡与表征冗余问题，且会自然表现出对时间步依赖的源偏好。通过引入可学习且时间步自适应的跨层聚合，该方法加快了收敛速度，并在生成质量上超越了仅依靠参数扩展或固定跳跃连接所能达到的水平。进一步的消融实验证实，该路由机制能够补充表征对齐技术，在中间块大小时表现最优，并在大规模模型蒸馏过程中有效保留细粒度视觉细节。最终结果表明，自适应的跨层信息流能够在不破坏 Transformer 同构结构的前提下，解决架构固有的僵化问题。

作者对比了结合 REPA 的 SiT 与 DAR 模型的性能，重点关注收敛速度与最终质量。结果显示，在所有训练迭代中，DAR 结合 REPA 均取得了比 SiT 结合 REPA 更低的 FID 分数，表明其性能更优且收敛更快。即使在考虑模型规模与替代架构设计的情况下，这些增益依然保持稳定。DAR 结合 REPA 在所有训练迭代中的 FID 表现均优于 SiT 结合 REPA。相较于 SiT 结合 REPA，DAR 结合 REPA 能以更少的训练迭代次数实现更低的 FID 分数。DAR 结合 REPA 的性能提升并非源于模型规模的增加或架构选择的改变。

作者分析了 DAR 不同变体在收敛速度与最终质量方面的表现，对比了注入与未注入时间步信息的模型。结果表明，动态变体在所有训练阶段均取得最低的指标值，展现出更优的性能；而未注入时间步信息的静态变体则持续表现较差。引入时间步注入后，静态与动态配置的结果均得到改善。动态变体在所有训练阶段均取得最佳性能。未注入时间步的静态变体表现劣于动态变体及注入时间步的静态变体。时间步注入对静态与动态变体的性能均有提升作用。

实验评估了不同块大小对模型性能的影响，通过 FID 与 IS 指标衡量生成质量。结果显示，块大小为 4 时取得最佳 FID 分数，而最高 IS 也出现在块大小为 4 时，表明不同配置下存在指标权衡。相较于块大小 1 和 8，块大小 4 的 FID 性能最优。块大小 4 实现了最高的 IS，暗示该设置下表征质量达到最优。性能随块大小变化呈现 U 型趋势，中间值表现优于极值。

作者将提出的 DAR 方法与标准残差及类 U-Net 路由基线在图像生成任务上进行对比，采用 FID 与 Inception Score 等指标。结果表明，DAR 以更少的训练迭代次数与参数实现了更优的性能，且与 REPA 结合时效果进一步增强。该方法在不同配置下均展现出强劲的收敛能力与质量提升，尤其在带引导的动态变体中表现突出。DAR 在减少训练迭代与参数量的同时，性能优于标准残差与类 U-Net 路由。DAR 的动态变体在质量与收敛速度上均超越基线，特别是在无分类器引导设置下。将 DAR 与 REPA 结合带来了显著改进，表明两者具有正交且叠加的增益效果。

该评估将提出的 DAR 方法与 SiT 及标准路由基线进行对比，以验证收敛速度、生成质量与参数效率。组件消融实验进一步证实，动态路由与时间步注入持续带来性能提升，而中间块大小则最优平衡了质量指标。综合来看，实验表明 DAR 以更少的参数实现了卓越的生成质量与更快的收敛速度，在与 REPA 集成时更是取得了叠加的性能提升。