Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning

📄 arXiv: 2412.12953v1 📥 PDF

作者: Moritz Reuss, Jyothish Pari, Pulkit Agrawal, Rudolf Lioutikov

分类: cs.LG, cs.RO

发布日期: 2024-12-17

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出MoDE：一种基于混合专家降噪器的参数高效扩散Transformer策略，用于多任务模仿学习。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散策略 模仿学习 混合专家模型 Transformer 机器人控制

📋 核心要点

  1. 现有扩散策略模型规模增大导致计算成本显著增加，阻碍了其在复杂任务中的应用。
  2. 提出混合降噪专家（MoDE），利用稀疏专家和噪声条件路由实现参数高效扩展，降低计算成本。
  3. MoDE在多个模仿学习基准测试中取得了SOTA性能，同时显著降低了FLOPs和活跃参数。

📝 摘要（中文）

扩散策略已广泛应用于模仿学习，它具有生成多模态和不连续行为等吸引人的特性。随着模型规模的增大以捕获更复杂的能力，其计算需求也随之增加。为了解决这一问题，我们提出混合降噪专家（MoDE）作为一种新的模仿学习策略。MoDE超越了当前最先进的基于Transformer的扩散策略，同时通过稀疏专家和噪声条件路由实现参数高效扩展，通过专家缓存将活跃参数减少40%，推理成本降低90%。我们的架构将这种高效扩展与噪声条件自注意力机制相结合，从而能够更有效地跨不同噪声水平进行降噪。MoDE在四个已建立的模仿学习基准（CALVIN和LIBERO）的134个任务中实现了最先进的性能。值得注意的是，通过在多样化的机器人数据上预训练MoDE，我们在CALVIN ABC上实现了4.01，在LIBERO-90上实现了0.95。与默认的扩散Transformer架构相比，它在4个基准测试中平均超过基于CNN和Transformer的扩散策略57%，同时使用减少90%的FLOPs和更少的活跃参数。此外，我们对MoDE的组件进行了全面的消融研究，为设计用于扩散策略的高效且可扩展的Transformer架构提供了见解。

🔬 方法详解

问题定义：论文旨在解决扩散策略在模仿学习中计算成本高昂的问题，尤其是在模型规模增大以处理更复杂任务时。现有基于Transformer的扩散策略虽然性能优异，但其计算需求随着模型规模的增长而迅速增加，这限制了它们在资源受限环境中的应用。

核心思路：论文的核心思路是利用混合专家模型（Mixture of Experts, MoE）来提高扩散策略的参数效率和计算效率。通过将模型分解为多个专家，并根据输入噪声水平动态地选择激活的专家子集，从而减少了活跃参数的数量和计算量。同时，噪声条件路由机制允许模型根据噪声水平自适应地调整专家组合，从而提高降噪性能。

技术框架：MoDE的整体架构包括一个噪声条件自注意力模块和一个混合专家降噪器。噪声条件自注意力模块用于提取输入状态和噪声水平的特征，并将这些特征传递给混合专家降噪器。混合专家降噪器由多个专家组成，每个专家都是一个小的Transformer模块。一个门控网络（Gating Network）根据噪声水平选择激活的专家子集，并将输入传递给这些专家进行降噪。最后，将激活专家的输出进行加权平均，得到最终的降噪结果。

关键创新：MoDE的关键创新在于将混合专家模型与扩散策略相结合，并引入了噪声条件路由机制。传统的混合专家模型通常根据输入数据的特征选择专家，而MoDE则根据噪声水平选择专家。这种噪声条件路由机制允许模型根据噪声水平自适应地调整专家组合，从而提高降噪性能。此外，MoDE还采用了专家缓存技术，进一步降低了推理成本。

关键设计：MoDE的关键设计包括：1) 稀疏专家网络结构，减少活跃参数数量；2) 噪声条件门控网络，根据噪声水平动态选择专家；3) 专家缓存机制，存储常用专家的输出，减少重复计算；4) 噪声条件自注意力机制，提高特征提取能力。损失函数采用标准的扩散模型损失函数，即预测噪声与真实噪声之间的均方误差。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MoDE在CALVIN ABC和LIBERO-90等四个模仿学习基准测试中取得了SOTA性能，平均超过基于CNN和Transformer的扩散策略57%。同时，MoDE将活跃参数减少40%，推理成本降低90%，FLOPs减少90%。在CALVIN ABC上实现了4.01，在LIBERO-90上实现了0.95。

🎯 应用场景

MoDE具有广泛的应用前景，尤其是在机器人控制、自动驾驶等需要高效模仿学习的领域。通过降低计算成本，MoDE使得扩散策略能够在资源受限的平台上部署，从而实现更智能、更灵活的机器人行为。此外，MoDE还可以应用于其他生成模型领域，例如图像生成、语音合成等，提高生成模型的效率和质量。

📄 摘要（原文）

Diffusion Policies have become widely used in Imitation Learning, offering several appealing properties, such as generating multimodal and discontinuous behavior. As models are becoming larger to capture more complex capabilities, their computational demands increase, as shown by recent scaling laws. Therefore, continuing with the current architectures will present a computational roadblock. To address this gap, we propose Mixture-of-Denoising Experts (MoDE) as a novel policy for Imitation Learning. MoDE surpasses current state-of-the-art Transformer-based Diffusion Policies while enabling parameter-efficient scaling through sparse experts and noise-conditioned routing, reducing both active parameters by 40% and inference costs by 90% via expert caching. Our architecture combines this efficient scaling with noise-conditioned self-attention mechanism, enabling more effective denoising across different noise levels. MoDE achieves state-of-the-art performance on 134 tasks in four established imitation learning benchmarks (CALVIN and LIBERO). Notably, by pretraining MoDE on diverse robotics data, we achieve 4.01 on CALVIN ABC and 0.95 on LIBERO-90. It surpasses both CNN-based and Transformer Diffusion Policies by an average of 57% across 4 benchmarks, while using 90% fewer FLOPs and fewer active parameters compared to default Diffusion Transformer architectures. Furthermore, we conduct comprehensive ablations on MoDE's components, providing insights for designing efficient and scalable Transformer architectures for Diffusion Policies. Code and demonstrations are available at https://mbreuss.github.io/MoDE_Diffusion_Policy/.