Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair
作者: Jincheng Ying, Li Wenlin, Minghui Xu, Yinhao Xiao
分类: cs.CV
发布日期: 2026-07-07
💡 一句话要点
提出教师对齐修复方法以连接扩散剪枝与步骤蒸馏
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 扩散模型 剪枝 蒸馏 教师对齐 图像生成 深度学习 模型压缩
📋 核心要点
- 现有的剪枝和蒸馏方法在处理扩散模型推理成本时存在不足，尤其是在剪枝后恢复性能方面。
- 论文提出了一种短期教师对齐修复阶段，旨在将剪枝后的生成器与教师模型对齐，以提高蒸馏效果。
- 实验结果表明，经过20%剪枝的模型参数减少至98.826M，推理时仅需一次网络评估，FID值达到3.12，性能显著提升。
📝 摘要（中文）
扩散模型能够生成高质量图像，但其推理成本主要来自于大型去噪网络和重复的去噪步骤。现有的压缩方法通常分别处理这些成本。剪枝方法虽然能够减少网络规模，但大多数仍需经过长时间的后剪枝再训练阶段以恢复多步采样器。而步骤蒸馏则减少去噪步骤，但通常假设学生模型已经足够好以接收有用的蒸馏梯度。本文探讨了是否可以用步骤蒸馏替代后剪枝再训练，发现直接替换失败。因此，提出了一种短期教师对齐修复阶段，作为剪枝与步骤蒸馏之间的桥梁。通过在噪声真实图像潜变量上对齐修复后的生成器，最终实现了一步蒸馏。
🔬 方法详解
问题定义：本文旨在解决扩散模型推理中的高成本问题，尤其是剪枝后如何有效恢复模型性能。现有方法在剪枝后通常需要长时间的再训练，而步骤蒸馏又对学生模型的性能有较高要求。
核心思路：论文提出通过教师对齐修复阶段，先将剪枝后的生成器与教师模型对齐，然后再进行一步蒸馏。这种方法旨在减少后剪枝再训练的需求，同时提高蒸馏的有效性。
技术框架：整体流程包括三个主要阶段：首先进行剪枝以减少网络规模，其次通过教师对齐修复阶段对剪枝后的生成器进行调整，最后进行一步蒸馏以进一步优化模型。
关键创新：最重要的创新在于引入了教师对齐修复阶段，作为剪枝与蒸馏之间的桥梁。这一设计使得模型在剪枝后能够更好地适应蒸馏过程，显著提高了生成质量。
关键设计：在实验中，选择了适当的蒸馏目标，并对剪枝比例进行了调整，20%剪枝后的模型参数为98.826M，推理时仅需一次网络评估，FID值为3.12，显示出良好的性能。实验还探讨了不同剪枝比例对模型性能的影响。
🖼️ 关键图片
📊 实验亮点
实验结果显示，经过20%剪枝的模型参数减少至98.826M，推理时仅需一次网络评估，FID值达到3.12，相较于原始模型的124.713M参数和63次评估，性能显著提升，展示了该方法的有效性。
🎯 应用场景
该研究的潜在应用领域包括图像生成、计算机视觉和深度学习模型的优化。通过有效减少模型推理成本，能够在资源受限的环境中实现高质量图像生成，具有重要的实际价值和广泛的应用前景。
📄 摘要（原文）
Diffusion models generate high-quality images, but their inference cost comes from two sources: large denoising networks and repeated denoising steps. Existing compression pipelines usually attack these costs separately. Pruning reduces the network, but most pruning methods still rely on a long post-pruning retraining stage to recover a many-step sampler. Step distillation reduces the number of denoising steps, but it usually assumes a student that can already follow the teacher well enough to receive useful distillation gradients. This paper asks whether post-pruning retraining can be replaced by step distillation. We find that the direct replacement fails: after pruning an EDM2-XS teacher, starting SiDA from the pruned checkpoint produces unusable samples. We introduce a short teacher-alignment repair stage as a bridge between pruning and step distillation. The bridge matches the pruned generator to the teacher on noisy real-image latents, then hands the repaired checkpoint to one-step distillation. On ImageNet-512, the original EDM2-XS baseline uses 124.713M parameters and 63 network evaluations, reaching an FID of 3.53. With a suitable distillation objective, our 20% pruned one-step generator uses 98.826M parameters and one network evaluation, reaching an FID of 3.12. With 30% pruning, the model uses 88.029M parameters and one network evaluation, with an FID of 4.26.