生成的表征对齐：训练扩散Transformer比你想象的更简单

paper是KAIST发表在ICLR 2025的工作

paper title:REPRESENTATION ALIGNMENT FOR GENERATION: TRAINING DIFFUSION TRANSFORMERS IS EASIER THAN YOU THINK

Code：链接

ABSTRACT

最近的研究表明，扩散模型中的去噪过程能够在模型内部产生有意义的（判别性）表示，尽管这些表示的质量仍然落后于通过最近的自监督学习方法学习的表示。我们认为，训练大规模扩散模型以进行生成的主要瓶颈之一在于如何有效地学习这些表示。此外，通过引入高质量的外部视觉表示，而不是仅依赖扩散模型独立学习这些表示，训练可以变得更加容易。我们通过引入一种简单的正则化方法——表示对齐（REPA），来研究这一点，该方法将去噪网络中噪声输入隐藏状态的投影与从外部预训练视觉编码器获得的干净图像表示对齐。结果非常显著：当应用于流行的扩散和基于流的Transformer模型（如DiTs和SiTs）时，我们的简单策略在训练效率和生成质量上都取得了显著的提升。例如，我们的方法可以将SiT训练速度提高超过17.5倍，并且在不到40万个步骤的训练中，达到与经过700万步骤训练的SiT-XL模型（无需分类器引导）的相同性能。在最终生成质量方面，我们的方法通过使用分类器引导和引导间隔，达到了FID=1.42的最新成果。

图1：表示对齐使扩散变压器训练变得明显容易。我们的框架REPA明确地将扩散模型表示与通过简单的正则化的强大预验证的视觉表示一致。值得注意的是，模型训练变得更加高效，高效，并且比vanilla模型的收敛速度> 17.5倍。

1 INTRODUCTION

基于去噪的生成模型，如扩散模型（Ho et al., 2020; Song et al., 2021）和基于流的模型（Albergo & Vanden-Eijnden, 2023; Lipman et al., 2022; Liu et al., 2023），已成为生成高维视觉数据的一种可扩展方法。它们在诸如零样本文本到图像（Podell et al., 2023; Saharia et al., 2022; Esser et al., 2024）或文本到视频（Polyak et al., 2024; Brooks et al., 2024）生成等具有挑战性的任务中取得了显著成功。

近期的研究探讨了将扩散模型用作表示学习方法（Li et al., 2023a; Xiang et al., 2023; Chen et al., 2024c; Mukhopadhyay et al., 2021），并表明这些模型能够在其隐藏状态中学习到判别性特征，同时更优的扩散模型能够学习到更好的表示（Xiang et al., 2023）。实际上，这一观察结果与早期采用去噪分数匹配（denoising score matching）（Vincent, 2011）的方法密切相关，该方法是一种自监督学习方法（Bengio et al., 2013），其隐式地学习一个表示 $h$，作为去噪自编码器 $s_{\theta }(\tilde{x})$ 的隐藏状态，通过从被破坏的输入 $ \tilde{x}$ 中重建 $x$（Yang & Wang, 2023）。然而，重建任务可能并不适用于学习优良的表示，因为它无法消除 $x$ 中不必要的细节以用于表示学习（LeCun, 2022; Assran et al., 2023）。

我们的方法。 在本研究中，我们确定了训练扩散模型的主要挑战在于需要学习一个高质量的内部表示 $h$。我们表明，对于生成型扩散模型而言，训练过程在借助外部表示 $y_{\ast }$ 进行支持时会变得显著更简单且更高效。具体而言，我们提出了一种简单的正则化技术，该技术利用了自监督视觉表示的最新进展，以 $y_{\ast }$ 作为辅助，从而在训练效率和扩散变换器的生成质量方面实现了显著提升。

我们首先对近期的扩散变换器（Peebles & Xie, 2023; Ma et al., 2024a）以及最先进的自监督视觉模型 DINOv2（Oquab et al., 2024）进行实证分析。与以往研究（Xiang et al., 2023）类似，我们首先观察到预训练的扩散模型确实能够学习有意义的判别性表示（如线性探测实验所示，见图 2a）。然而，这些表示明显逊色于 DINOv2 生成的表示。

随后，我们发现扩散模型所学习的表示与 DINOv2 的表示仍然较为不对齐（图 2b），我们通过测量表示对齐（representation alignment）（Huh et al., 2024）对此进行研究。最终，我们观察到扩散模型和 DINOv2 之间的这种对齐性在训练时间延长以及模型规模增大时会持续提升（图 2c）。

图2：预验证的SIT模型的比对行为。我们经验研究了Dinov2-G和原始的SIT-XL/2检查点之间的特征对齐，该检查点训练了7M迭代。 （a）虽然SIT学习了语义上有意义的表示，但与Dinov2相比，仍然存在显着的差距。 （b）使用cknna（Huh等，2024），我们观察到SIT已经显示出与Dinov2的一致性，尽管与其他视觉编码器相比，其绝对值较低。 （c）通过更大的模型和更长的训练来改善对齐方式，但是进度仍然缓慢且不足。

这些见解启发我们通过引入外部自监督表示来增强生成模型。然而，在使用现成的自监督视觉编码器（例如，通过微调编码器以适应生成任务）时，这种方法并不直观。第一个挑战是输入不匹配：扩散模型处理的是带有噪声的输入 $\tilde{x}$，而大多数自监督学习编码器是在干净的图像 $x$ 上进行训练的。这一问题在现代**潜变量扩散模型（latent diffusion models）中更加明显，因为它们以经过预训练的 VAE 编码器（Rombach et al., 2022）压缩后的潜变量图像 $z=E(x)$ 作为输入。此外，这些现成的视觉编码器并非专为重建或生成任务而设计。为了解决这些技术难题，我们引导扩散模型的特征学习，采用正则化（regularization）**技术，将预训练的自监督表示蒸馏到扩散表示中，从而提供了一种灵活的方式来整合高质量表示。

图3：桥接表示差距：（a）我们的方法，repa显着降低了扩散变压器和dinov2之间的“语义差距”，如成像网分类的线性探测结果所证明。 （b）使用REPA，即使仅几层（例如8）层，扩散变压器和DINOV2之间的比对也大大改善。 （c）值得注意的是，随着对齐方式的改善，我们可以推动SIT模型的一代代表性信封：在相同数量的训练迭代次数中，它既可以提供更好的生成质量，又提供更强的线性探测结果。我们使用一个在第8层训练了REPA的单个网络，并在不同层进行评估。

具体而言，我们提出了表示对齐（REPresentation Alignment, REPA），这是一种基于最新扩散变换器架构（Peebles & Xie, 2023）的简单正则化技术。本质上，REPA 通过蒸馏预训练的自监督视觉表示 $y_{\ast }$（源自干净图像 $x$），将其融入扩散变换器的表示 $h$，用于去噪过程中处理带噪输入 $\tilde{x}$。该正则化减少了表示 $h$ 和自监督目标表示 $y_{\ast }$ 之间的语义差距（如图 3a 所示），并使其更好地对齐（如图 3b 所示）。值得注意的是，这种增强的对齐显著提升了扩散变换器的生成性能（如图 3c 所示）。有趣的是，我们观察到，使用 REPA 时，仅需对齐前几个 Transformer 块即可实现充分的表示对齐。这使得扩散变换器的后续层可以专注于捕捉基于对齐表示的高频细节，从而进一步提高生成性能。

基于我们的分析，我们进行了系统级比较，以验证所提出方案的有效性，并将其应用于两个最新的扩散变换器：DiTs（Peebles & Xie, 2023）和 SiTs（Ma et al., 2024a）。在 SiT 训练中，我们的模型在类别条件的 ImageNet 生成任务（Deng et al., 2009）上，仅用 40 万次训练迭代（无需分类器引导）即可达到 FID=7.9（Ho & Salimans, 2022），训练速度比标准 SiTs 快 17.5 倍。此外，在使用分类器引导的情况下，我们的方法将最终的 FID 从 2.06 降至 1.80，并在使用引导区间（Kynkäänniemi et al., 2024）时达到了最新的最优结果 FID=1.42。

我们在本文中总结的主要贡献如下：

• 我们假设，在扩散变换器中学习高质量表示对于提高其生成性能至关重要。
• 我们提出了 REPA，这是一种用于对齐扩散变换器表示和强自监督视觉表示的简单正则化方法。
• 我们的框架提高了扩散变换器的生成性能，例如，在 SiTs 任务中，我们实现了 17.5 倍加速训练，并在 ImageNet 生成任务上提升了 FID 评分。

2 PRELIMINARIES

我们通过随机插值（stochastic interpolants）（Albergo et al., 2023; Ma et al., 2024a）的统一视角，对基于流（flow）和扩散（diffusion）的模型进行了简要概述；更多详细信息请参见附录 A。

我们考虑一个连续的时间相关过程，该过程的数据 ${\mathbf{x}}_{\ast }$ 服从分布 $p(\mathbf{x})$，并且高斯噪声 $\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$，其中 $t\in [0,T]$：
$${\mathbf{x}}_t={\alpha }_t{\mathbf{x}}_{\ast }+{\sigma }_t\mathbf{ϵ},{\alpha }_0={\sigma }_T=1,{\alpha }_T={\sigma }_0=0,$$
其中 ${\alpha }_t$ 和 ${\sigma }_t$ 分别是随时间 $t$ 递减和递增的函数。在这样的过程中，存在一个概率流常微分方程（probability flow ordinary differential equation, PF ODE），其速度场表示为：
$${\dot{\mathbf{x}}}_t=\mathbf{v}({\mathbf{x}}_t,t)。$$

其中，该 ODE 在 $t$ 时刻的分布等于边际分布 $p_t(\mathbf{x})$。因此，可以通过求解方程 (2) 中的概率流 ODE（PF ODE），利用现有的 ODE 采样器（例如 Euler 采样器）从随机高斯噪声 $\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$ 进行数据采样（Lipman et al., 2022; Ma et al., 2024a）。

该速度场 $\mathbf{v}(\mathbf{x},t)$ 由以下两个条件期望的和表示：
$$\mathbf{v}(\mathbf{x},t)=\mathbb{E}[{\dot{\mathbf{x}}}_t|{\mathbf{x}}_t=\mathbf{x}]={\dot{\alpha }}_t\mathbb{E}[{\mathbf{x}}_{\ast }|{\mathbf{x}}_t=\mathbf{x}]+{\dot{\sigma }}_t\mathbb{E}[\mathbf{ϵ}|{\mathbf{x}}_t=\mathbf{x}],$$
其中 ${\mathbf{v}}_{\theta }({\mathbf{x}}_t,t)$ 可用于近似该表达式，并通过最小化以下训练目标进行优化：
$${\mathcal{L}}_{\text{velocity}}(\theta ):={\mathbb{E}}_{{\mathbf{x}}_{\ast },\mathbf{ϵ},t}\left[\Vert {\mathbf{v}}_{\theta }({\mathbf{x}}_t,t)-{\dot{\alpha }}_t{\mathbf{x}}_{\ast }-{\dot{\sigma }}_t\mathbf{ϵ}{\Vert }^2\right]。$$

此外，存在一个逆向随机微分方程（stochastic differential equation, SDE），该 SDE 的边际分布 $p_t(\mathbf{x})$ 与方程 (2) 中的 PF ODE 的边际分布一致，扩散系数为 $w_t$（Ma et al., 2024a）：
$$d{\mathbf{x}}_t=\mathbf{v}({\mathbf{x}}_t,t)dt-\frac{1}{2}{w}_t\mathbf{s}({\mathbf{x}}_t,t)dt+\sqrt{w_t}d{\bar{\mathbf{w}}}_t,$$
其中分数 $\mathbf{s}(\mathbf{x},t)$ 为以下条件期望：
$$\mathbf{s}(\mathbf{x},t)=-{\sigma }_t^{-1}\mathbb{E}[\mathbf{ϵ}|{\mathbf{x}}_t=\mathbf{x}]。$$
它可以通过 $\mathbf{v}(\mathbf{x},t)$ 直接计算得到，对于 $t>0$：
$$\mathbf{s}(\mathbf{x},t)={\sigma }_t^{-1}\cdot \frac{{\alpha }_t\mathbf{v}(\mathbf{x},t)-{\dot{\alpha }}_t\mathbf{x}}{{\dot{\alpha }}_t{\sigma }_t-{\alpha }_t{\dot{\sigma }}_t}。$$
这意味着数据可以通过方程 (5) 使用 SDE 求解器生成。

根据 Ma et al. (2024a)，我们主要考虑一个简单的线性插值方案，并限制 $T=1$：${\alpha }_t=1-t$，${\sigma }_t=t$。然而，我们的方法同样适用于类似的变体（例如 DDPM; Ho et al. 2020），其采用离散化过程并使用不同的 ${\alpha }_t$ 和 ${\sigma }_t$，其中 $\mathcal{N}(0,\mathbf{I})$ 作为平衡分布（即如果 $t\to \infty$，则 ${\mathbf{x}}_t$ 收敛到 $\mathcal{N}(0,\mathbf{I})$）。

3 REPA: REGULARIZATION FOR REPRESENTATION ALIGNMENT

3.1 OVERVIEW

---

设 $p(\mathbf{x})$ 为数据 $\mathbf{x}\in \mathcal{X}$ 的未知目标分布。我们的目标是通过从 $p(\mathbf{x})$ 采样的数据集来近似 $p(\mathbf{x})$。为了降低计算成本，我们采用了当前流行的潜变量扩散模型（latent diffusion）（Rombach et al., 2022）。这涉及学习一个潜变量分布 $p(\mathbf{z})$，其中 $\mathbf{z}=E(\mathbf{x})$ 是通过预训练自编码器（例如 KL-VAE；Rombach et al., 2022）进行压缩得到的潜变量，而 $\mathbf{x}\sim p_{\text{data}}(\mathbf{x})$。

我们的目标是通过训练扩散模型 ${\mathbf{v}}_{\theta }({\mathbf{z}}_t,t)$ 来学习该分布，使用的目标函数包括速度预测等（详见第 2 节）。在此，我们从自监督表示学习（Bengio et al., 2013）的角度重新审视去噪得分匹配（denoising score matching）。从这一视角来看，可以将扩散模型 ${\mathbf{v}}_{\theta }({\mathbf{z}}_t,t)$ 视为两个函数的组合 $g_{\theta }\circ f_{\theta }$，其中 $f_{\theta }:\mathcal{Z}\to \mathcal{H}$ 是编码器，使得 $f_{\theta }({\mathbf{z}}_t)={\mathbf{h}}_t$，而 $g_{\theta }:\mathcal{H}\to \mathcal{Z}$ 是解码器，使得 $g_{\theta }({\mathbf{h}}_t)={\mathbf{v}}_t$，即编码器 $f_{\theta }$ 隐式地学习了一个表示 ${\mathbf{h}}_t$，用于重构目标 ${\mathbf{v}}_t$。

然而，通过预测输入空间（例如生成像素）来学习一个良好的表示可能是具有挑战性的，因为模型通常无法有效地去除无关细节，而这些细节的去除对于学习强大的表示至关重要（LeCun, 2022; Assran et al., 2023）。我们认为，大规模扩散模型在生成任务上的训练瓶颈之一就在于表示学习，这是当前扩散模型的弱点。此外，我们假设，通过高质量的外部视觉表示来引导模型，而不是单纯依靠扩散模型自行学习表示，可以使训练过程更加容易。

为了解决这一挑战，我们提出了一种简单的正则化方法，称为 表示对齐（REPresentation Alignment, REPA），该方法基于最新的扩散变换器架构（Peebles & Xie, 2023; Ma et al., 2024a）（详见附录 B 进行说明）。简而言之，我们的正则化方法能够将预训练的自监督视觉表示有效地蒸馏到扩散变换器中。这使得扩散模型能够利用这些具有丰富语义信息的外部表示来进行生成，从而显著提升性能。

3.2 OBSERVATIONS

---

为了更深入地研究这一问题，我们首先分析了预训练 SiT 模型（Ma et al., 2024a）在 ImageNet（Deng et al., 2009）上的逐层行为，该模型使用线性插值和速度预测进行训练。特别地，我们关注扩散变换器与当前最先进的自监督视觉模型 DINOv2（Oquab et al., 2024）之间的表示差距（representation gap）。我们从三个角度来研究这一问题：语义差距（semantic gap）、特征对齐的变化（feature alignment progression）及最终的特征对齐结果。

对于 语义差距（semantic gap），我们使用 DINOv2 的特征，并对其进行线性探测，与训练 700 万步的 SiT 模型进行对比，遵循 Xiang et al. (2023) 中相同的协议，即对扩散变换器隐藏状态进行全局池化后进行线性探测。接着，为了衡量 特征对齐（feature alignments），我们采用 CKRNA（Huh et al., 2024），这是一种基于最近邻匹配的核对齐度量方法，相关工作可参考 CKA（Kornblith et al., 2019）。这使得我们可以定量评估不同表示之间的对齐情况。具体结果见图 2，更多细节（例如 CKRNA 的定义）可见附录 C.1。

---

扩散变换器在最先进视觉编码器面前表现出显著的语义差距

如图 2a 所示，我们观察到预训练扩散变换器的隐藏状态表示在 layer 20 处达到了相对较高的线性探测峰值，这与先前研究（Xiang et al., 2023; Chen et al., 2024c）一致。然而，其表现仍然远低于 DINOv2，表明这两种表示之间存在显著的语义差距。此外，我们发现，在达到这一峰值后，线性探测性能迅速下降，这表明扩散变换器必须从单纯学习语义丰富的表示转变为生成包含高频细节的图像。

扩散模型的表示已经（较弱地）与其他视觉表示对齐

在图 2b 中，我们报告了 SiT 与 DINOv2 之间的表示对齐情况，使用了 CKNNA 进行评估。特别是，SiT 模型的表示已经比 MAE（He et al., 2022）表现出更好的对齐性，后者也是一种基于掩码补全的自监督学习方法。然而，其绝对对齐分数仍然低于其他自监督学习方法（例如 MoCov3（Chen et al., 2021）与 DINOv2 之间的对齐）。这些结果表明，尽管扩散变换器的表示在某种程度上与自监督视觉表示对齐，但这种对齐仍然较弱。

---

更大的模型和更长时间的训练能够改善对齐

我们还测量了不同模型规模和训练迭代次数下的 CKNNA 值。如图 2c 所示，我们观察到随着模型规模的增大和训练时间的延长，对齐性有所改善。然而，其绝对对齐分数仍然较低，甚至在经历了 700 万次迭代的长期训练后，仍未达到其他自监督视觉编码器（例如 MoCov3 和 DINOv2）之间的对齐水平。

这些发现并不是SIT模型独有的，而是在其他基于Denoising的生成变压器中观察到的。例如，在图2中，我们使用DIT模型（Peebles＆XIE，2023）提出了类似的分析，该模型在Imagenet上以DDPM物镜预测（Ho等，2020; Nichol＆Dhariwal，2021）。有关更多详细信息，请参见附录C.2。

3.3 REPRESENTATION ALIGNMENT WITH SELF-SUPERVISED REPRESENTATIONS

---

REPA 通过对齐模型隐藏状态的 patch 级投影与预训练的自监督视觉表示来进行优化。具体而言，我们使用干净图像的表示作为目标，并探索其影响。这种正则化的目标是让扩散变换器的隐藏状态能够预测无噪声的视觉表示，从而提取包含有用语义信息的噪声输入。这为后续层重构目标提供了有意义的指导。

形式化地，设 $f$ 是一个预训练编码器，并考虑一个干净图像 $x_{\ast }$。令 $y_{\ast }=f(x_{\ast })\in {\mathbb{R}}^{N\times D}$，其中 $N,D>0$ 分别表示 patch 的数量和 $f$ 的嵌入维度。REPA 通过 $h_{\varphi }(h_t)\in {\mathbb{R}}^{N\times D}$ 与 $y_{\ast }$ 对齐，其中 $h_{\varphi }(h_t)$ 是扩散变换器编码器输出 $h_t=f_{\theta }(z_t)$ 通过可训练投影头 $h_{\varphi }$ 进行投影后的结果。在实践中，我们通常使用多层感知机（MLP）来参数化 $h_{\varphi }$。

特别地，REPA 通过最大化预训练表示 $y_{\ast }$ 和隐藏状态 $h_t$ 之间的 patch 级相似度来实现对齐：

$${\mathcal{L}}_{\text{REPA}}(\theta ,\varphi ):=-{\mathbb{E}}_{x_{\ast },ϵ,t}\left[\frac{1}{N}\sum _{n=1}^N\text{sim}(y_{\ast }^{[n]},h_{\varphi }(h_t^{[m,n]}))\right],$$

其中 $n$ 表示 patch 索引，$m$ 是 $f_{\theta }$ 的深度，$\text{sim}(\cdot ,\cdot )$ 是一个预定义的相似性函数。

在实践中，我们将该项添加到扩散模型的目标函数中，如第 2 节和附录 A 所述。例如，对于 Eq. (4) 所描述的速度模型训练，优化目标变为：

$$\mathcal{L}:={\mathcal{L}}_{\text{velocity}}+\lambda {\mathcal{L}}_{\text{REPA}}$$

其中 $\lambda >0$ 是一个超参数，用于控制去噪和表示对齐之间的权衡。我们主要研究了这种正则化在两个流行目标上的影响：改进的 DDPM（Nichol & Dhariwal, 2021）用于 DiT（Peebles & Xie, 2023）和线性随机插值用于 SiT（Ma et al., 2024a），但其他目标也可以考虑。