[論文レビュー] TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts

Yu Xu, Hongbin Yan+12 more|arXiv (Cornell University)|Jan 12, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

TAG-MoE は、階層的なタスク意味付けと予測整合性正則化を用いて、拡散 MoE 専門家を高レベルのタスク意図に従ってルーティングすることで、統一画像生成/編集におけるタスク干渉を低減するタスク認識ゲーティング機構を導入します。

ABSTRACT

Unified image generation and editing models suffer from severe task interference in dense diffusion transformers architectures, where a shared parameter space must compromise between conflicting objectives (e.g., local editing v.s. subject-driven generation). While the sparse Mixture-of-Experts (MoE) paradigm is a promising solution, its gating networks remain task-agnostic, operating based on local features, unaware of global task intent. This task-agnostic nature prevents meaningful specialization and fails to resolve the underlying task interference. In this paper, we propose a novel framework to inject semantic intent into MoE routing. We introduce a Hierarchical Task Semantic Annotation scheme to create structured task descriptors (e.g., scope, type, preservation). We then design Predictive Alignment Regularization to align internal routing decisions with the task's high-level semantics. This regularization evolves the gating network from a task-agnostic executor to a dispatch center. Our model effectively mitigates task interference, outperforming dense baselines in fidelity and quality, and our analysis shows that experts naturally develop clear and semantically correlated specializations.

研究の動機と目的

• dense diffusion transformers における統一的な画像生成/編集での深いタスク干渉を動機づけ解決する。
• タスク認識ルーティングを保証しつつ、スパース Mixture-of-Experts (MoE) でモデル容量を拡大する。
• グローバルなタスク意図信号を提供する階層的タスク意味付けを導入する。
• MoE ルーティングをタスク意味と一致させる予測整合正則化を開発する。

提案手法

• 拡散トランスフォーマーの最後の10層の FFN を MoE 層に置換し、固定パラメータで容量を増加させる。
• 各タスクのスコープ・タイプ・保持を説明する階層的タスク意味付けを導入する。
• グローバルな意味埋め込みをタスクタグから、MoE ゲートスコアからの集約ルーティングシグネチャを計算する。
• ルーティングシグネチャを意味嵌め込みと整合させるためのコサイン類似度損失で整合ヘッドを訓練する。
• L_align を主フロー損失と MoE 負荷均衡損失と重み付き総目的関 mathematically で結合する。
• 統一的生成/編集タスクをカバーする大規模で多様なデータセットで訓練する。

実験結果

リサーチクエスチョン

• RQ1 タスク認識ゲーティング機構は統一的な画像生成と編集におけるタスク間干渉を緩和できるか。
• RQ2 階層的タスク意味と予測整合は MoE ルーティングを高位のタスク意図に対応させるのに役立つか。
• RQ3 TAG-MoE は統一ベンチマークにおいて、オープンソースとクローズドソースのベースライン間でどう性能を示すか。
• RQ4 専門家は実務的に意味的に有意な特化と空間的に認識可能なルーティングを発展させるか。

主な発見

• TAG-MoE は統一生成/編集ベンチマークでオープンソースベースライン中の最先端性能を達成する。
• 意味整合的なルーティングはタスク非依存の MoE ルーティングと比較して指示追従性と意味的整合を改善する。
• アブレーションにより密なモデルは劣化し、意味に基づくエキスパートのルーティングには L_align が決定的であることが示される。
• 予測整合を備えた MoE は、関連する編集領域に空間的に局在したエキスパート特化を明確に示す。
• ユーザ調査では、参照整合、プロンプト整合、全体的品質の観点で TAG-MoE を明確に好む傾向がある。