Экспертная гонка: Гибкая стратегия маршрутизации для масштабирования диффузионного трансформера с использованием смеси экспертов

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts

March 20, 2025

Авторы: Yike Yuan, Ziyu Wang, Zihao Huang, Defa Zhu, Xun Zhou, Jingyi Yu, Qiyang Min

cs.AI

Аннотация

Диффузионные модели стали основным фреймворком в области генерации изображений. Опираясь на этот успех, интеграция методов Mixture of Experts (MoE) показала потенциал в улучшении масштабируемости и производительности моделей. В данной статье мы представляем Race-DiT — новую MoE-модель для диффузионных трансформеров с гибкой стратегией маршрутизации, названной Expert Race. Позволяя токенам и экспертам конкурировать друг с другом и выбирать лучшие кандидаты, модель учится динамически назначать экспертов критически важным токенам. Кроме того, мы предлагаем регуляризацию на уровне слоев для решения проблем обучения на поверхностных слоях и функцию потерь на основе сходства маршрутизаторов, чтобы предотвратить коллапс мод и обеспечить лучшее использование экспертов. Многочисленные эксперименты на наборе данных ImageNet подтверждают эффективность нашего подхода, демонстрируя значительное улучшение производительности и перспективные свойства масштабирования.

English

Diffusion models have emerged as mainstream framework in visual generation. Building upon this success, the integration of Mixture of Experts (MoE) methods has shown promise in enhancing model scalability and performance. In this paper, we introduce Race-DiT, a novel MoE model for diffusion transformers with a flexible routing strategy, Expert Race. By allowing tokens and experts to compete together and select the top candidates, the model learns to dynamically assign experts to critical tokens. Additionally, we propose per-layer regularization to address challenges in shallow layer learning, and router similarity loss to prevent mode collapse, ensuring better expert utilization. Extensive experiments on ImageNet validate the effectiveness of our approach, showcasing significant performance gains while promising scaling properties.