Эффективные политики трансформации диффузии с смесью экспертных денойзеров для многозадачного обучения.

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

December 17, 2024

Авторы: Moritz Reuss, Jyothish Pari, Pulkit Agrawal, Rudolf Lioutikov

cs.AI

Аннотация

Политики Диффузии стали широко используемыми в Обучении по Имитации, предлагая несколько привлекательных свойств, таких как генерация мультимодального и дискретного поведения. Поскольку модели становятся все более крупными для охвата более сложных возможностей, их вычислительные требования увеличиваются, как показывают недавние законы масштабирования. Поэтому продолжение использования текущих архитектур представляет собой вычислительное препятствие. Для решения этого разрыва мы предлагаем Mixture-of-Denoising Experts (MoDE) в качестве новой политики для Обучения по Имитации. MoDE превосходит текущие передовые политики Диффузии на основе Трансформеров, обеспечивая эффективное масштабирование через разреженных экспертов и шумоусловленную маршрутизацию, снижая как активные параметры на 40%, так и затраты на вывод на 90% с помощью кэширования экспертов. Наша архитектура сочетает это эффективное масштабирование с механизмом шумоусловленного самовнимания, обеспечивая более эффективное подавление шума на разных уровнях шума. MoDE достигает передового уровня производительности на 134 задачах в четырех установленных бенчмарках по обучению по имитации (CALVIN и LIBERO). Значительно, предварительное обучение MoDE на различных данных робототехники позволяет достичь 4.01 на CALVIN ABC и 0.95 на LIBERO-90. Он превосходит как политики Диффузии на основе CNN, так и Трансформеры в среднем на 57% по 4 бенчмаркам, используя на 90% меньше операций с плавающей запятой и меньше активных параметров по сравнению с архитектурами Трансформеров по умолчанию для политик Диффузии. Кроме того, мы проводим всесторонние абляции компонентов MoDE, предоставляя идеи для разработки эффективных и масштабируемых архитектур Трансформеров для политик Диффузии. Код и демонстрации доступны по адресу 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/.