Reproduce “Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning” on OpenResearch

About this paper

Moritz Reuss, Jyothish Pari, Pulkit Agrawal, Rudolf Lioutikov

An architecture for diffusion policies, MoDE (Mixture-of-Denoising Experts), achieves state-of-the-art performance in multitask robotic imitation learning while operating with 90% fewer FLOPs and significantly faster inference compared to prior diffusion Transformers. Developed by researchers from KIT and MIT, MoDE integrates expert specialization for different noise levels within the denoising process and utilizes an expert caching mechanism to accelerate execution.

Key insights

• The denoising process in diffusion models is inherently multi-task, benefiting from specialized experts that are dynamically activated based on the noise level, which significantly improves both efficiency and performance.
• Decoupling expert routing from token content and conditioning it solely on the noise level enables a powerful expert caching strategy that drastically reduces inference-time computations.
• Both an explicit noise input token and noise-conditioned self-attention are crucial architectural components for optimal performance in diffusion Transformer policies, with an optimal balance found around 4 experts for efficiency and effectiveness.

Summary by alphaXiv.