arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2503.16057cs.CVcs.AIcs.LG

Expert Race:一种用于扩展混合专家扩散Transformer的灵活路由策略

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

ShanghaiTech University(上海科技大学)

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

AI总结:

针对扩散Transformer的扩展性问题,提出搭载Expert Race灵活路由策略的Race-DiT模型,结合逐层正则化与路由相似度损失优化专家利用,在ImageNet上实现显著性能提升。

摘要:

扩散模型已成为视觉生成领域的主流框架。在这一成功基础上,混合专家(Mixture of Experts, MoE)方法的融合在提升模型可扩展性与性能方面展现出潜力。本文提出Race-DiT,这是一种面向扩散Transformer的新型MoE模型,搭载灵活的路由策略Expert Race。通过让token与专家共同竞争并筛选最优候选,模型能够学习将专家动态分配给关键token。此外,我们提出逐层正则化以解决浅层学习难题,并提出路由相似度损失以防止模式崩溃,确保更优的专家利用率。在ImageNet上开展的大量实验验证了该方法的有效性,展现出显著的性能提升与良好的扩展潜力。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑