仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
机器之心报道
编辑:Panda
在建模复杂的数据分布方面,扩散生成模型表现出色,不过它的成果大体上与表征学习(representationlearning)领域关联不大。
通常来说,扩散模型的训练目标包含一个专注于重构(例如去噪)的回归项,但缺乏为生成学习到的表征的显式正则化项。这种图像生成范式与图像识别范式差异明显——过去十年来,图像识别领域的核心主题和驱动力一直是表征学习。
在表征学习领域,自监督学习常被用于学习适用于各种下游任务的通用表征。在这些方法中,对比学习提供了一个概念简单但有效的框架,可从样本对中学习表征。
直观地讲,这些方法会鼓励相似的样本对(正例对)之间相互吸引,而相异的样本对(负例对)之间相互排斥。研究已经证明,通过对比学习进行表征学习,可以有效地解决多种识别任务,包括分类、检测和分割。然而,还没有人探索过这些学习范式在生成模型中的有效性。
鉴于表征学习在生成模型中的潜力,谢赛宁团队提出了表征对齐(REPA)。该方法可以利用预训练得到的现成表征模型的能力。在训练生成模型的同时,该方法会鼓励其内部表征与外部预训练表征之间对齐。有关REPA的更多介绍可阅读我们之前的报道《扩散模型训练方法一直错了!谢赛宁:Representationmatters》。
REPA这项开创性的成果揭示了表征学习在生成模型中的重要性;然而,它的已有实例依赖于额外的预训练、额外的模型参数以及对外部数据的访问。
简而言之,REPA比较麻烦,要真正让基于表征的生成模型实用,必需一种独立且极简的方法。
这一次,MIT本科生RunqianWang与超70万引用的何恺明出手了。他们共同提出了DispersiveLoss,可译为「分散损失」。这是一种灵活且通用的即插即用正则化器,可将自监督学习集成到基于扩散的生成模型中。
论文标题:DiffuseandDisperse:ImageGenerationwithRepresentationRegularization
论文链接:https ://arxiv. org/abs/2506. 09027v1
分散损失的核心思想其实很简单:除了模型输出的标准回归损失之外,再引入了一个用于正则化模型的内部表征的目标(图1)。
直觉上看,分散损失会鼓励内部表征在隐藏空间中散开,类似于对比学习中的排斥效应。同时,原始的回归损失(去噪)则自然地充当了对齐机制,从而无需像对比学习那样手动定义正例对。
一言以蔽之:分散损失的行为类似于「没有正例对的对比损失」。
因此,与对比学习不同,它既不需要双视图采样、专门的数据增强,也不需要额外的编码器。训练流程完全可以遵循基于扩散的模型(及基于流的对应模型)中使用的标准做法,唯一的区别在于增加了一个开销可忽略不计的正则化损失。
与REPA机制相比,这种新方法无需预训练、无需额外的模型参数,也无需外部数据。凭借其独立且极简的设计,该方法清晰地证明:表征学习无需依赖外部信息源也可助益生成式建模。
带点数学的方法详解
分散损失
新方法的核心是通过鼓励生成模型的内部表征在隐藏空间中的分散来对其进行正则化。这里,将基于扩散的模型中的原始回归损失称为扩散损失(diffusionloss),将新引入的正则化项称为分散损失(DispersiveLoss)。
如果令X={x_i}为有噪声图像x_i构成的一批数据,则该数据批次的目标函数为:
其中,L_Diff(x_i)是一个样本的标准扩散损失,L_Disp(X)则是依赖于整个批次的分散损失项,λ是其加权项。
在实践中,该团队没有应用任何额外的层(如,投射头),而是直接将分散损失应用于中间表示,不增加额外的可学习参数。
该方法是自成一体且极简的。具体而言,它不会改变原始L_Diff项的实现:它不引入额外的采样视图,也不引入额外的数据增强,并且当λ为零时,它刚好就能约简为基线扩散模型。
这种设计之所以可行,是因为引入的分散损失L_Disp(X)仅依赖于同一输入批次中已经计算出的中间表示。这不同于标准对比学习——在标准对比学习中,额外的增强和视图可能会干扰每个样本的回归目标。
前面也说过,分散损失的行为类似于「没有正例对的对比损失」。在生成模型的背景下,这个公式是合理的,因为回归项提供了预先定义的训练目标,从而无需使用「正例对」。这与先前关于自监督学习的研究《Understandingcontrastiverepresentationlearningthroughalignmentanduniformityonthehypersphere》一致,其中正例项被解释为对齐目标,而负例项则被解释为正则化的形式。通过消除对正例对的需求,损失项可以定义在任何标准批次的(独立)图像上。
从概念上讲,可以通过适当移除正例项,从任何现有的对比损失中推导出分散损失。就此而言,「分散损失」一词并非指特定的实现,而是指一类鼓励实现分散的通用目标。下文将介绍分散损失函数的几种变体。
基于InfoNCE的分散损失变体
在自监督学习中,InfoNCE是被广泛使用且有效的对比损失变体。作为案例研究,该团队提出了与InfoNCE损失相对应的分散损失。
数学形式上,令z_i=f(x_i)表示输入样本x_i的生成模型的中间表示,其中f表示用于计算中间表示的层的子集。原始InfoNCE损失可以被解读为分类交叉熵目标,它会鼓励让正例对之间具有高相似度,而负样本对之间具有低相似度:
其中,表示一对正例(例如,通过对同一幅图像进行数据增强获得的数据),(z_i,z_j)表示包含正例对和所有负例对(即i≠j)的任意一对样本。D表示相异度函数(例如,距离),τ是一个称为温度的超参数。D的一个常用形式是负余弦相似度:
在(2)式的对数中,分子仅涉及正例对,而分母包含批次中的所有样本对。根据之前的一些研究,可以将公式(2)等效地重写为:
其中,第一项类似于回归目标,它最小化z_i与其目标之间的距离。另一方面,第二项则会鼓励任何一对(z_i,z_j)尽可能距离拉远。
为了构造对应的分散损失,这里只保留第二项:
该公式也可以被视为一种对比损失(公式(3)),其中每个正例对由两个相同的视图组成,使得为一个常数。等式(4)就等价于
只差一个常数项log(batchsize),而这个常数项不会影响优化过程。从概念上讲,此损失定义基于参考样本z_i。为了得到定义在一批样本Z={z_i}上的形式,这里按照之前的研究可将其重新定义为:
此损失函数对于批次内的所有样本具有相同的值,并且每个批次仅计算一次。在该团队的实验中,除了余弦相异度之外,我们还研究了平方ℓ₂距离:。使用这种ℓ₂形式时,只需几行代码即可轻松计算出分散损失,如算法1所示。
等式(6)中定义的基于InfoNCE的分散损失类似于前述先前关于自监督学习的论文中的均匀性损失(尽管这里没有对表示进行ℓ₂正则化)。在那篇论文中的对比表示学习,均匀性损失被应用于输出表示,并且必须与对齐损失(即正则项)配对。而这里的新公式则更进一步,移除了中间表示上的对齐项,从而仅关注正则化视角。
该团队注意到,当j=i时,就不需要明确排除项D(z_i,z_j)。由于不会在一个批次中使用同一图像的多个视图,因此该项始终对应于一个恒定且最小的差异度,例如在ℓ₂的情况下为0,在余弦情况下为-1。因此,当批次大小足够大时,这个项在那个对数中的作用是充当一个常数偏差,其贡献会变小。在实践中,无需排除该项,这也简化了实现。
分散损失的其他变体
分散损失的概念可以自然延伸到InfoNCE之外的一类对比损失函数。
任何鼓励排斥负例的目标都可以被视为分散目标,并实例化为分散损失的一种变体。基于其他类型的对比损失函数,该团队构建了另外两种变体。表1总结了所有三种变体,并比较了对比损失函数和分散损失函数。
铰链损失(HingeLoss)
在对比学习的经典公式中,损失函数定义为独立损失项之和,每个损失项对应一个正例对或负例对。正例对的损失项为;负例对的损失项公式化为平方铰链损失,即,其中ε>0为边界值。为了构造分散损失函数,只需舍弃正例对的损失项,仅计算负例对的损失项即可。见表1第2行。
协方差损失(CovarianceLoss)
另一类(广义)对比损失函数作用于表征的互协方差矩阵。这类损失函数可鼓励互协方差矩阵接近单位矩阵。
在这里的分散损失中,该团队只考虑了非对角线元素Cov_mn。由于不使用增强视图,因此互协方差就简化为基于单视图批次计算的协方差矩阵。在这种情况下,当表征经过ℓ₂正则化后,对角线元素Cov_mm自动等于1,因此无需在损失函数中显式地处理。最终的分散损失为。见表1第3行。
使用分散损失的扩散模型
如表1所示,所有分散损失的变体都比其对应的分散损失更简洁。更重要的是,所有分散损失函数都适用于单视图批次,这样就无需进行多视图数据增强。因此,分散损失可以在现有的生成模型中充当即插即用的正则化器,而无需修改回归损失的实现。
在实践中,引入分散损失只需进行少量调整:
指定应用正则化器的中间层;
计算该层的分散损失并将其添加到原始扩散损失中。
算法2给出了训练伪代码,其中包含算法1中定义的分散损失的具体形式。
该团队表示:「我们相信,这种简化可极大地促进我们方法的实际应用,使其能够应用于各种生成模型。」
分散损失的实际表现如何?
表2比较了分散损失的不同变体及相应的对比损失。
可以看到,在使用独立噪声时,对比损失在所有研究案例中均未能提高生成质量。该团队猜想对齐两个噪声水平差异很大的视图会损害学习效果。
而分散损失的表现总是比相应的对比损失好,而前者还避免了双视图采样带来的复杂性。
而在不同的变体中,采用ℓ₂距离的InfoNCE表现最佳。因此,在其它实验中,该团队默认使用基于ℓ₂的InfoNCE。
另外,该团队还研究了不同模块选择以及不同λ(控制正则化强度)和τ(InfoNCE中的温度)值的影响。详见原论文。
总体而言,这种趋势有力地证明了分散损失的主要作用在于正则化。由于规模更大、性能更强的模型更容易过拟合,因此有效的正则化往往会使它们受益更多。
图5展示了SiT-XL/2模型生成的一些示例图像。
当然,该团队也将新方法与REPA进行了比较。新方法的正则化器直接作用于模型的内部表示,而REPA会将其与外部模型的表示对齐。因此,为了公平起见,应同时考虑额外的计算开销和外部信息源,如表6所示。
REPA依赖于一个预训练的DINOv2模型,该模型本身是从已在1. 42亿张精选图像上训练过的11B参数主干网络中蒸馏出来的。
相比之下,新提出的方法完全不需要这些:无需预训练、外部数据和额外的模型参数。新方法在将训练扩展到更大的模型和数据集时非常适用,并且该团队预计在这种情况下正则化效果会非常好。
最后,新提出的方法可以直接泛化用于基于一步式扩散的生成模型。
在表7(左)中,该团队将分散损失应用于最新的MeanFlow模型,然后观察到了稳定持续的改进。表7(右)将这些结果与最新的一步扩散/基于流的模型进行了比较,表明新方法可增强MeanFlow的性能并达到了新的SOTA。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道:liyazhou@jiqizhixin. com