cvpr 2024｜注意力驱动的无需重训练的高效扩散模型增强

欢迎大家来到IT世界,在知识的湖畔探索吧!

欢迎大家来到IT世界,在知识的湖畔探索吧!

Attention-Driven Training-Free Efficiency Enhancement of Diffusion Models

研究背景

扩散模型（Diffusion Models, DMs）在生成高质量和多样化的图像方面表现出色，但其卓越的性能是以巨大的计算成本为代价的。尤其是在文本到图像的生成任务中，即使是最先进的潜在扩散模型（Latent Diffusion Models, LDMs），其推理过程仍然相对缓慢，这对于在资源受限的设备上部署这些模型构成了挑战。因此，提高扩散模型的效率成为了一个重要的研究方向。

研究意义

本研究旨在提出一种无需重训练即可提升扩散模型效率的方法，这对于加速扩散模型的推理过程、降低计算成本具有重要意义。通过提高模型效率，可以使扩散模型更加适用于移动设备和边缘计算场景，推动计算机视觉和生成式AI技术的广泛应用。

文献综述

现有提高扩散模型效率的方法主要分为两类：高效采样策略和高效模型架构。高效采样策略通过减少所需的去噪步骤来加速生成过程，但无法降低内存占用。而高效模型架构则通过减少每个步骤的成本来实现加速，但通常需要对模型进行重训练，这既耗时又昂贵。本研究关注于训练后的模型压缩，特别是利用注意力图来指导模型压缩，这是一个相对未被充分探索的领域。

具体方法

本文提出了注意力驱动的无需重训练的高效扩散模型（Attention-driven Training-free Efficient Diffusion Model, AT-EDM）框架。该框架通过以下两个关键步骤来提升扩散模型的效率：

单步去噪令牌修剪‌：

利用注意力图来识别不重要的令牌，并通过广义加权页面排名（Generalized Weighted Page Rank, G-WPR）算法为每个令牌分配重要性分数。

根据重要性分数生成修剪掩码，并应用掩码来修剪令牌，从而实现注意力块的加速。

为了解决修剪后的令牌与后续卷积操作不兼容的问题，提出了一种基于相似性的令牌复制方法来恢复修剪的令牌。

去噪步骤感知的修剪（Denoising-Steps-Aware Pruning, DSAP）‌：

分析不同去噪步骤中注意力图的方差，发现早期步骤中的注意力图较为混乱且信息量较少，因此提出在早期步骤中减少令牌修剪量，以提高生成图像的质量。

设计了一种去噪步骤感知的修剪计划，通过在不同去噪步骤中调整修剪预算来优化生成质量。

图1

• 标题：Examples of applying AT-EDM to SD-XL
• 内容：展示了应用AT-EDM框架到Stable Diffusion XL（SD-XL）模型上的效果。对比了原始全尺寸模型（上排）和加速后的模型（下排）在生成图像时的浮点运算数（FLOPs）。加速后的模型在FLOPs上有约40%的减少，同时保持了相近的生成质量。

图2

• 标题：U-Net FLOPs breakdown of SD-XL
• 内容：分析了SD-XL模型中U-Net组件的FLOPs分布。指出注意力块（Attention blocks）是生成图像时的主要工作负载，占据了大部分的计算资源。

图3

• 标题：Overview of our proposed framework AT-EDM
• 内容：概述了AT-EDM框架的两个主要部分：单去噪步骤的令牌剪枝方案（Token pruning scheme in a single denoising step）和去噪步骤感知的剪枝调度（Denoising-Steps-Aware Pruning schedule）。详细描述了从获取注意力图到令牌剪枝，再到恢复剪枝令牌以及跨去噪步骤的剪枝调度等步骤。

图4

• 标题：Our similarity-based copy method for token recovering
• 内容：展示了相似性基于的令牌复制方法，用于解决剪枝后令牌与ResNet层不兼容的问题。通过找到与剪枝令牌最相似的保留令牌来恢复剪枝令牌，从而保持图像质量。

图5

• 标题：Variance of attention maps in different denoising steps
• 内容：分析了在不同去噪步骤中注意力图的方差变化。将去噪步骤分为四个典型区域，并指出早期步骤的注意力图方差较小且增长迅速，而后期步骤的方差较大且相对稳定。这一观察结果支持了去噪步骤感知的剪枝调度设计。

图6: AT-EDM与ToMe的生成图像对比

• 描述：
• 比较了原始Stable Diffusion XL（SD-XL）模型、Token Merging（ToMe）方法以及提出的AT-EDM方法在相同FLOPs预算下的生成图像质量。
• AT-EDM在保持相近FID和CLIP分数的同时，实现了显著的FLOPs节省（约40%）。
• 相较于ToMe，AT-EDM生成的图像具有更清晰的物体、更锐利的细节以及更好的文本-图像对齐。
• 关键点：
• 效率：AT-EDM在减少计算量的同时保持了高质量的图像生成。
• 图像质量：AT-EDM生成的图像在物体保留、细节清晰度以及文本-图像对齐方面优于ToMe。

图7: FID-CLIP分数曲线

• 描述：
• 展示了不同Classifier-Free Guidance（CFG）尺度下，SD-XL、ToMe和AT-EDM的FID和CLIP分数。
• AT-EDM在大多数CFG尺度下，不仅降低了FID分数（提高图像质量），还提高了CLIP分数（改善文本-图像对齐）。
• 关键点：
• FID分数：衡量生成图像与真实图像分布的差异，分数越低表示图像质量越高。
• CLIP分数：衡量生成图像与文本提示之间的语义相似性，分数越高表示对齐越好。

表1展示了在不同浮点运算数（FLOPs）预算下，部署Token Merging（ToMe）和本文提出的AT-EDM方法在Stable Diffusion XL（SD-XL）模型上的性能对比。主要指标包括FID分数、CLIP分数和每个去噪步骤的平均FLOPs。

• FID分数：用于评估生成图像的质量，分数越低表示图像质量越好。从表中可以看出，在相同的FLOPs预算下，AT-EDM方法获得的FID分数明显低于ToMe方法，表明AT-EDM生成的图像质量更高。
• CLIP分数：用于评估生成图像与文本提示的对齐程度，分数越高表示对齐越好。AT-EDM方法在大多数FLOPs预算下都获得了更高的CLIP分数，说明其生成的图像与文本提示更加匹配。
• FLOPs：表示每个去噪步骤的平均浮点运算数，用于衡量计算成本。随着FLOPs预算的减少，两种方法的计算成本都相应降低，但AT-EDM在保持较低FLOPs的同时，仍能保持较好的图像质量和文本对齐性。

表2比较了SD-XL、ToMe和AT-EDM在不同实现情况下的采样延迟。

• w/o FO：表示在不使用Fused Operation（FO）库的情况下，各模型的采样延迟。AT-EDM在这种情况下的延迟略低于ToMe，但两者都显著高于SD-XL。
• w/FO under CI：表示在使用当前实现的FO库（不提供注意力图作为中间结果）的情况下，各模型的采样延迟。由于需要额外获取注意力图，AT-EDM的延迟有所增加，但仍低于ToMe。
• w/FO under DI：表示在理想情况下（FO库提供注意力图作为中间结果），各模型的采样延迟。在这种情况下，AT-EDM的延迟显著降低，甚至低于SD-XL，展示了其在实际应用中的潜力。

图8: G-WPR算法的不同实现对比

• 描述：
• 比较了基于自注意力（SA-based WPR）和基于交叉注意力（CA-based WPR）的G-WPR算法在图像生成和剪枝掩码上的差异。
• SA-based WPR生成的图像质量优于CA-based WPR，因为后者可能会移除过多的背景tokens，导致背景难以恢复。
• 关键点：
• 自注意力 vs 交叉注意力：自注意力在实现上更适合于保持图像质量，因为它能更准确地识别不重要的tokens。
• 剪枝掩码：展示了不同算法下哪些tokens被剪枝，反映了算法在识别冗余信息上的差异。

图9: 不同token恢复方法的对比

• 描述：
• 比较了零填充、双三次插值、直接复制和相似性复制四种方法在恢复剪枝tokens后的图像质量。
• 相似性复制方法通过利用注意力图选择最相似的保留tokens来填充剪枝位置，显著提高了图像质量。
• 关键点：
• 零填充与插值：这些方法通常会导致图像质量的明显下降，因为它们不能有效地恢复剪枝tokens的信息。
• 相似性复制：通过利用注意力图中的相关性信息，能够更准确地恢复剪枝tokens，从而保持图像质量。

图10: 不同数量prune-less步骤的对比

• 描述：
• 探索了在早期denoising步骤中设置不同数量的prune-less步骤对生成图像质量的影响。
• 当设置15个早期的prune-less步骤时，生成的图像质量最佳。
• 关键点：
• prune-less步骤：在早期denoising步骤中不进行剪枝，以保留对图像质量贡献较大的计算。
• 最优步骤数：通过消融实验发现，设置15个prune-less步骤能够在减少计算量和保持图像质量之间取得最佳平衡。