Políticas de Transformador de Difusión Eficiente con Mezcla de Expertos Des ruidificadores para Aprendizaje Multitarea

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

Resumen

Las políticas de difusión se han vuelto ampliamente utilizadas en el Aprendizaje por Imitación, ofreciendo varias propiedades atractivas, como la generación de comportamientos multimodales y discontinuos. A medida que los modelos se hacen más grandes para capturar capacidades más complejas, sus demandas computacionales aumentan, como lo muestran las recientes leyes de escalado. Por lo tanto, continuar con las arquitecturas actuales presentará un obstáculo computacional. Para abordar esta brecha, proponemos Mixture-of-Denoising Experts (MoDE) como una nueva política para el Aprendizaje por Imitación. MoDE supera a las actuales políticas de difusión basadas en Transformadores de última generación al permitir un escalado eficiente de parámetros a través de expertos dispersos y enrutamiento condicionado por ruido, reduciendo tanto los parámetros activos en un 40% como los costos de inferencia en un 90% mediante el almacenamiento en caché de expertos. Nuestra arquitectura combina este escalado eficiente con un mecanismo de autoatención condicionado por ruido, permitiendo una denoising más efectiva en diferentes niveles de ruido. MoDE logra un rendimiento de última generación en 134 tareas en cuatro benchmarks establecidos de aprendizaje por imitación (CALVIN y LIBERO). Es notable que, al preentrenar MoDE con datos robóticos diversos, logramos 4.01 en CALVIN ABC y 0.95 en LIBERO-90. Supera tanto a las políticas de difusión basadas en CNN como a las basadas en Transformadores en un promedio del 57% en los 4 benchmarks, al mismo tiempo que utiliza un 90% menos de FLOPs y menos parámetros activos en comparación con las arquitecturas predeterminadas de Transformadores de difusión. Además, realizamos abstracciones exhaustivas sobre los componentes de MoDE, proporcionando ideas para diseñar arquitecturas de Transformadores eficientes y escalables para políticas de difusión. El código y las demostraciones están disponibles en 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/.