训练扩散模型其实比你想象的更简单！何恺明团队新作Dispersive Loss：给扩散模型加正则化
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
↑点击蓝字关注极市平台
作者丨科技猛兽
编辑丨极市平台
极市导读
无需做额外预训练、无需额外模型参数、无需额外数据，即插即用的正则化技术超简单实现性能提升。>>加入极市CV技术交流群，走在计算机视觉的最前沿
训练扩散模型可能比你想象的更简单。
MIT何恺明老师团队最近的工作提出：
当我们训练扩散模型的时候，可以使用表征学习的LossFunction辅助。
本文的DispersiveLoss就是一种即插即用的正则化技术，用以增强扩散模型。
DispersiveLoss有点类似于Contrastive的自监督学习，鼓励内部的表征分散在隐空间中。
与RepresentationAlignment(REPA)方法相比，DispersiveLoss方法是一种极简实现：
无需做额外预训练。
无需额外模型参数。
无需额外数据。
DispersiveLoss也在多种常用模型上得到了验证，对性能有一致性地改进。
下面是对本文的详细介绍。
论文名称：DiffuseandDisperse:ImageGenerationwithRepresentationRegularization
论文地址：
https ://arxiv. org/pdf/2506. 09027
扩散模型作为一种生成式模型，可以建模复杂的数据分布。另一方面，扩散模型的发展与表征学习的进展脱节了。扩散模型的训练目标通常包括一个regressionterm(比如denoising)，却没有一个对表征(用于generation)的regularizationterm。
在表征学习(RepresentationLearning)领域，自监督学习在学习通用表征方面取得了很大进展。ContrastiveLearning[ 1]提供了一个简单有效的框架，从样本对学习表征。这些方法鼓励\“正样本\“(比较相似的样本)之间互相吸引，\“负样本\“(比较不相似的样本)之间互相排斥。使用ContrastiveLearning的方法做RepresentationLearning，已被证明在各种识别任务中很有用，包括分类、检测和分割。但是，这种学习范式对生成式建模的有效性仍是一个未被充分探索的问题。
对于这个问题，REPA[ 2]提出表征对齐(RepresentationAlignment)方法，利用预训练、现成的表征模型的能力做对齐。REPA训练生成模型时，鼓励其内部表征与外部预训练模型的表征对齐。REPA的缺点也很明显：依赖于额外的预训练、额外的模型参数以及外部数据。
虽然REPA在实践中产生了可观的收益，但其依赖于额外的预训练开销。更具体地说，在外部信息源上，很难讲清楚REPA所带来的提升是来自自监督的训练目标，还是主要来自对外部数据的计算和访问。基于以上，开发一种自给自足的极简方法，使得在做生成式建模时可以利用表征学习，仍是生成领域的一个重要的研究方向。
DispersiveLoss是一种灵活，即插即用的正则化项，把自监督学习引入到了diffusion-based的生成式模型中。DispersiveLoss的核心思想如图1所示，很简单，除了正常的回归训练目标之外，再额外地引入一个训练目标来规范模型的内部表征。
直观地讲，DispersiveLoss会鼓励内部表征在hiddenspace中发散开，与ContrastiveLearning中的\“排斥\“效果很像。
正式来讲，定义为一个batch的噪声图片，这个batch的训练目标可以写成：
式中，是单个sample的标准Diffusionloss，是依赖整个Batch的DispersiveLoss。
在实践中，不应用任何额外的层(例如projectionhead)，并且DispersiveLoss直接应用于中间表征，无需额外的可学习参数。
这个方法不会引入额外的采样，额外的数据增强，且权重超参数为0时，1式直接退化为DiffusionLoss。
因此从概念上讲，DispersiveLoss可以从任何现有的对比损失中，通过去除positive的项推导出来。从这个意义上讲，\“DispersiveLoss\“并非是指特定的实现，而是指鼓励分散的一般性的训练目标。作者在下面介绍了几种变体。
InfoNCE[ 3]是自监督学习中contrastiveloss的一种广泛使用的变体。
令表示输入样本在生成模型中的中间表征，其中表示模型计算中间层特征的那部分。原始的InfoNCELoss可以解释为一个分类交叉嫡目标，鼓励正对之间的相似性高，负对之间的相似性低：
式中，表示正样本对（同一张图片做数据增强得到），代表任何一对样本，包括正对和所有负对。代表距离函数，常用负的余弦相似度
式2的分子只涉及正样本对，分母包括batch中的所有对。上式2可被重写为：
式中，第一项类似于regressionobjective它最小化与其目标之间的距离。第二项鼓励任何一对尽可能远。
为了构建DispersiveLoss，只保留第二项：
式4可以写成：
其中，为batchsize。是定值，可以去掉，不影响优化。
从概念上讲，这种损失定义基于参考样本。为了有一个定义在一批样本上的形式，式5可以重新定义为：
式6对于batch中的所有样本具有相同的值，每个batch只计算一次。在本文实验中，除了余弦相似度的距离度量之外，作者还研究了距离：。当使用这种形式时，DispersiveLoss可以很容易地用几行代码计算，如Algorithm1所示。
注意，在式6中做优化时候，我们没必要明确排除当时的这一项。因为在每个batch中不会使用同一图像的多个视图，因此该项总是定值：
形式下。
余弦相似度情况下。
因此，该项可以视为一个固定bias，不影响式6的优化结果。因此在实践中，没有必要排除这一项，也简化了实现。
本文引入了两种基于其他类型的对比损失函数变体。下图2总结了3个变体并比较了contrastive和dispersive的对应关系。
损失函数被定义为独立损失项的和，每个损失项对应一个正负对。正对的损失项为，负对的损失项被表述为squaredhingeloss，即，其中是margin。这里在计算DispersiveLoss的时候，只计算负对的项。
将交叉协方差表示为Cov，元素由索引。损失鼓励对角元素为1（使用损失函数实现），鼓励非对角元素为0（使用损失函数实现）。
在DispersiveLoss中，只考虑非对角元素。在这种情况下，当表征被归一化后，对角元素自动变为1，因此loss函数中无需显式地定义，所以最后DispersiveLoss就是。
DispersiveLoss可以即插即用地放在生成式模型中，作为一种正则化器来使用，无需修改原有的回归损失。
Algorithm2给出了使用DispersiveLoss的方法，只需要选定需要计算的中间层，并且把DispersiveLoss加到原来的DiffusionLoss即可。
模型：DiT，SiT
VAEtokenizer生成32×32×4latentspace。
使用具有250步的基于ODE的HeunSampler做采样。
训练80Epoch的模型，并且没有(CFG)。
默认情况下，权重\lambda和温度\tau设置为0. 5。
ContrastiveLoss就是把正样本对的Loss也算进去，创建正样本对时候使用twoview。DispersiveLoss就不计算正样本对的。
图3显示，当使用独立噪声时，所有情况下，ContrastiveLoss都未能提高生成质量。作者认为是将两个视图与完全不同的噪声水平会损害学习。因为当把noise做限制时，如图3所示，会略有提升(3/4的情况下)。这些实验表明，对比学习对数据增强的选择是很敏感的。虽然对比学习会有一点点帮助，但是引入额外的view，以及twoviews直接的耦合可能限制其应用。
在这些变体中，距离的InfoNCE表现最好：将FID大幅提高4. 14，或相对11. 35%。这与自监督学习的常见做法形成对比，自监督学习通常首选余弦相似度。
作者注意到，在计算InfoNCELoss之前，没有对表征应用归一化，因此，两个样本之间的距离可以任意大。作者假设这种设计可以鼓励表征更加分散，从而导致更强的正则化。默认在其他实验中使用基于距离的InfoNCE。
图4研究了DispersiveLoss在不同层(即TransformerBlock)的影响。总体而言，在所研究的所有案例中，正则化器都大大优于baseline，显示了方法的通用性。将分散损失应用于所有Block会产生最好的结果，同时将其应用于任何单个Block的性能也几乎一样好。
为了仔细观察，图5显示了模型表征的范数，其中DispersiveLoss仅应用于Block3。值得注意的是，正则化器在Block3上产生更大的表征范数，并将这种影响传播到其他所有的Block。要知道在其他的Block并没有直接使用DispersiveLoss，但是还是会被影响了。
这有助于解释图4中观察到的现象，即无论在哪里使用DispersiveLoss，都会有增益。在其他实验中，作者把DispersiveLoss应用于第1/4位置的Block。
Loss权重控制正则化的强度。作者探索了不同Loss权重和温度的影响，结果如图6所示。所有的配置都比Baseline(FID36. 49)有所改进，进一步表明正则化器有效。
有趣的是，可以观察到当Baseline更强时，改进都更大。
对于每个特定的模型大小，SiT相对和绝对改进都大于DiT。SiT比DiT更强。
与B或S大小模型相比，L大小模型表现出更大的相对改进。
总体而言，这一趋势提供了强有力的证据，证明DispersiveLoss的主要作用在于正则化。更大的模型往往更趋向于过拟合，因此往往会从正则化中受益更多。
图8中作者训练了SiT-XL/2。在w/和w/oCFG的情况下训练模型更多Epoch。研究了基于ODE和基于SDE的Sampler。总体而言，DispersiveLoss在所有设置中都被证明是有益的，即使在Baseline变得更强时。该模型生成的一些示例图像如图9所示。
虽然本文的正则化器直接对模型的内部表示进行操作，但REPA将它们与来自外部模型的那些对齐。
因此，为了公平比较，应考虑额外的计算开销和外部信息源，如图10所示。
相比之下，本文方法是完全自给自足，无需预训练，无需外部数据，也无需额外的模型参数。当将训练扩展到更大的模型和数据集时，本文方法很适用。
本文方法可以直接推广到基于一步扩散的生成模型。在图11(左)中，作者将DispersiveLoss应用于最近的MeanFlow模型[ 4]，并观察到一致的改进。图11(右)将这些结果与最新的基于One-stepDiffusion-based模型或者Flow-based模型进行了比较。结果表明，本文方法增强了MeanFlow。
参考
Momentumcontrastforunsupervisedvisualrepresentationlearning
Representationalignmentforgeneration:Trainingdiffusiontransformersiseasierthanyouthink
Representationlearningwithcontrastivepredictivecoding
Meanflowsforone-stepgenerativemodeling
公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT
极市干货
技术专栏：多模态大模型超详细解读专栏｜搞懂Tranformer系列｜大视觉模型(LVM)解读｜扩散模型系列｜极市直播
技术综述：小目标检测那点事｜大模型面试八股含答案｜万字长文！人体姿态估计(HPE)入门教程
#极市平台签约作者#
科技猛兽
知乎：科技猛兽
清华大学自动化系19级硕士
研究领域：AI边缘计算(EfficientAIwithTinyResource)：专注模型压缩，搜索，量化，加速，加法网络，以及它们与其他任务的结合，更好地服务于端侧设备。
作品精选
搞懂VisionTransformer原理和代码，看这篇技术综述就够了
用Pytorch轻松实现28个视觉Transformer，开源库timm了解一下！（附代码解读）
轻量高效！清华智能计算实验室开源基于PyTorch的视频(图片)去模糊框架SimDeblur
投稿方式：
添加小编微信Fengcall（微信号：fengcall19），备注：姓名-投稿
△长按添加极市平台小编
觉得有用麻烦给个在看啦~
