把残差流从被动求和改成主动检索:镜像路由让 DiT-XL/2 的 FID 从 18.85 降到 15.07,并让 REPA-XL/2 从 5.9 降到 4.34
核心概要
该工作系统分析了扩散 Transformer 内部表征,发现模型偏好复用早期层特征并倾向对称层配对,据此提出把残差连接从被动求和改为主动检索的结构化连接设计:将 DiT 沿深度均分为编码器与解码器两阶段,把编码器侧表征作为可微残差源,每个解码器子层通过轻量 softmax 路由选择其镜像对应的编码器源进行自适应融合;在 ImageNet 类条件生成上,该方法以不到 0.1% 的额外参数让 DiT-S/2、DiT-B/2、DiT-XL/2 的 FID 分别从 69.81、44.21、18.85 降至 62.48、36.33、15.07,并把强基线 REPA-XL/2 从 5.9 降至 4.34(无引导),加分类器无关引导后达到 1.39 FI
深度剖析
论文对 DiT 内部路由做了系统分析,观察到两个偏离标准 Transformer 行为的现象:整个网络深度上都持续给早期层表征较高权重(编码器侧依赖),以及在获得跨层路由自由后模型自发偏好镜像/对称层配对。 此前 U-Net 的对称跳连多被视为手工设计的归纳偏置,而 U-ViT、U-DiTs 等把长跳连搬回 Transformer 时用的是静态投影融合;这里把对称性呈现为模型在可学习路由下自己寻找的结构倾向。 来自对全层路由权重的可视化(图 4)与 CKA 表征相似度分析(图 5(a)–(b)),属于内部表征层面的观察性证据,而非下游性能证据。
提出结构化残差路由:把 DiT 沿深度均分为编码器与解码器阶段,编码器侧表征(含 patch embedding 输出)作为不脱离计算图的可微残差源,每个解码器子层用 softmax 打分从候选源中选取镜像对应的单一源进行自适应融合,插入在自注意力与 MLP 两个子层之前。 与 LLM 中的 Attention Residuals 共享 softmax 打分机制,但把源池限制在编码器侧而非全部前序层输出,且每个解码器层只取一个阶段匹配的镜像源,而非稠密全连接路由;与 U-ViT/U-DiTs 的镜像配对相似,但融合权重随 token 表征、空间位置、子层与扩散时间步自适应变化,而非静态投影。 方法在 DiT-S/2 上做消融:相同镜像配对下,自适应融合(Fused-Mirror w/o Enc.,FID 64.83)优于静态跳连(U-ViT-style skip,FID 67.90,且多约 1.476M 参数);加入编码器内部路由后进一步到 62.48。
在 ImageNet 类条件生成上,该方法在相同训练协议下一致提升不同规模的 DiT,且相对 FID 降幅随模型增大而增大,额外参数始终低于 0.1%。 400K 迭代下 DiT-S/2 从 69.81 到 62.48、DiT-B/2 从 44.21 到 36.33、DiT-XL/2 从 18.85 到 15.07;在匹配 FID 的比较中最多可减少 1.73 倍训练迭代,且这一优势随训练推进而扩大。 表 2 的受控对比,三个规模使用同一训练协议;DiT-XL/2 在 800K 迭代下 11.40 对基线 13.71,1300K 基线为 11.76。
该方法可作为微调模块叠加在已训练好的强模型上:REPA-XL/2 从 5.9 FID 经 0.28M 步微调降至 4.94,再经 0.35M 步降至 4.34;而相同设置的持续微调基线几乎无增益(5.87、6.01)。 说明后期瓶颈并非单纯模型容量不足,而在连接与梯度结构上仍有空间;同时该方法在效率上以单一阶段匹配源替代稠密路由,把解码器侧源栈激活从 1.51 GB 降到 289 MB。 REPA 微调实验从公开的 4M 迭代检查点出发,仅加入路由模块并使用去噪目标、不含 REPA 投影损失;持续微调基线使用完全相同设置但不含路由。加分类器无关引导后达到 1.39 FID。
启示与展望
该结果面向以 ImageNet 类条件生成、潜空间扩散协议训练的扩散 Transformer,适用于从零训练的 DiT-S/2、DiT-B/2、DiT-XL/2,以及从公开检查点出发的 REPA-XL/2 微调场景。方法以轻量归一化与投影层实现,可作为一个模块叠加在已有模型上,因此对已有训练流程的改动成本低。论文把镜像路由定位为一种结构先验:它把模型在稠密路由下本已偏好的对称配对直接给定,从而降低路由歧义、简化优化。对读者而言,这提供了一条可复用的思路——在保持 DiT 单尺度架构不变的前提下,通过重新组织跨深度信息与梯度通路来释放已有模型的表达潜力,而不必增加模型容量。
论文自身指出,更广泛的验证——更大规模的文生图与文生视频模型——仍属未来工作,这些模型涉及更复杂的条件机制、更大规模数据集与不同训练配方,可能影响最优路由结构;在引入更强多模态条件或时间依赖时,镜像路由是否仍是最佳选择也不清楚。此外,镜像配对带来的梯度对齐部分由镜像连接本身诱导,论文将其定位为对所引入优化行为的刻画而非独立因果结论。读者还可留意:路由权重随扩散时间步变化的具体行为、镜像路由在更深或非均匀深度划分下的表现,以及该方法与其它表征增强策略组合时的相互作用,这些在正文中未展开。
