重新思考 Diffusion Transformers 中的跨层信息路由

TL;DR DiT 中残差流存在幅度膨胀、梯度衰减等问题，本文提出 DAR：一种可学习、自适应时间步的路由机制，替换传统残差连接，显著提升训练效率和生成质量，并与 REPA 等方法正交叠加。

问题

问题背景

扩散变换器（Diffusion Transformers，DiTs）已成为视觉生成的事实标准架构，研究者不断优化 token 化、注意力与条件注入等模块，但跨层信息流的传输机制却几乎未被重新审视。

现有方法局限

DiTs 直接继承原始 Transformer 的 残差连接，它在每一层简单地将输入与子层输出相加。系统性实证分析发现，这种经典设计在扩散模型中引发三个具体症状：

• 前向幅度单调膨胀：残差累积导致深层表示范数持续增大，造成数值不稳定。
• 反向梯度急剧衰减：梯度在反向传播中消失，阻碍浅层参数的有效训练。
• 显著的块间冗余：相邻层输出高度相似，浪费模型容量。 此外，标准残差连接与去噪时间步无关，无法根据不同噪声水平动态调整信息路由，限制了模型对扩散过程复杂动态的适应能力。

为什么这个问题难且重要

扩散模型需要处理从纯噪声到清晰图像的完整轨迹，不同时间步对层间信息流动的需求截然不同。固定残差路径无法灵活应对，而引入可学习路由机制不仅要保证训练稳定，还需兼顾推理效率。业界对 训练加速 和 生成质量 的追求日益迫切，跨层信息路由作为一个独立于表示对齐（如 REPA）的新设计轴，能与现有方法正交叠加，实现叠加增益。该研究揭示了扩散架构中一个长期被忽视的优化空间，有望成为下一代 DiTs 的基础组件。

行业类比

就像大语言模型中通过 Mixture-of-Experts 或动态路由精细控制 token 流动以提升效率，扩散模型也可通过 timestep 自适应的层间路由，在不显著增加参数的情况下加速训练并改善生成细节。

核心洞察

• 传统残差连接在扩散变换器中引发信息流病态，而现有工作几乎完全忽视了这一架构瓶颈。本文首次沿深度和去噪时间步联合诊断了 DiT 的跨层信息流动，揭示了正向幅度单调膨胀、反向梯度急剧衰减以及块间显著冗余三种症状，证明标准残差加法并不适合生成任务的动态需求。这一发现将信息路由重新定义为扩散模型的关键设计空间，与广泛探索的标记化、注意力等轴正交。
• 通过学习时间步自适应的非增量路由机制，DAR 大幅提升了训练效率，并表明扩散模型中跨层路由是独立于表示对齐的优化方向。与 REPA 等表示对齐损失叠加时能获得2倍早期加速，且无需额外损失函数，这提示研究社区：结构化的信息聚合策略能带来比单纯添加辅助目标更本质的改善。该模块可作为即插即用组件，从预训练到微调、蒸馏均有效，为工程实践提供了低成本的增效途径。

方法

输入与动机

DiT 的标准残差连接是将子层输出与输入直接相加：$h_{l+1} = h_l + F(h_l)$。这种做法导致三个问题：前向过程中特征幅度单调膨胀、反向传播时梯度急剧衰减、以及相邻层表示高度冗余。为打破这种刚性增量路径，DAR 重新设计跨层信息聚合逻辑，输入包括所有先前子层的输出历史与当前去噪时间步 $t$。

核心模块：扩散自适应路由

DAR 以非增量、可学习、时间步自适应的方式从历史中聚合特征，替代原始残差加法。

• Query 参数化：为每个待更新的位置生成一个与内容无关的可学习查询向量 $q_l$，使得信息检索不依赖当前层的输入偏差，鼓励模型主动搜索历史中有用的表示。
• 时间步注入：将去噪步长 $t$ 通过傅里叶嵌入后注入到查询或键中，使聚合权重随噪声水平变化。例如，在高噪声阶段可以聚焦深层结构信息，而在低噪声阶段偏好浅层细节。
• 分块聚合（Chunked Aggregation）：为控制计算开销，将历史序列按层分组为若干 chunk，在 chunk 内平均池化后再参与注意力计算，显著降低序列长度，同时保留宏观的跨层交互能力。

最终，第 $l$ 层的 DAR 输出为所有历史子层输出的加权和： $$ h_{l+1} = \text{DAR}_l({h_0, h_1, ..., h_l}, t) $$ 该模块可直接替换任何 Transformer 子层的残差连接，无需更改其他结构。

与同类方法的差异

不同于固定的 U-Net 跳跃连接（仅在对称层之间传递）或 REPA 对表示的全局对齐，DAR 关注的是 逐层、动态、跨深度的信息路由，其权重由模型自主决定并依赖时间步，与已有表征对齐目标正交，可叠加使用以获得复合收益。

实验

实验设计

在 ImageNet 256×256 上以 SiT-XL/2 为基线，将标准残差连接替换为 DAR（Diffusion-Adaptive Routing）。评估生成指标 FID 与训练迭代数，并验证与 REPA 的叠加效果。此外，在大规模 T2I 模型 微调及 Distribution Matching Distillation 场景中测试 DAR 的泛化性。

关键发现

• 质量飞跃：DAR 将 SiT-XL/2 的 FID 从 9.67 降至 7.56（↓2.11）。
• 训练效率倍增：仅需 8.75× 更少的训练迭代即可匹配基线的收敛质量。
• 时间步感知是核心：消融实验中移除时间步条件（timestep injection）导致性能大幅回退，证明扩散过程特有的噪声水平信息对路由至关重要。
• 与 REPA 正交互补：在 REPA 之上叠加 DAR 可获得约 2× 的早期训练加速，表明跨层信息路由是独立于表示对齐的全新设计维度。
• 分块聚合的灵活性：合理选择 chunk size 可在计算开销与模型容量间取得平衡。

与基线的深度对比

传统 残差流（residual stream）在 DiT 中机械地累加各层输出，导致前向幅度单调膨胀、反向梯度急速衰减以及显著的块间冗余。DAR 提出一种 非增量式、时间步自适应的聚合，利用可学习的查询向量与时间步嵌入，动态选择并加权历史子层输出。这种信息路由策略使得模型能更高效地传递梯度、减少冗余计算，从根本上改善了扩散 Transformer 的信息流。相比于仅优化注意力或 tokenization 的现有工作，DAR 解锁了“层间如何通信”这一尚未被充分探索的结构性轴，为扩散模型的高效缩放提供了新方向。

行业影响

落地场景

DAR (Diffusion-Adaptive Routing) 对基于 DiT 的生成模型具有即插即用特性，可无缝部署在文生图 (T2I)、文生视频、3D 资产生成等管线中。典型产品化场景包括：

• 内容平台：短视频特效生成、广告素材自动创建，通过加速模型迭代快速响应热点创意。
• 电商与虚拟拍摄：商品图合成、虚拟试穿等需要高频细节保留的业务，DAR 在蒸馏阶段能保持高频纹理，提升投产比。
• 影视与设计工具：高分辨率图像/视频编辑的底层模型微调，可缩短制图时间。

商业价值

DAR 直接作用于训练效率和生成质量两个维度，带来显著降本与体验提升。

• 训练降本：在 SiT-XL/2 上，达到相同 FID 需要的训练迭代数减少 8.75 倍；叠加 REPA 后早期训练再快 2 倍，大幅节省 GPU 算力成本。
• 推理体验：FID 从 9.67 降至 7.56，生成图像保真度提升，同时支持 Distribution Matching Distillation 时保留高频细节，有利于实时交互场景。
• 迭代加速：对于需要频繁更新模型的内容平台，DAR 让新模型更快上线，缩短从研发到落地的周期。

与现有工作流的接口

DAR 是传统残差连接的直接替代品，不改变 DiT 主干结构，兼容主流增强技术。

• 集成方式：只需替换 DiT 块内的残差聚合逻辑，可复用在任一 DiT-based 模型中（如 SD3、PixArt、Sora 等变体）。
• 协同增效：已证实与 REPA 等表示对齐方法正交，堆叠后训练加速效果叠加；未来也可与其他 token 压缩、高效注意力模块组合。
• 部署灵活：预训练阶段引入可加速收敛；也可仅在微调或蒸馏阶段插入 DAR，无需重新训练完整模型。

具体落地用例

1. 电商视觉生成平台：为全球商家提供 AI 商品图合成服务。使用基于 DAR 的 DiT 模型，在同等算力下可更快产出高清晰度的多角度商品展示图，且训练成本下降使小规模商家也能负担个性化模型微调。
2. 短视频创作工具：集成 DAR 的文生视频模型，支持创作者输入文本直接生成短片。DAR 带来的训练加速让平台能够每周更新风格化模型，快速响应用户流行趋势，同时生成质量提升减少人工后期修整。

局限

• - **泛化验证场景有限**：主要实验基于 ImageNet 256×256 的类别条件生成，大规模文本到图像模型仅在后训练阶段测试，未展示从零开始训练的完整流程。对于更高分辨率、无条件和更复杂的开放域生成任务，DAR 能否保持一致的增益尚不明确。此外，论文未涉及视频生成等时序生成场景，而残差流在这些任务中的行为可能完全不同。
• - **计算效率与工程落地成本**：DAR 引入了可学习的路由参数和分块聚合（chunked aggregation），增加了额外的前向计算和存储开销。论文未报告详细的推理延迟对比或内存峰值变化，对于实际部署而言，精度提升（FID 降低 2.11）是否足以抵消运行时开销需要进一步分析。尤其在实时生成或边缘设备场景中，DAR 的实用性可能受限。
• - **对模型深度和架构的敏感性**：DAR 的诊断基于特定深度（如 SiT-XL）的 DiT 残差流，其发现的单调幅值膨胀、梯度衰减等现象在更深的网络（如 40+ 层）或非标准 Transformer 变体中可能不成立。分块大小等超参数需要针对不同架构重新搜索，方法缺乏即插即用的彻底自动化，限制了其作为一种通用工具的适用范围。