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

arXiv: https://arxiv.org/abs/2601.08881 · HTML: https://mlsanity.com/paper/2601.08881/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 55 of its 67 references that are in the map.
- First posted on arXiv 12 Jan 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)
- Mixture-of-experts / [Vision mixture-of-experts](https://mlsanity.com/moe/vision-mixture-of-experts/)

## Builds on (most cited first, 30 of 55)

- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Qwen2.5-VL Technical Report](https://mlsanity.com/paper/2502.13923/) (2025)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [GPT-4o System Card](https://mlsanity.com/paper/2410.21276/) (2024)
- [Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution](https://mlsanity.com/paper/2409.12191/) (2024)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation](https://mlsanity.com/paper/2208.12242/) (2022)
- [InstructPix2Pix: Learning to Follow Image Editing Instructions](https://mlsanity.com/paper/2211.09800/) (2022)
- [Qwen-Image Technical Report](https://mlsanity.com/paper/2508.02324/) (2025)
- [FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space](https://mlsanity.com/paper/2506.15742/) (2025)
- [Emerging Properties in Unified Multimodal Pretraining](https://mlsanity.com/paper/2505.14683/) (2025)
- [MagicBrush: A Manually Annotated Dataset for Instruction-Guided Image Editing](https://mlsanity.com/paper/2306.10012/) (2023)
- [Step1X-Edit: A Practical Framework for General Image Editing](https://mlsanity.com/paper/2504.17761/) (2025)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [OmniGen: Unified Image Generation](https://mlsanity.com/paper/2409.11340/) (2024)
- [OminiControl: Minimal and Universal Control for Diffusion Transformer](https://mlsanity.com/paper/2411.15098/) (2024)
- [Emu Edit: Precise Image Editing via Recognition and Generation Tasks](https://mlsanity.com/paper/2311.10089/) (2023)
- [VITON-HD: High-Resolution Virtual Try-On via Misalignment-Aware Normalization](https://mlsanity.com/paper/2103.16874/) (2021)
- [UltraEdit: Instruction-based Fine-Grained Image Editing at Scale](https://mlsanity.com/paper/2407.05282/) (2024)
- [OmniGen2: Towards Instruction-Aligned Multimodal Generation](https://mlsanity.com/paper/2506.18871/) (2025)
- [In-Context Edit: Enabling Instructional Image Editing with In-Context Generation in Large Scale Diffusion Transformer](https://mlsanity.com/paper/2504.20690/) (2025)
- [Less-to-More Generalization: Unlocking More Controllability by In-Context Generation](https://mlsanity.com/paper/2504.02160/) (2025)
- [HunyuanImage 3.0 Technical Report](https://mlsanity.com/paper/2509.23951/) (2025)
- [OmniEdit: Building Image Editing Generalist Models Through Specialist Supervision](https://mlsanity.com/paper/2411.07199/) (2024)
- [DreamBench++: A Human-Aligned Benchmark for Personalized Image Generation](https://mlsanity.com/paper/2406.16855/) (2024)
- [Measuring Style Similarity in Diffusion Models](https://mlsanity.com/paper/2404.01292/) (2024)
- [UniReal: Universal Image Generation and Editing via Learning Real-world Dynamics](https://mlsanity.com/paper/2412.07774/) (2024)
- [ACE++: Instruction-Based Image Creation and Editing via Context-Aware Content Filling](https://mlsanity.com/paper/2501.02487/) (2025)
- [DreamO: A Unified Framework for Image Customization](https://mlsanity.com/paper/2504.16915/) (2025)

## Built on by (most cited first, 6 of 6)

- [Beyond Pixels: Visual Metaphor Transfer via Schema-Driven Agentic Reasoning](https://mlsanity.com/paper/2602.01335/) (2026)
- [Meta-CoT: Enhancing Granularity and Generalization in Image Editing](https://mlsanity.com/paper/2604.24625/) (2026)
- [GeodesicNVS: Probability Density Geodesic Flow Matching for Novel View Synthesis](https://mlsanity.com/paper/2603.01010/) (2026)
- [FontCrafter: High-Fidelity Element-Driven Artistic Font Creation with Visual In-Context Generation](https://mlsanity.com/paper/2603.22054/) (2026)
- [Symbiotic-MoE: Unlocking the Synergy between Generation and Understanding](https://mlsanity.com/paper/2604.07753/) (2026)
- [Breaking the Uniformity Trap: Scaling Video Diffusion Model via SplitMoE](https://mlsanity.com/paper/2609.38140/) (2026)

## Abstract (arXiv)

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.