Variational Distillation of Diffusion Policies into Mixture of Experts
作者: Hongyi Zhou, Denis Blessing, Ge Li, Onur Celik, Xiaogang Jia, Gerhard Neumann, Rudolf Lioutikov
分类: cs.LG, cs.AI, cs.RO
发布日期: 2024-06-18 (更新: 2024-10-18)
备注: Accepted by the 38th Annual Conference on Neural Information Processing Systems,
💡 一句话要点
提出VDD：通过变分推断将扩散策略蒸馏为混合专家模型，提升行为学习效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 扩散模型 混合专家模型 变分推断 蒸馏学习 行为学习 机器人控制 模仿学习
📋 核心要点
- 扩散模型在行为学习中表现出色，但推理速度慢，难以应用于实时控制。
- VDD通过变分推断将扩散模型蒸馏为混合专家模型，兼顾表达能力和推理效率。
- 实验表明，VDD在多个任务中超越了现有蒸馏方法和传统MoE训练方法。
📝 摘要（中文）
本研究提出了一种新颖的方法，即变分扩散蒸馏（VDD），它通过变分推断将去噪扩散策略蒸馏为混合专家模型（MoE）。扩散模型是当前生成建模领域最先进的技术，因为它们具有准确学习和表示复杂、多模态分布的卓越能力。这种能力使得扩散模型能够复制人类行为中固有的多样性，使其成为行为学习（如从人类演示中学习（LfD））的首选模型。然而，扩散模型也存在一些缺点，包括似然函数的难解性和由于其迭代采样过程导致的长推理时间。特别是推理时间对机器人控制等实时应用提出了重大挑战。相比之下，MoE有效地解决了上述问题，同时保留了表示复杂分布的能力，但众所周知难以训练。VDD是第一种将预训练的扩散模型蒸馏为MoE模型的方法，因此结合了扩散模型的表达能力和混合模型的优势。具体来说，VDD利用变分目标的分解上界，允许单独训练每个专家，从而为MoE提供了一个鲁棒的优化方案。VDD在九个复杂的行为学习任务中证明，它能够：i) 准确地蒸馏扩散模型学习的复杂分布，ii) 优于现有的最先进的蒸馏方法，以及iii) 超过传统的MoE训练方法。
🔬 方法详解
问题定义：论文旨在解决扩散模型推理速度慢的问题，尤其是在机器人控制等实时应用中。虽然扩散模型能很好地学习复杂分布，但其迭代采样过程导致推理时间过长。混合专家模型（MoE）虽然推理速度快，但训练困难。
核心思路：论文的核心思路是将预训练的扩散模型蒸馏为MoE模型，从而结合两者的优点。扩散模型负责学习复杂的行为分布，MoE负责快速推理。通过蒸馏，将扩散模型的知识迁移到MoE，避免了直接训练MoE的困难。
技术框架：VDD的整体框架包括以下步骤：1) 使用扩散模型预训练行为策略；2) 使用变分推断将预训练的扩散模型蒸馏为MoE模型；3) 使用分解的上界变分目标，独立训练每个专家。该框架的关键在于变分推断和分解的上界，使得MoE的训练更加稳定和高效。
关键创新：VDD最重要的创新点在于提出了一种新的蒸馏方法，可以将扩散模型有效地蒸馏为MoE模型。与传统的蒸馏方法不同，VDD利用变分推断和分解的上界，使得MoE的训练更加鲁棒，并且能够更好地保留扩散模型学习到的复杂分布。这是首次将扩散模型蒸馏到MoE模型的工作。
关键设计：VDD的关键设计包括：1) 使用变分推断来近似扩散模型的后验分布；2) 使用分解的上界来简化变分目标的计算，允许独立训练每个专家；3) 设计合适的损失函数，鼓励MoE模型学习扩散模型的行为策略。具体而言，损失函数包括重构损失和KL散度损失，用于衡量MoE模型与扩散模型之间的差异。
🖼️ 关键图片
📊 实验亮点
实验结果表明，VDD在九个复杂的行为学习任务中，显著优于现有的蒸馏方法和传统的MoE训练方法。具体来说，VDD能够更准确地蒸馏扩散模型学习到的复杂分布，并且能够更快地进行推理。在某些任务中，VDD的性能提升超过了10%。
🎯 应用场景
VDD在机器人控制、自动驾驶、游戏AI等领域具有广泛的应用前景。它可以用于学习复杂的人类行为，并将其应用于机器人控制，从而提高机器人的自主性和适应性。此外，VDD还可以用于生成逼真的人类行为，例如在游戏中生成更智能的NPC。
📄 摘要（原文）
This work introduces Variational Diffusion Distillation (VDD), a novel method that distills denoising diffusion policies into Mixtures of Experts (MoE) through variational inference. Diffusion Models are the current state-of-the-art in generative modeling due to their exceptional ability to accurately learn and represent complex, multi-modal distributions. This ability allows Diffusion Models to replicate the inherent diversity in human behavior, making them the preferred models in behavior learning such as Learning from Human Demonstrations (LfD). However, diffusion models come with some drawbacks, including the intractability of likelihoods and long inference times due to their iterative sampling process. The inference times, in particular, pose a significant challenge to real-time applications such as robot control. In contrast, MoEs effectively address the aforementioned issues while retaining the ability to represent complex distributions but are notoriously difficult to train. VDD is the first method that distills pre-trained diffusion models into MoE models, and hence, combines the expressiveness of Diffusion Models with the benefits of Mixture Models. Specifically, VDD leverages a decompositional upper bound of the variational objective that allows the training of each expert separately, resulting in a robust optimization scheme for MoEs. VDD demonstrates across nine complex behavior learning tasks, that it is able to: i) accurately distill complex distributions learned by the diffusion model, ii) outperform existing state-of-the-art distillation methods, and iii) surpass conventional methods for training MoE.