CVPR2024|Diffusion模型轻量化与计算效率优化

前言

做算法应该都有顶会梦吧，发不了顶会只能刷一刷顶会了哈哈，向顶会大佬学习
扩散模型的训练和推理都需要巨大的计算成本（显卡不足做DDPM的下游任务实在是太难受了），所以本文整理汇总了部分CVPR2024中关于扩散模型的轻量化与计算效率优化 的相关论文。

---

1、Diffusion Models Without Attention

Author：Jing Nathan Yan, Jiatao Gu, Alexander M. Rush
paper：https://arxiv.org/pdf/2311.18257

虽然去噪概率扩散模型（DDPMs）在图像生成方面取得了显著的进展，但在高分辨率应用中面临巨大的计算挑战，特别是依赖于自注意力机制导致计算复杂度呈二次方增长（训练过DDPM的朋友应该都对显存有着巨大渴望吧）。目前大部分解决办法都是通过加快Unet和Transformer架构中的进程，但是这样会牺牲模型的生成能力。

为了解决此问题，本文提出了扩散状态空间模型（DIFFUSSM），这是一种无需注意力机制的扩散架构，使用门控状态空间模型（SSM）作为扩散过程中的骨干网络。DIFFUSSM通过避免全局压缩来有效处理更高分辨率的图像，从而在整个扩散过程中保留详细的图像表示。

DIFFUSSM利用了长范围SSM核心和沙漏式前馈网络的交替层，不使用U-Nets或Transformers中的patchification或长范围块的缩放。采用了门控双向SSM作为核心组件，并通过hourglass架构提高效率。

论文在ImageNet和LSUN数据集上进行了评估，结果表明DIFFUSSM在FID和Inception Score指标上与现有带注意力模块的扩散模型相当或更优，同时显著减少了总的FLOP使用。

2、Fixed Point Diffusion Models

Author：Xingjian Bai, Luke Melas-Kyriazi
paper：https://arxiv.org/pdf/2401.08741
code：https://github.com/lukemelas/fixed-point-diffusion-models

本文提出了一种定点扩散模型（FPDM），在扩散模型中集成了固定点求解概念，通过在去噪网络中嵌入隐式固定点求解层，将扩散过程转化为一系列相关的固定点问题。该方法结合一种新的随机训练方法，显著减少了模型的尺寸和内存使用，并加快了训练速度。

此外，FPDM还开发了两种新技术来提高采样效率：在时间步长之间重新分配计算资源和在时间步长之间重复使用固定点解。实验证明，该方法在ImageNet、FFHQ、CelebA-HQ和LSUN-Church等数据集上显著提高了性能和效率。与最先进的DiT模型相比，FPDM参数减少了87%，训练时内存消耗减少了60%，在采样计算或时间受限的情况下，图像生成质量也得到了改善。

3、Towards More Accurate Diffusion Model Acceleration with A Timestep Tuner

Author：Mengfei Xia, Yujun Shen, Changsong Lei, Yu Zhou, Ran Yi, Deli Zhao, Wenping Wang, Yong-jin Liu
paper：https://arxiv.org/pdf/2310.09469
code：https://github.com/THU-LYJ-Lab/time-tuner

由于许多扩散模型都需要上千个去噪步骤来生成图像，所以这些模型的推理速度都非常慢，许多对应的改进方法就是减少其中的去噪步骤，但会引起较大的性能退化。作者认为生成图像的质量下降是由将不准确的积分方向应用于时间步长间隔引起的，并提出了时间步长调谐器（Timestep tuner），能够以最小的成本为特定间隔找到更准确的积分方向，即在每个去噪步骤通过在新的时间步长上调节网络来替换原始参数化，强制采样分布朝向真实分布。

实验表明TimeTuner在无条件生成、高阶采样器生成、标签条件生成和文本条件生成等任务中均表现出显著的性能提升。在极端函数评估次数情况下，TimeTuner也显示出对一致性蒸馏方法的有效性。

4、Attention-Driven Training-Free Efficiency Enhancement of Diffusion Models

Author：Hongjie Wang, Difan Liu, Yan Kang, Yijun Li, Zhe Lin, Niraj K. Jha, Yuchen Liu
paper：https://arxiv.org/pdf/2405.05252

扩散模型在生成高质量和多样化的图像方面表现出卓越的性能，但这种卓越的性能是以昂贵的计算为代价的，特别是大量使用注意力模块会导致计算资源需求的急剧增加。作者提出了注意力驱动的免训练高效扩散模型（AT-EDM）框架，该框架利用注意力图（attention maps）在运行时对冗余的token进行剪枝（pruning），从而加速DMs的推理过程，无需重新训练。

该论文

• 提出了一种新的基于图的算法，Generalized Weighted Page
  Rank（GWPR），用于识别冗余的token。
• 提出了一种基于相似度的恢复方法，用于在卷积操作中恢复被剪枝的token。
• 提出了一种Denoising-Steps-Aware Pruning（DSAP）方法，用于在不同的去噪时间步中调整剪枝预算，以提高生成质量。

AT-EDM在效率方面与先前的艺术作品相比表现出色，例如在FLOPs节省了38.8%，并且在速度上比Stable Diffusion XL快了1.53倍，同时保持了与完整模型几乎相同的FID和CLIP分数。

5、Improving Training Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architecture

Author：Huijie Zhang, Yifu Lu, Ismail Alkhouri, Saiprasad Ravishankar, Dogyoon Song, Qing Qu
paper：https://arxiv.org/pdf/2312.09181

扩散模型显著的性能受到缓慢的训练与采样的阻碍，作者认为这是由于需要跟踪广泛的正向和反向扩散轨迹，并采用跨多个时间步长（即噪声水平）具有大量参数的大型模型，所以作者提出了一个受经验发现启发的多阶段框架。该框架通过将时间间隔分割成多个阶段，并在每个阶段使用定制的多解码器U-Net架构，结合了时间依赖模型和通用共享编码器。

该论文中的工作有：

• 识别了导致训练扩散模型效率低下的两个关键因素：模型容量需求的显著变化和梯度的不一致性
• 提出了一个新的多阶段架构，该架构通过将时间间隔分割成多个阶段，并在每个阶段使用定制的多解码器U-Net架构，有效提高了训练和采样效率
• 在CIFAR-10和CelebA数据集上进行了广泛的数值实验，证明了所提出框架的有效性。

6、DeepCache: Accelerating Diffusion Models for Free

Author：Xinyin Ma, Gongfan Fang, Xinchao Wang
paper：https://arxiv.org/pdf/2312.00858
code：https://github.com/horseee/DeepCache

扩散模型由于其显著的生成能力，在图像合成领域获得了前所未有的关注。但这些模型通常会产生大量的计算成本，主要归因于连续的去噪过程和冗余的模型尺寸。而传统压缩扩散模型的方法通常涉及广泛的重新训练，这在成本和可行性上都存在挑战。
作者提出了一个简单有效的加速算法DeepCache，在运行时动态压缩扩散模型，提高图像生成速度，无需额外训练负担。具体而言，DeepCache利用在扩散模型的连续去噪步骤中观察到的固有时间冗余，跨相邻去噪阶段缓存和检索特征，从而减少冗余计算。利用U-Net的特性，重用高级特征，同时以非常廉价的方式更新低级特征。反过来，这种创新策略使Stable Diffusion v1.5的加速因子达到2.3倍，CLIP Score仅下降0.05，LDM-4-G的加速因子达到4.1倍，ImageNet上的FID仅下降0.22。此外，在相同的吞吐量下，DeepCache有效地实现了与DDIM或PLMS相当甚至略有改善的结果。

7、Accelerating Diffusion Sampling with Optimized Time Steps

Author：Shuchen Xue, Zhaoqiang Liu, Fei Chen, Shifeng Zhang, Tianyang Hu, Enze Xie, Zhenguo Li
paper：https://arxiv.org/pdf/2402.17376v1

扩散模型大量的采样步骤导致其效率低下。用于扩散模型的高阶数值常微分方程（ODE）求解器使得能够以更少的采样步骤生成高质量的图像，但大多数采样方法仍然采用统一的时间步长，这种方法对于少量采样步的情况是不理想的。为了解决这个问题，论文作者提出了一个设计优化问题的通用框架，为扩散模型的特定数值ODE求解器寻求更合适的时间步长。该优化问题旨在最小化ODE的真实解和对应于数值求解器的近似解之间的距离。优化问题可以使用约束信任域方法在不到15秒内解决。
该方法可以高效地结合最新的采样方法UniPC，显著提高图像生成性能。该方法也是可以作为一种即插即用的方式，与各种预训练的扩散模型结合使用，实现最先进的采样性能。

---

总结

本文是对2024CVPR发表的关于扩散模型效率优化的几篇论文的整理汇总，如有不足欢迎指正。

参考链接：https://github.com/52CV/CVPR-2024-Papers?tab=readme-ov-file#8