追赶蒸馏：仅需训练一次即可加速采样
机器学习
2024-12-16 v5 计算机视觉与模式识别
摘要
扩散概率模型（DPMs）在多个机器学习领域取得了令人印象深刻的进展。然而，获得高质量合成样本通常需要执行大量采样步数，这阻碍了实时样本合成的可能性。传统的基于知识蒸馏的加速采样算法依赖于预训练模型权重和离散时间步场景，需要额外的训练轮次来实现其目标。为解决这些问题，我们提出追赶蒸馏（Catch-Up Distillation, CUD），其鼓励速度估计模型的当前时刻输出“追赶”其上一步输出。具体而言，CUD调整原始的常微分方程（ODE）训练目标，使当前时刻输出同时对齐真实标签和上一步输出，利用基于Runge-Kutta的多步对齐蒸馏进行精确ODE估计，同时防止异步更新。此外，我们研究了连续时间步场景下CUD的设计空间，并分析了如何确定合适策略。为证明CUD的有效性，我们在CIFAR-10、MNIST和ImageNet-64上进行了充分的消融与对比实验。在CIFAR-10上，我们通过单次训练会话下15步采样获得2.80的FID，并通过额外训练下1步采样获得3.37的新的最先进（SOTA）FID。后一结果仅需62万次迭代、批大小128，而一致性蒸馏（Consistency Distillation）需要2100万次迭代、更大的批大小256。我们的代码发布于https://anonymous.4open.science/r/Catch-Up-Distillation-E31F。
引用
@article{arxiv.2305.10769,
title = {Catch-Up Distillation: You Only Need to Train Once for Accelerating Sampling},
author = {Shitong Shao and Xu Dai and Lujun Li and Huanran Chen and Yang Hu and Shouyi Yin},
journal= {arXiv preprint arXiv:2305.10769},
year = {2024}
}