【他山之石】CVPR 2024 | SD-DiT：判别式自监督学习范式，让DiT训练更智能！！！

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注！

Diffusion Transformer（DiT），作为文生视频模型SORA的主要网络结构，近几个月引起了人们的高度关注。DiT 给扩散模型带来了 Transformer 架构，并且使得模型的训练达到相对较大的规模，但是它的训练收敛速度相对较慢。之前的工作引入视觉自监督表征学习中的 Mask 策略，这一策略可以显著提升 DiT 的训练效率，并且促进了模型在图像内部情境学习方面的进一步发展。尽管如此， Mask 策略对于扩散生成模型来说存在一定的局限性，比如 mask 的重建任务与图像生成任务的关联性并不明确，限制了 DiT 训练过程的进一步优化。

本文将介绍一篇 CVPR 2024 的会议论文——《SD-DiT：使用判别式自监督学习范式来帮助DiT加快训练收敛》。本文的研究团队来自HiDream.ai智象未来，香港中文大学（深圳）和香港理工大学， 作者通过 teacher-student 网络构建了 SD-DiT 框架，而 teacher-student 网络的输入正样本对建立在沿同一个 PF-ODE 的扩散噪声上。除此之外，作者解耦 SD-DiT 为 DiT encoder 和 decoder ，来分别完成自监督判别和图像生成的学习目标，而不是在 DiT encoder 和 decoder 上都应用 mask 重建。具体来说，SD-DiT 利用自监督判别损失函数来完成特征空间中的图像间对齐，而后在 DiT decoder 中进行图像生成的学习任务。

在 ImageNet 256x256 数据集上进行了大量实验表明， SD-DiT 在训练效率和生成性能之间取得了很好的平衡。比起 DiT ， SD-DiT 不仅训练收敛速度提高了 5 倍，在生成性能上也表现出色。

论文标题：

SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer

论文链接：

https://arxiv.org/abs/2403.17004

视频链接：

https://www.youtube.com/watch?v=Ah2w4fewyR4

一、研究背景

DiT[1] 通过引入 Transformer 架构，为扩散模型提供了一种新的训练方式，使模型的训练达到相对较大的规模，但是它的训练收敛速度相对较慢。之前的工作引入视觉自监督表征学习中的 Mask 策略来帮助 DiT 的训练，如 MDT[2] 通过 Mask 策略的 Intra-view self-reconstruction 学习到上下文邻近 patch 的关系，可以显著提升 DiT 的训练收敛速度。而 MaskDiT[3] 则通过在 encoder 只使用 50% 的输入 patch 大幅提升了训练效率。

然而，这样的 Mask 重构方法对于 Diffusion 生成模型存在一定局限性，比如来自于自监督学习的 Mask 重构 loss 可能会使模型在表征学习和可学习的 mask token 上浪费过多的能力。为了克服这些问题，作者提出了使用视觉自监督表征学习的另一种方法——Inter-view 判别式方法，来帮助 DiT 的训练。

二、本文方法

SIMCLR:

https://github.com/google-research/simclr

判别式自监督表征学习的关键点在于从一张图片构造两个不同的视图，并将它们配对组成正样本对，之后在特征空间中拉近两个视图的距离，从而使模型学习到视图不变性。如上图所示，传统的表征学习依赖于数据增强来构造正样本对，这样的正样本对并没有连接数据分布和噪声分布，并不适合于 DiT 。而 SD-DiT 则是依据 Open AI 近期提出的一致性模型[4] 的概念，选择在同一个 PF-ODE 上对正样本添加不同程度的噪声：对于学生视图采用 PF-ODE 范围内的噪声，对于教师视图采用 PF-ODE 中最小的噪声。

在整体模型上，SD-DiT 采用了解耦的编码-解码结构。学生编码器和教师编码器采用 DINO 自蒸馏判别式的方法来进行判别式损失的学习。在这个过程中，教师编码器的参数是固定的，它不进行参数更新，而是代表了学生编码器的滑动指数平均。

对于对应的判别式损失函数，SD-DiT 使用了 DINO[5] 和 iBOT[6] 的损失函数，在没有 mask 策略的 token 和 CLS token 上，模型通过交叉熵损失来进行训练。

在 SD-DiT 模型中，解码器则负责执行扩散去噪损失。值得注意的是，SD-DiT 并没有使用 mask token 以及 mask 重建损失，为的是避免浪费 DiT 的模型的容量在表征学习上，但是 SD-DiT 还是保留了 Mask 策略，通过 mask 划分了编码器的可见区域，而那些不可见的区域则是在解码器的输入端与编码器的输出相结合，再进行去噪处理。这样使得 SD-DiT 保留了 Mask 策略学习局部上下文关系的优势的同时，也避免了扩散生成损失和 mask 重构损失之间可能存在的矛盾。

三、实验结果

本文的实验在 ImageNet-256x256 分辨率进行，测试采用 FID-50k 为主要指标。

在不同尺度的 DiT 网络下，SD-DiT 展现出了显著的性能提升，尤其是在评估图像生成质量的常用指标 FID 上。对于 DiT-XL 这样相对较大的模型尺度，SD-DiT 可以达到大约 5 倍的收敛速度提升，并且带来了更好的 FID 结果。

在更多的训练步数情况下，比起 DiT， SD-DiT 表现出更好的训练收敛性。

在和其他State-of-the-art模型对比上，SD-DiT在训练效率和生成性能之间取得了很好的平衡。

四、总结

本文提出了一种创新的模型—— SD-DiT，它成功地将自监督判别训练范式融入到 Diffusion Transformer 的训练中。

作者通过 teacher-student 网络构建了 SD-DiT 框架，而其中的输入的判别式正样本对建立在沿同一个 PF-ODE 的扩散噪声上。本文提及的表征学习与扩散生成模型结合的尝试，希望对后续的 DiT 训练以及表征模型的学习范式具有一定的启发作用。

参考文献

[1] Peebles，William，and Saining Xie. "Scalable diffusion models with transformers." ICCV 2023.[2] Gao，Shanghua，et al. "Masked diffusion transformer is a strong image synthesizer.“ ICCV 2023.[3] Zheng， Hongkai，et al. "Fast training of diffusion models with masked transformers.“ TMLR 2024.[4] Song，Yang，Dhariwal Prafulla，Chen Mark，Sutskever Ilya. "Consistency models." ICML 2023.[5]Caron，Mathilde，et al. "Emerging properties in self-supervised vision transformers." ICCV. 2021.[6] Zhou，Jinghao，et al. "ibot: Image bert pre-training with online tokenizer." ICLR 2022