本文目录

1 Representation Alignment：扩散模型与自监督方法的表征 "对齐"
(来自 KAIST，Korea University，Scaled Foundations，New York University)
1 REPA 论文解读
1.1 REPA 研究背景
1.2 DDPM 和 Flow-based Model 简介
1.3 本文观点总结
1.4 REPA 的三点观察

太长不看版

训练扩散模型可能比你想象的更简单。

纽约大学谢赛宁老师团队最近的工作提出：

当我们训练扩散模型时，把 扩散模型与自监督方法的表征"对齐" (Representation Alignment) ，使得扩散模型的训练比你想象的还简单。
因为最近有些研究表明：扩散模型中包含的 (生成式) 去噪过程，可以诱导得到一些 (判别式) 的特征。虽然这些判别式的表征的质量落后于自监督学习 (DINO v2) 得到的表征。

作者认为：扩散模型训练的一个瓶颈是学习这些表征的过程没那么有效，并坦率直言，表征的对齐真的很重要！之前我们训练扩散模型的路可能是错误的。
并进一步指明一条训练扩散模型更容易的道路：集成一些高质量的外部视觉表征，而不是仅仅靠扩散模型本身自己学习这些表征。

本文通过提出一个叫表征对齐的方法研究这个问题，将 "去噪网络中噪声输入隐藏状态的投影" 与 "来自外部预训练的视觉编码器获得的干净图像表征" 对齐。

更进一步，他也给出了一些核心观察：
1) 扩散模型能够给出一些合理的表征，越好的扩散模型的 Representation 越强。
2) 即便如此，这些 Representation 与自监督方法得到的视觉架构 (如 DINOv2, MAE) 的 Representation 相比还有所不足。
3) 当把扩散模型和 DINOv2 的表征对齐时，扩散模型在训练时可以稳步提升。

这个方法带来的效益是：
训练主流扩散模型 DiTs 和 SiTs 变得 明显容易 ，在 训练效率 和 生成质量 方面都有显著的改进：ImageNet 256×256 实现了最先进的 FID=1.42。将 SiT 训练速度提高了 17.5 倍以上，本方法 400K 步训练的模型与常规条件下训练了 7M 步的 SiT-XL 模型的性能匹配。此外，这种方法还具有良好的扩展性，即对于更大的模型，改进幅度更大。

下面是对本文的详细介绍。

1 Representation Alignment：扩散模型与自监督方法的表征 "对齐"

论文名称：Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

论文地址：

http://arxiv.org/pdf/2410.06940

代码链接：

http://github.com/sihyun-yu/REPA

项目主页：

http://sihyun.me/REPA/

1.1 REPA 研究背景

基于去噪的生成模型，例如 Diffusion Model[1][2]和 Flow-based Model[3][4]已成为生成高维视觉数据的可扩展方法。这些模型在文生图 (SDXL, SD3) 等有挑战性的任务中取得了不错的结果。

最近的工作探索了使用扩散模型做表征学习，比如[5][6]和恺明的 l-DAE[7]，并表明扩散模型的 hidden state 可以学习到判别式的表征，而且更好的扩散模型往往可以学习到更好的表征。

1.2 DDPM 和 Flow-based Model 简介

Diffusion Model

Diffusion Model 通过学习从高斯分布 到 的逐渐去噪过程来建模目标分布 。在形式上, 扩散模型的正向过程 : 从 开始, 对于 , 逐渐添加高斯噪声。扩散模型学习反向过程 。

对于给定的 可以被建模为 。其中 是预定义超参数。DDPM 表明，如果反向过程 ( , for 表示为:

在 DDPM 中， 。

即：

其中，均值满足：

其中， 可以使用由简单去噪自编码器目标进行训练：

IDDPM[8]进一步展示出如果模型通过下面的目标函数同步学习 ，则可以进一步提高性能：

其中, 表示每个维度的变量, 且有 。

在足够大的 和合适的 的调度下, 分布 几乎变成各向同性高斯分布。因此, 可以从随机噪声开始生成样本并执行迭代反向过程 来获得数据样本 。

Flow-based Model

Flow-based Model 处理连续时间相关过程 , 使用数据 和高斯噪声

其中, 和 分别为 的递减和递增函数。存在一个速度场的概率流常微分方程 (Probability Flow Ordinary Differential Equation, PF ODE)：

其中这个 ODE 在 处的分布等于边际 。

速度 表示为两个条件期望的总和：

通过最小化以下训练目标，可以用模型 近似：

注意这也对应于以下反向随机微分方程 (Stochastic Differential Equation, SDE)：

其中 score 类似地变为条件期望：

与 类似, 可以用模型 近似, 目标如下:

这里, score 可以使用 的速度 直接计算为:

因此，只估计两个向量中的一个就足够了。

随机插值 (Stochastic interpolants) 显示任何 和 都满足 3 个条件：

2. and are differentiable,

导致在 和 之间进行插值而不产生偏差的过程。因此，可以通过在训练和推理期间将它们定义为一个简单的函数来使用一个简单的插值，例如：

• 线性插值器: 。
• 方差保持 (Variance-Preserving, VP) 插值器: 。

随机插值的另一个优点是扩散系数 在训练任何分数或速度模型时是独立的。因此, 当使用反向 SDE 进行采样时, 也可以在训练后明确选择 。

注意现有的基于分数的扩散模型，包括 DDPM，同样可以解释为 SDE。它们的前向扩散过程可以解释为预定义的

离散化的前向 SDE, 其平衡分布为 , 其中训练是在 上进行的 ( 足够大比如为 1000), 其中 几乎为各向同性高斯分布。生成是通过假设 , 并从随机高斯噪声开始求解相应的反向 SDE 来完成的, 其中 和扩散系数 是从前向扩散过程中隐式选择的, 这可能导致 Score-based Diffusion Model 的设计空间过于复杂。

1.3 本文观点总结

本文观点是： 训练扩散模型的主要挑战和主要瓶颈是需要学习高质量的内部表征h。

本文证明了： 在生成式扩散模型的训练过程中，当有外部表征 支持时，训练会变得更加简单，更加有效。

本文贡献是： 提出了一种简单的正则化技术，该技术利用自监督视觉表征 ，提高了训练效率和扩散模型的生成质量。

本文的探索过程：

发现预训练的扩散模型的确会学习到有意义的判别式表征 → 但是，这些表征明显不如 Dinov2 的表征 → 发现扩散模型学习的表征与 DINOv2 的表征之间的对齐仍然很弱(相比于 Dinov2 与自监督模型比如 MoCov3 这种表征之间的对齐) → 观察到扩散模型和 Dinov2 之间的对齐，随训练更长和模型更大而不断提高

这些发现启发本文 通过结合外部自监督模型的表征来增强生成式扩散模型 。

然而，使用现成的自监督视觉编码器 (比如通过对生成任务的编码器进行微调) 时，这种方法并不直接。

1. 挑战1： 输入不匹配，扩散模型的输入是有噪声的 ，而大多数自监督学习编码器都是在干净的图像 上训练的。
2. 挑战2： 这些现成的视觉编码器不是为重建或生成等任务设计的。

为了克服这些技术问题，本文使用一种正则项技术 表征对齐 (REPresentation Alignment, REPA) 来指导扩散模型的表征学习，将预训练的自监督表征蒸馏到扩散模型的表征中。

本质上, REPA 将干净图像 的预训练自监督视觉表征 蒸馏为噪声输入 的扩散模型的表征 。这种正则化减少了表征 中的语义差距，并更好地将其与目标自监督视觉表征 对齐。这种增强的对齐显着提高了扩散 Transformer 的生成性能。有趣的是, 对于 REPA, 作者观察到仅通过对齐前几个 Transformer Block 就可以实现足够的表征对齐。反过来, 这允许 Diffusion Transformer的后续层专注于基于对齐的表征来捕获高频细节，以进一步提高生成性能。

1.4 REPA 的三点观察

假设扩散模型为 , 其中 是 latent 变量, 满足 。作者把扩散模型 视为 2 个函数的组合： 。其中 Encoder 是 , Decoder 是 。其中编码器 隐式地学习 , 来重建 。

作者首先研究了 ImageNet 上预训练的 SiT[10]模型的逐层行为，该模型使用线性插值和速度预测进行训练。作者专注于测量 Diffusion Transformer 与最先进的自监督 DINOv2[11]模型之间的表征的差距。

作者从 3 个角度检查这一点：语义差距、特征对齐进展及其最终特征对齐。

对于语义差距，作者使用 DINOv2 特征与为 7M training iterations 训练的 SiT 模型的 linear probing 的结果进行比较。

对于特征对齐，作者使用 CKNNA[12]，这是一种与 CKA 相关的内核对齐度量，但基于相互最近邻。这允许定量评估不同表示之间的对齐。

观察1：Diffusion Transformer 与最先进的 Visual Encoder 表现出显著的语义差距。

如图 4(a) 所示，与之前的工作[6][7]一致，作者观察到预训练的 Diffusion Transformer 的隐藏状态表征在第 20 层实现了相当高的 Linear Probing 峰值。然而，它的性能仍然远低于 DINOv2，这表明两种表征之间存在实质性的语义差距。此外，本文发现，在达到峰值后，Linear Probing 性能迅速下降，这表明 Diffusion Transformer 必须摆脱只关注学习语义丰富的表征，以生成高频细节的图像。

观察2：扩散表征与其他视觉表示已经较弱地完成对齐。

在图 4(b) 中，作者使用 CKNNA 报告了 SiT 和 DINOv2 之间的表征对齐。SiT 模型表示已经显示出比 MAE 更好的对齐。然而，绝对的对齐分数仍然低于在其他自监督学习方法 (例如 MoCov3 与 DINOv2）之间观察到的分数。这些结果表明，虽然 Diffusion Transformer 的表征表现出与自监督视觉表示的一些对齐，但对齐仍然很弱。

观察3：随着模型规模扩大，训练更多，对齐也会更好。

作者还测量了不同模型大小和训练迭代的 CKNNA 值。如图 4(c) 所示，作者观察到与模型更大，训练更长可以改进对齐水平。然而，绝对的对齐数值仍然很低，并且没有达到其他自监督视觉编码器 (例如 MoCov3 和 DINOv2) 的水平。

这些发现不是只对 SiT 模型有效，也对其他的扩散模型有用。比如作者也在 ImageNet 上预训练的 DiT 模型进行了类似的分析，也在图 4 中观察到了类似的结论。

1.6 特征对齐

REPA 将模型隐藏状态的 Patch 投影与预训练的自监督视觉表征对齐。作者使用干净的图像表征作为目标，研究这个影响。这个正则项的目的是为了 Diffusion Transformer 的隐藏层从 带噪声 的图片中预测 干净的 视觉表征，这些表征包含有用的语义信息。这为后续层重建目标提供了有意义的指导。

如图 6 所示, 设 为预训练的编码器, 考虑干净的图像 。设