Lightweight Diffusion Models with Distillation-Based Block Neural Architecture Search

📄 arXiv: 2311.04950v2 📥 PDF

作者: Siao Tang, Xin Wang, Hong Chen, Chaoyu Guan, Yansong Tang, Wenwu zhu

分类: cs.CV, cs.LG

发布日期: 2023-11-08 (更新: 2023-11-15)

💡 一句话要点

提出DiffNAS以解决扩散模型高计算成本问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 神经架构搜索 蒸馏训练 计算效率 块级结构 动态损失 模型压缩

📋 核心要点

  1. 现有的扩散模型在生成任务中表现优异，但其高计算成本限制了实际应用。
  2. 本文提出DiffNAS，通过块级神经架构搜索自动去除扩散模型中的冗余结构，优化计算效率。
  3. 实验表明，DiffNAS在潜在扩散模型上实现了约50%的计算量和参数减少，显著提升了模型性能。

📝 摘要（中文）

扩散模型在生成任务中表现出色，但其高计算成本仍然是一个挑战。为了解决这一问题，本文提出了一种基于蒸馏的块级神经架构搜索方法DiffNAS，旨在自动去除扩散模型中的结构冗余。通过利用一个大型预训练教师模型，DiffNAS能够搜索出最小的架构，并在性能上与教师模型持平或更优。该方法在每个块内独立进行神经架构搜索，显著减少了搜索空间，并引入了块级局部搜索策略和动态联合损失的重训练策略，提升了模型性能。实验结果表明，DiffNAS在潜在扩散模型上实现了约50%的MACs和参数减少。

🔬 方法详解

问题定义：本文旨在解决扩散模型的高计算成本问题，现有方法在结构上存在冗余，导致计算效率低下。

核心思路：通过Diffusion Distillation-based Block-wise Neural Architecture Search (DiffNAS)，在每个块内独立进行架构搜索，寻找最小化的有效架构，从而减少冗余。

技术框架：DiffNAS的整体流程包括：首先，利用大型预训练教师模型进行架构搜索；其次，在每个块内进行局部搜索以选择最佳子网络；最后，采用动态联合损失进行重训练，以保持超网络训练与子网络重训练的一致性。

关键创新：DiffNAS的创新在于引入了块级局部搜索策略和动态联合损失，这与以往的全局搜索策略形成鲜明对比，能够有效避免不公平性并提升模型性能。

关键设计：在损失函数设计上，采用动态联合损失以提供每个块的有用目标，并缩短梯度传播路径，从而提升训练效率和模型性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，DiffNAS在潜在扩散模型上实现了约50%的MACs和参数减少，相较于基线方法，模型性能显著提升，验证了动态联合损失的有效性。

🎯 应用场景

该研究的潜在应用领域包括图像生成、视频生成和其他需要高效生成模型的任务。通过降低计算成本，DiffNAS可以使扩散模型更易于在资源受限的环境中部署，推动其在实际应用中的广泛使用。

📄 摘要（原文）

Diffusion models have recently shown remarkable generation ability, achieving state-of-the-art performance in many tasks. However, the high computational cost is still a troubling problem for diffusion models. To tackle this problem, we propose to automatically remove the structural redundancy in diffusion models with our proposed Diffusion Distillation-based Block-wise Neural Architecture Search (DiffNAS). Specifically, given a larger pretrained teacher, we leverage DiffNAS to search for the smallest architecture which can achieve on-par or even better performance than the teacher. Considering current diffusion models are based on UNet which naturally has a block-wise structure, we perform neural architecture search independently in each block, which largely reduces the search space. Different from previous block-wise NAS methods, DiffNAS contains a block-wise local search strategy and a retraining strategy with a joint dynamic loss. Concretely, during the search process, we block-wisely select the best subnet to avoid the unfairness brought by the global search strategy used in previous works. When retraining the searched architecture, we adopt a dynamic joint loss to maintain the consistency between supernet training and subnet retraining, which also provides informative objectives for each block and shortens the paths of gradient propagation. We demonstrate this joint loss can effectively improve model performance. We also prove the necessity of the dynamic adjustment of this loss. The experiments show that our method can achieve significant computational reduction, especially on latent diffusion models with about 50\% MACs and Parameter reduction.