dorsal/arxiv
View SchemaTAG-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 |
|---|---|
| Categories | |
| ArXiv ID | 2601.08881vv1 |
| URL | https://arxiv.org/abs/2601.08881 |
| License | http://arxiv.org/licenses/nonexclusive-distrib/1.0/ |
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.
{
"annotation_id": "f0a4991f-13f4-499a-a2f6-4b475d5e7537",
"date_created": "2026-02-17T05:53:19.289000Z",
"date_modified": "2026-02-17T05:53:19.289000Z",
"file_hash": "df4f9bb91ffe2cc9a7b08330fca6fefbd30d559e1906017f8d7a30d35236099e",
"private": false,
"record": {
"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\u0027s 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.",
"arxiv_id": "2601.08881",
"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"
],
"categories": [
"cs.CV",
"cs.AI"
],
"license": "http://arxiv.org/licenses/nonexclusive-distrib/1.0/",
"title": "TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts",
"url": "https://arxiv.org/abs/2601.08881",
"version": "v1"
},
"schema_id": "dorsal/arxiv",
"source": {
"execution_id": "ae1a28b5-cfbe-4999-9fb8-8e243efc5da8",
"id": "arXiv Dataset",
"type": "Model",
"variant": "snapshot-2026-01-17",
"version": "0.1.0"
},
"user_id": 1000002
}