【深度学习】解决Diffusion Transformers优化困境!南京大学与字节等提出DDT:解耦扩散模型
仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
作者丨科技猛兽编辑丨极市平台
导读
解耦Encoder-DecoderTransformer能否加速收敛并增强样本质量?
一个类似MAR+REPA的技术。
本文认为,这个过程带来了一个优化的困境:Encode的过程需要减少高频分量,而Decode过程又需要增加高频分量,就会带来优化的难度。
DDT的结果:对于ImageNet256×256,DDT-XL/2实现了1.31FID的SOTA性能。对于ImageNet512×512,DDT-XL/2实现了1.28的SOTAFID。
DDT还有一个额外的好处:可以增加推理速度。原因是对不同denosingsteps可以共享condition。还专门提出了一种基于统计动态规划的方法来识别出最佳的共享策略。
专栏目录:https ://zhuanlan.zhihu.com/p/687092760
1DDT:解耦扩散模型(来自南京大学,字节SeedVision)1DDT论文解读1.1DDT研究背景1.2DiffusionTransformer困境:低频语义编码能力有限1.3DDT方法:解耦编解码器1.4采样加速1.5实验结果
论文名称:DDT:DecoupledDiffusionTransformer
论文地址:
http ://arxiv.org/pdf/2504.05741
代码链接:
http ://github.com/MCG-NJU/DDT
图像生成是计算机视觉研究的一项基本任务,其目的是捕捉原始图像数据集的固有数据分布,并通过采样合适的分布来生成高质量的图像。
Diffusion的前向过程按照SDE前向过程逐渐向原始数据添加高斯噪声。去噪过程从这个损坏过程中学习分数估计。一旦准确地学习了得分函数,就可以通过数值求解反向SDE来合成数据样本。
DiffusionTransformer将Transformer架构引入扩散模型中,以取代传统的主流的基于UNet的模型。在给定足够的训练迭代的情况下,即使在不依赖长残差连接的情况下,DiffusionTransformer也优于传统方法。然而,由于成本高,它们的收敛速度慢仍然给开发新的模型带来了巨大的挑战。
本文希望从模型设计的角度解决上述缺点。传统模型常使用Encoder-Decoder架构,本文探索的是解耦Encoder-DecoderTransformer能否解锁加速收敛和增强样本质量,希望回答这个问题:解耦Encoder-DecoderTransformer能否加速收敛并增强样本质量?
本文的结论是:DiffusionTransformer在”提取抽象结构信息”以及”恢复详细细节信息”之间存在优化的困境。此外,由于原始像素监督,DiT在提取语义表征方面受限。
本文分析的主题是Flow-basedModel。代表纯噪声。如图3所示,扩散模型其实相当于是对谱分量Autoregressive的细化[1]。DiffusionTransformer对noisylatent进行编码,以捕获低频语义,然后再解码高频细节。然而,这种语义编码过程不可避免地会衰减高频信息,造成优化困境。
因此,DDT觉得应该将传统的Decoder-Only的扩散模型架构解耦为显式的Encoder-Decoder架构。
引理1:对于一个第步的flow-matchingnoisescheduler,定义为干净数据的最高频率。noisylatent的最大保持频率为:
代表纯噪声。当增加,噪声降低,此时语义编码变得更容易了(因为噪声降低),而解码变得更加困难(因为从低频的输入预测高频输出,这步需要预测的频率差大了)。
考虑第去噪步骤的最坏情况,DiffusionTransformer将频率编码为,直到第步,它需要解码至少的残差频率。如果没能在第步解码此残差频率,就会为后续步骤创造瓶颈。
如果将更多的计算分配给噪声较大的时间步能够获得提升,意味着DiffusionTransformer难以编码低频信息,以提供语义信息。
如果将更多的计算分配给噪声较小的时间步能够获得提升,意味着Flow-matchingTransformer难以解码出更高的频率,来提供更精细的细节。
如下图4所示,证明了与uniformscheduling相比,将更多的计算分配给早期时间步可以提高最终性能。这表明扩散模型面对的挑战,更多的是在比较noisy的时间步。
这就引出了一个关键的结论:DiffusionTransformer从根本上受到低频语义编码能力的限制。因此,本文探索了Encoder-Decoder的架构。
之前的研究MAR中,轻量级的diffusionMLPhead的解码能力受限,因此MAR通过Backbone得到的语义latent做辅助,以完成高质量图像生成。
之前的研究REPA中,通过预训练的visionfoundations的表征对齐,来增强低频编码。
DecoupledDiffusionTransformer包括一个conditionencoder和一个velocitydecoder。
DDT的整个框架是linearflowdiffusion。
conditionencoder的输入包括noisyinput,时间步,以及classlabel。它提取的低频分量self-condition作为velocitydecoder的输入条件。
其结构按照DiT和SiT设计,一系列Attention和FFN的堆叠。
为了在相邻时间步中保持的局部一致性,还采用了REPA的表征对齐技术。如式3所示,该方法将Encoder第层的中间特征与DINOv2表征对齐。是可学习投影MLP:
这种简单的正则化加速了训练收敛,促进了相邻步骤之间的局部一致性。它允许在相邻步骤之间共享编码器的输出。实验表明,这种编码器共享策略显着提高了推理效率,且性能下降可忽略不计。
VelocityDecoder采用与ConditionEncoder相同的架构设计,由堆叠的Attention和FFN组成,类似于DiT/SiT。
为了提高相邻步骤之间condition的一致性,通过AdaLN-Zero把condition引入Decoder特征。
Decoder使用Flow-matching损失进行训练,如式5所示:
通过将显式表示对齐合并到Encoder中,并将隐式Self-condition注入Decoder中,在训练期间跨相邻步骤实现的局部一致性。这使我们能够在适当的局部范围内共享,从而减少Encoder的计算量。
给定总推理步骤和Encoder的计算预算,共享率为,定义为计算Self-condition的时间步,如式6所示。如果当前时间步不在中,将先前计算的重用为。否则,使用Encoder和当前noiselatent重新计算:
一种简单的方法UniformEncoderSharing是每步计算一次。以前的工作,例如DeepCache,使用这种范式来加速UNet模型。
本文还提出一种统计动态规划的方法计算最优集。具体细节不展开了。
首先是ImageNet256×256实验。batchsize是256,使用的VAE是VAE-ft-EMA,下采样8倍,channel为4(SDf8d4)。
架构改善:SwiGLU,RoPE和RMSNorm。
训练加速:lognormsampling。
通过结合这些先进技术开发了改进的基线模型。
这些改进基线的性能如图6所示。为了验证我们实现的可靠性,作者还复现了REPA-B/2的结果,实现了略高于REPA中最初报告的指标。
图6中本文方法的结果,在没有REPA的情况下始终优于之前方法。但在使用REPA时,性能迅速接近饱和点。这在XL模型大小中尤为明显,其中这些增加的技术带来的收益逐渐下降了。
作者在图6中展示了不同大小的模型在400K训练步骤的性能。DDT在各种模型大小上表现出一致且显着的改进。DDT-B/2(8En4De)模型超过Improved-REPA-B/22.8的FID。DDT-XL/2(22En6De)超过REPA-XL/21.3的FID。虽然decoder-only的DiT使用REPA接近性能饱和,但DDT继续提供更好的结果。架构改进以及其他训练技术在模型尺寸增加时的收益递减,但DDT模型保持了显著的性能优势。
图7报告了DDTXL/2(22En6De)和DDT-L/2(20En4De)的最终指标。
DDT展示了卓越的效率,与REPA和其他DiT模型相比,使用近1/4的Epochs即实现收敛。
图8是512×512分辨率的结果。为了验证DDT模型的优越性,作者在256Epoch下在ImageNet256×256上训练的DDT-XL/2作为初始化,在ImageNet512×512上微调DDT-XL/2100K步。本文实现了1.90的FID,比REPA高出0.28。在将更多的训练步骤分配给DDT-XL/2时,在500K步上实现了1.28个FID,时间间隔为CFG3.0[0.3,1.0]。
如图5所示,ConditionEncoder中得到的condition具有很强的局部一致性。即使与的相似度高于0.8。这种一致性提供了一个机会,通过共享相邻步骤之间的Encoder来加速推理。
作者采用uniformencodersharing策略和统计动态规划策略。具体来说,对于uniform策略,每步重新计算。对于统计动态规划,通过动态规划求解相似度矩阵上的最小求和路径,并根据求解的策略重新计算。如图9所示,当小于6时,可以在实现显著推理加速的情况下,几乎没有视觉质量损失。如图10所示,指标损失微不足道,推理却显著加速。新提的统计动态规划策略略微优于原始的uniformencodersharing策略。
作者系统地探索了不同模型大小下从2:1到5:1的Encoder-Decoder比例。图11和图12显示了对架构优化的关键见解。
作者观察到,随着模型尺寸的增加,更大的Encoder有利于进一步提高性能。对于图11中的基本模型,最佳配置显示为8Encoder+4Decoder,提供卓越的性能和收敛速度。值得注意的是,图12中更大的模型表现出明显的偏好,使用20Encoder+4Decoder实现了峰值性能,这个Encoder-Decoder的比例更激进。这种意想不到的发现促使作者将DDT-XL/2比例缩放到22Encoder+6Decoder来探索DiT架构的性能上限。
参考
^Diffusionisspectralautoregression
公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT
极市干货
技术专栏:多模态大模型超详细解读专栏|搞懂Tranformer系列|大视觉模型(LVM)解读|扩散模型系列|极市直播
技术综述:小目标检测那点事|大模型面试八股含答案|万字长文!人体姿态估计(HPE)入门教程
#极市平台签约作者#
科技猛兽
知乎:科技猛兽
清华大学自动化系19级硕士
研究领域:AI边缘计算(EfficientAIwithTinyResource):专注模型压缩,搜索,量化,加速,加法网络,以及它们与其他任务的结合,更好地服务于端侧设备。
作品精选
搞懂VisionTransformer原理和代码,看这篇技术综述就够了
用Pytorch轻松实现28个视觉Transformer,开源库timm了解一下!(附代码解读)
轻量高效!清华智能计算实验室开源基于PyTorch的视频(图片)去模糊框架SimDeblur
投稿方式:
添加小编微信Fengcall(微信号:fengcall19),备注:姓名-投稿
△长按添加极市平台小编
觉得有用麻烦给个在看啦~