TAG-MoE: Gate Task-Aware per un Modello Generativo Unificato a Mistura di Esperti

Abstract

I modelli unificati di generazione e modifica di immagini soffrono di una grave interferenza tra compiti nelle architetture dense di transformer diffusivi, dove uno spazio parametrico condiviso deve scendere a compromessi tra obiettivi conflittuali (ad esempio, modifica locale rispetto a generazione guidata dal soggetto). Sebbene il paradigma sparso della Mistura di Esperti (MoE) sia una soluzione promettente, le sue reti di gating rimangono task-agnostic, operando sulla base di caratteristiche locali, inconsapevoli dell'intento globale del compito. Questa natura task-agnostic impedisce una specializzazione significativa e non risolve l'interferenza sottostante tra i compiti. In questo articolo, proponiamo una nuova architettura per iniettare l'intento semantico nell'instradamento MoE. Introduciamo uno schema di Annotazione Semantica Gerarchica dei Compiti per creare descrittori strutturati dei compiti (ad esempio, ambito, tipo, preservazione). Progettiamo quindi una Regolarizzazione per Allineamento Predittivo per allineare le decisioni di instradamento interne con la semantica di alto livello del compito. Questa regolarizzazione evolve la rete di gating da un esecutore task-agnostic a un centro di smistamento. Il nostro modello mitiga efficacemente l'interferenza tra compiti, superando i baseline densi in fedeltà e qualità, e la nostra analisi mostra che gli esperti sviluppano naturalmente specializzazioni chiare e semanticamente correlate.

English

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.

TAG-MoE: Gate Task-Aware per un Modello Generativo Unificato a Mistura di Esperti

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

Abstract

Support