Carrera de Expertos: Una Estrategia de Enrutamiento Flexible para Escalar Transformadores de Difusión con Mezcla de Expertos

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

March 20, 2025

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

cs.AI

Resumen

Los modelos de difusión se han consolidado como un marco principal en la generación visual. Basándose en este éxito, la integración de métodos de Mezcla de Expertos (MoE) ha mostrado potencial para mejorar la escalabilidad y el rendimiento de los modelos. En este artículo, presentamos Race-DiT, un novedoso modelo MoE para transformadores de difusión con una estrategia de enrutamiento flexible llamada Expert Race. Al permitir que los tokens y los expertos compitan juntos y seleccionen a los mejores candidatos, el modelo aprende a asignar dinámicamente expertos a los tokens críticos. Además, proponemos una regularización por capa para abordar los desafíos en el aprendizaje de capas superficiales, y una pérdida de similitud del enrutador para evitar el colapso de modos, garantizando una mejor utilización de los expertos. Experimentos exhaustivos en ImageNet validan la efectividad de nuestro enfoque, mostrando mejoras significativas en el rendimiento junto con propiedades prometedoras de escalabilidad.

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.