[论文解读] TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts

Yu Xu, Hongbin Yan+12 more|arXiv (Cornell University)|Jan 12, 2026

Generative Adversarial Networks and Image Synthesis被引用 0

一句话总结

TAG-MoE 引入一个面向任务的门控机制，用于统一图像生成/编辑，配合分层的任务语义注解和预测对齐正则化，根据高层次任务意图来路由扩散 MoE 专家，降低任务干扰。

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.

研究动机与目标

• 用以应对统一图像生成/编辑中的严重任务干扰问题，结合密集扩散变换器。
• 在固定参数下通过稀疏混合专家 MoE 提升模型容量，同时确保面向任务的路由。
• 引入分层的任务语义注解，为全局任务意图信号提供描述。
• 开发预测对齐正则化，使 MoE 路由与任务语义对齐。

提出的方法

• 将扩散变换器最后 10 层的前馈网络（FFN）替换为 MoE 层，在固定参数量下增加容量。
• 引入分层任务语义注解，描述每个任务的范围、类型与保留要素。
• 从任务标签中计算全局语义嵌入，以及从 MoE 门控分数聚合得到路由签名。
• 训练一个预测对齐头，通过余弦相似性损失使路由签名与语义嵌入对齐。
• 在总目标中用加权的 L_align、主流损失和 MoE 负载均衡损失组合成最终目标。
• 在覆盖统一生成/编辑任务的大规模、多样数据集上进行训练。

实验结果

研究问题

• RQ1 面向任务的门控机制能否缓解统一图像生成与编辑中的跨任务干扰？
• RQ2 分层任务语义和预测对齐是否能改善 MoE 路由，使之更匹配高层次的任务意图？
• RQ3 在统一基线的开源与闭源对比中，TAG-MoE 的表现如何？
• RQ4 专家是否能够在实践中发展出具有语义意义的专业化和空间感知的路由？

主要发现

• 在统一生成/编辑基准上，TAG-MoE 达到开源基线中的最先进性能。
• 语义对齐的路由相对任务无关的 MoE 路由，在遵循指令和语义对齐方面有改进。
• 消融研究表明密集模型容易失效，且 L_align 对语义引导的专家路由至关重要。
• 具备预测对齐的 MoE 显示出清晰的专家专业化，并在相关编辑区域具有空间定位性。
• 用户研究表明，在参考对齐、提示对齐和整体质量方面，用户普遍更偏好 TAG-MoE。