MoEにおける経路選択の重要性：明示的経路制御によるDiffusion Transformerのスケーリング

要旨

Mixture-of-Experts (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