Politiche efficienti del trasformatore di diffusione con miscela di esperti denoiser per l'apprendimento multi-task

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

December 17, 2024

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

cs.AI

Abstract

Le politiche di diffusione sono diventate ampiamente utilizzate nell'Apprendimento per Imitazione, offrendo diverse proprietà interessanti, come la generazione di comportamenti multimodali e discontinui. Poiché i modelli stanno diventando sempre più grandi per catturare capacità più complesse, le loro richieste computazionali aumentano, come dimostrato dalle recenti leggi di scalabilità. Pertanto, continuare con le attuali architetture presenterà un ostacolo computazionale. Per affrontare questa lacuna, proponiamo Mixture-of-Denoising Experts (MoDE) come una nuova politica per l'Apprendimento per Imitazione. MoDE supera le attuali politiche di diffusione basate su Transformer all'avanguardia consentendo una scalabilità efficiente dei parametri attraverso esperti sparsi e routing condizionato al rumore, riducendo sia i parametri attivi del 40% che i costi di inferenza del 90% tramite la memorizzazione nella cache degli esperti. La nostra architettura combina questa scalabilità efficiente con un meccanismo di autoattenzione condizionato al rumore, consentendo una denoising più efficace attraverso diversi livelli di rumore. MoDE raggiunge prestazioni all'avanguardia su 134 compiti in quattro benchmark consolidati di apprendimento per imitazione (CALVIN e LIBERO). In particolare, preaddestrando MoDE su dati robotici diversi, otteniamo 4.01 su CALVIN ABC e 0.95 su LIBERO-90. Supera sia le politiche di diffusione basate su CNN che su Transformer di default di un 57% in media su 4 benchmark, utilizzando il 90% in meno di FLOP e meno parametri attivi rispetto alle architetture di default dei Transformer di diffusione. Inoltre, conduciamo approfondite analisi sulle componenti di MoDE, fornendo spunti per progettare architetture di Transformer efficienti e scalabili per le politiche di diffusione. Il codice e le dimostrazioni sono disponibili su 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/.