2025年5月7日星期三

无需外挂!DiT自蒸馏神技SRA横空出世 :偷偷自学的表征对齐术让生成质量暴涨!

点击下方卡片,关注"AI生成未来"

如您有工作需要分享,欢迎联系:aigc_to_future

作者:Dengyang Jiang

解读:AI生成未来

图片

文章:https://arxiv.org/pdf/2505.02831 
代码:https://github.com/vvvvvjdy/SRA  

图片

亮点直击

  • 分析发现:文章深入分析diffusion transformer的潜在表征趋势,发现当只执行生成训练时,其会拥有一个大致从粗糙到细致,从差到好的表征(判别)过程。
  • 方法提出:文章引入SRA方法,通过将早期层高噪声条件下的潜在表征与后期层低噪声条件下的潜在表征对齐,实现自我表征引导,提升diffusion transforme的训练速度和生成表现。
  • 实验验证:文章通过一系列实验验证了SRA对DiTs(diffusion-based)和SiTs(flow-based)有显著性能提升,并与依赖复杂表征训练框架或强大外部表征先验的方法进行对比,证明了SRA的简单有效性。

背景介绍

在生成式预训练中引入表征指导

Diffusion transformer在视觉生成领域展现出了有效的可扩展性和可迁移性。最近的研究发现高质量的内部表征能够加速其生成训练,并提升生成样本的质量然而,现有方法在提供表征指导时,要么引入复杂训练框架,要么依赖大规模预训练模型,这限制了其应用场景和资源效率。

先前表征引入方法的局限性

  • 复杂训练框架:如MaskDiT和SD-DiT采用MAE或IBOT的训练范式,需要设计额外的训练任务和框架,增加了训练的复杂度和资源消耗。
  • 外部表征先验依赖:如REPA利用DINOv2,CLIP等大规模预训练模型作为表征指导模型,虽然能大幅提升性能,但依赖于大量数据和计算资源训练得到的外部先验,限制了方法的灵活性和普适性。
图片

灵感和发现

灵感来源

不同于将干净的图像作为输入的表征模型然后输出语义丰富的特征,扩散模型通常以latent噪声作为输入并一步一步清理出干净图像。换句话说,扩散模型的生成机制操作宏观上可以被认为是一个coarse-to-fine 过程。因此,我们推测,扩散模型中的表征也遵循这一趋势。

图片

发现

文章对普通的SiT和DiT进行了实证分析。我们首先发现其中的潜在特征随着块层数的增加和噪声水平的降低而逐步细化,呈现出大致从粗糙到精细的过程。接下来,我们观察到仅通过生成式预训练的diffusion transformer已经能够学习到有意义的判别性表征。

同时,尽管在大约第20层达到峰值后,由于模型需要转移注意力以生成高频细节的图像,准确率有所下降,但通过增加块层数和降低噪声水平, 表征的质量总体上是从差到好逐渐转变的。

方法(SRA)

SRA的核心思想是利用扩散Transformer自身在不同训练阶段和噪声条件下的表征差异,通过自我蒸馏的方式进行表征对齐。具体来说,SRA将早期层(高噪声条件下)的潜在表征与后期层(低噪声条件下)的潜在表征进行对齐,从而实现自我表征引导。这种方法无需额外的表征组件或外部预训练模型,仅在diffusion transformer自身的生成训练过程中引入指导。

图片

实现细节

  • 表征对齐:将学生模型(可训练模型)早期层的潜在表征与教师模型(通过指数移动平均更新权重的模型)后期层的潜在表征进行对齐。学生模型的输出通过轻量级可训练MLP头进行非线性变换后,与教师模型的输出进行对齐。
  • 损失函数:SRA通过最小化教师模型输出和学生模型输出变体之间的patch-wise距离来实现自我表征对齐。同时,将该目标函数与原始的扩散模型目标函数结合,进行联合学习。
  • EMA教师网络:为了避免直接使用同模型输出作为监督信号导致的训练不稳定性,SRA采用指数移动平均(EMA)策略构建教师模型。教师模型的权重基于学生模型的权重进行加权移动平均更新,从而提供稳定的表征引导。
  • 超参数设置:在实验中,我们对超参数进行了细致的调整,包括表征对齐的块层选择、时间间隔、对齐目标函数以及EMA教师网络的动量系数等。这些超参数的选择对SRA的性能有着显著的影响。

实验

实验设置

实验中,文章采用了ImageNet 1K数据集,并遵循DiT和SiT的训练设置,包括使用AdamW优化器、恒定学习率1e-4、批量大小256等。评估指标包括FID、sFID、IS、Precision和Recall,使用相同的评价框架和参考图像以确保与先前方法的公平比较。

元素级分析

图片

文章通过详细的组件分析,研究了SRA中不同设计选择对性能的影响。实验结果表明:

  • 块层选择:将教师模型的后期层与学生模型的早期层进行对齐能取得最佳性能,因为早期层需要更多的指导以捕捉语义丰富的表征。
  • 时间间隔:使用教师模型比学生模型低噪声输入的特征能提升性能,且时间间隔值在0.1左右时效果最佳。动态时间间隔表现略优于固定时间间隔。
  • 对齐目标函数:smooth-ℓ1、ℓ1和ℓ2三种回归训练目标函数均能有效提升性能。
  • 教师网络更新策略:采用固定动量系数0.9999的EMA策略更新教师网络效果最佳,其他更新策略并不适合SRA。
  • 投影头的作用:使用轻量级投影头对学生的输出进行后处理,能显著提升性能,因为它有助于让模型在进行非线性变换前捕捉更有效的隐藏表征以便对齐。

系统级比较

图片
图片

文章将SRA与最新扩散模型方法进行了系统级比较,结果表明:

  • 性能提升:SRA在不同模型大小和类型的DiT和SiT上均能显著提升性能。特别是在大型模型中,SRA的效果更为显著,这可能是因为大型模型能提供更丰富的表征引导。
  • 与现有方法对比:SRA在400个epoch内取得了优于原始SiT-XL模型的性能,并且在800个epoch时进一步提升。与依赖复杂表征训练框架的方法(如MaskDiT)相比,SRA在性能上有显著优势,并且与依赖强大外部表征先验的方法(如REPA)相当。

消融研究

图片

由于SRA是隐式地引入表征监督,文章的消融研究旨在验证表征在SRA中的重要性。实验结果表明:

  • 表征能力提升:SRA显著提升了diffusion transformer的潜在表征质量,这通过linear probing结果在不同层和时间步上的表现得以验证。
  • 生成质量与表征引导的紧密关联:随着用于对齐的教师网络层的变化,linear probing准确率与FID分数之间显示出强烈的关联性,证明了SRA中生成能力与自我表征引导机制的紧密联系。

结论

研究的核心观点是扩散Transformer自身具备提供表征引导的能力,无需引入额外的表征组件。通过提出SRA方法,证明了在仅进行生成训练的过程中,通过自我表征对齐可以有效提升扩散Transformer的表征学习能力和生成性能。实验结果表明,SRA在多个模型和数据集上均能带来显著的性能提升,并且在与依赖复杂表征训练框架或强大外部表征先验的方法对比中展现出优越性。

未来研究展望

尽管SRA取得了显著的成果,但仍有一些开放问题值得进一步研究:

  • 模型规模和数据扩展:研究SRA在更大模型规模和更多数据集上的可扩展性,特别是在文本到图像等更复杂的生成任务中的应用潜力。
  • 理论洞察:探索表征学习对生成任务有益的理论依据,以期为diffusion model的训练提供更深入的理解和指导。
  • 方法的普适性:验证SRA在其他生成任务中的应用效果,探索其作为通用表征增强方法的潜力。
图片
图片
图片
图片
图片
图片
图片
图片

参考文献

[1] No Other Representation Component Is Needed:  Diffusion Transformers Can Provide Representation Guidance by Themselves

致谢

如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。同时,欢迎给个星标⭐,以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进,共同探索未知,见证一个充满希望和伟大的未来!



技术交流

加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!

图片

没有评论:

发表评论

一键生成电影级分镜,轻松实现人物一致性,多图融合爽到爆,谷歌Gemini强大功能加持,非常适合制作AI广告和电商宣传

不需要本地部署,不需要会员,只需要一段简单的提示词,就可以生成完整的视频故事板,还可以用文字轻松编辑图片,保持人物一致性,多张图片轻松融合到一个场景,制作电影质感的AI视频,呈现多个稳定角色和产品,非常适合制作AI广告和电商宣传。 不需要本地部署,不需要会员,只需要一段简单的...