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

TL;DR

TAG-MoE tackles task interference in unified image generation and editing by injecting global task semantics into sparse Mixture-of-Experts routing. It introduces a hierarchical task semantic annotation and a predictive alignment regularization that forces the gating network to align routing with macroscopic task intents, enabling meaningful expert specialization. The approach uses a MM-DiT backbone with MoE layers in the deeper blocks, an explicit semantic routing signal, and a large-scale mixed dataset of public and in-house samples (~11M). Empirical results across ICE-Bench, EmuEdit, GEdit, DreamBench++, and OmniContext show state-of-the-art open-source performance and robust task adherence, with qualitative evidence of spatially localized, semantically guided routing. The work offers a practical path to scalable, unified generative models that respect diverse user intents while maintaining high fidelity and editing accuracy.

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.

Figures (7)

• Figure 1: We present TAG-MoE, by injecting high-level task semantic intent into the local routing decisions of the MoE gating network, we enabling the diffusion transformer model to handle diverse generative tasks.
• Figure 2: Pipeline of our method. TAG-MoE consists of: (1) A MM-DiT with MoE layers; (2) A Hierarchical Task Semantic Annotation that labels training data with atomic task descriptors; (3) A novel Semantic-Aligned Router explicitly aligns MoE routing behavior with task semantics through Predictive Alignment Regularization.
• Figure 3: Qualitative comparison on diverse tasks. Our model successfully resolves complex task conflicts where baselines fail.
• Figure 4: Comprehensive scores on different image editing and generation tasks.
• Figure 5: Compare with specialized image editing models and subject-driven generation models.