CVPR 2024 Highlight | 北航等发布「时间特征维护」:无需训练,极致压缩加速Diffusion

共 3360字,需浏览 7分钟

 ·

2024-06-18 11:59



  新智元报道  

编辑:LRST 好困
【新智元导读】拯救4bit扩散模型精度,仅需时间特征维护——以超低精度量化技术重塑图像内容生成!

近日,北航、莫纳什、UT Austin联合推出了时间特征维护的扩散模型低精度无损量化方法TFMQ-DM。

不仅以4bit的权重大小实现了目前无损条件下最极限的扩散模型训练后压缩,同时还实现了超过2.38倍真实硬件加速。
这一发现再次将Diffusion压缩推向全新的高度。
目前,该工作已被CVPR 2024高分接收,并被接收为Highlight Poster (Top 2.8%)。

论文地址:https://arxiv.org/pdf/2311.16503
代码地址:https://github.com/ModelTC/TFMQ-DM


扩散模型由于引入了时间变量t,于是在模型中注入了时序信息,这正是扩散模型不同于以往传统视觉模型的一大特征。


同时该变量也通过将时序特征融入模型去控制去噪过程。研究人员首次定义了时间特征emb,同时发现现有量化算法对于这些特征将产生严重扰动,从而破坏图片生成质量:


1. 时间特征扰动:研究人员发现量化导致了明显的时间特征误差,并将这种特征错误内的扰动现象称为时间特征扰动;


2. 时间信息失配:时间特征扰动改变了原始嵌入的时间信息。具体来说,旨在对应于时间步长。然而,由于存在显著的误差,量化模型的不再准确地与相关联,倾向于与对应的时间特征更为接近,这导致了该方法所说的时间信息不匹配;


3. 去噪轨迹偏离:时间信息不匹配传递了错误的时间信息,因此导致图像在去噪轨迹中对应的时间位置发生了偏差,最终导致图片不再按原轨迹去噪:



(Diffusion中的时间特征扰动)


扰动诱因分析


研究人员发现该扰动主要由以下两个原因造成:


1. 不合适的重建对象:已有量化重建方法并未直接优化时间特征,同时时间特征将会受到有限的校准数据影响产生过拟合现象,如下图Prev所示,其中Freeze代表冻结相关量化参数;



2. 忽略了时间特征相关模块中的有限激活:由于输入是有限整数,因此产生时间特征的模块将仅产生有限且随时间变化的激活,而已有量化策略均考虑分布层级优化,忽略了对于此类有限激活的拟合近似。


(Diffusion中时间特征相关模块有限激活)


时间特征维护­——TFMQ-DM


基于以上的诱因分析研究人员提出了如下时间特征维护策略,在低bit量化下完美的保证了Diffusion模型的时间特征精度与图像生成质量。


(TFMQDM整体压缩框架)


1. 时间信息块:将时间特征生成相关模块进行整合得到时间信息块(见框架图):



2. 时间信息感知重建:基于时间信息块,研究人员提出了时间信息感知重建 (TIAR) 来应对第一个诱因。在重构过程中,该块的优化目标如下:



3. 有限集校准:为了解决第二个诱因中有限集内激活范围宽泛的挑战,研究人员提出了有限集校准 (FSC) 用于激活量化。这个策略为所有时间信息块内的每个激活使用T组量化参数,例如激活x的量化参数可为。在时间步长为t时,x的量化函数可以表示为:



其中sT, zt分别为量化缩放因子和零偏移。


实验结果


研究团队在DDIM,LDM以及Stable Diffusion系列模型上验证了无条件生成、分类条件生成、文本条件以及多种先进采样器生成下TFMQ-DM框架性能。


TFMQ-DM在平均4或8 bit权重,8或32 bit激活时,所有评价指标均超过Q-Diffusion,PTQD等已有最先进方法,且在大部分场景下,该方案使用4bit权重量化超越了已有技术在8-bit权重甚至于全精度权重的模型的性能。


结果表明,TFMQ-DM率先在4bit权重、8bit权重下实现了对于 Diffusion的无损压缩加速,推动了Diffusion无训练量化的边界。


(LDM系列无条件生成对比结果)


(左:LDM系列分类条件生成对比结果;右:DDIM系列无条件生成对比结果)


具体来说,在 CelebA-HQ 256 × 256 数据集上,与当前最先进的方法相比,该团队的方法在 w4a8 设置下将FID与sFID分别降低了 6.71和 6.60(越低越好)。


值得注意的是,现有方法,无论是4 bit还是8bit权重量化,在人脸数据集如 CelebA-HQ 256 × 256 和 FFHQ 256 × 256 上与 FP 模型相比都显示出明显的性能下降,而 TFMQ-DM 与全精度模型相比几乎没有性能损失。


(Stable Diffusion系列文本条件生成对比结果)


(左:PLMS系列先进采样器无条件生成对比结果右:DPM++系列先进采样器无条件生成对比结果)


此外,对于当下最流行的文本条件生成类模型Stable Diffusion,TFMQ-DM在 w8a8 设置下的 FID 和在 w4a32 设置下的 sFID 甚至略低于全精度模型。


然而,虽然现有的指标无法充分评估生成图像的语义一致性以及物体细节,该团队提出的方法产生了更高质量的图像(见后文),具有更真实的细节,更好地展示了语义信息。


由于现有指标并不能完全反映生成图像效果的优劣,该团队研究人员提供了大量可视化效果对比图,体现出更加细腻和准确的生成质量:


(LDM上无条件图像生成效果图)


(LDM上分类条件图像生成效果图)


(Stable Diffusion上文本条件图像生成效果图,左提示词:“A digital illustration of the Babel tower, detailed, trending in artstation, fantasy vivid colors”。右提示词:“A beautiful castle beside a waterfall in the woods”。)


除测试集精度和可视化效果外,团队人员还在Intel® Xeon® Gold 6248R 处理器上验证框架的推理加速效果,相比原始浮点模型实现了 2.38 倍的显著速度提升。


(Stable Diffusion在CPU上真实加速)


总结


基于时间特征维护的校准量化可确保 Diffusion的量化参数准确保留原始时间信息。广泛的实验证明,TFMQ-DM 在 DDIM、LDM及Stable-Diffusion 系列中实现了令人信服的精度提升,尤其是在w4a8等极低比特设置下表现出了明显优势;同时该方案也实现了真实硬件部署加速。


其次TFMQ-DM 具有良好的兼容性,可与各种Diffusion量化框架无缝集成。总的来说,TFMQ-DM 提供的显著量化精度提升与其对于硬件的高度友好,有助于在资源受限的情况下进行实际部署,进一步促进了文生图模型的更广泛普及和应用。


TFMQ-DM 具有显著的兼容性,可与各种Diffusion量化框架无缝集成。TFMQ-DM 提供的显著量化精度,有助于在资源受限的情况下进行实际部署。


参考资料:
https://arxiv.org/pdf/2311.16503




浏览 20
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报