论文
桥接扩散修剪和步骤 蒸馏 与教师对齐修复
Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair
摘要
扩散模型生成高质量的图像,但其推理成本来自两个来源:大型去噪网络和重复的去噪步骤。现有的压缩管道通常会单独解决这些成本。剪枝减少了网络,但大多数剪枝方法仍然依赖于较长的剪枝后重新训练阶段来恢复多步采样器。步骤 蒸馏 减少了去噪步骤的数量,但它通常假设学生已经能够很好地跟随老师来接收有用的 蒸馏 梯度。本文询问是否可以用步骤蒸馏来代替后剪枝再训练。我们发现直接替换失败:修剪 EDM2-XS 教师后,从修剪的检查点启动 SiDA 会产生不可用的样本。我们引入了一个简短的教师对齐修复阶段作为修剪和步骤 蒸馏 之间的桥梁。该桥将修剪后的生成器与嘈杂的实像潜伏上的教师进行匹配,然后将修复的检查点交给单步 蒸馏。在ImageNet-512上,原始EDM2-XS基线使用124.713M参数和63次网络评估,达到FID为3.53。借助合适的 蒸馏 目标,我们的 20% 修剪一步生成器使用 98.826M 个参数和一次网络评估,达到 3.12 的 FID。剪枝率为 30% 时,模型使用 88.029M 个参数和 1 个网络评估,FID 为 4.26。