概要

本リポジトリ「ProMoE」は、ICLR 2026 採択論文「Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance」の公式コード実装です。拡散モデルにトランスフォーマ＋Mixture-of-Experts（MoE）を適用する際、単純なルーティング設計が性能・スケーラビリティ・学習安定性に大きく影響することを示し、明示的なルーティング指導（explicit routing guidance）を導入してこれらを改善する手法を提供します。リポジトリには実験の再現に必要となる設定やスクリプト、コンフィグ類が含まれており、研究者や実務家が手元で試すための出発点になります。

リポジトリの統計情報

• スター数: 7
• フォーク数: 1
• ウォッチャー数: 7
• コミット数: 3
• ファイル数: 12
• メインの言語: Python

主な特徴

• 明示的ルーティング指導を導入したMoEを拡散トランスフォーマに適用する公式実装。
• ルータ設計、負荷分散、容量制約に関する実装と設定が含まれ、拡張性と再現性を重視。
• 設定ファイル（configs）で実験のハイパーパラメータを管理し、実験の追試が容易。
• 軽量で読みやすい構成により、研究者が各構成要素（router, expert, gating）を解析・改良しやすい。

技術的なポイント

本研究の技術的要点は「ルーティング（routing）の設計がMoEの性能に与える影響を系統的に評価し、明示的にガイダンスを与えることで拡張と安定化を図る」点にあります。従来のMoEではルータ（ゲート）が各トークンを動的に専門家（expert）へ割当てるが、無作為性や負荷不均衡、容量オーバーフローに起因する性能劣化や学習不安定が問題でした。ProMoEでは以下の技術要素が注目されます。

• ルーティングガイダンス: 追加の損失項やヒューリスティックを導入してルータの決定を制御し、特定の専門家への偏りを抑制しつつ意味的類似性に基づく割当てを促進します。
• 負荷分散と容量制御: 各エキスパートのキャパシティを定義し、割当て数が閾値を超えた場合のリダイレクションや再スケジュールのロジックを実装。これにより一部の専門家に負荷が集中するのを防ぎます。
• 拡散トランスフォーマ統合: 拡散モデルのステップごとに専門家を活用する設計で、時間的・空間的な特徴分解を行い効率的に計算を配分。トレーニング中の勾配の流れや安定化を考慮した実装が含まれます。
• 実験再現性: 設定ファイルによるハイパーパラメータ管理、モデル構成の明確化、評価プロトコルの記載により、他者が結果を追試しやすく設計されています。

これらを組み合わせることで、同等の計算リソース下でもより大規模なMoE構成を安定して学習させられる点が本実装の核です。ルーティングに関する設計選択（soft/hard gating、top-k 選択、auxiliary loss の重み付けなど）がモデルの性能や効率に直接影響するため、実装は実験的な比較が容易なモジュール化を意識して作られています。

プロジェクトの構成

主要なファイルとディレクトリ：

• .gitignore: file — 無視するファイル・ディレクトリの定義
• LICENSE: file — ライセンス情報（リポジトリ全体の利用条件）
• README.md: file — 論文概要、実験設定、導入方法などの説明（本リポジトリの入口）
• config.py: file — 実験全体で使われる共通設定やユーティリティ（ハイパーパラメータのデフォルト定義）
• configs: dir — 個別実験のコンフィグ群（モデルサイズ、学習率、ルーティング設定等を分離管理）
• …他 7 ファイル

補足: ファイル数が少数にまとまっており、コードベースは研究実装として簡潔に保たれています。主要なコンポーネントはモジュール化され、ルータ（routing）、エキスパート（experts）、拡散トランスフォーマの統合ロジック、評価スクリプトや設定の読み込み部分が分かれている想定です。configsディレクトリを編集することで、実験の再現・改良を行いやすくなっています。

まとめ

ルーティング設計に注目したMoE拡張で拡散トランスフォーマのスケーリングを目指す、再現性の高い公式実装です。

リポジトリ情報：

• 名前: ProMoE
• 説明: [ICLR2026] The official code of “Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance”
• スター数: 7
• 言語: Python
• URL: https://github.com/ali-vilab/ProMoE
• オーナー: ali-vilab
• アバター: https://avatars.githubusercontent.com/u/125999365?v=4

READMEの抜粋：

[ICLR 2026] Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance