Expert Race: Een flexibel routeringsstrategie voor het schalen van Diffusion Transformers met een Mixture of 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

Samenvatting

Diffusiemodellen zijn uitgegroeid tot een mainstream raamwerk in visuele generatie. Voortbouwend op dit succes heeft de integratie van Mixture of Experts (MoE)-methoden belofte getoond in het verbeteren van de schaalbaarheid en prestaties van modellen. In dit artikel introduceren we Race-DiT, een nieuw MoE-model voor diffusie-transformers met een flexibele routeringsstrategie, Expert Race. Door tokens en experts samen te laten concurreren en de beste kandidaten te selecteren, leert het model experts dynamisch toe te wijzen aan kritieke tokens. Daarnaast stellen we per-laag regularisatie voor om uitdagingen in het leren van ondiepe lagen aan te pakken, en router similarity loss om mode-collaps te voorkomen, wat een betere benutting van experts garandeert. Uitgebreide experimenten op ImageNet valideren de effectiviteit van onze aanpak, waarbij aanzienlijke prestatieverbeteringen worden getoond terwijl schaaleigenschappen worden beloofd.

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.