Efficient Diffusion Transformer Policies with Mixture of Expert
  Denoisers for Multitask Learning
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/.
6763c3475f32892ca1133cc2
Moritz Reuss
claimed_verified
2024-12-19T18:01:10.717Z
650af93422ce64f22b619549
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/RBybzGRphbiV2MXYHrDoc.png
Moritz Reuss
user
mbreuss
6763c3475f32892ca1133cc3
Jyothish Pari
6763c3475f32892ca1133cc4
Pulkit Agrawal
6763c3475f32892ca1133cc5
Rudolf Lioutikov
admin_assigned
2024-12-20T09:48:20.326Z
6606fab1d81d611249506cb1
/avatars/334faf35bd96aee3240a1622c1325459.svg
Rudolf Lioutikov
user
lioutikov
Efficient Diffusion Transformer Policies with Mixture of Expert
  Denoisers for Multitask Learning
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/.
Mixture-of-Denoising Experts (MoDE) enhances imitation learning by providing efficient scaling and state-of-the-art performance with reduced computational requirements.
Diffusion Policies Mixture-of-Denoising Experts MoDE Transformer-based Diffusion Policies parameter-efficient scaling sparse experts noise-conditioned routing noise-conditioned self-attention imitation learning CALVIN LIBERO CNN-based Diffusion Policies Diffusion Transformer architectures FLOPs