Efficiënte Diffusion Transformer-beleidsregels met een mengeling van expertdenoiser voor multitask-leren.

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

Samenvatting

Diffusiebeleid is wijdverbreid geworden in Imitatie Leren, met verschillende aantrekkelijke eigenschappen, zoals het genereren van multimodale en discontinu gedrag. Aangezien modellen steeds groter worden om complexere mogelijkheden vast te leggen, nemen hun computationele eisen toe, zoals blijkt uit recente schalingswetten. Daarom zal doorgaan met de huidige architecturen een computationele hindernis vormen. Om deze kloof aan te pakken, stellen we Mixture-of-Denoising Experts (MoDE) voor als een nieuw beleid voor Imitatie Leren. MoDE overtreft de huidige state-of-the-art Transformer-gebaseerde Diffusiebeleidslijnen en maakt schaalvergroting mogelijk door middel van schaarse experts en ruis-geconditioneerde routering, waarbij zowel actieve parameters met 40% worden verminderd als inferentiekosten met 90% via expertcaching. Onze architectuur combineert deze efficiënte schaling met een ruis-geconditioneerd zelfaandachtsmechanisme, waardoor effectievere denoising mogelijk is bij verschillende geluidsniveaus. MoDE behaalt state-of-the-art prestaties op 134 taken in vier gevestigde imitatie-leerbenchmarks (CALVIN en LIBERO). Opmerkelijk is dat door MoDE vooraf te trainen op diverse robotica-gegevens, we 4.01 behalen op CALVIN ABC en 0.95 op LIBERO-90. Het overtreft zowel op CNN gebaseerde als Transformer Diffusiebeleidslijnen met gemiddeld 57% over 4 benchmarks, terwijl het 90% minder FLOPs en minder actieve parameters gebruikt in vergelijking met standaard Diffusie Transformer-architecturen. Bovendien voeren we uitgebreide ablaties uit op de componenten van MoDE, die inzichten bieden voor het ontwerpen van efficiënte en schaalbare Transformer-architecturen voor Diffusiebeleid. Code en demonstraties zijn beschikbaar op 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/.