【深度学习】无需借助DINOv2！DiT自表征对齐：用自蒸馏取代 REPA，高效提高生成性能
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
作者丨科技猛兽编辑丨极市平台
导读
本文提出了一种新的自表征对齐方法，通过自蒸馏的方式就可获得表征指导，实验结果表明，将SRA应用于DiTs和SiTs会产生一致的性能改进。SRA还实现了与依赖外部表征先验的方法相当的性能。
DiT训练使用自蒸馏(Self-Distillation)取代REPA。
最近关于表征对齐REPA的研究表明：在扩散Transformer中，学习有意义的内部表征既可以加速生成训练，也可以提高扩散Transformer的生成质量。但是，REPA方法需要引入额外的大规模预训练的视觉基础模型比如DINOv2，CLIP等，在DiT的训练过程中提供指导。
SRA将早期层(噪声程度较高)与末期层(噪声程度较低)的输出的latentrepresentation对齐，辅助生成模型的训练。实验结果表明，将SRA应用于DiTs和SiTs会产生一致的性能改进。SRA还实现了与依赖外部表征先验的方法相当的性能。
论文地址：
https ://arxiv. org/pdf/2505. 02831
代码链接：
https ://github. com/vvvvvjdy/SRA
DiffusionTransformer和VisionTransformer因其预训练的可扩展性以及下游任务的泛化能力，在视觉生成领域很常用。最近，许多工作表明：对于DiT而言，学习高质量的内部表征不仅可以加快训练进度，还可以提高生成质量。
这些工作主要有图1的2种范式，在DiT训练期间提供表征的指导：
如图1(a)所示，利用表征学习的方法(例如MAE、IBOT)。这种方法需要复杂的训练框架设计。
如图1(b)所示，利用大规模预训练视觉基础模型(例如DINOv2、CLIP)的表征。这种方法依赖一个强大的先验，该先验是在数千个GPU的海量数据上训练的。
因此，本文提出了一个重要的问题：在训练生成模型的时候，能不能不用额外的组件，获得类似的表征指导？
表征模型通常输入干净的图像，输出语义丰富的特征。而扩散模型通常输入的是噪声，输出的每一次都比上一次更加清晰的图片。换句话说，扩散模型运行的生成机制通常可以被认为是从粗到细的过程。受这种行为的启发，本文假设DiT的表征也遵循这种趋势。为了验证这一点，作者对DiT进行了实证分析。
如图2(a)所示，作者首先发现随着block的增加和noiselevel降低，DiT的latent特征逐渐细化，从粗糙逐渐变得精细。
如图2(b)的ImageNetLinearProbing结果所示，作者观察到扩散转换器已经学习了有意义的判别式的表征。LinearProbing的精度在大约20层达到峰值后下降，因为模型需要偏移以专注于生成具有高频细节的图像。这说明增加block，降低noiselevel，表征质量基本上在慢慢变好。
从以上结果可以得出结论：DiT预训练时，DiT的表征大致从粗到细。且DiT学习到了有意义的判别式表征。
基于这种趋势，本文在生成模型DiT的训练中将DiT中的较弱表征与更好表征对齐，从而增强模型的表征学习，且不使用任何外部组件。
Denoise-based的模型学习通过逐步去噪过程将Gaussiannoise转换为datasample。给定一个逐渐添加噪声的前向过程，这些模型学习反向过程来恢复原始数据。
对于分布的数据点，正向过程如下：。模型学习使用神经网络反转这个过程，该网络预测每一步添加的噪声。该网络使用简单的均方误差目标进行训练，该目标衡量它可以预测噪声的程度：
与Denoise-based的模型不同，Flow-basedmodel学习概率流常微分方程(PFODE)的速度场。PFODE允许模型通过流向数据分布，对数据采样。PFODE的前向过程定义为：
其中，是数据，是高斯噪声，和分别是的单调递减和递增函数。PFODE由下式给出：
其中，该ODE在时间的边际分布与正向过程的边际分布相匹配。为了学习速度场，训练模型以最小化以下损失函数：
如图1(c)所示，SRA不需要任何外部组件；本质上，SRA将早期层的输出潜在表征(噪声程度高)与末期层的输出潜在表征(噪声程度低)对齐。
同时，为了使训练过程更加稳定，SRA从另一个与可训练模型共享相同架构的模型(通过EMA更新权重)获得目标特征。也就是说，学生模型的输出latent特征首先通过投影层，然后与教师输出的目标特征对齐。
SRA可以提供一种灵活的方式来获得表征的指导，且无需外部组件需求和架构修改。
如图3所示是SRA方法的框架。
设是可训练的学生模型，是教师模型。输入noiselatent，timestep，以及condition分别为。
学生编码器latent输出。
教师latent输出：。
其中，是batchsize，numberofpatches和embeddingdimension。表示中第层输出。
SRA中设置：。使用教师输出和学生的输出变换进行自我表征对齐。其中，是学生编码器输出的投影，通过轻量级MLP。
轻量级MLP可以在训练后可选地被丢弃。
SRA通过最小化教师输出和学生输出变体之间的patch－wisedistance来实现Self－alignment：
其中，是一个patchindex，是预定义的距离函数，是学生DiT和投影头的参数。
把生成模型的训练目标与上述Self－alignment结合起来，一起学习：
其中，是超参数。
在SRA中，不需要现成的教师模型给出先验指导。SRA使用学生权重的EMA从学生网络的过去迭代中构建教师模型。EMA：，其中是动量系数。
实验部分关注以下几个问题：
SRA中的每个设计选择和组件如何影响性能？
SRA对于不同baselines和modelsizes是不是work？
SRA与其他使用外部组件(比如使用额外的表征学习范式，或者额外的视觉基础模型)的方法相比，性能如何？
SRA是否真正增强了基线模型的表示能力，并且生成能力确实与表征指导密切相关？
除非另有说明，训练细节严格遵循DiT和SiT中的设置，无权值衰减，BatchSize为256，使用SDVAE提取latent。对于模型配置，使用DiT和SiT论文中引入的B/2、L/2和XL/2架构，处理PatchSize为2的输入。
对于DiT，使用DDPMSampler并将函数评估(NFE)的数量默认设置为250。对于SiT，使用SDEEuler-Maruyama采样器（对于的SDE），默认情况下将NEF设置为250。
下面是对每个组件的影响进行的分析。使用SiT-B/2以及用SRA训练400Kiterations进行评估，如图4所示。
首先，作者分析了使用不同学生和教师模型的层对齐的影响。这里的结论是：观察到使用教师模型相对后期的层，但是不是最后一层(比如第8层)，来监督学生模型相对早期的层(比如第3层)，可以得到最优的结果。作者认为前几层需要更多的指导，因此它们可以捕获语义上有意义的表征以进行后续生成。同时，教师层的表示质量与相应对齐学生的表现之间存在很强的相关性。基于这些结果，将B、L和XL模型的对齐层分别设置为3→8、6→16和8→20。
然后作者研究了用于对齐的Timeinterval(第1. 4节中的k)。这里作者研究了固定和动态间隔。这里的结论是：观察到教师模型相比学生模型，输入的特征噪声更低时，可以提高性能，间隔为0. 1或平均值为0.1是最优的。
作者认为较低的噪声水平可以提供更好的表示指导，但过大的时间间隔会阻碍模型的学习过程，导致只关注优化对齐损失，而忽略了生成方面。动态区间表现出稍好的性能，实验中将时间间隔设置为0-0. 2。
作者比较了3种简单的对齐回归目标，包括smooth－和。发现3个目标都可以带来良好的性能并且在训练期间是稳定的。实验中使用smooth－。
EMA模型通常仅用于评估。然而，由于SRA需要它在训练期间提供指导，因此作者研究了不同的更新方法。
作者研究了几种不同的策略来构建教师模型。首先，发现教师模型如果直接从学生模型复制，会损害性能。接下来，使用自监督学习中使用的策略，在训练过程中将动量系数从0. 996更新为1。然而，发现效果不佳。最后，使用0. 9999的动量系数不变，发现比较适合。
作者观察到使用简单的head来对学生的输出进行后处理，比直接使用学生的输出进行对齐要好得多。
作者认为这种轻量的操作使模型能够使得模型学习到更有效的hiddenrepresentation，来进行最终对齐。这个就不用显式地对齐整个latent特征，因为这些特征可能会破坏每一个layer和timestep负责的原始生成内容。
首先，作者比较了vanillaDiT或SiT与使用SRA训练的模型之间的FID值。如图5所示。
使用SRA训练的DiT在不同类型的每个trainingstep以及各种大小上的性能都有显着提高。此外，与一些SSL工作的观察结果类似，SRA在更大的尺寸模型中的影响更为显著，这可能是因为较大的模型倾向于提供更丰富的指导。而且，即使模型已经获得了较低的FID分数，SRA的优势也不会饱和。
作者认为这可能是由于教师在整个训练过程中的能力也不断提高，允许在训练时为学生提供更好的表征指导。
最后，使用CFG提供了SiT-XL与SRA和其他方法之间的定量比较，如图6所示。400Epoch的SiT+SRA超过了SiT方法。随着训练的延长，还会进一步提升。在800Epoch时，SRA的SiT-XL的FID为1. 58，IS为311. 4。这一结果远远优于依赖于MaskDiT，并且与REPA相当。此外，由于在整个训练过程中，教师模型的质量越来越高，作者发现与REPA相比，SRA方法不太可能遇到饱和。
