MoE에서 라우팅의 중요성: 명시적 라우팅 지도를 통한 Diffusion Transformer 확장

초록

혼합 전문가(MoE) 모델은 계산 효율성을 유지하면서 모델 용량을 확장하기 위한 강력한 패러다임으로 부상했습니다. 대규모 언어 모델(LLM)에서의 두드러진 성공에도 불구하고, MoE를 Diffusion Transformer(DiT)에 적용하려는 기존 시도들은 제한된 성능 향상만을 보여왔습니다. 우리는 이러한 격차가 언어 토큰과 시각 토큰 간의 근본적인 차이에서 비롯된다고 분석합니다. 언어 토큰은 의미론적으로 밀집되어 있고 토큰 간 변이가 뚜렷한 반면, 시각 토큰은 공간적 중복성과 기능적 이질성을 나타내어 시각 MoE에서의 전문가 특화를 저해합니다. 이를 해결하기 위해 우리는 명시적 라우팅 지도를 통한 2단계 라우터를 특징으로 하여 전문가 특화를 촉진하는 MoE 프레임워크인 ProMoE를 제안합니다. 구체적으로, 이 지도는 라우터가 이미지 토큰을 기능적 역할에 따라 조건부 라우팅을 통해 조건부 집합과 무조건부 집합으로 분할하도록 장려하고, 학습 가능한 프로토타입을 기반으로 한 의미론적 내용에 따른 프로토타입 라우팅을 통해 조건부 이미지 토큰의 할당을 정제합니다. 더 나아가, 프로토타입 라우팅에 의해 가능해진 잠재 공간 내 유사도 기반 전문가 할당은 명시적 의미론적 지도를 통합하는 자연스러운 메커니즘을 제공하며, 우리는 이러한 지도가 시각 MoE에 중요함을 입증합니다. 이를 기반으로, 우리는 프로토타입 라우팅 과정을 명시적으로 향상시켜 전문가 내 일관성과 전문가 간 다양성을 촉진하는 라우팅 대조 손실을 제안합니다. ImageNet 벤치마크에 대한 광범위한 실험을 통해 ProMoE가 Rectified Flow 및 DDPM 훈련 목표 모두에서 최첨단 방법들을 능가함을 입증합니다. 코드와 모델은 공개될 예정입니다.

English

Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fundamental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert specialization in vision MoE. To this end, we present ProMoE, an MoE framework featuring a two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and refine the assignments of conditional image tokens through prototypical routing with learnable prototypes based on semantic content. Moreover, the similarity-based expert allocation in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and we validate that such guidance is crucial for vision MoE. Building on this, we propose a routing contrastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on ImageNet benchmark demonstrate that ProMoE surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code and models will be made publicly available.

MoE에서 라우팅의 중요성: 명시적 라우팅 지도를 통한 Diffusion Transformer 확장

Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

초록

Support