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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Xu, Yu, Yan, Hongbin, Cao, Juan, Cheng, Yiji, Hang, Tiankai, He, Runze, Yin, Zijin, Zhang, Shiyi, Zhang, Yuxin, Li, Jintao, Wang, Chunyu, Lu, Qinglin, Lee, Tong-Yee, Tang, Fan
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915891445235712
author Xu, Yu
Yan, Hongbin
Cao, Juan
Cheng, Yiji
Hang, Tiankai
He, Runze
Yin, Zijin
Zhang, Shiyi
Zhang, Yuxin
Li, Jintao
Wang, Chunyu
Lu, Qinglin
Lee, Tong-Yee
Tang, Fan
author_facet Xu, Yu
Yan, Hongbin
Cao, Juan
Cheng, Yiji
Hang, Tiankai
He, Runze
Yin, Zijin
Zhang, Shiyi
Zhang, Yuxin
Li, Jintao
Wang, Chunyu
Lu, Qinglin
Lee, Tong-Yee
Tang, Fan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_08881
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts
Xu, Yu
Yan, Hongbin
Cao, Juan
Cheng, Yiji
Hang, Tiankai
He, Runze
Yin, Zijin
Zhang, Shiyi
Zhang, Yuxin
Li, Jintao
Wang, Chunyu
Lu, Qinglin
Lee, Tong-Yee
Tang, Fan
Computer Vision and Pattern Recognition
Artificial Intelligence
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.
title TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts
topic Computer Vision and Pattern Recognition
Artificial Intelligence
url https://arxiv.org/abs/2601.08881