AI 摘要 · 由英文经 AI 翻译

Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, Saining Xie

arXiv:2410.06940v4cs.CVcs.LG

TL;DR

Diffusion transformer 难以学到高质量表征,限制了生成训练的效率。REPA 将其含噪隐藏状态与外部自监督视觉表征对齐,从而加快收敛并提升生成质量;采用 guidance interval 时,FID=1.42。

  • 问题

    扩散模型学到的表征落后于自监督视觉表征,使表征学习成为大规模生成训练的一项核心挑战。

    出处:§1 · p. 2, §3.1 · p. 4
  • 方法

    REPA 通过将噪声隐状态的投影与预训练自监督视觉编码器提取的干净表征对齐,对扩散 Transformer 进行正则化。

    出处:§1 · pp. 2–3
  • 结果

    REPA 持续提升 diffusion transformer 的生成质量并加快收敛,在 400K 次 SiT 迭代后达到 FID=7.9,采用 guidance interval 时达到 FID=1.42。

    出处:§1 · p. 3, §6 · p. 10
  • 启示与局限

    外部视觉表征能够大幅加速扩散 Transformer 训练,同时提升 DiT 和 SiT 模型的生成性能。

    出处:§1 · pp. 2–3, §6 · p. 10

    通过实例判别和去噪学到的表征之间的理论关系,仍是未来值得探索的方向。

    出处:M LIMITATIONS AND FUTURE WORK · p. 43

Abstract

from arXiv · show

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$\times$, 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.

1 引言

论文指出,高质量的内部表征学习是扩散 Transformer 训练的核心瓶颈,并提出 REPA,将带噪隐藏状态与外部自监督视觉表征对齐。REPA 缩小了表征差距,显著提升了 DiTs 和 SiTs 的训练效率与生成质量。

贡献

  • 论文提出假设:学习高质量表征对于提升扩散 Transformer 的生成性能至关重要。主要贡献包括表征学习假设、REPA 正则化方法,以及提升 ImageNet 生成效率和 FID 分数。出处:§1 · pp. 2–3

动机

  • 扩散 Transformer 能够学习有意义的判别性表征,但其质量显著不及 DINOv2 学到的表征。这一实证差距促使研究者利用外部自监督视觉表征来辅助生成模型训练。出处:§1 · p. 2

方法

  • REPA 将干净图像的自监督表征蒸馏到扩散 Transformer 对带噪输入形成的表征中。该正则项在缩小语义差距的同时,解决扩散模型带噪输入与编码器干净图像输入之间的不匹配。出处:§1 · pp. 2–3

结果

  • 将 REPA 应用于 DiTs 和 SiTs 后,生成性能和表征对齐效果均有所提升。研究通过对这两种近期提出的扩散 Transformer 架构进行系统级比较来评估该方法。出处:§1 · p. 3
  • SiT 在无 classifier-free guidance 的类别条件 ImageNet 生成任务上,经过 400K 次训练迭代达到 FID=7.9。据报告,这比原始 SiT 快 >17.5×。出处:§1 · p. 3

2 预备知识

预备知识通过随机插值统一了基于流和基于扩散的生成方法。随机插值定义了含噪状态、概率流 ODE,以及等价的 SDE 采样过程。本文主要采用线性插值,同时指出该方法可扩展到 DDPM 等相关变体。

随机插值

  • 随机插值根据数据 x* ∼ p(x) 和高斯噪声 ϵ ∼ N(0, I),定义连续过程 x_t = α_t x* + σ_t ϵ。边界条件为 α_0 = σ_T = 1 且 α_T = σ_0 = 0。出处:§2 · p. 3

概率流 ODE

  • 该插值对应一个概率流 ODE,其在时刻 t 的分布等于边缘分布 p_t(x),因此可通过 ODE 求解器从高斯噪声中采样。文中以 Euler 采样器为例。出处:§2 · pp. 3–4

速度建模

  • ODE 速度为 v(x, t) = E[ẋ_t | x_t = x] = ẋα_t E[x* | x_t = x] + ẋσ_t E[ϵ | x_t = x]。模型 v_θ(x_t, t) 通过训练目标来近似该速度。出处:§2 · p. 4

反向 SDE

  • 等价的反向 SDE 与概率流 ODE 具有相同的边缘分布,并提供了另一种利用 SDE 求解器生成数据的途径。其形式中包含扩散系数 w_t 和得分函数 s(x_t, t)。出处:§2 · p. 4

插值选择

  • 主要实验采用 T = 1,并使用线性插值 α_t = 1 − t 和 σ_t = t;该方法也适用于 DDPM 等变体。DDPM 使用离散化过程和不同的 α_t、σ_t,并以 N(0, I) 作为平衡分布。出处:§2 · p. 4

3 REPA:表征对齐正则化

REPA 将预训练自监督视觉表征蒸馏到扩散 Transformer 中,以解决扩散模型表征学习的瓶颈。它将噪声输入产生的隐状态投影与干净图像的逐 patch 视觉表征对齐,并将这一对齐项加入标准扩散目标。

方法

  • REPA 将扩散 Transformer 的投影隐状态与预训练自监督视觉编码器提取的干净图像表征对齐。该方法引导噪声输入产生的隐状态趋向于不受噪声影响、且对后续重建具有语义价值的表征。出处:§3.1 · p. 4, §3.3 · p. 6

动机

  • 尽管扩散 Transformer 的线性探测峰值表现相当不错,其语义表征与最先进视觉编码器仍存在显著差距。该差距通过预训练 SiT 模型在 ImageNet 上与 DINOv2 的比较来衡量。出处:§3.2 · p. 5
  • 随着模型规模增大、训练时间延长,表征对齐效果有所提升,但即使经过 7M 次迭代,仍不及自监督视觉编码器。该趋势通过 CKNNA 在不同模型规模和训练迭代次数下进行报告,其中包括与 MoCov3 和 DINOv2 的比较。出处:§3.2 · p. 5

方法

  • 该方法在 patch 级别进行对齐,通过最大化扩散编码器输出经 MLP 投影后的结果与预训练编码器输出之间的相似度来实现。对于干净图像 x*,预训练表征 y*=f(x*) 包含 N 个 patch,嵌入维度为 D;投影后的隐状态具有相同维度。出处:§3.3 · p. 6
  • REPA 将对齐项加入原始扩散目标,其中 λ > 0 用于控制去噪与表征对齐之间的权衡。论文主要研究 DiT 中的 Improved DDPM,以及 SiT 中的线性随机插值。出处:§3.3 · p. 6

4 实验

实验表明,REPA 能显著提升不同模型规模和视觉编码器下的扩散Transformer生成质量、表征质量与收敛速度。这些收益在不同对齐方案下依然成立,最终带来当前最佳的生成质量。

组件分析

  • 在不同扩散Transformer设计方案中,REPA都能稳定提升生成性能,FID显著优于原始模型。出处:§4.2 · p. 7

目标表征

  • 预训练目标表征越强,生成质量和线性探测性能就越好;编码器质量与对齐模型的性能高度相关。出处:§4.2 · pp. 7–8, Figure 5 · p. 8

对齐深度

  • 仅对前几层施加正则化就已足够;在第6或8层接入REPA还能进一步提升生成质量,并将后续层留给高频细节。出处:§4.2 · p. 8

对齐目标

  • NT-Xent在50-100K次迭代的早期阶段表现更好,但其优势会随训练推进而减弱,因此后续实验采用余弦相似度。出处:§4.2 · p. 8

可扩展性

  • 随着扩散Transformer规模增大,REPA带来的收敛加速效果愈发显著;在目标编码器固定时,更大的模型也能更快达到更优性能。出处:§4.2 · p. 8, Figure 5 · p. 8

5 相关工作

相关工作通过兼顾生成与理解的混合模型,以及对扩散表征的分析或应用,将扩散模型与表征学习联系起来。另有研究借助预训练视觉编码器提升扩散模型的效率和性能,包括两阶段文生图框架。

连接扩散模型与表征学习

  • 混合方法同时训练模型执行基于扩散的生成任务和分类或分割等理解任务。Yang et al. (2022) 和 Deja et al. (2023) 将分类与生成相结合,而 Tian et al. (2024) 将分割与生成相结合。出处:§5 · p. 10
  • 近期研究也分析并利用扩散模型学到的表征。这些研究属于更广泛的工作范畴,旨在利用或改进扩散模型的表征。出处:§5 · p. 10

结合外部表征的扩散模型

  • 预训练视觉编码器已被用于提升扩散模型的效率和性能。文中列举的研究包括 Pernias et al. (2024) 和 Li et al. (2024)。出处:§5 · p. 10
  • Würstchen 采用两阶段文生图框架:文本条件模型先生成语义图,再进行图像合成。随后,第二个扩散模型基于生成的语义图合成图像。出处:§5 · p. 10

6 结论 … C 分析细节

论文总结指出,REPA 将扩散 Transformer 的表征与自监督表征对齐,从而提升生成性能并加快收敛。相关章节还介绍了基于扩散的生成过程,以及所用的 DiT/SiT 架构。

6 结论

  • REPA 将扩散 Transformer 的表征与近期自监督表征对齐,在提升生成性能的同时加快收敛。作者将 REPA 描述为一种简单的正则化方法,并指出它既能加快收敛,也能提升生成质量。出处:§6 · p. 10

A 基于扩散的模型描述

A.1 去噪扩散概率模型

  • 扩散模型从高斯噪声 N(0, I) 出发,学习反向去噪过程以逼近目标分布 p(x),从而逆转正向加噪过程。生成从随机噪声开始,通过迭代应用反向过程得到数据样本。出处:§A.1 · p. 18
  • DDPM 使用简单的去噪自编码器目标,在预先定义的正向和反向过程中训练去噪网络。正向过程逐步加入高斯噪声,反向过程则逐步将其去除。出处:§A.1 · p. 18

A.2 随机插值

  • 随机插值在数据 x∗∼p(x) 与高斯噪声 ϵ∼N(0, I) 之间定义连续时间路径,从而支持基于流的建模。该路径使用时间相关函数 α_t 和 σ_t;其概率流常微分方程与边缘分布相匹配。出处:§A.2 · p. 19
  • 插值速度通过条件期望表示,并由使用相应目标训练的模型进行近似。得分也可以用模型建模,但对于 t > 0,也可以直接由速度计算。出处:§A.2 · p. 19
  • 随机插值要求使用正值可微的调度函数,并满足 α_1=σ_0=0 和 α_0=σ_1=1,从而实现数据与噪声之间的无偏插值。线性调度和方差保持调度是其中的例子;训练后还可选择扩散系数,用于反向 SDE 采样。出处:§A.2 · p. 19

B 扩散 Transformer 架构

  • 模型沿用 DiT 和 SiT 的做法,将输入序列切分为 patch,并基于 stable-diffusion VAE 得到的下采样潜在图像 z=E(x) 进行运算。该架构被描述为与视觉 Transformer 十分相似。出处:B DIFFUSION TRANSFORMER ARCHITECTURE · p. 20

C.1 评估细节 · C.2 DiT 分析 · C.3 预训练视觉编码器说明 •

附录定义了 CKNNA 评估方法,分析了 DiT 表征与 DINOv2 的对比,并介绍了预训练视觉编码器及其数据集影响。总体而言,DiT 表征具有意义,但对齐较弱,与 DINOv2 相比仍有明显差距。

C.1 评估细节

  • CKNNA 通过 k 近邻嵌入而非所有数据集样本对来衡量对齐程度,从而放宽了 CKA 的要求。CKA 使用根据网络表征计算得到的核矩阵来衡量模型的整体相似度。出处:§C.1 · p. 20
  • 随机抽取 10,000 张 ImageNet 验证图像;由于较小的 k 显示出更好的对齐效果,CKNNA 报告采用 k = 10。该评估遵循 Huh et al. (2024) 的观察结果。出处:§C.1 · p. 21
  • 线性探测训练一个线性层,共 90 epochs,batch size 为 16,384;采用无参数的 batch normalization、Adam、余弦衰减和 0.001 的初始学习率。该设置遵循 DAE (Chen et al., 2024c)。出处:§C.1 · p. 21

C.2 DiT 分析

  • DiT 学到了有意义的判别表征,但线性探测结果显示,其与 DINOv2 相比存在较大差距。分析使用的是经过 7M 次迭代训练的预训练 DiT-XL/2 模型。出处:§C.2 · p. 21, Figure 10 · p. 21
  • DiT 与 DINOv2 的对齐较弱,且其 CKNNA 绝对值仍然较小。Figure 10 使用 CKNNA 衡量 DINOv2-g 与 DiT-XL/2 之间的特征对齐程度。出处:Figure 10 · p. 21, §C.2 · p. 21

C.3 预训练视觉编码器说明 •

  • 预训练视觉编码器包括 MAE、DINO、MoCov3、CLIP、DINOv2、I-JEPA 和 SigLIP,涵盖重建、自蒸馏、对比学习、联合嵌入和图文目标。MAE 重建被遮挡的图像块;DINO 使用动量教师;MoCov3 将 MoCo 扩展到视觉 Transformer;CLIP 使用图文对比学习;DINOv2 结合像素级和图像块级判别;I-JEPA 预测图像中缺失的部分;SigLIP 使用成对 sigmoid 损失。出处:§C.3 · p. 21
  • 在大规模图像数据集上训练出的视觉表征越好,带来的提升就越大,无论其训练数据是否包含 ImageNet。在其他数据集上训练时,SigLIP 和 CLIP 等编码器不会在 ImageNet 上微调,从而区分数据集泄漏的影响。出处:§C.3 · p. 22

D 超参数及更多实现细节 · E 评估细节 · F 基线方法

论文说明了训练、预训练编码器、采样和计算资源方面的实现选择,随后定义评估流程并概述基线方法。评估遵循 ADM 的参考设置,并采用成熟的生成指标和扩散模型基线。

D 超参数及更多实现细节

  • 模型遵循原始 SiT 实现及 DiT/SiT 结构,使用 AdamW,学习率恒定为 1e-4,β1=0.9,β2=0.999,不使用权重衰减,采用 fp16、梯度裁剪和预先计算的压缩潜变量。所有实验均采用这些设置以加快训练。出处:D HYPERPARAMETER AND MORE… · p. 23
  • 预训练编码器检查点来自 RCG 的 MoCov3-B/-L 实现或官方发布版本;当 patch 数量不同时,会对位置嵌入进行插值。插值使预训练编码器能够适配 patch 数量不同的扩散 Transformer。出处:D HYPERPARAMETER AND MORE… · p. 23
  • 采样器采用 Euler-Maruyama 方法和 Eq. (5) 中的 SDE,扩散系数为 w_t=σ_t,最后一个 SDE 步长设为 0.04,这显著提升了性能。报告的提升与原始 SiT 论文中的结果相近。出处:D HYPERPARAMETER AND MORE… · p. 23
  • 使用 8 张 NVIDIA H100 80GB GPU,可在批量大小为 256 时实现约 5.4 step/s 的训练速度;通过预先计算预训练特征等工程优化,还能进一步提速。通过额外优化可以提高所述训练速度。出处:D HYPERPARAMETER AND MORE… · p. 23

E 评估细节

  • 评估严格遵循 ADM 的设置和参考批次,使用 H100 80GB 或 4090Ti GPU,并采用 tf32 精度;与 fp32 相比,其生成结果差异可以忽略。为保证评估一致性,保留参考批次和官方实现。出处:E EVALUATION DETAILS · p. 24
  • FID 衡量真实图像分布与生成图像分布之间的 Inception-v3 特征距离;sFID 则使用中间层空间特征来捕捉空间分布。FID 假设特征分布服从多元高斯分布;所提供的段落还将 IS 介绍为一种使用 Inception-v3 的评估指标。出处:E EVALUATION DETAILS · p. 24

F 基线方法

  • 评估将方法与 ADM、VDM++、Simple diffusion 和 CDM 进行比较,分别代表涉及分类器引导采样、自适应噪声调度、简化扩散设计及其他扩散基线的方法。所提供的基线段落将 ADM、VDM++、Simple diffusion 和 CDM 列为基线方法,但对 CDM 的描述不完整。出处:F BASELINES · p. 24

G 详细定量结果 · H 更多定性结果 · I 关于相关工作的补充讨论

详细结果表明,REPA 在不同设置下都能提升 SiT 的训练速度和性能;定性结果图则展示了未经筛选的类别条件生成样例。相关工作将 REPA 置于预训练编码器方法、去噪 Transformer、辅助自监督目标和基于去噪的表征学习等研究脉络中。

G 详细定量结果

  • SiT 模型使用 λ = 0.5 和负余弦相似度与 DINOv2-B 表征对齐,所选层随模型规模而异。基础模型使用第 4 层隐状态,而 large 和 xlarge 模型使用第 8 层隐状态。出处:G DETAILED QUANTITATIVE RESULTS · p. 25
  • 在 4M 次迭代时评估 SiT-XL/2+REPA,采用不同的 classifier-free guidance 强度,包括使用 guidance interval 的设置。详细评估涵盖不同 guidance 强度 w,以及 guidance interval 的应用情况。出处:G DETAILED QUANTITATIVE RESULTS · p. 25, Table 9 · p. 25, Table 10 · p. 25
  • 当使用 classifier-guidance 强度 w = 1.35 测量 FID 时,相比原始 SiT-XL/2 模型,REPA 大幅提升训练速度和性能。Figure 11 展示训练迭代次数与 FID 的关系,数值是在 classifier-free guidance 下测得的。出处:G DETAILED QUANTITATIVE RESULTS · p. 26, Figure 11 · p. 26

H 更多定性结果

  • 更多图例展示了 SiT-XL/2+REPA 在 classifier-free guidance w = 4.0 下,对动物、物体、食物和风景类别进行的未经筛选的生成结果。示例包括红海龟、金刚鹦鹉、金毛寻回犬、原声吉他、芝士汉堡、珊瑚礁和火山。出处:Figure 12 · p. 27, Figure 13 · p. 27, Figure 15 · p. 28

I 关于相关工作的补充讨论

  • 预训练视觉编码器已用于 GAN 训练和蒸馏,通过中间特征充当判别器来加速收敛。引用的方法包括 Sauer et al.、Kumari et al. 和 Kang et al.。出处:I MORE DISCUSSION ON RELATED WORK · p. 39
  • U-ViT、MDT、DiffiT 和 DiT 等去噪 Transformer 表明,Transformer 主干网络适用于扩散模型或基于流的模型训练。文中指出,跳跃连接可能有所帮助,但 DiT 表明纯 Transformer 架构并非必须采用跳跃连接。出处:I MORE DISCUSSION ON RELATED WORK · p. 39
  • 其他生成模型方法通过增加辅助自监督任务来加速训练,包括掩码重建、判别损失和掩码图像建模。示例包括 MaskDiT、SD-DiT 和 MAGE。出处:I MORE DISCUSSION ON RELATED WORK · p. 39
  • 先前工作也将去噪视为自监督学习任务,以改进表征学习或生成紧凑的潜在表征。示例包括扩展扩散目标、扩散引导编码器,以及对基于去噪的表征学习的分析。出处:I MORE DISCUSSION ON RELATED WORK · p. 39

J ImageNet 512×512 实验

ImageNet 512×512 实验通过将 SiT 输入调整为压缩的 64 × 64 × 4 潜变量,并将 DINOv2 输入缩放至 448 × 448,检验 REPA 的可扩展性。训练迭代次数减少超过三倍时,采用 REPA 的 SiT-XL/2 在 FID、sFID、IS 和 Prec 上仍优于原版 SiT-XL/2。

实验设置

  • 该 512×512 设置将 512 × 512 × 3 像素图像压缩为 64 × 64 × 4 潜变量,作为 SiT 输入。这些潜变量由 stable diffusion VAE 生成。出处:J IMAGENET 512×512 EXPERIMENT · p. 40
  • 在 ImageNet 512×512 实验中,输入 DINOv2 的图像缩放至 448 × 448。除此之外,该实验沿用 ImageNet 256×256 设置。出处:J IMAGENET 512×512 EXPERIMENT · p. 40

定量结果

  • 训练迭代次数减少 >3×,采用 REPA 的 SiT-XL/2 仍优于原版 SiT-XL/2,指标包括 FID、sFID、IS 和 Prec。比较结果见 Table 11。出处:J IMAGENET 512×512 EXPERIMENT · p. 40

定性结果

  • 图中展示了 SiT-XL/2+REPA 在 ImageNet 512×512 上生成的样本,使用 classifier-free guidance,w = 4.0。定性结果见 Figure 36。出处:Figure 36 · p. 40

定量结果

  • ImageNet 512×512 上的系统级比较使用 classifier-free guidance,w = 1.35。定量结果见 Table 11。出处:Table 11 · p. 40

K 文本到图像生成实验 · L 特征图可视化

在 MS-COCO 上训练的 MMDiT 模型中,REPA 提升了文本到图像生成效果,无论定性还是定量表现都优于原始模型。特征可视化进一步表明,REPA 生成由粗到细的表征,而原始模型的特征在 t 较大时会变得嘈杂。

K 文本到图像生成实验

L 特征图可视化

M 局限性与未来工作

论文指出,未来可进一步理解并改进 REPA 的对齐深度,将其扩展到不同数据类型和文本到图像任务,发展相关理论,并探索随时间变化的权重设置。

对齐深度

  • 在第8层应用 REPA 比对齐更靠后的 Transformer 层嵌入更有益,因此有必要深入分析对齐深度。作者将这项分析列为进一步改进 REPA 的方向。出处:M LIMITATIONS AND FUTURE WORK · p. 43

不同输入数据类型

  • 应将 REPA 的研究从潜空间图像扩散拓展到其他领域,包括像素级扩散、视频以及大规模文本到图像扩散模型。文本到图像方向的动机来自论文报告的 MS-COCO 生成结果。出处:M LIMITATIONS AND FUTURE WORK · p. 43

理论分析

  • 未来的理论研究应解释 REPA 为何有效,并考察通过实例判别目标与去噪目标学习到的表征之间的关系。作者提出将这两类目标进行比较,作为理论分析的具体切入点。出处:M LIMITATIONS AND FUTURE WORK · p. 43

随时间变化的 REPA

  • 随时间变化的 REPA 方案可以采用基于扩散过程噪声调度的权重函数;作者在本文中并未探索这一方案。作者将这一设计留作未来研究方向,本文则聚焦于目标表征、对齐深度和可扩展性等其他分析。出处:M LIMITATIONS AND FUTURE WORK · p. 43
Loading 2410.06940v4…