改进扩散模型中跨层信息传递,提升生成质量和训练效率。
Rethinking Cross-Layer Information Routing in Diffusion Transformers

- 提出可学习、时间自适应的非增量信息聚合机制
- 在ImageNet上将FID降低2.11,训练迭代减少8.75倍
- 兼容主流优化方法,适合大规模文本到图像模型微调
扩散Transformer(DiTs)已成为现代视觉生成的骨干架构,其几乎每个设计维度——包括标记化、注意力、条件输入、目标函数和潜在自编码器——都已被深入研究。然而,控制信息跨层累积的残差流却直接继承自原始Transformer。本文对DiTs中跨层信息流进行系统性实证分析,同时考虑深度与去噪时间步,发现传统残差相加存在三个具体问题:前向幅度单调膨胀、反向梯度急剧衰减、块级显著冗余。受此诊断启发,我们提出扩散自适应路由( extsc{DAR}),一种即插即用的残差替换方案,实现对子层输出历史的可学习、时间自适应且非增量的聚合。此外, extsc{DAR}兼容多种现代Transformer增强技术,如REPA。在ImageNet $256\times256$ 上, extsc{DAR}使SiT-XL/2的FID从9.67提升至7.56,训练迭代次数减少8.75倍即可达到基线收敛质量。叠加于REPA之上,早期阶段训练速度提升2倍,表明跨层信息路由是扩散建模中尚未充分探索的设计轴线,与现有表征对齐目标正交。该方法不仅适用于预训练,也可用于大规模文本到图像模型的微调阶段,并在分布匹配蒸馏中有效保留高频细节。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have become a de facto backbone of modern visual generation, and nearly every major axis of their design -- tokenization, attention, conditioning, objectives, and latent autoencoders -- has been extensively revisited. The residual stream that governs how information accumulates across layers, however, has been directly inherited from the original Transformer. In this paper, we present a systematic empirical analysis of cross-layer information flow in DiTs, jointly along depth and denoising timestep, and identify three concrete symptoms of traditional residual addition, namely monotonic forward magnitude inflation, sharp backward gradient decay, and pronounced block-wise redundancy. Motivated by this diagnosis, we propose Diffusion-Adaptive Routing (\textsc{DAR}), a drop-in residual replacement that performs \emph{learnable, timestep-adaptive, and non-incremental} aggregation over the history of sublayer outputs. Moreover, the proposed \textsc{DAR} is compatible with many modern Transformer enhancement methods, such as REPA. On ImageNet $256\times256$, \textsc{DAR} improves SiT-XL/2 by $2.11$ FID ($7.56$ vs.\ $9.67$) and matches the baseline's converged quality with $8.75\times$ fewer training iterations. Stacked on top of REPA, it yields a $2\times$ training acceleration in the early stage, suggesting cross-layer information routing as an underexplored design axis in diffusion modeling, one that operates orthogonally to existing representation-alignment objectives. Beyond pretraining, \textsc{DAR} can also be applied during the fine-tuning stage of large-scale T2I models and preserves high-frequency details during Distribution Matching Distillation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。