Corrida de Especialistas: Uma Estratégia de Roteamento Flexível para Escalonar Transformadores de Difusão com Mistura de Especialistas

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

Resumo

Os modelos de difusão emergiram como um framework predominante na geração visual. Com base nesse sucesso, a integração de métodos de Mistura de Especialistas (MoE) tem mostrado potencial para melhorar a escalabilidade e o desempenho dos modelos. Neste artigo, apresentamos o Race-DiT, um novo modelo MoE para transformadores de difusão com uma estratégia de roteamento flexível, chamada Expert Race. Ao permitir que tokens e especialistas compitam juntos e selecionem os melhores candidatos, o modelo aprende a atribuir dinamicamente especialistas aos tokens críticos. Além disso, propomos uma regularização por camada para abordar desafios no aprendizado de camadas superficiais e uma perda de similaridade do roteador para evitar o colapso de modos, garantindo uma melhor utilização dos especialistas. Experimentos extensivos no ImageNet validam a eficácia da nossa abordagem, demonstrando ganhos significativos de desempenho e prometendo propriedades de escalabilidade.

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.