Politiques de transformation de diffusion efficaces avec un mélange d'experts débruiteurs pour l'apprentissage multitâche

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

December 17, 2024

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

cs.AI

Résumé

Les politiques de diffusion sont devenues largement utilisées en Apprentissage par Imitation, offrant plusieurs propriétés attrayantes, telles que la génération de comportements multimodaux et discontinus. Alors que les modèles deviennent de plus en plus grands pour capturer des capacités plus complexes, leurs exigences computationnelles augmentent, comme le montrent les récentes lois d'échelle. Par conséquent, poursuivre avec les architectures actuelles présentera un obstacle computationnel. Pour combler cette lacune, nous proposons Mélange d'Experts de Débruitage (MoDE) comme une nouvelle politique pour l'Apprentissage par Imitation. MoDE surpasse les politiques de diffusion basées sur les Transformers de pointe actuelles tout en permettant une mise à l'échelle efficace des paramètres grâce à des experts clairsemés et un routage conditionné par le bruit, réduisant à la fois les paramètres actifs de 40 % et les coûts d'inférence de 90 % via la mise en cache des experts. Notre architecture combine cette mise à l'échelle efficace avec un mécanisme d'auto-attention conditionné par le bruit, permettant un débruitage plus efficace à travers différents niveaux de bruit. MoDE atteint des performances de pointe sur 134 tâches dans quatre benchmarks d'apprentissage par imitation établis (CALVIN et LIBERO). Notamment, en pré-entraînant MoDE sur des données robotiques diverses, nous obtenons 4,01 sur CALVIN ABC et 0,95 sur LIBERO-90. Il surpasse à la fois les politiques de diffusion basées sur les CNN et les Transformers de pointe de 57 % en moyenne sur 4 benchmarks, tout en utilisant 90 % de FLOPs en moins et moins de paramètres actifs par rapport aux architectures de Transformers de diffusion par défaut. De plus, nous menons des ablations approfondies sur les composants de MoDE, fournissant des informations pour la conception d'architectures de Transformers efficaces et évolutives pour les politiques de diffusion. Le code et les démonstrations sont disponibles sur 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/.