Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning (2412.12953v1)

Abstract: 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/.

• The paper introduces MoDE, a Mixture-of-Denoising Experts model that cuts active parameters by 40% and inference costs by 90% compared to standard diffusion policies.
• The method integrates noise-conditioned self-attention and expert caching to optimize denoising performance across 134 tasks on benchmarks like CALVIN and LIBERO.
• Comprehensive ablation studies show that MoDE achieves a 57% performance increase, offering practical guidelines for scalable imitation learning in robotics.

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

The paper "Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning" introduces a novel approach to improving the efficiency and scalability of Diffusion Policies in Imitation Learning. This work is particularly relevant in the context of Imitation Learning, which leverages expert demonstrations to enable agents to learn versatile skills, but often suffers from high computational costs due to the size and complexity of models required to capture sophisticated behavior patterns.

Key Contributions

1. Introduction of Mixture-of-Denoising Experts (MoDE): The authors propose MoDE as a novel policy framework that advances the current state-of-the-art in Diffusion Policies, especially in computational efficiency. The architecture employs sparse experts and noise-conditioned routing, which significantly reduces active parameters by 40% and inference costs by 90% compared to baseline Transformer-based Diffusion Policies.
2. Architectural Innovation: MoDE integrates a noise-conditioned self-attention mechanism that enhances denoising performance across varying noise levels. This is coupled with expert caching mechanisms, which optimize inference by precomputing and storing expert decisions.
3. Performance and Efficiency: The paper highlights how MoDE achieves state-of-the-art performance on 134 tasks across four established benchmarks, including CALVIN and LIBERO. Notably, MoDE offers an average of 57% performance increase over existing CNN and Transformer-based Diffusion Policies while drastically reducing required computational resources (90% fewer FLOPs).
4. Comprehensive Ablation Studies: The authors conduct detailed ablation studies examining different components of MoDE, which provide insights for designing efficient and scalable Transformer architectures. This includes testing variations in noise-conditioned routing and expert distribution strategies, offering practical guidelines for future developments in AI and robotics.

Numerical Results and Experimental Validation

MoDE demonstrates strong performance metrics on established benchmarks, achieving significant gains in efficiency and task success rates. On the CALVIN benchmark, MoDE recorded a significant improvement by achieving an average success rate of 4.01 on CALVIN ABCD tasks and 0.95 on LIBERO-90 tasks, surpassing traditional Diffusion Policies by considerable margins.

Implications and Theoretical Impacts

The introduction of MoDE has important implications for the development of scalable AI systems, particularly in robotics and complex multitask environments where computational resources are limited. The approach provides a scalable method to deploy large-scale learning models without prohibitive resource consumption, promoting more accessible deployment in real-time applications.

Future Directions

The integration of Mixture-of-Experts with denoising Transformer Polices paves the way for future work focused on refining expert routing mechanisms and further reducing computational overhead. Potential developments could explore more dynamic expert activation strategies and adaptive architectures that can respond robustly to an even wider array of tasks and environmental variances.

In conclusion, this paper makes a substantial contribution to the field of AI and Imitation Learning by providing a more efficient alternative to conventional Diffusion Policies, ultimately enhancing both scalability and task performance in a computationally efficient manner.