EM Distillation for One-step Diffusion Models

虽然扩散模型可以学习复杂的分布，但抽样需要计算密集的迭代过程。现有的蒸馏方法可以实现高效的抽样，但有明显的限制，比如在非常少的抽样步骤中性能下降，依赖于训练数据访问，或者是寻找模式的优化可能无法捕捉到完整的分布。我们提出了EM蒸馏（EMD），这是一种基于最大似然的方法，将扩散模型蒸馏为一步生成模型，并且在保持感知质量最小损失的情况下。我们的方法是通过期望最大化（EM）的视角推导出来的，其中生成器参数使用来自扩散教师先验和推断的生成器潜在变量的联合分布的样本进行更新。我们开发了重新参数化的抽样方案和噪声消除技术，共同稳定了蒸馏过程。我们进一步揭示了我们的方法与现有的最小化寻找模式KL的方法之间的有趣联系。在ImageNet-64和ImageNet-128的FID分数方面，EMD优于现有的一步生成方法，并且在蒸馏文本到图像扩散模型的先前工作中表现良好。

本篇论文提出了一种名为EM Distillation (EMD)的方法，以最小化感知质量的损失，将扩散模型转化为一步生成器模型，解决了现有扩散模型采样过程中计算量大的问题。

EMD方法通过期望最大化算法，利用来自扩散教师先验和推断生成器潜在变量的联合分布的样本更新生成器参数，同时使用重新参数化的采样方案和噪声抵消技术来稳定蒸馏过程。

论文的实验结果表明，EMD在ImageNet-64和ImageNet-128上的FID分数优于现有的一步生成方法，并且在蒸馏文本到图像扩散模型方面也表现出色。值得关注的是，EMD方法与现有的最小化模式寻求KL的方法之间存在有趣的联系。

与该论文相关的研究包括：《Density Matching for Large-Scale Subset Selection》、《Distilling Non-Autoregressive Transformers via Matrix Factorization》等。