ARIA: 基于自适应区域重要性分配的条件扩散蒸馏

# ARIA: 基于区域的自适应重要性分配用于条件扩散蒸馏 来源：https://arxiv.org/html/2606.23898 Loay Mualem<sup>1,2</sup> Vinh Tong<sup>1</sup> Samir Darouich<sup>1,3</sup> Mathias Niepert<sup>1</sup>  
<sup>1</sup>斯图加特大学人工智能研究所 <sup>2</sup>IMPRS-IS <sup>3</sup>斯图加特大学理论化学研究所。如有任何问题，请发送至：[email protected] (https://arxiv.org/html/2606.23898v1/mailto:[email protected])  

###### 摘要  
条件扩散模型的蒸馏旨在将大型教师模型的行为迁移到较小的学生模型，同时保持对不同条件输入的语义对齐。与识别任务不同，条件扩散中的知识蒸馏往往难以将知识迁移到训练分布之外，因为预测的噪声高度依赖于条件信号。因此，有效的蒸馏需要对大规模条件空间进行探索。在实际环境中，这构成了一个主要瓶颈。成对的图像–条件数据可能有限，为每个可用条件生成合成图像通常在计算上不可行，而条件（如文本提示）的池子可能极其庞大。最近的研究通过在训练过程中切换条件来应对这一问题，在不改变蒸馏目标的前提下，使学生模型暴露于更广泛的条件空间。然而，这引出了一个互补的问题：一旦拥有了大规模的条件语料库，训练努力应该如何分配？在本文中，我们提出了ARIA，一个自适应地在粗粒度条件区域之间分配训练努力的框架。通过在线维护区域级别的教师–学生差异估计，ARIA在保持原始蒸馏目标的同时，将更新聚焦于持续存在不对齐的区域。实验上，ARIA在大多数架构和设置下优于随机条件（RC），在未见和代表性不足的场景中取得了最明显的改进。我们还提供了理论分析，表明所提出的追踪机制在有限方差和漂移假设下能够跟随训练过程中的演变差异。

## 1 引言  
扩散模型在多个领域展现出强大的生成性能，包括图像[21 (https://arxiv.org/html/2606.23898#bib.bib2), 51 (https://arxiv.org/html/2606.23898#bib.bib3), 12 (https://arxiv.org/html/2606.23898#bib.bib4), 2 (https://arxiv.org/html/2606.23898#bib.bib7), 42 (https://arxiv.org/html/2606.23898#bib.bib5), 38 (https://arxiv.org/html/2606.23898#bib.bib6)]、视频[20 (https://arxiv.org/html/2606.23898#bib.bib8), 15 (https://arxiv.org/html/2606.23898#bib.bib9), 3 (https://arxiv.org/html/2606.23898#bib.bib10)]、音频[27 (https://arxiv.org/html/2606.23898#bib.bib11), 34 (https://arxiv.org/html/2606.23898#bib.bib12)]和机器人[9 (https://arxiv.org/html/2606.23898#bib.bib13), 6 (https://arxiv.org/html/2606.23898#bib.bib43), 9 (https://arxiv.org/html/2606.23898#bib.bib13)]。这些模型迭代地将噪声转化为目标分布的样本。大规模文本到图像系统如Stable Diffusion[41 (https://arxiv.org/html/2606.23898#bib.bib34)]能够根据自然语言提示生成图像，但其强大性能通常需要大模型和大量采样步骤，这促使研究者开发更高效的变体。知识蒸馏（KD）是一种广泛使用的模型压缩方法，通过训练较小的学生模型模仿教师模型的预测或中间表示[19 (https://arxiv.org/html/2606.23898#bib.bib18), 60 (https://arxiv.org/html/2606.23898#bib.bib19), 63 (https://arxiv.org/html/2606.23898#bib.bib20), 66 (https://arxiv.org/html/2606.23898#bib.bib21), 61 (https://arxiv.org/html/2606.23898#bib.bib49), 36 (https://arxiv.org/html/2606.23898#bib.bib50), 59 (https://arxiv.org/html/2606.23898#bib.bib51), 67 (https://arxiv.org/html/2606.23898#bib.bib52), 44 (https://arxiv.org/html/2606.23898#bib.bib53), 37 (https://arxiv.org/html/2606.23898#bib.bib54), 11 (https://arxiv.org/html/2606.23898#bib.bib55)]。在文本到图像扩散模型中，这通常涉及采样图像–文本对，沿着扩散轨迹生成噪声中间变量，并训练学生模型在每一步匹配教师模型的输出[19 (https://arxiv.org/html/2606.23898#bib.bib18), 60 (https://arxiv.org/html/2606.23898#bib.bib19), 63 (https://arxiv.org/html/2606.23898#bib.bib20), 66 (https://arxiv.org/html/2606.23898#bib.bib21)]或特征表示[60 (https://arxiv.org/html/2606.23898#bib.bib19), 63 (https://arxiv.org/html/2606.23898#bib.bib20), 66 (https://arxiv.org/html/2606.23898#bib.bib21), 8 (https://arxiv.org/html/2606.23898#bib.bib22), 30 (https://arxiv.org/html/2606.23898#bib.bib23)]。虽然KD能够在识别模型中将知识迁移到训练分布之外[8 (https://arxiv.org/html/2606.23898#bib.bib22), 30 (https://arxiv.org/html/2606.23898#bib.bib23)]，但在条件扩散模型中这种效果较弱[26 (https://arxiv.org/html/2606.23898#bib.bib17)]，因为预测的噪声高度依赖于条件。因此，有效蒸馏教师模型通常需要探索条件空间的大部分。在实践中，由于版权、隐私和许可限制，大规模成对图像–文本数据集的获取往往受限。使用教师模型生成合成图像是一种可能的替代方案，但大规模扩散采样计算成本高且存储密集[67 (https://arxiv.org/html/2606.23898#bib.bib52), 44 (https://arxiv.org/html/2606.23898#bib.bib53)]。相比之下，文本数据是丰富的：大量提示和标注可以以极低成本获得。这造成了根本性的不平衡——虽然条件空间（文本）可以非常庞大，但可用于蒸馏的图像数量通常受到计算资源的限制。为了解决这种不匹配，Kim等人[26 (https://arxiv.org/html/2606.23898#bib.bib17)]提出了随机条件（RC），它将噪声图像与随机采样的文本条件配对，从而无需为每个提示生成图像即可探索条件空间。RC回答了一个重要的可行性问题：当成对图像–提示数据有限时，辅助文本能否改进扩散蒸馏？我们的工作提出了一个互补的分配问题：一旦拥有了大规模的辅助文本池，训练预算应如何在其上分配？静态覆盖策略将辅助条件空间视为固定资源并均匀采样。然而，一个区域的有用性依赖于学生模型的状态，并且随着训练而变化——有些区域学习得很快，而其他区域则持续存在较大的教师–学生差距。为此，我们提出了ARIA（基于区域的自适应重要性分配），它不仅利用辅助文本来扩展覆盖范围，还提供自适应反馈。ARIA将辅助条件分组为粗粒度区域，在线追踪其教师–学生差异，并将采样重新分配给当前最需要监督的区域。  
参考图注  
图1：ARIA、RC和教师模型在三种设置下生成图像的并排比较：训练时排除动物、车辆或食物；使用25%的训练检查点；以及在整个数据集上完全训练。

具体来说，ARIA是RC中条件选择步骤的轻量级替代方案。成对图像缓存、教师–学生目标和模型架构保持不变，而辅助提示采样器变得差异感知。我们在有限成对图像和丰富辅助文本的文本条件扩散模型蒸馏中实例化了ARIA。ARIA保持底层蒸馏目标不变，仅修改条件的采样方式。我们提供了区域级差异追踪机制的理论分析，并证明ARIA在大多数架构和数据设置下优于随机条件，在未见和代表性不足的场景中取得了最强的改进，且收敛更快。

**贡献。** 我们的主要贡献如下：(1) 我们提出了ARIA，一个基于区域的自适应重要性分配框架，优先处理差异较大的粗粒度条件区域。与样本级重要性采样不同，ARIA在粗粒度条件区域级别追踪差异，从而无需逐样本评分即可实现对极大规模输入空间的可扩展自适应分配。(2) 我们在有限成对图像和丰富辅助文本的文本到图像扩散蒸馏中实例化了ARIA，在随机采样基础上实现了一致的改进，特别是对于未见提示和不平衡的条件分布。(3) 我们进行了广泛的实证评估，展示了稳定的训练行为以及跨架构和区域构建的鲁棒性，包括SD 1.4、SD 2.1和SDXL教师模型，以及块剪枝、通道剪枝和KOALA学生模型。(4) 我们提供了ARIA基于EMA的区域评分规则的理论分析，建立了在噪声和时间漂移下追踪演化差异的有限时间保证。

**组织。** 第2节 (https://arxiv.org/html/2606.23898#S2) 回顾相关工作，第3节 (https://arxiv.org/html/2606.23898#S3) 形式化蒸馏设置和RC基线。第4节 (https://arxiv.org/html/2606.23898#S4) 介绍ARIA及其算法公式和追踪分析。第5节 (https://arxiv.org/html/2606.23898#S5) 将ARIA应用于两种实验设置下的文本到图像蒸馏，展示性能改进和鲁棒性。我们在第6节 (https://arxiv.org/html/2606.23898#S6) 进行讨论并展望未来方向。

## 2 相关工作

**压缩扩散模型。** 现代扩散模型通常规模较大，在训练和推理时需要大量的计算和内存资源。多项工作旨在通过模型压缩来降低扩散模型的计算和内存成本。现有方法包括量化[48 (https://arxiv.org/html/2606.23898#bib.bib33)]、架构优化[32 (https://arxiv.org/html/2606.23898#bib.bib37)]、剪枝[14 (https://arxiv.org/html/2606.23898#bib.bib44), 28 (https://arxiv.org/html/2606.23898#bib.bib45), 68 (https://arxiv.org/html/2606.23898#bib.bib46), 58 (https://arxiv.org/html/2606.23898#bib.bib47), 7 (https://arxiv.org/html/2606.23898#bib.bib48)]和知识蒸馏[57 (https://arxiv.org/html/2606.23898#bib.bib30), 26 (https://arxiv.org/html/2606.23898#bib.bib17), 25 (https://arxiv.org/html/2606.23898#bib.bib1)]。例如，BK-SDM[25 (https://arxiv.org/html/2606.23898#bib.bib1)]将知识蒸馏与块剪枝相结合来压缩Stable Diffusion[41 (https://arxiv.org/html/2606.23898#bib.bib34)]，而KOALA[29 (https://arxiv.org/html/2606.23898#bib.bib35)]对SDXL[39 (https://arxiv.org/html/2606.23898#bib.bib36)]进行逐层压缩和知识蒸馏。

**用于模型压缩的知识蒸馏。** 知识蒸馏（KD）[19 (https://arxiv.org/html/2606.23898#bib.bib18)]被广泛用于通过从较大教师模型迁移软预测或中间表示来训练紧凑模型[60 (https://arxiv.org/html/2606.23898#bib.bib19), 63 (https://arxiv.org/html/2606.23898#bib.bib20), 66 (https://arxiv.org/html/2606.23898#bib.bib21), 8 (https://arxiv.org/html/2606.23898#bib.bib22), 30 (https://arxiv.org/html/2606.23898#bib.bib23)]。KD已成功应用于多个领域，包括语言模型、视觉变换器和扩散模型[52 (https://arxiv.org/html/2606.23898#bib.bib26), 23 (https://arxiv.org/html/2606.23898#bib.bib27), 16 (https://arxiv.org/html/2606.23898#bib.bib28), 55 (https://arxiv.org/html/2606.23898#bib.bib29), 25 (https://arxiv.org/html/2606.23898#bib.bib1), 26 (https://arxiv.org/html/2606.23898#bib.bib17), 57 (https://arxiv.org/html/2606.23898#bib.bib30)]。在扩散模型中，KD常用于训练压缩架构。然而，Kim等人[26 (https://arxiv.org/html/2606.23898#bib.bib17)]表明，传统的蒸馏策略可能难以将知识迁移到代表性不足或未覆盖的概念，这促使在训练过程中改进数据分配策略。

**用于减少采样步数的知识蒸馏。** 多项工作专注于在不重新训练的情况下加速扩散模型的去噪过程[35 (https://arxiv.org/html/2606.23898#bib.bib58), 64 (https://arxiv.org/html/2606.23898#bib.bib59), 62 (https://arxiv.org/html/2606.23898#bib.bib60), 65 (https://arxiv.org/html/2606.23898#bib.bib61)]，将采样步数从几千步减少到10-25步。然而，进一步减少通常会导致显著的性能下降。基于蒸馏的加速方法[43 (https://arxiv.org/html/2606.23898#bib.bib14), 50 (https://arxiv.org/html/2606.23898#bib.bib15), 54 (https://arxiv.org/html/2606.23898#bib.bib16), 11 (https://arxiv.org/html/2606.23898#bib.bib55), 37 (https://arxiv.org/html/2606.23898#bib.bib54), 67 (https://arxiv.org/html/2606.23898#bib.bib52), 36 (https://arxiv.org/html/2606.23898#bib.bib50), 59 (https://arxiv.org/html/2606.23898#bib.bib51), 61 (https://arxiv.org/html/2606.23898#bib.bib49)]通过训练学生模型将多步去噪轨迹压缩为更少的步数（有时甚至一步）来应对这一问题。重要的是，这些方法侧重于减少采样步数，而非压缩模型容量。相比之下，我们的工作针对基础扩散模型本身的压缩，提供了一个紧凑的基础，可以自然地补充和增强步数加速技术。

## 3 问题设置  
我们研究将预训练的文本到图像扩散模型（教师模型）蒸馏为更小的学生模型，如[26 (https://arxiv.org/html/2606.23898#bib.bib17), 25 (https://arxiv.org/html/2606.23898#bib.bib1)]所述。为清晰起见，我们使用噪声预测目标，尽管相同的推导适用于分数、速度或数据预测。令 ε<sub>T</sub>(x<sub>t</sub>, t, c) 表示教师模型，ε<sub>S</sub>(x<sub>t</sub>, t, c) 表示学生模型，其中 x<sub>t</sub> 是时间步 t 处的噪声潜变量，c 是文本条件。

**条件蒸馏损失。** 一种直接的无图像蒸馏策略是首先合成以文本提示为条件的图像，并构建成对数据集 D = {(x<sup>n</sup>, c<sup>n</sup>)}<sub>n=1</sub><sup>N</sup>，其中 x<sup>n</sup> 表示对应于文本条件 c<sup>n</sup> 的生成图像。生成的图像作为原始干净样本 x<sub>0</sub>，我们可以从中为任何时间步 t 在条件 c<sup>n</sup> 下生成噪声输入 x<sub>t</sub>。由于扩散模型在图像生成时需要大量的计算成本，这些合成图像通常预先生成并存储，形成训练数据集。随后，教师模型可以通过以下目标蒸馏为学生模型：

L<sub>out</sub> = E<sub>(x<sub>t</sub>, c) ∈ D, t</sub> [ ‖ ε<sub>T</sub>(x<sub>t</sub>, c, t) - ε<sub>S</sub>(x<sub>t</sub>, c, t) ‖<sub>2</sub><sup>2</sup> ]   (1)

其中 ε<sub>T</sub> 和 ε<sub>S</sub> 分别表示教师和学生网络的噪声预测。在这里，(x<sub>t</sub>, c) 是 samp