Expert Race: Eine flexible Routing-Strategie zur Skalierung von Diffusionstransformatoren mit Mixture of Experts

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

March 20, 2025

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

cs.AI

Zusammenfassung

Diffusionsmodelle haben sich als Mainstream-Framework in der visuellen Generierung etabliert. Aufbauend auf diesem Erfolg hat die Integration von Mixture of Experts (MoE)-Methoden vielversprechende Ergebnisse bei der Verbesserung der Skalierbarkeit und Leistung von Modellen gezeigt. In diesem Artikel stellen wir Race-DiT vor, ein neuartiges MoE-Modell für Diffusionstransformer mit einer flexiblen Routing-Strategie, Expert Race. Indem Token und Experten gemeinsam konkurrieren und die besten Kandidaten auswählen, lernt das Modell, Experten dynamisch kritischen Token zuzuweisen. Zusätzlich schlagen wir eine Regularisierung pro Schicht vor, um Herausforderungen beim Lernen in flachen Schichten zu bewältigen, sowie einen Router-Ähnlichkeitsverlust, um Modus-Kollaps zu verhindern und eine bessere Nutzung der Experten zu gewährleisten. Umfangreiche Experimente auf ImageNet bestätigen die Wirksamkeit unseres Ansatzes und zeigen signifikante Leistungssteigerungen bei vielversprechenden Skalierungseigenschaften.

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.