Эффективные политики диффузии с помощью смеси экспертов для многозадачного обучения
В последние годы модели диффузии стали важным инструментом в области имитационного обучения (Imitation Learning, IL), благодаря своим выдающимся свойствам, таким как способность генерировать разнообразное и многомодальное поведение. Однако с увеличением размеров моделей и их вычислительных требований возникают значительные проблемы, связанные с производительностью и эффективностью. В данной статье мы рассмотрим новую архитектуру, предложенную в работе "Mixture-of-Denoising Experts (MoDE)", которая направлена на решение этих проблем.
Модели диффузии учат обратному процессу, добавляя гауссовский шум к образцам данных. После обучения они могут генерировать новые образцы, основанные на заданных целях, таких как инструкции или изображения. Однако высокий вычислительный расход современных моделей ограничивает их применение в реальных задачах, особенно в сценариях, где ресурсы ограничены, например, в мобильных роботах.
Для преодоления этих ограничений исследователи предлагают использовать архитектуру Mixture-of-Experts (MoE). Основная идея заключается в том, чтобы использовать только подмножество параметров модели на каждом шаге, что позволяет значительно сократить вычислительные затраты.
Проблема и решение
Формулировка проблемы
Целью IL является обучение агентов на основе демонстраций экспертов, чтобы они могли развивать универсальные навыки. Политика, которую мы обучаем, предсказывает последовательность действий, основываясь на истории состояний и заданной цели. При этом необходимо максимизировать логарифмическую вероятность последовательности действий с учетом контекста.
Архитектура MoDE
MoDE использует непрерывную модель диффузии как представление политики. Эта модель включает в себя несколько ключевых компонентов:
-
Шумовая маршрутизация: новый механизм маршрутизации, который распределяет токены к экспертам на основе текущего уровня шума, что позволяет использовать специализированные эксперты для различных фаз денойзинга.
-
Шумовая самовнимательность: механизм самовнимательности, который адаптируется к уровню шума, улучшая качество денойзинга.
-
Кэширование экспертов: механизм, позволяющий заранее вычислить и объединить необходимые эксперты для более эффективного вывода.
Эти компоненты позволяют MoDE достигать выдающихся результатов при значительно меньших вычислительных затратах по сравнению с традиционными архитектурами.
Связанные работы
Диффузия в робототехнике
Модели диффузии нашли широкое применение в робототехнике, включая представление политики для IL и оффлайн обучения с подкреплением. Наиболее распространенной архитектурой для использования моделей диффузии в робототехнике являются сверточные нейронные сети (CNN), однако в последнее время трансформеры становятся стандартом благодаря своей способности масштабироваться с увеличением объема данных.
Смесь экспертов
MoE представляет собой класс моделей, которые позволяют избирательно маршрутизировать информацию через модель. Современная версия MoE была предложена в 2017 году и с тех пор активно используется в различных архитектурах, включая трансформеры. Однако ни одна из предыдущих работ не рассматривала применение MoE для повышения вычислительной эффективности и скорости вывода в контексте диффузионных политик.
Многозадачное обучение в моделях диффузии
Показано, что процесс денойзинга является многозадачным. Некоторые работы уже используют архитектуры, подходящие для многозадачного обучения, однако MoDE предлагает новый уровень адаптации и специализации для различных фаз денойзинга.
Методология
Архитектура MoDE
MoDE состоит из L блоков трансформеров, каждый из которых специализирован для различных фаз денойзинга. Основные компоненты включают:
- Шумовая маршрутизация: на каждом уровне осуществляется выбор экспертов на основе уровня шума, что позволяет эффективно использовать ресурсы.
- Шумовая самовнимательность: обеспечивает адаптацию внимания к текущему уровню шума, улучшая качество денойзинга.
Кэширование экспертов
Кэширование экспертов позволяет заранее определить, какие эксперты будут использоваться на каждом уровне шума, что значительно снижает вычислительные затраты и время вывода.
Предобучение
MoDE предобучается на разнообразных данных робототехники, что позволяет улучшить его обобщающие способности и производительность на сложных задачах.
Оценка
Сравнение с другими политиками
MoDE был протестирован на нескольких задачах и показал выдающиеся результаты по сравнению с другими архитектурами, такими как CNN и традиционные трансформеры. Например, в тестах на 134 задачах MoDE показал среднее увеличение производительности на 57% при использовании на 90% меньшего количества FLOPs.
Эффективность и скорость
MoDE демонстрирует значительное улучшение вычислительной эффективности, требуя в среднем на 40% меньше времени вывода по сравнению с плотными трансформерами при аналогичном количестве параметров.
Ablation Studies
Ablation studies показали, что ключевыми факторами, влияющими на производительность MoDE, являются механизмы маршрутизации и шумовой инъекции. Например, использование шумовых токенов и шумовой самовнимательности продемонстрировало явное преимущество по сравнению с традиционными методами.
Заключение
В данной работе была представлена новая архитектура MoDE, которая использует смесь экспертов для повышения производительности и эффективности диффузионных политик. Результаты экспериментов показывают, что MoDE не только превосходит предыдущие модели по производительности, но и значительно снижает вычислительные затраты, что делает его идеальным решением для многозадачных сценариев в робототехнике.
Будущие исследования будут направлены на дальнейшее улучшение и адаптацию архитектуры для более широкого спектра задач и условий.