[논문 리뷰] TAG-MoE: Task-Aware Gating for Unified Generative Mixture-of-Experts

Yu Xu, Hongbin Yan외 12명|arXiv (Cornell University)|2026. 01. 12.

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.

연구 동기 및 목표

• 밀집 확산 트랜스포머를 사용한 통합 이미지 생성/편집에서 심각한 작업 간 간섭을 동기 부여하고 해결한다.
• 작업 인식 라우팅을 보장하면서 희소 Mixture-of-Experts (MoE)를 사용해 모델 용량을 확장한다.
• 전역 작업 의도 신호를 제공하기 위해 계층적 작업 의미 주석을 도입한다.
• 작업 의미와 MoE 라우팅을 정렬하기 위한 예측 정렬 정규화를 개발한다.

제안 방법

• 확산 트랜스포머의 마지막 10개 계층에서 FFN을 MoE 계층으로 바꿔 고정 파라미터에서 용량을 늘린다.
• 각 작업의 범위, 유형, 보존 여부를 설명하기 위해 계층적 작업 의미 주석을 도입한다.
• 작업 태그로부터 글로벌 시맨틱 임베딩을 계산하고 MoE 게이트 점수로부터 집계된 라우팅 시그니처를 얻는다.
• 코사인 유사도 손실을 통해 라우팅 시그니처를 시맨틱 임베딩과 맞추도록 예측 정렬 헤드를 학습한다.
• 주 흐름 손실과 MoE 부하 균형 손실에 L_align를 가중 합산한 총 목적 함수를 구성한다.
• 통합 생성/편집 작업을 포괄하는 대규모 다변량 데이터셋에 대해 학습한다.

실험 결과

연구 질문

• RQ1 작업 인식 게이팅 메커니즘이 통합 이미지 생성/편집에서 작업 간 간섭을 완화할 수 있는가?
• RQ2 계층적 작업 의미와 예측 정렬이 고수준 작업 의도에 맞는 MoE 라우팅을 개선하는가?
• RQ3 통합 벤치마크에서 오픈 소스 대비 클로즈드 소스 기준선에서 TAG-MoE의 성능은 어떠한가?
• RQ4 실무에서 전문가들이 의미적으로 의미 있는 특화 및 공간적으로 의식된 라우팅을 개발하는가?

주요 결과

• TAG-MoE는 통합 생성/편집 벤치마크에서 오픈 소스 기준선들 중 최첨단 성능을 달성한다.
• 의미 정렬된 라우팅은 작업 비의존적 MoE 라우팅에 비해 지시 사항 준수와 의미 정렬을 향상시킨다.
• 삭제실험 결과 밀집 모델은 실패하고 L_align가 의미 가이드된 전문가 라우팅에 결정적임이 드러난다.
• MoE와 예측 정렬은 관련 편집 영역에 공간적으로 국한된 명확한 전문가 특화를 보여준다.
• 사용자 연구는 참조 정렬, 프롬프트 정렬, 전반적 품질에서 TAG-MoE에 대한 명확한 선호를 나타낸다.