Source-linked AI summary
TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts
Yu Xu, Hongbin Yan, Juan Cao, Yiji Cheng, Tiankai Hang, Runze He, Zijin Yin, Shiyi Zhang, Yuxin Zhang, Jintao Li, Chunyu Wang, Qinglin Lu, Tong-Yee Lee, Fan Tang
TL;DR
Unified image generation and editing models suffer from task interference because shared parameters and task-agnostic routing must handle conflicting objectives. TAG-MoE adds hierarchical task semantics and predictive alignment regularization to guide MoE routing, and it reports improved performance over dense and task-agnostic MoE baselines while exposing a scope boundary in visual understanding.
Problem
Shared parameters and task-agnostic MoE gating struggle to handle conflicting requirements across unified generation and editing tasks.
Method
TAG-MoE uses hierarchical task semantic annotations and predictive alignment regularization to align MoE routing with global task intent.
Results
TAG-MoE mitigates task interference and outperforms dense models and task-agnostic MoE baselines in quantitative metrics and qualitative fidelity.
Takeaways & Limitations
The framework yields meaningful expert specialization by injecting global task intent into local routing decisions.
Takeaways & Limitations
The framework relies on pre-processed intent and cannot jointly reason over instructions and source-image content, limiting integrated semantic-perceptual tasks.
Abstract
from arXiv · showhide
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.
1. Introduction
Unified generation and editing models consolidate diverse image-manipulation tasks but face severe interference from conflicting objectives and task-agnostic routing. TAG-MoE injects structured task semantics into MoE routing to promote specialization.
- Unified models combine subject customization, style transfer, inpainting, and instruction-based editing in a single dense diffusion-transformer system.
- Conflicting objectives force shared parameters toward a mediocre compromise, degrading performance across user intents.
- Standard MoE expands capacity but conventional gating remains task-agnostic and fails to account for heterogeneous tasks in unified generation.
- TAG-MoE represents each task through scope, semantic type, and preservation constraints using hierarchical semantic annotations.
- Predictive alignment regularization correlates task semantics with routing decisions, turning local gating into an intent-aware dispatch mechanism.
2. Related Work
Related work has advanced unified image generation and editing, but existing approaches do not explicitly model the conflicting requirements of different tasks. MoE improves capacity, yet purely data-driven routing remains limited.
- Unified models evolved from sequence-to-sequence token concatenation toward improved multimodal conditioning and specialized embeddings.
- Editing requires regional preservation while modification occurs elsewhere, whereas customization requires identity consistency across new contexts.
- Without explicit modeling of competing requirements, unified approaches struggle to adapt across user intents and generalize robustly.
- Sparse MoE architectures increase capacity through expert routing without proportional per-sample computation, motivating their use in visual generation.
- Existing MoE applications use data-driven routing, leaving limitations when task diversity and conflicting objectives must be handled jointly.
3. Method
TAG-MoE combines an MM-DiT with hierarchical task semantics and a semantic-aligned router. The router aligns expert selection with explicit task meaning to guide specialization interpretably.
- The framework uses an MM-DiT with MoE layers for efficient, dynamic handling of diverse generation and editing tasks.
- Hierarchical task semantic annotation supplies explicit task descriptors for routing.
- A semantic-aligned router aligns routing decisions with task semantics to guide interpretable expert specialization.
3.1. MoE-based Multimodal Diffusion Transformer
The method processes text and image information as one MM-DiT token sequence, replacing later image-stream FFNs with sparse MoE layers. This increases capacity while keeping activation costs comparable and uses flow matching for end-to-end training.
- Text embeddings, conditional-image latents, noisy target latents, and timestep information are concatenated into one MM-DiT input sequence.
- MoE layers replace image-stream FFNs in later transformer blocks, where deeper semantic synthesis benefits from increased capacity.
- A gating network maps each token to expert probabilities, selects top-k experts, and combines their outputs as a weighted sum.
- Sparse activation increases model capacity at a fixed activation-parameter budget compared with dense models.
- The MoE-enhanced architecture is trained end-to-end with a Flow Matching objective.
3.2. Hierarchical Task Semantic Annotation
TAG-MoE represents each unified generation or editing task with structured semantic descriptors rather than a single coarse label. The annotation covers operational scope, semantic type, and preservation constraints, and is used during training to supervise task-aware routing.
- Hierarchical annotation: The three-tier annotation scheme decomposes each task into scope, type, and preservation descriptors.Examples include local or global editing, attribute or action editing, and identity or style preservation.
- Annotation pipeline: An automated Qwen-VL pipeline analyzes training triplets and outputs atomic tags under continuously refined rules.The triplets contain a source image, instruction, and target image.
- Annotation example: A sunglasses-editing example receives tags for local editing, object editing, identity preservation, background preservation, and style preservation.
- Inference: The hierarchical tags are used only during training; at inference, a VLM rewrites the user instruction into a detailed prompt for MM-DiT.The enriched prompt is encoded as text embedding C and fed into the MM-DiT.
3.3. Semantic-Aligned Gating Network
The semantic-aligned gating network converts task annotations and expert usage into vectors, then trains routing to predict global task semantics. Its auxiliary alignment loss is combined with generation and load-balancing objectives.
- Router overview: The router aligns a semantic embedding s with an aggregated routing signature g through predictive alignment.The semantic embedding represents task intent, while the routing signature summarizes the model’s expert usage.
- Global semantic embedding: The global semantic embedding s sums learned embeddings for the sample’s atomic tags and is permutation-invariant.It provides the structured supervisory signal for the alignment loss.
- Aggregated routing signature: Routing scores are generated for every token across all MoE layers, then averaged across layers and the token dimension to form g.The resulting vector captures which experts are activated on average for the sample.
- Predictive alignment: A two-layer prediction head projects g from expert space R^N into semantic space R^D, producing predicted semantics ŝ.The alignment loss compares ŝ with the target semantic embedding s using cosine similarity.
- Predictive alignment: Minimizing the alignment loss backpropagates through g to all gating networks, encouraging routing signatures that contain information about global task semantics.
- Overall training objective: The total objective combines generation loss L_flow, load-balancing loss L_lbl, and semantic alignment loss L_align with balancing hyperparameters.The alignment term is auxiliary to the primary generation objective.
3.4. Dataset Construction
TAG-MoE is trained on a hybrid dataset combining public benchmarks with proprietary in-house data. The construction pipeline targets broad coverage of unified generation and editing tasks.
- Dataset composition: The training set contains over 11 million samples from public and proprietary in-house sources.The public portion contains 2.2 million samples from established editing, virtual try-on, and subject-driven generation datasets.
- In-house construction: The in-house dataset uses staged instruction and target-image generation to cover specialized tasks.Images are sourced from public datasets, language models generate diverse instructions, and specialist and generalist models produce target images.
4. Experiments
TAG-MoE is evaluated across unified generation, editing, subject-driven generation, qualitative comparisons, ablations, expert specialization, and user preference. It achieves strong instruction adherence, preservation, and robustness while showing that sparse semantically guided routing mitigates task interference.
- Implementation: The model replaces final-layer FFNs with four-expert MoE layers using a two-layer MLP gate and top-1 routing.The MoE layers are inserted into the image stream’s final 10 diffusion-transformer layers.
- Evaluation Setup: Evaluation combines ICE-Bench with specialized editing and subject-driven benchmarks, using metrics for aesthetics, alignment, instruction execution, and preservation.ICE-Bench covers diverse unified tasks; EmuEdit-Bench, GEdit-Bench, DreamBench++, and OmniContext provide finer-grained evaluation.
- Unified Generation Evaluation: The model leads open-source baselines on ICE-Bench aesthetic quality, CLIP-cap, and vllmqa, with CLIP-cap also exceeding GPT-4o and Gemini-2.5-flash.Its results provide a more favorable balance of instruction adherence and semantic alignment than baselines with higher source fidelity.
- Unified Generation Evaluation: Across 26 ICE-Bench task categories, the model achieves state-of-the-art performance in most categories with robust and well-balanced capability.The analysis cautions that DreamOmni2’s high reference-generation scores can be inflated by copy-paste behavior.
- Image Editing Evaluation: On EmuEdit-Bench and GEdit-Bench, the model achieves the highest vllmqa scores on both benchmarks despite not leading every metric.The vllmqa measure evaluates whether the instruction was correctly executed from the source image, instruction, and output.
- Subject-Driven Evaluation: For subject-driven generation, the model attains SOTA Face-ref scores on both benchmarks, the highest DreamBench++ Style-ref, and the top OmniContext DINO-ref.These results indicate strong subject, identity, and style preservation while remaining competitive on other preservation measures.
- Qualitative Comparison: Qualitative comparisons show improved handling of interfering intents, including copy-paste artifacts, stylistic dissonance, incomplete compositional edits, and conflicting local-global requirements.Predictive Alignment Regularization routes conflicting subtasks to specialized experts.
- Qualitative Comparison: Against specialized models, TAG-MoE handles complex 3D-aware edits and preserves identity while following demanding motion instructions.Specialized editing baselines struggle with structural changes, while subject-driven baselines compromise identity, context, or pose accuracy.
5. Limitations and Future Work
TAG-MoE’s key limitation is that it relies on pre-processed intent separately from visual content, restricting integrated semantic and perceptual understanding.
- The framework cannot jointly reason over editing intent and the source image’s visual content.It understands task attributes such as scope and type but not contextual information contained in pixels.
- This separation restricts tasks requiring integrated semantic and perceptual understanding.
- The model fails at content-based reasoning, such as solving a math problem shown in an image.
6. Conclusion
TAG-MoE addresses task-agnostic routing in unified image generation and editing with semantic task annotations and predictive alignment regularization. It mitigates task interference and outperforms dense models and task-agnostic MoE baselines in quantitative metrics and qualitative fidelity.
- TAG-MoE identifies task-agnostic routing as the core bottleneck for applying MoE to diverse, conflicting tasks.
- The framework injects global task intent into local routing through Hierarchical Task Semantic Annotation and Predictive Alignment regularization.
- TAG-MoE mitigates task interference while outperforming dense models and task-agnostic MoE baselines in quantitative metrics and qualitative fidelity.