DiT模型缺陷探讨：TREAD策略下的新发现与挑战

「警惕！DiT或存在架构缺陷」

近日，一篇在X平台上的帖子引发了广泛讨论，有博主指出DiT（Diffusion Transformer）存在设计上的不足，并附上了一份论文截图作为佐证。

图 1. 我们引入了TREAD，这是一种能够显著提升基于token的扩散模型骨干网络训练效率的训练策略。当应用于标准的DiT骨干网络时，我们在无引导FID指标上实现了14/37倍的训练速度提升，同时生成质量也显著提升。

图中横轴代表训练时间（以A100 GPU的小时数为单位，采用对数刻度，范围从100小时到10000小时），纵轴代表FID分数（数值越低表示生成图像质量越高）。

博主认为，此图的核心并非TREAD的速度优势，而是DiT的FID过早稳定，这可能暗示DiT存在「隐性架构缺陷」，导致它无法从数据中持续学习。

博主提到的论文今年1月发表（3月更新v2），介绍了一种名为TREAD的新方法。该工作通过创新的「令牌路由」（token routing）机制，在不改变模型架构的情况下，极大提升了训练效率和生成图像的质量，从而在速度和性能上都显著超越了DiT模型。

具体来讲，TREAD在训练过程中采用「部分令牌集」vs「完整令牌集」，通过预定义路由保存信息并重新引入到更深层，跳过部分计算以降低成本，且此做法仅用于训练阶段，推理时仍采用标准设置。这与MaskDiT等方法类似，但更高效。

论文标题：TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

论文地址：https://arxiv.org/abs/2501.04765

代码：https://github.com/CompVis/tread

在后续回复中，博主逐步展开了对DiT的批判，并解释TREAD如何揭示这些问题。

博主指出，该论文揭示了DiT模型的设计缺陷。具体来说，研究发现，在训练过程中，如果将模型中的部分计算单元替换为「恒等函数」（即让这些单元不执行任何计算，仅「直通」数据，相当于被临时禁用），模型的最终评估分数反而会提高。

接着，博主指出了DiT的两个「可疑」设计：

• 整个架构都使用后层归一化（Post-LayerNorm）

博主认为DiT使用了一种已知不太稳定的技术（后层归一化）来处理一个数值范围变化极其剧烈的任务（扩散过程）。

• adaLN-zero

博主认为，尽管该模型自称是「Transformer」架构，但在处理最关键的「指导信息」（即条件数据）时，并未使用强大的Transformer，而是用了一个非常简单的MLP网络（多层感知机）。

更具体地，adaLN-zero通过完全覆盖注意力单元的输入并注入任意偏置来覆盖输出，这限制了模型的表达能力，相当于「讨厌注意力操作」（hate the attention operation），从而削弱了DiT的整体潜力。

博主还提到了与早期论文相关的LayerNorm研究，指出LayerNorm的偏置和增益参数可能对梯度调整影响更大，而非真正改善模型性能。他认为adaLN-zero正是利用这一点，名为「梯度调节」，实则像是在「给小模型偷偷注入过拟合的偏置」。

论文标题：Understanding and Improving Layer Normalization

论文地址：https://arxiv.org/abs/1911.07...

看了这篇帖子后，DiT的作者——纽约大学计算机科学助理教授谢赛宁有些坐不住了。

在2022年，谢赛宁发表了DiT的论文，这是扩散模式首次和Transformer相结合。

论文标题：Scalable Diffusion Models with Transformers

论文链接： 免费服务器云服务器服务器教程