Expert Race : Une stratégie de routage flexible pour la mise à l'échelle des Transformers de diffusion avec un mélange d'experts

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

March 20, 2025

Auteurs: Yike Yuan, Ziyu Wang, Zihao Huang, Defa Zhu, Xun Zhou, Jingyi Yu, Qiyang Min

cs.AI

Résumé

Les modèles de diffusion sont devenus un cadre de référence dominant dans la génération visuelle. S'appuyant sur ce succès, l'intégration des méthodes de Mélange d'Experts (MoE) a montré des perspectives prometteuses pour améliorer l'évolutivité et les performances des modèles. Dans cet article, nous présentons Race-DiT, un nouveau modèle MoE pour les transformeurs de diffusion doté d'une stratégie de routage flexible, Expert Race. En permettant aux tokens et aux experts de concourir ensemble et de sélectionner les meilleurs candidats, le modèle apprend à attribuer dynamiquement les experts aux tokens critiques. De plus, nous proposons une régularisation par couche pour résoudre les défis liés à l'apprentissage des couches superficielles, ainsi qu'une perte de similarité du routeur pour éviter l'effondrement des modes, garantissant ainsi une meilleure utilisation des experts. Des expériences approfondies sur ImageNet valident l'efficacité de notre approche, démontrant des gains de performance significatifs tout en promettant des propriétés d'évolutivité.

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.