CVPR 2025 扩散模型炼出新绝技！注意力蒸馏技术：图像生成效果全面碾压传统方法

论文链接：https://arxiv.org/pdf/2502.20235

项目链接：https://xugao97.github.io/AttentionDistillation/

亮点直击

• 分析了之前即插即用注意力特征方法的局限性，并提出了一种新颖的注意力蒸馏损失，用于重现参考图像的视觉特征，取得了显著优越的结果。
• 开发了attention distillation引导采样，这是一种改进的分类器引导方法，将注意力蒸馏损失整合到去噪过程中，大大加快了合成速度，并支持广泛的视觉特征转移和合成应用。

生成效果一览

给定参考图，文生图

艺术风格迁移

纹理合成

总结速览

解决的问题

• 现有生成扩散模型在图像风格和语义理解方面虽然有进展，但在将参考图像的视觉特征转移到生成图像中时，使用即插即用注意力特征的方法存在局限性。

提出的方案

• 提出了一种新颖的注意力蒸馏损失，用于在理想和当前风格化结果之间计算损失，并在隐空间中通过反向传播优化合成图像。
• 开发了一种改进的分类器引导方法，即注意力蒸馏引导采样，将注意力蒸馏损失整合到去噪采样过程中。

应用的技术

• 利用预训练扩散网络的自注意力特征。
• 通过注意力蒸馏损失优化生成图像。
• 将注意力蒸馏损失整合到去噪采样过程中，以改进分类器引导方法。

达到的效果

• 大大加快了合成速度。
• 支持广泛的视觉特征转移和合成应用。
• 在将示例的风格、外观和纹理转移到新图像的合成中表现出色。

方法

预备知识

隐空间扩散模型（LDM），如Stable Diffusion，由于其对复杂数据分布的强大建模能力，在图像生成方面达到了最先进的性能。在LDM中，首先使用预训练的VAE 将图像 压缩到一个学习到的隐空间中。随后，基于UNet的去噪网络 被训练用于在扩散过程中预测噪声，通过最小化预测噪声与实际添加噪声 之间的均方误差来实现。

其中 表示条件， 表示时间步长。去噪 UNet 通常由一系列卷积块和自注意力／交叉注意力模块组成，所有这些都集成在残差架构的预测分支中。

KV注入在图像编辑 [2, 7, 58]、风格迁移 [10, 25, 29] 和纹理合成 [70] 中被广泛使用。它建立在自注意力机制之上，并将扩散模型中的自注意力特征用作即插即用的属性。自注意力机制的公式为：

在注意力机制的核心，是基于查询 和键（ ）之间的相似性计算权重矩阵，该矩阵用于对值（ ）进行加权聚合。KV注入通过在不同的合成分支之间复制或共享KV特征来扩展这一机制。其关键假设是 特征代表图像的视觉外观。在采样过程中，将合成分支中的 特征替换为示例的相应时间步长的 特征，可以实现从源图像到合成目标的外观转移。

注意力蒸馏损失

尽管KV注入取得了显著的效果，但由于残差机制的影响，它在保留参考的风格或纹理细节方面表现不足；例如，参见下图3（a）。KV注入仅作用于残差，这意味着信息流（红色箭头）随后受到恒等连接的影响，导致信息传递不完整。因此，采样输出无法完全再现所需的视觉细节。

本项工作提出了一种新的损失函数，通过在自注意力机制中重新聚合特征来提取视觉元素，因此称之为注意力蒸馏（AD）损失。利用预训练的T2I扩散模型Stable Diffusion的UNet，从自注意力模块中提取图像特征。如下图2（a）所示，首先根据目标分支的 ，从参考分支重新聚合KV特征（ 和 ）的视觉信息，这与KV注入相同。

将此注意力输出视为理想的风格化。然后，我们计算目标分支的注意力输出，并计算相对于理想注意力输出的L1损失，这定义了AD损失：

可以使用提出的AD损失通过梯度下降来优化随机隐空间噪声，从而在输出中实现生动的纹理或风格再现；例如，参见上图3（b）。这归因于优化中的反向传播，它不仅允许信息在（残差）自注意力模块中流动，还通过恒等连接流动。通过持续优化， 和 之间的差距逐渐缩小，使得注意力越来越准确，最终特征被正确聚合以产生所需的视觉细节。

根据最近的实验分析 [7, 29, 58]，经验性地选择UNet的最后6个自注意力层来计算AD损失。此外，在优化过程中，通过线性减少输入到UNet的时间步长 从 到 0来模拟扩散模型的采样过程。从不同的随机隐空间噪声开始，并在100步内优化它们。请注意，在整个优化过程中，UNet预测的噪声完全被丢弃，不断更新相同的隐空间变量。

为了更好地理解我们的AD损失，展示了多次运行的优化结果，如下图4所示。这些结果表明：i）AD损失有效地蒸馏出高质量的风格和纹理视觉特征；ii）AD损失自适应于不同的空间结构，展示了多次运行中的多样性。

内容保留优化

通过AD损失蒸馏的纹理和风格，可以进一步使用内容损失将合成内容与另一个参考图像对齐。这种优化允许图像合成在保留目标内容的同时转换一个图像的视觉元素，实现风格迁移、外观迁移等任务。

如上图2（b）所示，定义内容损失与AD损失类似，也基于自注意力机制，充分利用扩散模型中对图像的深刻理解。特别是，计算目标查询 和参考查询 之间的L1损失，构成了内容损失：

在实现中，同样选择最后6个自注意力层来计算内容损失，这与AD损失保持一致。内容保留优化的目标是：

优化后，经过预训练的VAE将优化后的隐空间code解码为图像空间。

注意力蒸馏引导采样

上述合成是使用反向传播优化的。在本节中，介绍如何将注意力蒸馏损失以改进的分类器引导方式纳入扩散模型的采样过程中。

根据[13]，分类器引导在去噪过程中改变去噪方向，从而生成来自的样本，其公式可以表示为：

其中， 是时间步长， 表示提示， 和 分别指去噪网络和LDM中的隐空间变量。 控制引导强度。受 ［17］启发，使用基于注意力蒸馏损失的能量函数来引导扩散采样过程。

具体来说，在DDIM采样［55］过程中，时间步长 的隐空间变量 根据从去噪网络估计的更新方向 被转换为 ：

其中， 。将方程（6）中的 替换为 ：

其中， 是与 相关的常数。通过用 替换为 来定义能量函数。然而，我们发现配置引导强度是一个棘手的问题；详见下文的验分析。为了解决这个问题，我们引入了Adam优化器［35］来自动管理强度并计算梯度。为简化起见，将 近似为 ，从而能够对 进行初始DDIM采样，然后进行简单的优化更新。AD损失以隐空间的 和噪声扰动的参考隐空间 作为输入，如上图2（c）所示。通过AD损失优化更新 为：

通过AD损失的引导，可以在时间步条件下计算隐空间的损失，而不是将隐空间转换为图像空间并计算图像级损失，如最近的工作所做的那样 [4, 41, 49]。使用Adam优化器能够建立一个通用的学习率，从而减轻设定引导强度的挑战。请注意，方程 (4) 中提出的内容损失也可以与AD损失一起添加到采样过程中，以进一步保留内容参考图像的结构。

改进的VAE解码

实验证据 表明，隐空间扩散模型中使用的VAE在感知上是有损的。对于需要高频局部细节的任务，例如纹理合成，可以选择在示例图像 上微调VAE解码器 的权重 ，以使用 L 1 损失来增强其重建质量，按照［3］的方法：

其中，表示VAE编码器。下图5展示了一些重建和采样的结果，显示了经过微调的VAE在感知质量上的提升。

实验

应用和对比

接下来，将注意力蒸馏损失应用于各种视觉特征转移任务，并将结果与每个应用中的最新方法进行比较。

风格和外观转移。 遵循Gatys等人 [22] 的著名工作精神，通过前文中描述的优化方法实现风格和外观转移。方法与CSGO 、StyleShot、StyleID、StyTR2和NST进行风格转移的比较，以及与Cross-Image Attention和SpliceViT进行外观转移的比较。下图6展示了定性比较结果。在风格转移中，本文方法有效地捕捉到高质量、一致的风格特征，同时保留了内容图像的语义结构。这在第3和第4行的素描风格中尤为明显。相比之下，尽管基线方法保留了原始结构，但在风格上存在显著差异。在外观转移中，本文的方法也表现出优越性，避免了Cross-image Attention中颜色过饱和的问题。

特定风格的文本到图像生成。 如前文所述，可以在扩散采样中应用我们的AD损失，从而实现特定风格的文本到图像生成。将参考图像设置为所需的风格图像。下图7展示了一些生成结果，并与其他方法进行比较，包括Visual Style Prompting、InstantStyle和RB-Modulation。下图7中的结果表明，本文的方法在文本语义上与现有方法相当，同时在风格一致性上与参考图像相比具有显著优势。除了上述方法之外，进一步结合ControlNet，以在各种模态（例如深度和Canny边缘）上进行额外条件的特定风格文本到图像生成。下图8展示了一些生成的例子。

受控纹理合成。 本文方法可以应用于纹理优化，如前文所示。受［7］的启发，在计算注意力蒸馏损失时进一步引入了mask引导，从而限制了由 查询的值，实现了受控的纹理合成。具体来说，给定一个源纹理 ，其对应的源分割图 ，以及目标分割图 ，首先将 和 展平，并将它们下采样以匹配注意力特征的分辨率，得到 和 。然后，计算引导 mask ；通过这个 mask，在注意力计算中限制了由 聚合的视觉信息，使其仅关注相应的区域：

为了进一步确保与目标mask的对齐，用从源图像对应标记区域随机抽取的像素填充目标mask作为初始化。然后，将此初始化视为内容参考，并基于查询特征添加内容损失，如前文所述。上图1和下图9展示了我们的受控纹理合成结果。

与基于补丁的神经纹理优化方法GCD相比，本文结果在保持可比纹理细节的同时展现了更平滑的物体边缘。相比之下，GCD存在颜色混叠的伪影问题；参见图9的第二行。最近，Self-Rectification引入了一种“懒编辑”控制，用于生成非平稳纹理。为了实现相同的目标，利用SDEdit来保留用户编辑的布局图像的结构。然后，将提出的AD损失和内容损失整合到采样过程中。如图9所示，Self-Rectification输出的纹理过渡更为平滑，而我们的结果更好地遵循了原始纹理示例。

纹理扩展。 由于补丁来源有限，使用传统方法合成超高分辨率纹理非常困难。在此，将注意力蒸馏引导的采样应用于MultiDiffusion模型，使纹理扩展到任意分辨率。尽管SD-1.5是在尺寸为512×512的图像上训练的，但令人惊讶的是，当结合注意力蒸馏时，它在大尺寸纹理合成中表现出了强大的能力。下图10展示了将纹理扩展到512×1536的尺寸与GCD和GPDM的比较。本文方法在如此具有挑战性的任务中显示出显著的优势。

消融研究

在本节中，展示了关于我们方法的两个方面的消融研究结果：i) 内容保持优化中内容损失权重的影响，以及ii) 在注意力蒸馏引导采样中管理引导强度的优化器。

内容损失权重。 如下图11所示，变化内容损失权重对转移结果有着有趣的影响。例如，在风格转移中，左侧展示的抽象风格示例说明了的调整如何导致不同程度的抽象化，为艺术创作提供了灵活性。在外观转移中，由于扩散网络的精确语义理解，右侧展示的人脸图像转移随着的变化表现出平滑的身份过渡。

优化器。 下图12展示了优化器在管理引导强度中的重要性。实验测试了一种手动设置引导强度的简单策略，以控制对隐空间变量的梯度更新规模。然而，手动变化这个强度通常无法产生合理的结果：示例中的纹理或外观特征通常会丢失，如下图12的最后三列所示。相比之下，引入Adam优化器来管理隐空间优化产生的结果与输入示例的视觉特征非常匹配（下图12的第2到4列）。此外，在每个采样时间步中增加优化迭代次数通常可以提高生成结果的质量，尽管这也带来了额外的计算时间。在实践中，将迭代次数设置为2，以实现高效平衡，提供高质量的结果。

用户偏好研究

为了验证定性分析，在三个转移任务上进行了包含30个问题的用户研究（每个任务针对6个选定的竞争对手各设置5个问题）。在每个问题中，展示两个结果：一个来自我们的方法，一个来自竞争对手。用户被要求根据提供的说明和标准选择更好的一个。从50名参与者那里收集了1500个回应，整体偏好得分总结在下图13中。我们的方法在很大程度上持续优于其他选项。

结论

本文提出了一种统一的方法来处理各种视觉特征转移任务，包括风格/外观转移、特定风格的图像生成和纹理合成。该方法的关键是一种新颖的注意力蒸馏损失，它计算理想风格化与当前风格化之间的差异，并逐步修改合成。本文方法克服了以往工作的局限性，实验验证了其优越性。

参考文献

[1] Attention Distillation: A Unified Approach to Visual Characteristics Transfer