Effiziente Diffusionstransformator-Richtlinien mit einer Mischung von Experten-Denoisierern für Multitask-Lernen.

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

December 17, 2024

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

cs.AI

Zusammenfassung

Diffusionsrichtlinien sind in der Imitationslernen weit verbreitet und bieten verschiedene attraktive Eigenschaften, wie die Erzeugung multimodaler und diskontinuierlicher Verhaltensweisen. Da Modelle immer größer werden, um komplexere Fähigkeiten zu erfassen, steigen ihre Rechenanforderungen, wie durch aktuelle Skalierungsgesetze gezeigt. Daher wird die Fortführung der aktuellen Architekturen zu einem Rechenhindernis führen. Um diese Lücke zu schließen, schlagen wir Mixture-of-Denoising Experts (MoDE) als neuartige Richtlinie für das Imitationslernen vor. MoDE übertrifft aktuelle State-of-the-Art Transformer-basierte Diffusionsrichtlinien und ermöglicht eine parameter-effiziente Skalierung durch spärliche Experten und rauschbedingtes Routing, wodurch sowohl die aktiven Parameter um 40 % als auch die Inferenzkosten um 90 % durch Expertencaching reduziert werden. Unsere Architektur kombiniert diese effiziente Skalierung mit einem rauschbedingten Self-Attention-Mechanismus, der eine effektivere Rauschunterdrückung bei verschiedenen Rauschniveaus ermöglicht. MoDE erzielt State-of-the-Art-Leistung in 134 Aufgaben in vier etablierten Imitationslern-Benchmarks (CALVIN und LIBERO). Bemerkenswert ist, dass wir durch das Vortrainieren von MoDE auf vielfältigen Robotikdaten 4,01 bei CALVIN ABC und 0,95 bei LIBERO-90 erreichen. Es übertrifft sowohl CNN-basierte als auch Transformer-Diffusionsrichtlinien im Durchschnitt um 57 % über 4 Benchmarks hinweg, während es 90 % weniger FLOPs und weniger aktive Parameter im Vergleich zu Standard-Diffusionstransformer-Architekturen verwendet. Darüber hinaus führen wir umfassende Ablationen an MoDE-Komponenten durch, um Einblicke in die Gestaltung effizienter und skalierbarer Transformer-Architekturen für Diffusionsrichtlinien zu geben. Der Code und Demonstrationen sind verfügbar unter 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/.