Expert Race: Una strategia di routing flessibile per scalare i Transformer a diffusione con una miscela di esperti

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

March 20, 2025

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

cs.AI

Abstract

I modelli di diffusione sono emersi come framework principale nella generazione visiva. Basandosi su questo successo, l'integrazione dei metodi Mixture of Experts (MoE) ha dimostrato potenzialità nel migliorare la scalabilità e le prestazioni del modello. In questo articolo, introduciamo Race-DiT, un nuovo modello MoE per trasformatori di diffusione con una strategia di routing flessibile, Expert Race. Consentendo ai token e agli esperti di competere insieme e selezionare i candidati migliori, il modello impara ad assegnare dinamicamente gli esperti ai token critici. Inoltre, proponiamo una regolarizzazione per strato per affrontare le sfide nell'apprendimento degli strati superficiali e una perdita di similarità del router per prevenire il collasso modale, garantendo un migliore utilizzo degli esperti. Esperimenti estensivi su ImageNet convalidano l'efficacia del nostro approccio, mostrando significativi miglioramenti nelle prestazioni e promettenti proprietà di scalabilità.

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.