[论文解读] Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

Sihyun Yu, Sangkyung Kwak+5 more|arXiv (Cornell University)|Oct 9, 2024

Reinforcement Learning in Robotics被引用 4

一句话总结

本文提出Representation Alignment (REPA)，一种简单的正则化技术，通过将扩散Transformer中的噪声潜在表征与来自预训练自监督编码器（如DINOv2）的高质量外部视觉特征对齐，实现高效训练。通过将这些外部表征蒸馏到扩散模型的早期层，REPA使训练速度提升超过17.5倍，并在无分类器指导设置下实现SOTA的FID=1.42，显著提升了训练效率与生成质量。

ABSTRACT

Recent studies have shown that the denoising process in (generative) diffusion models can induce meaningful (discriminative) representations inside the model, though the quality of these representations still lags behind those learned through recent self-supervised learning methods. We argue that one main bottleneck in training large-scale diffusion models for generation lies in effectively learning these representations. Moreover, training can be made easier by incorporating high-quality external visual representations, rather than relying solely on the diffusion models to learn them independently. We study this by introducing a straightforward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders. The results are striking: our simple strategy yields significant improvements in both training efficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiTs and SiTs. For instance, our method can speed up SiT training by over 17.5$ imes$, matching the performance (without classifier-free guidance) of a SiT-XL model trained for 7M steps in less than 400K steps. In terms of final generation quality, our approach achieves state-of-the-art results of FID=1.42 using classifier-free guidance with the guidance interval.

研究动机与目标

• 为解决大规模扩散Transformer因内部表征学习能力差而导致的训练缓慢与低效问题。
• 探究引入高质量外部视觉表征是否能显著提升扩散模型的训练效率与生成质量。
• 通过一种简单而有效的正则化策略，克服扩散模型输入的噪声性与视觉编码器预训练数据中清晰图像之间的输入不匹配问题。
• 证明与强自监督编码器（如DINOv2）进行表征对齐可实现更快收敛与更优的生成性能。
• 提供一种实用、即插即用的方法，在不改变网络架构的前提下提升扩散Transformer的训练效率。

提出的方法

• REPA引入一种正则化损失，将扩散Transformer中噪声潜在隐藏状态的投影与应用于干净图像的预训练视觉编码器（如DINOv2）的对应特征对齐。
• 该方法使用线性投影头将扩散模型的隐藏状态映射到与外部视觉编码器相同的特征空间，从而实现特征匹配。
• 对齐损失在训练过程中计算，并仅应用于扩散Transformer的早期层，确保计算开销最小化。
• 该方法兼容基于扩散和基于流的Transformer架构，包括DiT与SiT结构，并适用于使用VAE嵌入输入的潜在扩散框架。
• 该方法无需微调外部编码器，也无需修改扩散模型的架构，具有良好的可部署性。
• 训练过程采用无分类器指导，并在标准文本到图像基准上通过FID与用户研究进行评估。

Figure 1: Representation alignment makes diffusion transformer training significantly easier. Our framework, REPA, explicitly aligns the diffusion model representation with powerful pretrained visual representation through a simple regularization. Notably, model training becomes significantly more e

实验结果

研究问题

• RQ1 将扩散Transformer表征与高质量外部视觉特征对齐，是否能显著加速训练收敛？
• RQ2 与标准扩散训练相比，使用DINOv2等预训练自监督编码器进行表征对齐是否能提升生成图像的质量？
• RQ3 扩散模型特征与外部视觉特征之间的对齐在训练过程中如何演化？能否通过显式正则化加速这一过程？
• RQ4 所提出的REPA方法在不同扩散Transformer架构（如DiT与SiT）中是否均有效？
• RQ5 REPA是否能在远少于基线模型的训练步数下实现SOTA的生成质量？

主要发现

• REPA将SiT-XL/2达到基线700万步性能所需训练步数减少至40万以下，实现17.5倍的收敛速度提升。
• 在无分类器指导、指导尺度为4.0的设置下，该方法实现了SOTA的FID得分1.42，优于此前所有方法。
• 在下游分类任务中进行线性探测时，使用REPA的模型准确率显著提升，表明其学习到的语义表征更优。
• 与DINOv2的表征对齐效果随训练时间延长与模型规模增大而持续提升，证实了外部监督的价值。
• 对齐损失有效缩小了扩散模型特征与强自监督编码器特征之间的语义差距，即使在输入为噪声的情况下亦然。
• 定性结果表明，无论在“消防车”、“笔记本电脑”、“航天飞机”还是“珊瑚礁”等多样化类别中，生成结果均具有高保真度、多样性与提示对齐性。

(a) Semantic gap: Linear probing