何恺明团队的这篇文章提出了一种名为「Dispersive Loss」的即插即用正则化方法,用来弥合扩散模型与表示学习之间长期存在的鸿沟。
- 当前扩散模型主要依赖回归目标进行训练,普遍缺乏对内部表示的显式正则化。
- Dispersive Loss 鼓励模型内部的特征表示在隐藏空间中「分散」开来,类似于对比自监督学习中的「斥力」机制,但其关键优势在于无需正样本对,因此不干扰扩散模型的采样过程。
- 相较于现有方法 REPA,Dispersive Loss 具有自包含、极简主义的特点,无需预训练、额外参数或外部数据。
- 在 ImageNet 数据集上,该方法在 DiT 和 SiT 等主流扩散模型上均取得了持续且显著的性能提升,甚至在一步生成模型 MeanFlow 上达到了新的最先进水平。
- 这项研究为生成模型引入表示学习提供了简洁高效的途径,让扩散模型在生成图像的同时,也能更好地「理解」图像。
论文: Diffuse and Disperse: Image Generation with Representation Regularization
一、扩散模型概述
1.1 从噪声到图像的「逆向修复」
扩散模型的工作原理,可以被形象地理解为「先涂鸦,再修复」。
- 「涂鸦」过程(前向扩散): 我们会逐步地、一点点地给这张清晰的照片添加随机噪声,就像用橡皮擦一点点抹掉细节,直到这张照片完全变成一堆毫无意义的随机噪声。这个过程是逐渐失去信息的过程。
- 「修复」过程(逆向去噪): 模型的训练目标,就是学习如何逆转这个「涂鸦」过程。给定一张被噪声污染的图片,模型的目标是预测出其中的噪声,然后从噪声图片中「减去」预测到的噪声,从而让图片变得更清晰一点。这个过程会重复很多次(通常是几百到几千步),每一步都让图片离原始清晰状态更近一点。最终,从一堆纯粹的随机噪声开始,模型就能一步步地「修复」出全新的、逼真的图像。
1.2 生成 Vs 理解
扩散模型的成功,主要归功于其回归(Regression) 式的训练目标。简单来说,它被训练成一个强大的「预测器」:给定一个带噪声的输入,它就预测出原始的噪声或者原始的图像。这个过程就像一个学生,我们给他大量的习题(加噪图像),并告诉他正确答案(原始噪声或图像),他通过不断练习来提高「做题」的准确率。
然而,这里隐藏着一个小秘密:这种训练方式,虽然能让模型「做对题」,但它没有显式地(explicitly)要求模型去学习和优化其内部的「知识结构」或「思维方式」。换句话说,模型可能只是擅长「死记硬背」和「模式匹配」,而没有真正地「理解」图像的深层语义和结构。
这就引出了我们今天的主角——表示学习(Representation Learning)。
二、表示学习
2.1 把像素变成「有意义的词汇」
在计算机世界里,一张图片最初只是一堆像素值,对机器来说,它们是冰冷的数据。表示学习的目的,就是把这些原始的、高维的、难以直接处理的像素数据,转换成一种更「抽象」、「压缩」且「有意义」的低维向量,我们称之为「表示」(Representation)或「特征」(Feature)。
你可以把这个过程想象成:
- 原始像素: 就像一堆混乱的字母和符号。
- 表示学习: 就像一个翻译器,把这些混乱的符号翻译成一个个「词汇」或「句子」,这些「词汇」和「句子」能够捕捉到图像的核心语义信息。比如,一张猫的图片,它的「表示」可能包含「有胡须」、「尖耳朵」、「毛茸茸」等「词汇」,而不是仅仅记录每个像素的颜色值。
一个好的表示对 AI 来说非常重要:
- 更高效: 机器不需要处理海量的原始像素,只需处理精炼过的特征。
- 更通用: 学到的特征可以迁移到各种不同的任务中,比如识别、检测、分割,甚至生成。
- 更鲁棒: 好的表示能够过滤掉无关紧要的细节,抓住数据的本质,对噪声和变化不那么敏感。
2.2 图像识别 Vs 图像生成
在图像识别领域,表示学习可谓过去十年的「发动机」。无论是监督学习还是自监督学习(尤其是对比学习),都在努力学习判别性(Discriminative) 的表示。也就是说,让「猫」和「狗」的表示在特征空间中离得远远的,这样分类器就能轻松区分它们。
在图像生成领域, 传统的生成模型,包括扩散模型,主要关注生成图像的逼真度(Fidelity) 和多样性(Diversity)。它们往往通过直接优化输出像素来达到这些目标,而对模型内部的表示质量关注较少。这就好比一个画家,他可能只关注画出来的画像不像,而不太关注自己「心里」对事物的理解是否深刻。
2.3 扩散模型 + 表示学习 = ?
以扩散模型为例,它在生成图像方面非常强大,但其内部表示可能不够「精炼」;而表示学习则专注于学习高质量的内部表示。理论上,如果能把表示学习的优势引入扩散模型,岂不是能让 AI 绘图更上一层楼?
但现实是,这两个领域的发展在很长一段时间里是相对独立的:
- 扩散模型侧重于如何有效地「去噪」和「生成」。
- 表示学习侧重于如何有效地「理解」和「区分」。
这种「擦肩而过」导致了一个遗憾:扩散模型虽然生成能力惊人,但其内部的「知识结构」可能没有得到充分的优化和正则化。这就可能导致模型在处理复杂场景、生成多样化内容时,仍然存在一些「盲点」或「不稳定性」。
三、前人探索:REPA
3.1 REPA 的思路
REPA (Representation Alignment) 是一个重要的先行者。它的核心思想是:将生成模型的内部表示,与一个外部的、已经预训练好的、高质量的表示模型(比如 DINOv2)的表示进行「对齐」(Alignment)。
你可以把这想象成:扩散模型这个学生在「做题」的同时,我们给他请了一个「顶级家教」(预训练的 DINOv2 模型)。这个家教已经对「知识点」理解得非常透彻,所以我们要求学生在做题的过程中,也要努力让自己的「思维方式」和「知识结构」向家教看齐。通过这种「对齐」,希望学生能更好地理解知识,从而提高做题质量。
3.2 REPA 的局限
REPA 确实在一定程度上提升了扩散模型的性能,展现了表示学习的潜力。但它也存在一些明显的局限性:
- 依赖外部「家教」: REPA 严重依赖一个预训练好的、高容量的外部编码器。这意味着你需要先花巨大的计算资源和时间去训练这个「家教」模型。比如,DINOv2 是在 1.42 亿张图片上训练了相当于 1500 个 ImageNet Epoch 的时间,并拥有 1.1 亿参数。这笔「学费」可不便宜。
- 额外的计算开销和参数: 除了训练外部模型,REPA 在训练扩散模型时,还需要额外加载和运行这个「家教」,甚至可能需要额外的模型参数来处理对齐,这会增加训练的复杂性和资源消耗。
- 不够「纯粹」: 很难说 REPA 的性能提升,究竟是来源于表示对齐本身,还是仅仅因为投入了更多的计算资源和使用了更多的外部数据。就像学生成绩提高,到底是家教教得好,还是因为请了家教后学生学得更久、做了更多题?
有没有一种更简洁、更自包含、成本更低的方法,也能让扩散模型实现「认知升级」呢?
四、Dispersive Loss 的设计思路
4.1 核心思想:鼓励「分散」
Dispersive Loss 的核心理念非常直观:在扩散模型训练的同时,额外施加一个约束,鼓励模型内部的特征表示在隐藏空间中尽可能地「分散」开来。
这种「分散」意味着不同的输入样本,即使它们在原始像素层面非常相似(比如两张猫的图片,但猫的姿态略有不同),它们在隐藏空间中也应该被映射到不同的、可区分的位置。这有助于模型学习到更丰富、更具判别力的特征,避免特征坍缩(Feature Collapse) 乃至 模式坍缩(Mode Collapse)。
PS:特征坍缩指所有特征都挤在一起;模式坍缩指模型只能生成少数几种类型的图像。
4.2 「无正样本对」的对比学习
这是 Dispersive Loss 最巧妙的地方,也是它与传统对比学习的本质区别。
传统对比学习 的哲学是「拉近正样本,推远负样本」。
- 正样本: 同一张图片的不同增强版本,被认为是「相似」的,它们的表示应该被「拉近」。
- 负样本: 不同图片,被认为是「不相似」的,它们的表示应该被「推远」。 这种方法需要精心设计数据增强,并确保正负样本的正确配对。
而 Dispersive Loss 的作者认为,在扩散模型中,原始的「去噪回归损失」本身就天然地充当了「正样本对」的对齐作用。
回忆一下扩散模型训练过程:模型接收一个加了噪声的图像 xtx_t,并被要求去预测原始的清晰图像 x0x_0。这意味着模型在学习一个映射 xt→x0x_t \rightarrow x_0。
你可以把 xtx_t 及其对应的真实 x0x_0 看作是「语义上对齐」的一对。模型的目标就是让 xtx_t 经过网络处理后的输出,能够与 x0x_0 完美匹配。这种内在的「对齐」机制,实际上就起到了传统对比学习中「拉近正样本」的作用。
因此,我们不需要再显式地定义正样本对,只需要关注如何「推远负样本」——即鼓励批次内所有样本的表示相互分散。
4.3 「即插即用」与「极简主义」
Dispersive Loss 的另一个巨大优势是其即插即用(Plug-and-play) 和极简主义(Minimalist) 的设计理念:
- 不干扰原始训练: 它只作用于模型内部已经计算出的中间表示,不涉及额外的输入采样、特殊数据增强或修改原始扩散模型的复杂去噪计算。这意味着它不会对扩散模型本身复杂的训练流程造成任何干扰。
- 无额外开销: 你不需要额外预训练一个巨大的外部模型(像 REPA 那样),也不需要引入新的可学习参数。它仅仅是总损失函数上的一个额外项,计算开销极小。
这种设计大大降低了应用门槛,使得任何正在使用扩散模型进行研究或开发的团队,都可以轻松地集成 Dispersive Loss,从而在不增加复杂性的前提下,提升模型性能。
五、Dispersive Loss 的具体实现
5.1 损失函数
论文将 Dispersive Loss 作为标准扩散模型训练总目标函数的一个附加项。形式上,对于一个批次(batch)的噪声图像 X={xi}X = \{ x _ { i } \},总的目标函数表示为:
L(X)=Exi∈X[LDiff(xi)]+λLDisp(X)\mathcal { L } ( X ) { = } \mathbb { E } _ { x _ { i } \in X } [ \mathcal { L } _ { \mathrm { D i f f } } ( x _ { i } ) ] + \lambda \mathcal { L } _ { \mathrm { D i s p } } ( X ) \\
- LDiff(xi)\mathcal { L } _ { \mathrm { D i f f } } ( x _ { i } ):这是标准的扩散损失,通常是每个样本 xix_i 的去噪回归损失。例如,模型预测的噪声与真实噪声之间的 ℓ2\ell_2 距离。这是模型「做题」的部分。
- LDisp(X)\mathcal { L } _ { \mathrm { D i s p } } ( X ):这就是 Dispersive Loss,它作用于整个批次 XX 的中间表示。这是模型「提升认知」的部分。
- λ\lambda:一个权重超参数,用于平衡扩散损失和 Dispersive Loss 的重要性。你可以把它看作是「做题」和「认知提升」这两种训练目标之间的优先级。
注意,当 λ=0\lambda=0 时,整个模型就完全退化为标准的扩散模型,也就是说 Dispersive Loss 是一种纯粹的「增强」或「正则化」手段。
论文主要聚焦于 InfoNCE-based Dispersive Loss,因为它在实验中表现最佳。
5.2 传统 InfoNCE
在自监督学习中,InfoNCE 是一种非常流行的对比损失。对于一个样本 xix_i,我们从中提取出它的中间表示 zi=f(xi)z_i = f(x_i)。如果 zi+z_i^+ 是 xix_i 的另一个增强视图的表示(正样本),那么原始 InfoNCE 损失通常是这样定义的:
LContrast=−logexp(−D(zi,zi+)/τ)∑jexp(−D(zi,zj)/τ)\mathcal { L } _ { \mathrm { C o n t r a s t } } { = } { - } \log \frac { \exp \left( - \mathcal { D } \left( z _ { i } , z _ { i } ^ { + } \right) / \tau \right) } { \sum _ { j } \exp \left( - \mathcal { D } \left( z _ { i } , z _ { j } \right) / \tau \right) } \\
其中:
- D(zi,zj)\mathcal { D } ( z _ { i } , z _ { j } ) 是衡量 ziz_i 和 zjz_j 之间不相似度的函数(例如距离)。
- τ\tau 是一个超参数,称为温度(temperature),它控制着相似度度量中的「软化」程度。
这个公式可以等价地被重写为两项之和(这也是论文的关键洞察之一):
LContrast=D(zi,zi+)/τ+log∑jexp(−D(zi,zj)/τ)\mathcal { L } _ { \mathrm { C o n t r a s t } } { = } \mathcal { D } \big ( z _ { i } , z _ { i } ^ { + } \big ) / \tau + \log \sum _ { j } { \exp \big ( { - \mathcal { D } \big ( z _ { i } , z _ { j } \big ) / \tau } \big ) } \\
- 第一项 D(zi,zi+)/τ\mathcal { D } \big ( z _ { i } , z _ { i } ^ { + } \big ) / \tau: 这一项鼓励正样本对 (zi,zi+)(z_i, z_i^+) 之间的距离最小化,即相互「吸引」或「对齐」。这是传统对比学习中的「吸引项」。
- 第二项 log∑jexp(−D(zi,zj)/τ)\log \sum _ { j } { \exp \big ( { - \mathcal { D } \big ( z _ { i } , z _ { j } \big ) / \tau } \big ) }: 这一项鼓励所有样本对 (zi,zj)(z_i, z_j) 之间的距离最大化,即相互「排斥」或「分散」。这是传统对比学习中的「斥力项」。
5.3 构建 Dispersive InfoNCE 损失
正如我们之前所说,扩散模型的回归损失已经承担了「吸引」或「对齐」的作用。因此,为了构建 Dispersive Loss,我们只需要保留对比学习中的斥力项。为了在整个批次(batch)上统一计算,论文将其定义为:
LDisp=logEi,j[exp(−D(zi,zj)/τ)]\mathcal { L } _ { \mathrm { D i s p } } { = } { \log } \mathbb { E } _ { i , j } \left[ { \exp \left( { - \mathcal { D } \left( z _ { i } , z _ { j } \right) / \tau } \right) } \right] \\
这个公式表示,我们计算批次中任意两个样本 ziz_i 和 zjz_j 之间的不相似度,然后取指数、求平均,最后取对数。这个损失函数的目标就是让批次中所有样本的表示尽可能地「分散」开来。
5.4 不相似度函数的选择
在选择不相似度函数 D\mathcal{D} 时,常见的有:
- 负余弦相似度: D(zi,zj)=−zi⊤zj/(‖zi‖‖zj‖)\mathcal { D } ( z _ { i } , z _ { j } ) = - z _ { i } ^ { \top } z _ { j } / ( \| z _ { i } \| \| z _ { j } \| )。通常在对比学习中,会先对 zi,zjz_i, z_j 进行 ℓ2\ell_2 归一化,然后使用余弦相似度。
- 平方 ℓ2\ell _ { 2 } 距离: D(zi,zj)=‖zi−zj‖22\mathcal { D } ( z _ { i } , z _ { j } ) = \| z _ { i } - { z } _ { j } \| _ { 2 } ^ { 2 }。
论文在实验中发现了一个有趣的现象:当使用平方 ℓ2\ell _ { 2 } 距离,并且不进行 ℓ2\ell_2 归一化时,InfoNCE-based Dispersive Loss 的效果最好。这与传统对比学习中常使用归一化余弦相似度的经验是相反的。
为什么会这样?
- 传统对比学习的目的: 是在超球面上学习判别性特征。归一化确保所有特征向量都在单位球面上,余弦相似度衡量它们在球面上方向的差异。
- Dispersive Loss 的目的: 是最大化表示的分散性。当不进行 ℓ2\ell_2 归一化时,特征向量的范数(长度) 是被允许自由变化的。
如果模型为了最小化 Dispersive Loss,它会倾向于让特征向量的范数变得非常大。因为向量越长,它们之间的 ℓ2\ell_2 距离就越容易变得非常大,从而实现更强的「分散」效果。
这种设计提供了一种额外的「推力」,不仅仅是让向量的方向不同,更是让它们在整个高维空间中「扩散」开来,从而实现更强大的正则化。下图也证实了这一点:应用 Dispersive Loss 后,模型内部特征的范数普遍增大了。
5.5 训练流程
Dispersive Loss 作为一个附加项「即插即用」地集成到标准的扩散模型训练中:
六、实验效果
6.1 实验设置概览
- 数据集: 主要在业界标准的大规模图像数据集 ImageNet 256×256256 \times 256 分辨率上进行。
- 基线模型: 选择了当前最先进且广泛使用的扩散模型架构——DiT (Diffusion Transformers) 和 SiT (Scalable interpolant Transformers) 。这些模型都是基于 Transformer 架构,代表了扩散模型的前沿水平。
- 评估指标: 主要使用 FID-50k (Fréchet Inception Distance)。FID 是衡量生成图像质量和多样性的常用指标,值越低表示生成图像越真实、多样性越好。也报告了 Inception Score (IS),越高越好。
- 训练细节: 默认在 VAE 编码的潜空间上训练,使用 Heun 采样器,训练 80 或更多个 Epoch。
6.2 核心发现
主题 | 关键发现 & 数据 | 意义/结论 |
---|---|---|
对比学习对比 | 传统对比学习对噪声极其敏感,FID 恶化 14-19%;Dispersive Loss 始终提升,InfoNCE-ℓ₂降 11.35% | 「无正样本对」策略更适合扩散模型 |
Dispersive Loss 变体 | 所有变体均优于基线,InfoNCE-ℓ₂最强 | 「分散表示」机制普适且有效 |
正则化层位置与涟漪效应 | 仅在 Block 3 加损失,全模型多层表示范数同步提升 | 效果可辐射全网,适用不同层 |
超参数鲁棒性 | λ∈[0.25,1.0]、τ∈[0.25,2.0] 均有效,对τ极为不敏感 | 易于调参,应用省心 |
规模与模型通用性 | DiT/SiT 全尺寸均提升,模型越大提升越明显 | 大模型更受益,适应性强 |
与 REPA 对比 | REPA 需超大预训练和数据,Dispersive Loss 无需额外资源,FID 1.97(优于大多数方法) | 极简高效,资源友好,实用性强 |
一步生成模型通用性 | MeanFlow-XL/2 FID 3.21(原版为 3.43) | 不挑架构,生成步数通用 |
七、总结
7.1 主要贡献
- 正则化新视角: 强调了对生成模型内部表示进行显式正则化的重要性,尤其对大型模型而言,这可能成为未来提升生成质量、泛化能力和训练稳定性的关键方向。
- 新自监督学习范式: 「无正样本对」的对比学习思想可能推广到其他自监督学习场景,尤其是在数据本身具有某种「对齐」或「回归目标」的领域。
- 跨模态融合: 鼓励研究者进一步探索表示学习和生成模型之间更深层次的协同效应,例如在文本到图像、文本到视频等跨模态生成任务中,如何优化内部表示以获得更好的语义一致性和生成质量。
7.2 局限与未来
不足/挑战 | 具体描述 | 未来方向/改进建议 |
---|---|---|
理论解释深度 | 缺少对为何ℓ₂距离效果更好等问题的深入理论分析,目前多为实验观察 | 从信息论、流形学习等角度推导理论基础,指导设计更优正则方式 |
超参数自动化 | λ、τ 虽鲁棒但仍需人工设定,遇新任务需重新调优 | 发展自适应调整机制,如元学习、梯度自调优等 |
评估指标局限 | FID 等客观指标不全面,不能完全反映人类主观感知和多样性等 | 引入主观评价、更多能捕捉语义和结构质量的多维评估指标 |
计算开销 | InfoNCE 需 O(N²) 距离计算,大批次/高维时计算瓶颈可能显现 | 采样负样本、近似搜索等降复杂度方法 |
模态/任务泛化 | 现集中于图像生成,跨模态/其他生成任务有效性尚待验证 | 拓展到文本、音频、视频等生成及跨模态任务 |
与扩散模型深度集成 | 目前为附加正则项,未与扩散原理深度结合 | 探索与去噪、采样等过程联合优化,简化模型结构 |
作者:tomsheep
链接:https://www.zhihu.com/question/1916513058118939910/answer/1916828390629025489
来源:知乎