Políticas Eficientes de Transformador de Difusão com Mistura de Especialistas Desnitrificadores para Aprendizado Multitarefa

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

December 17, 2024

Autores: Moritz Reuss, Jyothish Pari, Pulkit Agrawal, Rudolf Lioutikov

cs.AI

Resumo

As Políticas de Difusão tornaram-se amplamente utilizadas em Aprendizado por Imitação, oferecendo várias propriedades atrativas, como a geração de comportamentos multimodais e descontínuos. À medida que os modelos estão se tornando maiores para capturar capacidades mais complexas, suas demandas computacionais aumentam, conforme demonstrado pelas recentes leis de escalonamento. Portanto, prosseguir com as arquiteturas atuais representará um obstáculo computacional. Para lidar com essa lacuna, propomos o Modelo de Especialistas em Desnublamento (MoDE) como uma nova política para Aprendizado por Imitação. O MoDE supera as atuais Políticas de Difusão baseadas em Transformadores de última geração, permitindo escalonamento eficiente de parâmetros por meio de especialistas esparsos e roteamento condicionado a ruído, reduzindo tanto os parâmetros ativos em 40% quanto os custos de inferência em 90% por meio de armazenamento em cache de especialistas. Nossa arquitetura combina esse escalonamento eficiente com um mecanismo de autoatenção condicionado a ruído, possibilitando uma desanuviamento mais eficaz em diferentes níveis de ruído. O MoDE alcança desempenho de última geração em 134 tarefas em quatro benchmarks estabelecidos de aprendizado por imitação (CALVIN e LIBERO). Notavelmente, ao pré-treinar o MoDE em dados robóticos diversos, alcançamos 4.01 no CALVIN ABC e 0.95 no LIBERO-90. Ele supera tanto as Políticas de Difusão baseadas em CNN quanto as baseadas em Transformadores em uma média de 57% em quatro benchmarks, enquanto utiliza 90% menos FLOPs e menos parâmetros ativos em comparação com as arquiteturas padrão de Transformadores de Difusão. Além disso, realizamos ablações abrangentes nos componentes do MoDE, fornecendo insights para o projeto de arquiteturas de Transformadores eficientes e escalonáveis para Políticas de Difusão. O código e demonstrações estão disponíveis em https://mbreuss.github.io/MoDE_Diffusion_Policy/.

English

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