TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts

2026-01-12

AI Summary

This work addresses the challenge of task interference in unified image generation and editing models based on dense diffusion Transformers, where shared parameters struggle to reconcile conflicting objectives such as localized editing and subject-driven generation. To mitigate this, the authors propose a task-aware Mixture-of-Experts (MoE) routing mechanism that leverages hierarchical task semantic annotations and prediction alignment regularization to guide the gating network in dispatching experts according to high-level semantic intent. This approach transforms the gate from a task-agnostic executor into a semantics-driven scheduler, enabling semantically grounded expert specialization. While preserving sparse activation, the method significantly outperforms dense baselines, achieving notable improvements in generation fidelity, editing controllability, and expert specialization.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• task interference
• Mixture-of-Experts
• image generation
• image editing
• task-aware gating

Innovation

• Task-Aware Gating
• Mixture-of-Experts
• Task Interference
• Semantic Routing
• Diffusion Transformers