Variational Distillation of Diffusion Policies into Mixture of Experts

这项工作介绍了一种新方法，称为变分扩散蒸馏（Variational Diffusion Distillation，简称VDD），它利用变分推断将去噪扩散策略蒸馏成专家混合模型（Mixtures of Experts，简称MoE）。扩散模型是生成建模领域的最新技术，由于其精确学习和表示复杂的多模态分布的能力而备受赞誉。这种能力使得扩散模型能够复制人类行为中的内在多样性，因此成为学习人类示范（Learning from Human Demonstrations，简称LfD）等行为学习领域首选的模型。然而，扩散模型也存在一些缺点，包括难以计算的似然度和由于迭代采样过程而导致的长推断时间。特别是推断时间对于机器人控制等实时应用来说是一个重大挑战。相比之下，MoE有效地解决了上述问题，同时保持了表示复杂分布的能力，但是训练难度极大。VDD是第一种将预训练的扩散模型蒸馏成MoE模型的方法，从而结合了扩散模型的表达能力和混合模型的优点。具体而言，VDD利用变分目标的分解上界，使得每个专家的训练可以分别进行，从而实现了MoE的鲁棒优化方案。在九个复杂的行为学习任务中，VDD展示了以下能力：i）准确地蒸馏扩散模型学习到的复杂分布；ii）优于现有的最先进蒸馏方法；iii）超越传统的MoE训练方法。

本文旨在将Diffusion Models中学习到的复杂分布提取到MoE模型中，以解决Diffusion Models中似然函数不可解以及长时间推断的问题，同时保留MoE模型的表达复杂分布的能力。

本文提出了一种新的方法，即Variational Diffusion Distillation (VDD)，通过变分推断将Diffusion Models转化为MoE模型，从而结合了Diffusion Models的表达能力和MoE模型的优点。VDD使用一种分解的上界变分目标来训练每个专家，从而实现了MoE的稳健优化方案。

本文在九个复杂的行为学习任务中展示了VDD的优越性能，包括准确地提取Diffusion Models学习到的复杂分布、优于现有的最先进的蒸馏方法以及超越传统的MoE训练方法。实验结果表明，VDD在行为学习任务中具有广泛的适用性和可扩展性。

近年来，关于Diffusion Models的研究逐渐增多，如DDPM、DALL-E等。同时，MoE模型也是深度学习领域的热门研究方向，例如MoCo、DeepMoD等。