表征对齐在训练DiT模型中的重要性

Diffusion Models专栏文章汇总：入门与实战

前言：训练过DiT模型的读者们肯定有所体会，相比于UNet模型训练难度大了很多，模型不仅很难收敛，而且非常容易训崩，其中一个很重要的原因是没有进行表征对齐！这篇博客详细介绍表征对齐在训练DiT模型中的重要性。

 

目录

训练过程的问题

训练DiT模型为什么慢？

仅仅依赖"渲染"损失是不够的

多层DiT之间表征能力的区别

表征对齐的具体方法

相关资料

---

训练过程的问题

扩散模型虽然强大，但训练起来却是个噩梦。训练一个DiT/SiT模型需要700万次迭代，要跑好几周才能达到论文里报告的FID分数。

这简直是在考验研究者的耐心和计算资源！

训练DiT模型为什么慢？

• 扩散transformer 确实能学到不错的表征，而且生成效果越好，表征质量也越高。

• 但是，这些表征还是比不上像DINOv2、JEPA或MAE这些自监督学习方法得到的表征。

• 当我们测量扩散模型特征和DINOv2特征的对齐程度时，发现扩散模型在训练过程中确实在慢慢进步，但这个过程实在是太慢了。

仅仅依赖"渲染"损失是不够的

如果你只关注重建漂亮的像素，就无法过滤掉输入中的无关细节——而这恰恰是学习强大表征的关键。

看来，即使你的目标是生成好看的图像，你也需要先学习一个强大的表征，然后才能处理那些让图像看起来漂亮的细节。

多层DiT之间表征能力的区别

仅通过对齐前几个DiT块可以实现足够的表示对齐。反过来，这允许DiT的后续层专注于基于对齐表示捕获高频细节，进一步提高生成性能。

如下图，学习能力有一个峰值，在20层之后显著下降。

表征对齐的具体方法

REPA 通过最大化预训练表示 y∗ 和隐藏状态 ht 之间的补丁相似性来实现对齐：

最终的loss如下：

相关资料

- Paper: https://arxiv.org/abs/2410.06940  

- Project page: https://sihyun.me/REPA/  

- Code: https://github.com/sihyun-yu/REPA