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

Computer Science > Computer Vision and Pattern Recognition arXiv:2601.08881 (cs) [Submitted on 12 Jan 2026 (v1), last revised 26 Mar 2026 (this version, v2)] Title:TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts Authors: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 View a PDF of the paper titled TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts, by Yu Xu and 13 other authors View PDF HTML (experimental) 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...