전문가 경주: 전문가 혼합을 통한 확산 트랜스포머 확장을 위한 유연한 라우팅 전략

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

March 20, 2025

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

cs.AI

초록

디퓨전 모델은 시각적 생성 분야에서 주류 프레임워크로 자리 잡았습니다. 이러한 성공을 바탕으로, 전문가 혼합(Mixture of Experts, MoE) 방법의 통합은 모델의 확장성과 성능 향상에 있어 유망한 가능성을 보여주었습니다. 본 논문에서는 디퓨전 트랜스포머를 위한 새로운 MoE 모델인 Race-DiT를 소개합니다. 이 모델은 유연한 라우팅 전략인 Expert Race를 통해 토큰과 전문가가 함께 경쟁하며 최상위 후보를 선택함으로써, 모델이 중요한 토큰에 전문가를 동적으로 할당하는 방법을 학습합니다. 또한, 얕은 층 학습의 문제를 해결하기 위해 층별 정규화를 제안하고, 모드 붕괴를 방지하기 위한 라우터 유사성 손실을 도입하여 전문가 활용도를 개선합니다. ImageNet에 대한 광범위한 실험을 통해 우리의 접근 방식의 효과를 검증하였으며, 상당한 성능 향상과 함께 확장 가능성을 확인하였습니다.

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.