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

🎯 Pitch

Standard mixture-of-experts routers are blind to what task you’re actually asking for, so editing fidelity and creative generation end up sabotaging each other. TAG-MoE fixes this by injecting hierarchical task semantics directly into the routing signal, forcing experts to specialize by intent rather than by local token patterns. The result: the right experts automatically concentrate on the right image regions for each operation, enabling a single unified model to beat specialized editors and generators alike.

1. Executive Summary

This paper proposes TAG-MoE, a task-aware Mixture-of-Experts framework that injects high-level task semantics into the local routing decisions of diffusion transformer gating networks for unified image generation and editing. Built on a Qwen-Image MM-DiT backbone and evaluated across ICE-Bench, EmuEdit, GEdit, DreamBench++, and OmniContext, TAG-MoE introduces a Hierarchical Task Semantic Annotation scheme (decomposing each task into Scope, Type, and Preservation tags) paired with Predictive Alignment Regularization (a cosine similarity loss that forces the aggregated routing signature to predict the task's global semantic embedding). This alignment enables experts to develop spatially-aware, task-specific specializations—different expert combinations activate for "change material" versus "change color," and active experts concentrate computation precisely on the relevant image region—achieving state-of-the-art overall performance among open-source baselines on ICE-Bench while surpassing specialized editing and subject-driven models on the vllmqa instruction-following metric, establishing that task-aware routing resolves the interference between conflicting objectives (editing fidelity versus generative diversity) that cripples standard task-agnostic MoE gating in unified models.

2. Context and Motivation

The Core Problem: Task Interference in Unified Image Generation Models

The fundamental problem this paper tackles is deceptively simple to state but devastating in practice: when you train a single model to perform diverse image generation and editing tasks, the objectives of those tasks actively fight each other. The shared parameter space must simultaneously learn to do two fundamentally contradictory things — preserve content faithfully during editing while also generating novel, diverse content during subject-driven generation — and the result is a model that does neither particularly well.

The paper characterizes this failure mode precisely in Section 1:

"The shared parameter space must simultaneously execute inherently contradictory objectives: local editing demands precise content preservation, while subject-driven generation requires expressive diversity and novel synthesis. This fundamental conflict forces the network toward a 'mediocre compromise solution,' preventing the necessary representational specialization and ultimately degrading performance across the spectrum of user intents."

This is not merely an efficiency concern. The "mediocre compromise" manifests as concrete, observable failures visible in the qualitative results (Figure 3). Unified baselines produce "copy-paste" artifacts in subject-driven generation — the model simply pastes the source subject into the output without adapting it to the new context, because learning strong identity preservation (useful for editing) conflicts with learning flexible generation. In inpainting tasks, baselines exhibit stylistic dissonance between the inpainted region and the background. In compositional editing, baselines show incomplete execution where only part of the instruction is followed. These are not subtle degradations — they represent fundamental functional breakdowns caused by the interference.

Why This Problem Matters Now

The urgency of this problem is driven by a clear trend in the field: the industry is rapidly converging toward unified image generation and editing models, as documented in Section 1 with references to GPT-4o, Gemini 2.5, and Flux-Kontext. The vision is a single model that handles everything — object removal, style transfer, identity-preserving customization, instruction-based editing, subject insertion, and so on. This unification promises substantial practical benefits: reduced infrastructure complexity (one deployment instead of task-specific models), simpler product interfaces (one input channel for all image manipulation), and better utilization of large-scale pretraining data.

But the paper argues that this unification is "critically bottlenecked by severe task interference" (Section 1). Without solving the interference problem, the unified model vision cannot deliver on its promise. The key insight is that the bottleneck is architectural, not merely a matter of more data or more parameters: a dense model with a shared parameter space fundamentally cannot represent the distinct processing pathways needed for conflicting task families.

The paper's own ablation (Figure 6 left, Table 4) provides direct evidence: a dense model with equivalent activated parameters to the MoE model shows "a severe performance drop on ICE-Bench metrics" and slower convergence. This validates that the problem isn't just capacity — it's the inability of dense architectures to structurally disentangle the conflicting objectives.

Prior Approaches and Where They Fall Short

The paper situates its work against three streams of prior approaches, each with critical limitations:

1. Unified Architectures Without Task-Aware Routing (Section 2.1)

The first wave of unified models treated image generation and editing as a sequence-to-sequence problem, concatenating text embeddings with source and target image tokens into a single input for large transformers (OmniGen, ACE, UniVG). Subsequent works iterated on architectural details: trainable index and condition embeddings for better multimodal alignment (UniReal, RealGeneral), 3D rotary positional encodings to distinguish source from target (Flux-Kontext), dual-branch designs that decouple subject and background processing (BlobCtrl), and integrated MLLMs for improved scene understanding (BAGEL, OmniGen2, Query-Kontext).

The critical failure, as the paper diagnoses it in Section 2.1, is that none of these works model the inherent conflict between task objectives:

"Without explicitly modeling these distinct—and often competing—requirements, existing approaches struggle to adaptively serve the full spectrum of user intents, limiting their practical robustness and generalization."

In other words, these models rely on the same dense parameters to handle every task. The model learns some task-specific behaviors implicitly through data, but there is no structural mechanism to prevent editing-specific adaptations from interfering with generation-specific adaptations. The result is the "mediocre compromise" that degrades performance everywhere.

2. MoE in Image Generation Without Task Awareness (Section 2.2)

The Mixture-of-Experts paradigm has been successfully adopted in diffusion transformers for scaling capacity — pioneering works like DiT-MoE scaled diffusion transformers to 16 billion parameters by replacing FFN layers with sparse MoE layers, and subsequent variants (HunyuanImage-3.0, Dense2MoE) showed that sparse expert architectures enhance expressiveness. ICEdit took the next step by integrating LoRA-based MoE modules into attention blocks for image editing specifically.

However, the paper identifies a critical architectural failure when applying standard MoE to the unified multi-task setting:

"Standard routers rely solely on local token features, remaining entirely oblivious to the high-level, global task intent (e.g., 'identity preservation' or 'style modification'). This profound information gap between the local gate and the global objective leads to spontaneous, inefficient expert specialization, fundamentally failing to structurally disentangle multi-task interference."

This is the paper's core diagnosis. In standard MoE, the gating network looks at individual token representations to decide which expert to route them to. But a token's local features — say, a pixel belonging to a person's face — don't encode whether the model is supposed to preserve that face (editing task) or adapt that face to a new context (subject-driven generation). The gate makes its decision in a task-agnostic vacuum.

The consequence, demonstrated in the ablation (Table 4, "MoE w/o ℒ_align" variant), is that a task-agnostic MoE does benefit from its larger effective capacity (it still outperforms the dense baseline), but it leaves substantial performance on the table compared to task-aware routing. The experts do not develop meaningful task-specific specialization because the gate has no signal about what task it should be specializing for.

3. Conflicting Findings Across Prior Work

The paper doesn't explicitly frame prior work as contradictory, but the tension is implicit: Task-specific editing models (InstructPix2Pix, MagicBrush, UltraEdit, ICEdit, EmuEdit) perform well on their target tasks but cannot generalize to subject-driven generation. Task-specific subject-driven models (DreamO, OmniControl, UNO) excel at identity preservation but cannot handle general editing. And existing unified models (ACE++, Flux-Kontext, BAGEL, OmniGen2, DreamOmni2) attempt both but degrade relative to specialists, especially on complex tasks where the objectives conflict most severely.

This sets up a clear research question: can a unified model match or exceed specialized models across both editing and generation tasks by structurally resolving the interference between them?

How This Paper Positions Itself

The paper positions TAG-MoE as a third way that is fundamentally different from both dense unified models and task-agnostic MoE. The intellectual move is elegantly summarized in Section 2.2:

"How to inject the high-level, global task semantics into the local MoE routing mechanism to enable task-aware specialization remains an open challenge."

The paper's answer has two components that work together:

First, the Hierarchical Task Semantic Annotation scheme (Section 3.2) creates a structured representation of what the model should be doing. Rather than relying on a coarse task label ("edit this image"), it decomposes each training instance into Scope (local/global editing, content customization), Type (object editing, style transfer, attribute editing), and Preservation constraints (identity, background, structure). This three-tier annotation captures the semantic intent that the routing mechanism needs to be aware of — it tells the model, in a structured form, what facets of the task matter.

Critically, the authors note in Section 3.2 that "a single coarse label (e.g., 'edit') cannot capture user intent." The same surface label ("edit") covers "change the background to a beach" (global editing, low identity preservation needed) and "make the person smile" (local editing, high identity preservation needed). The hierarchical annotation disambiguates these cases.

Second, the Predictive Alignment Regularization (Section 3.3) creates the architectural bridge. The key principle is stated upfront:

"semantically similar generation tasks evoke similar expert usage patterns."

The alignment loss enforces that the model's actual routing behavior — aggregated into a routing signature vector $\mathbf{g}$ — must be predictive of the task's semantic embedding $\mathbf{s}$ . A lightweight prediction head maps from routing space to semantic space, and the cosine similarity loss backpropagates through the gating network parameters. This means the gate cannot just route based on local token features — it must route in a way that, when aggregated across the entire sample, encodes global information about what task is being performed.

The paper frames this as transforming the gate's functional role: it "evolves from a task-agnostic executor into a semantic-aware scheduler" (Section 3.3.3). This is not an incremental tweak to MoE — it redefines what the routing decision means. The gate's output now carries task-level semantic information, which forces experts to develop specializations aligned with task semantics rather than low-level feature patterns.

Importantly, the paper positions this as a training-time mechanism only. During inference, "ground-truth tags are no longer required" (Section 3.2). Instead, a lightweight preprocessing step uses a VLM to rewrite the user's instruction into a more detailed prompt, which implicitly carries the task semantics that the model learned to route on. This means the method adds no inference-time overhead from the annotation pipeline — the semantic-awareness is baked into the trained routing weights.

The paper's positioning relative to prior work can be summarized as: prior unified models fail because they lack any mechanism to structurally resolve task interference; prior MoE models fail because their task-agnostic routing cannot develop meaningful task specialization; TAG-MoE succeeds by making the routing mechanism task-aware, which allows experts to specialize along semantically meaningful task boundaries. The evidence for this claim comes from both the quantitative results (Tables 1-3, where TAG-MoE leads open-source baselines) and the expert activation analysis (Section 4.5 and Figure 6), which shows experts developing task-specific and spatially-aware activation patterns that correspond precisely to the hierarchical task tags.

3. Technical Approach

3.1 Reader Orientation

TAG-MoE is a diffusion transformer with MoE layers whose gating network is trained to be aware of what task it is performing, so that different experts can specialize for different types of image generation and editing operations. The system solves the problem of task interference — the fact that editing (which demands preservation) and generation (which demands creative diversity) require fundamentally different computational behaviors that a shared dense model cannot simultaneously accommodate — by structurally decoupling these conflicting objectives into distinct expert pathways that are activated based on the semantic intent of the current task, not just local pixel features.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components, operating in two phases (training with semantic supervision, inference without it):

MM-DiT Backbone with MoE Layers — a Multimodal Diffusion Transformer (built on Qwen-Image) where the final 10 transformer blocks have their image-stream FFN layers replaced by sparse MoE layers, each containing 4 experts with top-1 routing.
Hierarchical Task Semantic Annotation Pipeline — an automated system that uses Qwen-VL to label each training sample (source image, instruction, target image) with structured tags across three dimensions: Scope, Type, and Preservation. This produces the ground-truth semantic embedding $\mathbf{s}$ .
Semantic-Aligned Gating Network — the standard MoE router, but trained with an additional objective that forces its aggregated routing decisions (the vector $\mathbf{g}$ ) to predict the task's semantic embedding $\mathbf{s}$ . This makes the gate task-aware.
Predictive Alignment Regularization — the loss function $\mathcal{L}_{\text{align}}$ (cosine similarity between predicted semantics and ground-truth semantics) plus a lightweight 2-layer MLP prediction head $\mathcal{H}_{\text{pred}}$ that maps from routing space to semantic space.
VLM-Based Instruction Rewriting (Inference Only) — at inference time, Qwen-VL rewrites the user's raw instruction into a detailed prompt, which carries the implicit task semantics that the trained routing weights have learned to recognize.

Information flow during training: A training triplet (source image, instruction, target image) enters → the VLM annotation pipeline produces hierarchical tags → tags are embedded via a learned tag embedding matrix $\mathbf{W}_{\text{tag}}$ and summed into the global semantic embedding $\mathbf{s}$ → the MM-DiT processes the input tokens through MoE layers, producing per-token routing scores → scores are averaged across layers and tokens into the aggregated routing signature $\mathbf{g}$ → a prediction head maps $\mathbf{g}$ to a predicted semantic embedding $\hat{\mathbf{s}}$ → cosine similarity loss between $\mathbf{s}$ and $\hat{\mathbf{s}}$ backpropagates through the gating networks → the gate learns to route tokens in a way that encodes global task information.

Information flow during inference: The user provides a raw instruction + source image → a VLM rewrites the instruction into a detailed prompt → the prompt is encoded into text embeddings → the MM-DiT processes the full token sequence through the learned, semantically-aware gating networks → experts that specialized for the task's semantic profile are preferentially activated → the model generates the output image.

3.3 Roadmap for the Deep Dive

First, the MoE-based MM-DiT architecture itself: how the base model is structured, where MoE layers are inserted, the routing mechanism, and why these design choices were made — this is the substrate everything else builds on.
Second, the Hierarchical Task Semantic Annotation scheme: the three-tier taxonomy, how the automated annotation pipeline works, and how tags are converted into the ground-truth semantic embedding $\mathbf{s}$ — this provides the supervisory signal that makes task-aware routing possible.
Third, the Semantic-Aligned Gating Network: how the aggregated routing signature $\mathbf{g}$ is constructed from per-token gating scores, how the prediction head maps routing to semantics, and the predictive alignment loss that connects them — this is the core innovation that bridges local routing with global task intent.
Fourth, the overall training objective: how the alignment loss is combined with the flow matching loss and load balancing loss, including the balancing hyperparameters — this completes the training recipe.
Fifth, the dataset construction and inference procedure: how training data is assembled and how the system operates at test time without ground-truth task tags — this closes the loop from training to deployment.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a method paper whose core idea is that MoE gating networks in unified image generation models must be made task-aware through semantic supervision, and that a predictive alignment objective between routing decisions and task semantics can force experts to develop meaningful, task-specific specializations that resolve the interference problem.

MoE-Based Multimodal Diffusion Transformer Architecture

The base architecture is a Multimodal Diffusion Transformer (MM-DiT) built on Qwen-Image, a text-to-image diffusion model. The paper takes this pretrained model and adapts it for the unified image generation and editing setting by modifying both its input processing pipeline and its internal architecture.

Input processing. The model accepts diverse inputs and unifies them into a single token sequence. The process works as follows:

A pretrained Multimodal Large Language Model (MLLM) encodes the text instruction $c_{\text{text}}$ into a sequence of text embeddings $C$ . The paper does not specify which MLLM is used for encoding during training, but during inference, Qwen-VL serves this role through instruction rewriting (§3.2).
A pretrained VAE encoder $\mathcal{E}$ maps both the conditional (source) image $I_c$ and the target image $I_0$ into latent representations $z_c$ and $z_0$ . During training, Gaussian noise is added to $z_0$ to produce a noisy latent $z_t$ at timestep $t$ .
Both $z_c$ and $z_t$ are patchified into sequences of visual tokens — flattened 2D patches of the latent representations, following the standard Vision Transformer convention established by Peebles and Xie (2023).
The complete input sequence is formed by concatenating: the text embeddings $C$ , the conditional image tokens from $z_c$ , the noisy target image tokens from $z_t$ , and a timestep embedding. This unified sequence is what the MM-DiT processes, and the architecture must learn to attend appropriately across these different token types to understand the relationship between the instruction, the source image, and the desired output.

MoE layer insertion strategy. The paper does not convert the entire model to MoE. Instead, it makes a targeted architectural modification:

"We replace the feed-forward networks (FFNs) of the image stream in diffusion transformer blocks with MoE layers... We only implement MoE layers in the later transformer blocks as high-level semantic synthesis in these deeper layers benefits most from the increased capacity."

Specifically, the MoE layers replace the FFNs of the image stream in the final 10 layers of the diffusion transformer (Section 4.1). This is a deliberate design choice grounded in observations from the MoE literature (DeepSpeed-MoE, DeepSeek-V3): earlier layers tend to learn low-level features (edges, textures, basic shapes) that are largely task-invariant, while deeper layers learn high-level semantic representations where task-specific processing matters most. By concentrating MoE capacity in these deeper layers, the model gets the benefits of specialization where it needs them while avoiding unnecessary overhead in layers where dense computation suffices.

The text stream (the MLLM-encoded instruction tokens) is not MoE-ified — only the image stream. This makes sense because the text tokens encode the task intent and remain relatively compact, while the image tokens are far more numerous and carry the representational complexity that benefits from expert specialization. The text stream can influence image processing through cross-attention without needing its own expert pathways.

MoE layer structure and routing. Each MoE layer consists of:

A set of $N$ expert networks $E = \{E_1, E_2, \ldots, E_N\}$ , where each expert has an architecture identical to the original FFN it replaces. This is a critical detail: switching from dense to MoE does not change the per-expert computation — it changes how many parallel experts are available and how inputs are distributed among them. Section 4.1 specifies that each MoE layer has four experts ( $N = 4$ ), and a top-1 routing strategy is used, meaning each token is sent to exactly one expert.
A gating network $\mathcal{G}$ (also called a router) that maps each input token to a probability distribution over the $N$ experts. The gating network is implemented as a two-layer MLP (Section 4.1). Given an input token representation $x \in \mathbb{R}^D$ (where $D$ is the model's hidden dimension), the gate produces scores $\mathcal{G}(x) \in \mathbb{R}^N$ — one logit per expert. The top- $k$ selection ( $k = 1$ ) picks the expert with the highest score.

The MoE forward pass for a single token is:

$\text{MoE}(x) = \sum_{E_i \in \mathcal{T}(x)} \mathcal{G}(x)_i \cdot E_i(x)$

where $\mathcal{T}(x) \subseteq \{E_1, \ldots, E_N\}$ is the set of top- $k$ selected experts for token $x$ , $\mathcal{G}(x)_i$ is the gating score (routing weight) assigned to expert $E_i$ , and $E_i(x)$ is the output of expert $E_i$ when processing token $x$ .

What it computes: For each image token, the gate evaluates which of the four experts is most appropriate, routes the token exclusively to that expert (since $k=1$ ), and weights the expert's output by the gate's confidence score. The output is a single vector — the expert's transformed representation — scaled by the routing weight.

Why this form: The sparse top- $k$ routing is the standard MoE formulation because it achieves the key property that makes MoE attractive: the total activated parameters per token is only $1/N$ of the total expert capacity (for top-1 with $N$ experts). This means the model can have $4\times$ the FFN parameters of a dense model while paying only approximately the same per-token inference cost (plus the small overhead of the gating network). The weighted sum formulation allows the gate to express confidence — if the gate is uncertain between experts, the selected expert's output is downweighted, providing a soft signal about routing ambiguity that gradients can flow through.

Training objective for the MM-DiT. The entire MoE-enhanced architecture is trained end-to-end using a Flow Matching objective (Section 3.1). Flow matching is a modern alternative to diffusion that learns a continuous normalizing flow between noise and data distributions rather than a score function. The paper inherits this objective from the Qwen-Image base model and does not modify it — the flow matching loss $\mathcal{L}_{\text{flow}}$ is the primary generation loss.

Load balancing. Without explicit intervention, MoE training can collapse to a state where the gate always routes to the same one or two experts (the "rich get richer" problem), defeating the purpose of having multiple experts. The paper includes a standard MoE load balancing loss $\mathcal{L}_{\text{lbl}}$ in the total objective (Equation 5), though the specific formulation (auxiliary loss encouraging uniform expert utilization) is not detailed — this is a well-established technique from the MoE literature that the paper treats as standard infrastructure.

Design choice: why MoE layers only in the image stream of later blocks. Beyond the general observation about deep-versus-shallow layer properties, this choice has a specific motivation in the unified generation setting. Early layers process raw latent features that are shared across all tasks — the VAE latent representation of an image of a person is similar regardless of whether the model will edit that person or use them as a subject for novel generation. But in deeper layers, the model needs to make task-specific decisions: should it preserve or transform this feature? Should it generate novel content or maintain consistency? By concentrating MoE in these deeper layers, the model deploys expert specialization precisely where task identity matters most, while avoiding the risk that early-layer MoE would learn spurious task-unrelated specializations.

Hierarchical Task Semantic Annotation

The goal of the annotation scheme is to provide a structured, machine-readable representation of what a given training example requires the model to do, so that this information can serve as a supervisory signal for training the gating network.

The three-tier taxonomy. The paper defines three orthogonal dimensions for characterizing any image generation or editing task:

Scope — the operational nature and spatial extent of the task. Examples given include: global editing, local editing, content customization. This dimension captures where in the image the model should operate. "Global editing" (e.g., changing the entire image to a painting style) requires broad attention across all tokens; "local editing" (e.g., adding sunglasses to a person) requires focused attention on a specific region while leaving the rest untouched.
Type — the semantic category of the manipulation. Examples: object editing, style transfer, attribute editing. This dimension captures what kind of change is being made. "Object editing" (adding, removing, or modifying objects) is fundamentally different from "style transfer" (changing the aesthetic presentation while preserving content), and the model needs to apply different computational strategies for each.
Preservation — the invariants that must remain unchanged. Examples: identity preservation, background preservation, structure preservation. This dimension captures what must NOT change. This is particularly critical because the conflict between preservation and generation is the core interference problem. A task with "identity preservation" (subject-driven generation) requires the model to rigidly maintain face/object features while adapting the context; a task with "background preservation" (local editing) requires the opposite — maintain the context while changing the foreground.

Why these three dimensions? The taxonomy is not arbitrary — it directly mirrors the dimensions along which task interference manifests. Two tasks that share the same Scope and Type but differ in Preservation constraints (e.g., "change the background" versus "change the background while preserving the person's style") will conflict on how strongly to weight identity features. Two tasks that share Preservation but differ in Scope (e.g., local versus global style transfer) will conflict on how broadly to attend. The taxonomy decomposes each task into atoms that, collectively, describe the computational behavior the model should exhibit.

The annotation pipeline. The paper describes an automated process using Qwen-VL (Qwen2.5-VL, referenced as [1] in Section 3.2):

The system provides Qwen-VL with the definitions of the three-tier system — what each dimension means, what the possible tag values are, and how to assign them based on analyzing the training triplet (source image, instruction, target image).
For each training triplet, Qwen-VL analyzes the source image, the text instruction, and the target image, then outputs a set of atomic tags across the three dimensions.
The rule set defining valid tag assignments is "continuously refined to maintain consistency and semantic quality" — this is an iterative process where human inspection of Qwen-VL's outputs leads to adjustments in the prompt/definitions to improve annotation reliability.

The paper provides a concrete example to illustrate the scheme: for the task "Make the person in the photo wear sunglasses," the annotation would produce tags such as:

Scope: local editing (the change affects only a specific region — the face area — not the entire image)
Type: object editing (the manipulation involves adding an object — sunglasses — to the scene)
Preservation: identity preservation, background preservation, style preservation (the person's identity should remain recognizable, the background should stay the same, and the overall style should be maintained)

This example demonstrates how a single instruction decomposes into multiple semantic facets, each imposing different (and potentially conflicting) requirements on the model.

Global vocabulary and tag embedding. The paper defines a fixed global vocabulary $\mathcal{V}$ containing all $K$ possible atomic tags across the three dimensions. The total number $K$ is not explicitly stated, but the vocabulary must cover the full range of Scope, Type, and Preservation tags needed to describe all tasks in the training set.

A learnable tag embedding matrix $\mathbf{W}_{\text{tag}} \in \mathbb{R}^{K \times D}$ is instantiated, where $D$ is the model's hidden dimension (the same dimension as the transformer's internal representations). Each row of this matrix is a learned $D$ -dimensional embedding vector for one specific atomic tag (e.g., "local editing," "identity preservation," "style transfer").

Constructing the global semantic embedding $\mathbf{s}$ . For a given training sample, let $T_p \subseteq \mathcal{V}$ be the set of atomic tags assigned to that sample by the annotation pipeline. The paper uses element-wise summation to aggregate these tag embeddings into a single fixed-dimensional vector:

$\mathbf{s} = \sum_{t \in T_p} \mathbf{W}_{\text{tag}}[\text{index}(t)]$

where $\mathbf{W}_{\text{tag}}[\text{index}(t)]$ retrieves the learned embedding vector $\mathbf{e}_t \in \mathbb{R}^D$ for tag $t$ , and the summation runs over all tags in the sample's tag set $T_p$ . The result is $\mathbf{s} \in \mathbb{R}^D$ , the global semantic embedding that serves as the ground-truth target for the alignment loss.

What it computes: For each training sample, the system takes the discrete set of atomic tags (variable in number across samples), retrieves their learned embedding vectors, and sums them element-wise. The output is a single $D$ -dimensional vector that encodes the composite semantic identity of the task — a "macro-level semantic ground truth" that resides in the same vector space as the model's internal representations.

Why this form (summation rather than concatenation or averaging): The element-wise summation has a crucial property that makes it appropriate for this setting: it is permutation-invariant. The order of tags doesn't affect the final embedding — "Scope: local editing; Type: object editing; Preservation: identity preservation" produces the same $\mathbf{s}$ regardless of which tag is listed first. This is correct because task semantics are unordered — there is no natural sequence among Scope, Type, and Preservation; they are orthogonal dimensions that jointly describe the task. If the paper had used concatenation (which is order-dependent) or a sequence model (which imposes sequential structure), it would have introduced spurious ordering dependencies that don't reflect the semantic structure.

Additionally, summation naturally handles the variable-sized tag set problem. Different training samples have different numbers of tags (a simple task might have 3 tags; a complex compositional task might have 8+ tags). Summing embeddings produces a fixed $D$ -dimensional vector regardless of tag count, and the magnitude of the resulting vector roughly reflects the number of active semantic constraints, which is a reasonable inductive bias — more complex tasks (more tags) produce larger-magnitude semantic embeddings, which will require the routing signature to carry more information to predict them successfully.

The summation also means that the tag embeddings are additive: $\mathbf{s}$ is the vector sum of all active tag embeddings, making the space compositional. The embedding for "local editing AND identity preservation" is the vector sum of the individual embeddings for "local editing" and "identity preservation." This compositionality is important because the model needs to handle novel combinations of tags at test time — it may encounter a task that combines a Scope tag and a Type tag it has seen separately but never together, and the additive representation naturally generalizes.

Inference-time handling. The annotation pipeline is exclusively a training mechanism. The paper explicitly states:

"This hierarchical annotation scheme is exclusively used for training. During the inference stage, these ground-truth tags are no longer required."

Instead, at inference time, the user's raw instruction text $c_{\text{text}}$ and source image $I_c$ are passed to Qwen-VL, which performs instruction rewriting: it analyzes the image and text to generate a more detailed, descriptive prompt. This enriched prompt is encoded as the text embedding $C$ and fed into the MM-DiT. The intuition is that the rewritten prompt implicitly carries the same task semantics that the hierarchical tags explicitly encoded during training — "make the person smile" gets rewritten to something like "Edit the person's facial expression to add a natural smile while preserving their identity, the background, and overall lighting conditions," which contains the same Scope/Type/Preservation information that the tags would have provided. The trained gating network, having learned to route based on task semantics, can extract this information from the richer text representation.

Semantic-Aligned Gating Network

This is the paper's core technical contribution: a mechanism that forces the MoE gating network's routing decisions to become predictive of the global task semantics, thereby injecting high-level task intent into what would otherwise be purely local, task-agnostic routing.

The mechanism has three sub-components that work together: (1) the aggregated routing signature $\mathbf{g}$ , which summarizes the model's actual routing behavior for a given sample; (2) a prediction head $\mathcal{H}_{\text{pred}}$ that maps from routing space to semantic space; and (3) the predictive alignment loss $\mathcal{L}_{\text{align}}$ that enforces the mapping.

Aggregated Routing Signature $\mathbf{g}$ .

Each MoE layer's gating network $\mathcal{G}$ produces routing scores for every image token. For a model with $L$ MoE layers (here, $L = 10$ , the final 10 transformer blocks) and $T$ image tokens in the sequence, the routing scores form a tensor of shape $(L, T, N)$ , where $N = 4$ is the number of experts.

The paper aggregates these scores into a single vector $\mathbf{g} \in \mathbb{R}^N$ through a two-stage averaging procedure:

Stage 1 — average across layers: For each token $t$ , compute the average routing score across all $L$ MoE layers:

$\bar{S}_t = \frac{1}{L} \sum_{l=1}^{L} S_{l,t}$

where $S_{l,t} \in \mathbb{R}^N$ is the vector of gating scores for token $t$ in layer $l$ , and the sum is over the $L = 10$ MoE layers. This produces a per-token average score vector $\bar{S}_t \in \mathbb{R}^N$ that represents, across all MoE layers, how strongly token $t$ tends to activate each expert.

Stage 2 — average across tokens: Compute the mean of these per-token average scores across all $T$ image tokens:

$\mathbf{g} = \frac{1}{T} \sum_{t=1}^{T} \bar{S}_t = \frac{1}{T \cdot L} \sum_{t=1}^{T} \sum_{l=1}^{L} S_{l,t}$

The result is a single vector $\mathbf{g} \in \mathbb{R}^N$ (where $N = 4$ , one value per expert) that encodes the overall expert usage pattern for the entire sample.

What it computes: For each of the four experts, $\mathbf{g}$ contains a scalar representing the average routing score assigned to that expert, averaged across all 10 MoE layers and all image tokens in the sequence. If expert 3 is consistently chosen with high confidence across many tokens, $\mathbf{g}_3$ will be high; if expert 1 is rarely activated, $\mathbf{g}_1$ will be near zero.

Why this form (layer-averaging then token-averaging): The two-stage averaging serves several purposes:

Dimensionality reduction: It collapses the $(L, T, N)$ tensor into a compact $N$ -dimensional vector ( $N = 4$ ), which is small enough that a lightweight prediction head can map it to the $D$ -dimensional semantic space without introducing too many parameters.
Layer averaging before token averaging ensures that experts active in different layers contribute equally to the signature. Without this, a late-layer expert (in the 10th MoE block) and an early-layer expert (in the first MoE block) would contribute identically, despite processing different levels of abstraction. The layer averaging makes the signature a summary of expert usage across all representational depths, not just the final layer.
Token averaging captures the overall distribution of routing behavior — is expert 2 typically chosen for most tokens, or only for a few? — without being sensitive to sequence length. This is critical because different images produce different numbers of tokens.
Gradient flow: The averaging operations are differentiable, so gradients from the alignment loss flow through $\mathbf{g}$ back to every gating network in every MoE layer. This means the alignment signal influences routing decisions at all levels, not just a single layer.

Prediction Head $\mathcal{H}_{\text{pred}}$ .

The aggregated routing signature $\mathbf{g} \in \mathbb{R}^N$ (dimension $N = 4$ ) and the global semantic embedding $\mathbf{s} \in \mathbb{R}^D$ (dimension $D$ , the model's hidden dimension, which is likely hundreds or thousands) live in different spaces with dramatically different dimensionalities. The prediction head bridges this gap:

$\hat{\mathbf{s}} = \mathcal{H}_{\text{pred}}(\mathbf{g})$

where $\mathcal{H}_{\text{pred}}$ is implemented as a two-layer MLP (Section 3.3.3). The input dimension is $N = 4$ , the output dimension is $D$ , and there is one hidden layer of unspecified size between them. The MLP learns a non-linear mapping from the low-dimensional routing space to the high-dimensional semantic space.

What it computes: Given the 4-dimensional routing signature $\mathbf{g}$ , the MLP produces a predicted semantic embedding $\hat{\mathbf{s}} \in \mathbb{R}^D$ — the model's best guess, based solely on how it chose to route tokens, of what the task's semantic embedding should be. If the routing signature carries sufficient information to predict the task semantics, this prediction will be close to the ground-truth $\mathbf{s}$ .

Why a two-layer MLP: The architecture is deliberately simple and lightweight. A deeper or more complex prediction head could learn to extract more information from $\mathbf{g}$ , but that would partially defeat the purpose: the goal is to force $\mathbf{g}$ to encode task semantics in a way that a simple function can extract, not to build a powerful prediction model that compensates for poor routing. A shallow MLP acts as a bottleneck — it can only succeed at predicting $\mathbf{s}$ if $\mathbf{g}$ already contains rich, linearly-accessible information about the task. This design choice creates pressure on the gating network to produce routing signatures that cleanly encode task semantics, rather than relying on the prediction head to disentangle messy representations.

Predictive Alignment Loss $\mathcal{L}_{\text{align}}$ .

The loss function that ties everything together measures the cosine similarity between the predicted semantic embedding $\hat{\mathbf{s}}$ and the ground-truth semantic embedding $\mathbf{s}$ :

$\mathcal{L}_{\text{align}} = 1 - \text{sim}(\hat{\mathbf{s}}, \mathbf{s}) = 1 - \frac{\hat{\mathbf{s}} \cdot \mathbf{s}}{|\hat{\mathbf{s}}| |\mathbf{s}|}$

where $\hat{\mathbf{s}} \cdot \mathbf{s}$ is the dot product of the predicted and ground-truth semantic embeddings, and $|\hat{\mathbf{s}}|$ and $|\mathbf{s}|$ are their L2 norms (Euclidean lengths). The cosine similarity ranges from $-1$ (perfectly opposite directions) to $+1$ (perfectly aligned), so $\mathcal{L}_{\text{align}}$ ranges from $0$ (perfect alignment) to $2$ (perfect opposition).

What it computes: The loss penalizes angular deviation between the predicted semantic vector $\hat{\mathbf{s}}$ (which is a function of the routing decisions $\mathbf{g}$ , which are in turn functions of the gating network parameters) and the ground-truth semantic vector $\mathbf{s}$ . Minimizing this loss encourages the routing signature to contain information that allows accurate prediction of the task's semantic embedding — that is, the model's actual computational behavior (as summarized by $\mathbf{g}$ ) should be strongly correlated with the task's semantic identity.

Why cosine similarity rather than MSE or L1 loss: Cosine similarity is scale-invariant — it only cares about the direction of the vectors, not their magnitudes. This is the right inductive bias for this setting because:

The magnitude of $\mathbf{s}$ depends on how many tags are active (summing more embeddings produces a larger vector), which is an artifact of the annotation scheme, not a meaningful property of task semantics. A task with 3 tags should not be considered "closer" to its prediction than a task with 8 tags simply because the vectors have smaller magnitudes — what matters is whether the directions align.
The magnitude of $\hat{\mathbf{s}}$ is controlled by the prediction head's weights and could drift during training. If the loss penalized magnitude differences (as MSE would), the optimization might focus on scaling $\hat{\mathbf{s}}$ to match $\mathbf{s}$ 's norm rather than actually learning meaningful semantic representations.
Cosine similarity directly measures whether the routing signature encodes which semantic properties are active (the direction), independent of how many are active (the magnitude). This aligns with the paper's stated principle: "semantically similar generation tasks evoke similar expert usage patterns." Two tasks with similar semantic profiles (tags pointing in similar directions in embedding space) should produce similar routing signatures, and cosine similarity is the natural metric for this.

Gradient flow and the alignment mechanism. The key insight is that minimizing $\mathcal{L}_{\text{align}}$ updates two sets of parameters:

The parameters of the prediction head $\mathcal{H}_{\text{pred}}$ — making it better at extracting semantic information from whatever routing signature it receives.
More importantly, the parameters of all gating networks $\mathcal{G}$ across all $L = 10$ MoE layers — because the gradient of $\mathcal{L}_{\text{align}}$ with respect to $\hat{\mathbf{s}}$ flows through $\mathcal{H}_{\text{pred}}$ into $\mathbf{g}$ , and from $\mathbf{g}$ through the averaging operations back to every individual gating score $S_{l,t}$ produced by every gate in every layer.

This means the gating networks are trained with a dual objective: they must route tokens in a way that produces good generation results (minimizing $\mathcal{L}_{\text{flow}}$ ), AND they must route in a way that, when aggregated, allows prediction of the task semantics (minimizing $\mathcal{L}_{\text{align}}$ ). The alignment loss provides a global, task-level signal that is otherwise absent from the purely local, per-token routing decisions — it tells the gate "whatever you're doing for individual tokens, the aggregate pattern must encode what task we're performing."

The paper characterizes this transformation elegantly: the gating network "evolves from a task-agnostic executor into a semantic-aware scheduler" (Section 3.3.3). The word "scheduler" is apt — rather than just reactively dispatching tokens based on their local features, the gate learns to proactively allocate tokens to experts in patterns that are systematically related to the overall task, as if it were a dispatch center that knows what job is being done and assigns workers accordingly.

Overall Training Objective

The predictive alignment loss is an auxiliary objective that complements, rather than replaces, the model's primary training signal. The total loss is a weighted sum of three terms:

$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{flow}} + \lambda_{\text{lbl}} \mathcal{L}_{\text{lbl}} + \lambda_{\text{align}} \mathcal{L}_{\text{align}}$

where:

$\mathcal{L}_{\text{flow}}$ is the primary flow matching loss for the diffusion generative objective. This is the standard loss used to train the Qwen-Image base model, and it ensures the model actually learns to generate and edit images correctly. It is the dominant term — without it, the model would produce meaningless outputs regardless of how well its routing aligned with task semantics.
$\mathcal{L}_{\text{lbl}}$ is the standard MoE load balancing loss that encourages uniform utilization across experts. Without this term, the gate could collapse to always selecting the same one or two experts, defeating the purpose of having multiple experts. The paper does not provide the explicit formulation, but standard load balancing losses penalize deviation from uniform expert utilization across the batch.
$\mathcal{L}_{\text{align}}$ is the proposed predictive alignment loss described in the previous section.
$\lambda_{\text{lbl}}$ and $\lambda_{\text{align}}$ are scalar hyperparameters that control the relative weight of each auxiliary loss relative to the primary generation loss. The paper does not report the specific values used in experiments.

Why this weighting scheme: The auxiliary losses are explicitly supplementary — they shape the how of computation (which experts handle which tasks, and ensuring all experts get used) without determining the what (the actual image generation quality). The flow matching loss handles the what. The alignment loss handles the how. The load balancing loss prevents degenerate solutions within the how. The hyperparameters $\lambda_{\text{lbl}}$ and $\lambda_{\text{align}}$ must be tuned to ensure the auxiliary objectives guide without dominating — too large, and the model optimizes for routing patterns at the expense of generation quality; too small, and the routing remains task-agnostic. The paper's ablation study (Table 4) confirms that adding $\mathcal{L}_{\text{align}}$ substantially improves performance over the MoE baseline without it, validating that the chosen weighting achieves the intended effect.

Dataset Construction and Inference Procedure

Training data composition. The model is trained on a large-scale hybrid dataset totaling over 11 million samples, combining publicly available data (2.2M samples) with proprietary in-house data (the remaining ~8.8M+ samples). The public portion draws from established benchmarks chosen for their coverage of the unified task space:

Universal instructive editing: InstructP2P, UltraEdit, and OmniEdit — these provide diverse instruction-based image editing examples.
Virtual try-on: VTON-HD — provides specialized data for clothing transfer tasks, which involve a specific type of local editing with strong preservation constraints.
Subject-driven generation: OmniControl — provides data for conditioning generation on control signals, relevant for subject-driven tasks.

The proprietary in-house dataset is constructed through a multi-stage pipeline described in Section 3.4:

Source image curation: Pristine images are sourced from large-scale public datasets — the paper does not specify which ones, but standard practice would involve datasets like LAION, CC, or other web-scale image collections.
Instruction generation: Large language models (specifically GPT-4o) generate a vast array of diverse editing and generation instructions for these images. This step creates the text portion of the training triplets.
Target image generation: To obtain high-quality target images that correspond to the generated instructions, the pipeline uses a combination of specialist and generalist models:
- Specialist models (e.g., ControlNet) are used for specific task types like "Control generation."
- Generalist models (Flux-Kontext, Qwen-Edit, SeedEdit) are employed for a broad range of edits. This diverse set of "teacher" models ensures the training data covers multiple task families with reasonable quality.
Video frame processing: Following the methodology of UniReal, the pipeline processes video frames to create dynamic editing datasets — for example, extracting frames showing pose changes or viewpoint shifts, which provide data for tasks involving 3D-aware editing.
Data augmentation via inverse tasks: To enhance robustness and quality, the pipeline systematically augments the data by constructing corresponding inverse tasks and instructions. The example given is pairing "object addition" with "object removal." This teaches the model bidirectionality — if it learns to add an object, it also learns to remove it, which should improve its understanding of object semantics and spatial reasoning. The paper states this "significantly improves generative fidelity."

Inference procedure. The deployment pipeline differs from training in one key respect: ground-truth task tags are not available. The paper describes the inference workflow in Section 3.2:

The user provides a raw instruction $c_{\text{text}}$ and a source image $I_c$ .
A Vision-Language Model (specifically Qwen-VL) performs instruction rewriting: it analyzes the source image and raw instruction jointly, then generates a more detailed, descriptive prompt. This rewritten prompt serves as the text input to the MM-DiT.
The MM-DiT processes the text embeddings (from the rewritten instruction), the conditional image tokens (from $I_c$ ), and the noisy latent tokens through its 10 semantically-aware MoE layers, generating the output image.

The rewritten instruction carries implicit task semantics that the trained gating networks have learned to respond to. Because the gating networks were trained to produce routing signatures predictive of the hierarchical task tags, and because the rewritten instruction contains the same semantic information that those tags captured, the gates can route tokens to the appropriate experts based on the enriched text alone. The semantic awareness is baked into the gating weights — it doesn't require explicit tag inputs at inference time.

This design means the system adds minimal inference overhead from the TAG-MoE contributions: the VLM-based instruction rewriting happens once per user query (a preprocessing step), and the MoE routing itself adds only the cost of the gating MLP (a 2-layer network producing 4 scores) plus the cost of executing one of four FFN experts per token — which is comparable to the FFN cost in the original dense model, since each expert is architecturally identical to the FFN it replaced.

4. Key Insights and Innovations

Innovation 1: Diagnostic Reframing of MoE Routing Failure as an Information Gap, Not a Capacity Problem

The paper's most fundamental intellectual move is its diagnosis of why standard MoE fails in unified generation models. Prior work on MoE for image generation (DiT-MoE, Dense2MoE, HunyuanImage-3.0) implicitly assumed that sparse expert architectures work because they expand capacity — more parameters under a fixed activation budget allow the model to learn more complex functions. The standard narrative was: dense models hit a capacity ceiling; MoE breaks through it by routing different inputs to different sub-networks, and the routing mechanism will naturally discover useful specializations through gradient descent. ICEdit extended this logic to image editing by adding LoRA-based MoE modules, but still trusted purely data-driven routing to figure out which experts should handle which editing operations.

TAG-MoE fundamentally reinterprets the failure mode. The problem is not that standard MoE lacks capacity — the ablation in Table 4 shows that task-agnostic MoE ("MoE w/o ℒ_align") still outperforms the dense baseline, confirming that raw capacity helps. The problem is that the gating network operates with an information deficit: it sees only local token features and has zero access to the global task identity. As Section 2.2 states, standard routers are "task-agnostic, operating based on local features, unaware of global task intent." This means the gate cannot distinguish between a token from a face that needs preservation (editing task) and an identical token from a face that needs adaptation (subject-driven generation task). The local features are the same; the required expert is different; the gate has no signal to make the correct choice.

This is a diagnostic reframing, not an incremental improvement. The field had been treating MoE routing as a representation-learning problem — give the gate enough data and it will learn to route correctly. TAG-MoE treats it as an information architecture problem — the gate lacks access to a critical variable (task identity) that determines the correct routing decision, and no amount of capacity or data can compensate for that missing information. This reframing matters because it changes what solution approaches are viable: if routing failure is a capacity problem, you add more experts or more layers; if it's an information problem, you need to inject the missing signal into the routing mechanism. TAG-MoE chooses the latter path.

The evidence for this diagnosis comes not only from the performance gains (Table 4, where the alignment loss provides substantial improvement over task-agnostic MoE) but from the expert activation analysis in Section 4.5 and Figure 6. Without semantic alignment, experts might still develop some specialization (the task-agnostic MoE outperforms dense), but it is "spontaneous" and "inefficient" — the specialization patterns are driven by low-level feature correlations rather than task structure. With semantic alignment, specializations become semantically coherent: different experts activate for "change material" versus "change color," and active experts concentrate computation precisely on the task-relevant image region (the backpack pixels). This demonstrates that the information gap was the bottleneck — once task semantics are injected, meaningful specialization emerges.

Innovation 2: Hierarchical Task Semantic Annotation as a Structured Supervisory Signal for Routing

Before TAG-MoE, unified generation models either used no explicit task representation (treating all tasks as undifferentiated sequence-to-sequence problems, as in OmniGen, ACE, and UniVG) or used coarse, flat task labels ("editing" vs. "generation") that were insufficient to capture the nuanced differences between task instances. The paper argues that a single coarse label like "edit" cannot capture user intent because the same label covers tasks with fundamentally different computational requirements — "change the background to a beach" (global editing, low identity preservation) and "make the person smile" (local editing, high identity preservation) are both "editing" but demand different routing patterns.

The Hierarchical Task Semantic Annotation scheme (Section 3.2) introduces a three-dimensional decomposition of task semantics — Scope, Type, and Preservation — that is genuinely novel in its design principles:

It maps directly to the dimensions of task interference. The three tiers are not arbitrary — they correspond to the axes along which task objectives conflict. Scope captures where the model operates (local vs. global attention), which determines whether the model should attend narrowly or broadly. Type captures what change is being made, which determines the nature of the transformation. Preservation captures what must not change, which is the dimension that most directly conflicts with generative diversity. By decomposing tasks along these axes, the annotation scheme provides a supervisory signal that tells the gate exactly which facets of a task create routing-relevant distinctions.

It is compositional and generalizes. The tag embeddings are summed element-wise to form the semantic embedding vector $\mathbf{s}$ , making the representation additive: the embedding for "local editing with identity preservation" is the vector sum of "local editing" and "identity preservation." This means the model can handle novel combinations of tags at test time that it never saw during training — if it learned to route for "local editing" and "identity preservation" separately, the additive representation ensures it can route for their combination. This compositionality is essential because the space of possible task specifications (combinations of Scope, Type, and Preservation tags) is combinatorially larger than any training set.

It is automated and scalable. Rather than requiring expensive human annotation, the scheme uses Qwen-VL to automatically tag training triplets. This makes the approach feasible at the 11M-sample scale of the training set. The paper notes that the rule set is "continuously refined to maintain consistency and semantic quality" — this is an iterative human-in-the-loop process where human inspection of VLM outputs leads to improved annotation prompts, but the per-sample cost is dramatically lower than manual labeling.

This contribution is significant beyond TAG-MoE because it establishes a general framework for task-aware training in multi-task generative models. The specific three-tier taxonomy might not transfer to other domains (video generation, 3D synthesis, audio processing), but the design principle — decompose task semantics into orthogonal dimensions that correspond to the axes of computational conflict — is transferable. Any domain where different task instances require different computational strategies can benefit from a structured annotation scheme that captures those distinctions.

The evidence that this annotation scheme captures meaningful semantic structure comes from the expert activation analysis: the experts develop specializations that align with the tag dimensions (Section 4.5, Figure 6). Different task types ("change material" vs. "change color") activate different expert combinations, confirming that the tags capture routing-relevant distinctions.

Innovation 3: Predictive Alignment as a Mechanism for Injecting Global Intent into Local Routing

The core architectural innovation — and the mechanism that distinguishes TAG-MoE from all prior task-aware or task-conditioned approaches — is the Predictive Alignment Regularization (Section 3.3). This is not simply "conditioning the gate on task labels" (which would be a standard approach: concatenate a task embedding to the gate's input). Instead, it is a bidirectional alignment mechanism that forces the gate's internal behavior to encode task semantics, rather than injecting task information as an additional gate input.

The distinction is critical. A standard task-conditioned gate would receive the task embedding as input: $\mathcal{G}(x, \text{task\_embedding})$ . This tells the gate what task it should be performing, but it places a single burden — the gate must learn to USE that information appropriately. The gate could still ignore the task embedding, or use it in superficial ways that don't produce meaningful expert specialization.

TAG-MoE's predictive alignment works in the opposite direction: the gate's OUTPUT (routing decisions, aggregated into $\mathbf{g}$ ) must PREDICT the task semantics (via the prediction head $\mathcal{H}_{\text{pred}}$ and the cosine similarity loss). This means the gate cannot simply consume task information — it must PRODUCE routing behavior that CONTAINS task information. The alignment loss creates pressure on the gate to route in patterns that are systematically correlated with task identity, which in turn forces experts to develop task-specific specializations (because different tasks produce different routing patterns, and experts receive different input distributions).

This is a fundamentally different inductive bias from standard task conditioning. Standard conditioning says: "Here's what task we're doing; now route accordingly." Predictive alignment says: "Route in whatever way you want, but your aggregate routing pattern must encode what task we're doing." The latter is more powerful because it doesn't prescribe how the gate should use task information — it only requires that the routing behavior be informative about the task. The gate can discover whatever routing strategy best satisfies both the generation objective and the alignment objective, potentially finding specializations that a human-designed task-conditioning scheme would not anticipate.

This innovation draws a conceptual parallel to self-supervised representation learning, where models are trained to predict pretext tasks (e.g., rotation angle, patch position) from their internal representations, forcing those representations to encode meaningful structure. Here, the "pretext task" is predicting the global semantic embedding from the routing signature, and the "representation" being shaped is the gating network's routing behavior. The prediction head acts as an auxiliary probe that provides a learning signal without constraining the form of the routing strategy.

The evidence that predictive alignment produces qualitatively different behavior from standard task conditioning is visible in the expert activation analysis (Figure 6): experts develop spatially-aware specializations — they concentrate computation on precisely the image region relevant to the current edit, not just on the task type in abstract. This spatially-aware behavior is unlikely to emerge from simple task-conditioned routing (which would activate the same expert for all tokens of a given task type) but is a natural consequence of predictive alignment, which rewards the gate for producing routing patterns that, when aggregated, are task-informative — and spatially focused routing is more informative than diffuse routing.

Innovation 4: Demonstration That Task-Aware MoE Resolves Cross-Task Interference, Not Just Scales Capacity

The paper makes a critical empirical claim that goes beyond architectural novelty: task-aware MoE does not merely increase capacity — it structurally resolves the interference between conflicting task objectives, enabling a unified model to match or exceed specialized models on their own benchmarks. This claim is significant because it distinguishes TAG-MoE from prior MoE-for-generation work (DiT-MoE, Dense2MoE) that used MoE purely as a capacity-scaling technique for single-task text-to-image generation.

The evidence for this claim is multi-faceted:

Quantitative comparison with specialized models (Tables 2 and 3). A unified model that only benefited from increased capacity would be expected to underperform task-specific specialists, which can dedicate their entire parameter budget to a narrow objective. Yet TAG-MoE achieves the highest vllmqa scores (a VLM-based measure of instruction following) on both EmuEdit-bench and GEdit-bench (Table 2), surpassing specialized editing models. On subject-driven generation (Table 3), it achieves state-of-the-art Face-ref scores on both DreamBench++ and OmniContext, and the highest Style-ref score on DreamBench++. It also obtains the top DINO-ref score on OmniContext. These results are not marginal — they show a unified model outperforming models specifically designed and trained for individual task families.

Per-category analysis on ICE-Bench (Figure 4). The radar chart covering 26 task types shows TAG-MoE achieving "state-of-the-art performance in the vast majority of categories." This breadth of leadership — rather than excelling in one category at the expense of others — suggests that the model is not simply shifting interference around (improving editing at the cost of generation, or vice versa) but genuinely resolving it.

The DreamOmni2 contrast. The paper explicitly calls out DreamOmni2's behavior as diagnostic: "DreamOmni2's high reference-generation scores largely stem from copy-paste behavior on source subjects, which artificially inflates similarity metrics" (Section 4.3). This is a concrete example of the "mediocre compromise" — DreamOmni2 resolves the editing-vs-generation tension by defaulting to copy-paste, which achieves high source similarity scores but fails as actual generation. TAG-MoE avoids this degenerate solution because its task-aware routing allows it to deploy different computational strategies for tasks that require preservation versus tasks that require generation, rather than finding a single compromise strategy.

Qualitative evidence of cross-task benefit (Figure 5). The paper argues that the framework's resolution of interference actually produces positive transfer between task families. The qualitative comparison shows that TAG-MoE handles 3D-aware edits (car turning around, shelf structure from a new angle) that specialized editing baselines fail at, "benefiting from the structural diversity and geometric awareness implicitly learned from subject-driven data." Conversely, in subject-driven tasks, TAG-MoE maintains better identity preservation than specialized models, with fidelity "attributed to the high consistency derived from editing alignment data during unified training." This bidirectional cross-task benefit is direct evidence that task interference has been resolved — rather than editing and generation objectives pulling the model in opposite directions, they now contribute complementary skills that enhance each other.

This innovation matters for the field because it establishes that unified models are not inherently limited by task interference if architectures provide structural mechanisms for disentangling conflicting objectives. The implication is that future unified generation models should invest in task-aware routing rather than simply scaling dense parameters or relying on data-driven specialization to implicitly resolve conflicts. The paper provides both a mechanism (predictive alignment) and a diagnostic tool (expert activation visualization) for achieving and verifying this disentanglement.

5. Experimental Analysis

Evaluation Methodology

Dataset. For training, the model uses a hybrid dataset totaling over 11 million samples, comprising 2.2M public samples from InstructP2P, UltraEdit, OmniEdit, VTON-HD, and OmniControl, plus a proprietary in-house dataset constructed using GPT-4o for instruction generation and specialist/generalist models (ControlNet, Flux-Kontext, Qwen-Edit, SeedEdit) for target image generation, with video frame processing and inverse task augmentation (Section 3.4). For evaluation, five benchmarks are used: ICE-Bench (the primary unified benchmark, spanning both editing and subject-driven generation across 26 task types), EmuEdit-Bench and GEdit-Bench (specialized editing benchmarks), and DreamBench++ and OmniContext (subject-driven generation benchmarks), all described in Section 4.2.
Base model(s). TAG-MoE is built on the Qwen-Image T2I model (a text-to-image diffusion transformer), as stated in Section 4.1. The MoE layers replace the FFNs of the image stream in the final 10 transformer blocks. The paper does not report the total parameter count of the base model or the MoE-augmented variant, which is a notable omission — the scale of the model is left unspecified beyond the fact that each MoE layer contains 4 experts with architectures identical to the original FFN, and top-1 routing is used. For pretraining baselines in the dense comparison, a "dense model of an equivalent activated parameter count" is constructed (Section 4.5).
Metrics. A comprehensive metric suite is employed (Section 4.2). Aesthetic quality is measured via a SigLip-based predictor. Source consistency is captured by CLIP-src (for editing tasks, measuring similarity between output and source image) and CLIP-ref (for subject-driven generation). Text alignment is measured by CLIP-cap (similarity between output image and text instruction). For editing evaluation, a vllmqa score uses Qwen2-VL-72B to judge whether the instruction was correctly executed, providing an intelligent measure of editing success that goes beyond static similarity. For subject-driven tasks, three preservation dimensions are assessed: Face-ref (facial identity preservation using the buffalo model from InsightFace), DINO-ref (subject similarity via DINO features), and Style-ref (style fidelity via CSD). All metrics are normalized so higher values indicate better performance, though the specific normalization procedure (from what original ranges to what final range) is not detailed.
Baselines. The paper compares against three categories (Section 4.2). Unified generation and editing baselines: ACE++, Flux.1 Kontext, BAGEL, OmniGen2, Qwen-Edit, and DreamOmni2 — all open-source unified models. Closed-source product-level models (GPT-4o, Gemini-2.5-flash) are included for context but primary claims are benchmarked against open-source baselines. Specialized zero-shot instruction-based editing baselines: InstructPix2Pix, EmuEdit, MagicBrush, UltraEdit, ICEdit, and Step1X-Edit. Specialized zero-shot subject-driven generation baselines: DreamO, OminiControl, and UNO. This three-category comparison is designed to test whether a unified model with task-aware routing can match or exceed specialists on their own terrain.
Generation budget / compute accounting. For the dense-versus-MoE comparison, the dense baseline is constructed with "an equivalent activated parameter count" to the MoE model (Section 4.5, Table 4). This means per-token inference FLOPs are roughly matched between the two architectures — the MoE model has more total parameters (due to having 4 experts) but only activates one expert per token, resulting in comparable per-token computation. The paper does not report exact parameter counts or FLOP numbers to verify this equivalence, which makes the comparison somewhat qualitative. For training, the convergence curves in Figure 6 (left) show loss versus training steps, but no wall-clock time or total FLOP comparison is provided.
Cross-validation / statistical protocol. No cross-validation or statistical significance testing is reported. The paper presents point estimates on benchmark test sets without confidence intervals or error bars. The user study (Section 4.6) involves 65 participants evaluating 50 cases across 350 sets total, but this provides only a qualitative preference signal rather than statistical validation of the quantitative benchmark results. The ablation study (Table 4) reports single numbers per configuration without variance estimates, making it impossible to assess whether the observed improvements are statistically reliable. The per-category radar chart (Figure 4) aggregates scores across 26 task types but does not report per-category sample sizes, which may vary substantially.

Main Quantitative Results

Unified Generation Evaluation on ICE-Bench

The headline result appears in Table 1, which reports performance on ICE-Bench across four metrics: Aesthetic, CLIP-src, CLIP-cap, and vllmqa. TAG-MoE achieves the highest scores among all open-source baselines on three of the four metrics:

Aesthetic: TAG-MoE achieves the top score (bolded), though the exact numerical value requires reading the table directly — the paper text states it is highest among open-source models but does not quote the number.
CLIP-cap: TAG-MoE's score "not only surpasses all open-source competitors but also exceeds closed-source, product-level models such as GPT-4o and Gemini-2.5-flash, indicating stronger alignment with user instructions across diverse generation and editing tasks." This is the standout metric because it measures instruction following, which is the core challenge in unified generation.
vllmqa: TAG-MoE again achieves the highest score among open-source baselines.
CLIP-src: The paper acknowledges that "some baselines exhibit high source fidelity (e.g., DreamOmni2 on CLIP-src)" but frames this as an acceptable trade-off — TAG-MoE achieves "a more favorable overall balance by excelling in instruction adherence and semantic alignment."

An important interpretive note accompanies the CLIP-src result: the paper states that DreamOmni2's high source fidelity scores "largely stem from copy-paste behavior on source subjects, which artificially inflates similarity metrics." This is a critical methodological observation — a model that simply copies the source image unchanged would achieve perfect CLIP-src but zero task completion, highlighting why the vllmqa score (which explicitly evaluates instruction execution) is a more meaningful metric for unified models.

The per-category breakdown is visualized in the radar chart in Figure 4, covering all 26 task types on ICE-Bench. The paper states that TAG-MoE "achieves state-of-the-art performance in the vast majority of categories, demonstrating robust and well-balanced capability." The radar chart shows TAG-MoE's polygon (presumably in a distinct color, though this cannot be verified from text alone) encompassing most of the outer boundary of the plot, indicating leadership across most axes. The specific per-category scores are not tabulated, so the exact number of categories where TAG-MoE leads and the margins of leadership cannot be quantified from the paper text. The claim of "vast majority" leadership is qualitative judgment from the visualization rather than a counted statistic.

Image Editing Evaluation on Specialized Benchmarks

Table 2 reports results on EmuEdit-Bench and GEdit-Bench, comparing TAG-MoE against specialized editing baselines. The key finding is that TAG-MoE achieves the highest vllmqa score on both benchmarks, which the paper identifies as "the most important indicator" because it "uses a powerful VLLM to evaluate the correctness of the executed instruction, offering a more intelligent and reliable measure of editing success."

The paper explicitly notes that TAG-MoE "does not achieve top-1 performance on every metric" in Table 2 — it trails some specialists on static similarity metrics like CLIP-src or CLIP-cap in specific cases — but frames the vllmqa leadership as the decisive factor. The specific numerical values for each metric and baseline are presented in the table structure, which requires direct visual inspection. The paper does not report the magnitude of the vllmqa advantage (how many points ahead of the next-best specialist).

EmuEdit is a special case in this comparison: since it is not open-source and only provides pre-generated outputs on its own benchmark, its performance on GEdit-Bench is unavailable. This creates an incomplete comparison — the strongest specialist on EmuEdit-Bench cannot be evaluated on the second benchmark, making the cross-benchmark comparison somewhat asymmetric.

Subject-Driven Generation Evaluation on Specialized Benchmarks

Table 3 reports results on DreamBench++ and OmniContext for subject-driven generation metrics. The key findings:

Face-ref: TAG-MoE achieves state-of-the-art scores on both benchmarks, indicating superior facial identity preservation.
Style-ref: TAG-MoE achieves the highest score on DreamBench++ (OmniContext does not include style-related tasks).
DINO-ref: TAG-MoE achieves the top score on OmniContext and is "highly competitive on DreamBench++."

The paper interprets these results as demonstrating that the "unified model can match or surpass specialized models, effectively mitigating the typical tension between subject fidelity and generative diversity." The fact that TAG-MoE leads on both Face-ref (which measures preservation) and on vllmqa for editing (which measures instruction-following correctness) is presented as direct evidence that task interference has been resolved — the model is not trading off preservation quality for editing capability or vice versa.

Qualitative Results

Figure 3 presents qualitative comparisons with unified baselines across "complex tasks characterized by interfering intents." The paper describes three failure modes of baseline unified models that TAG-MoE avoids:

"Copy-paste" artifacts in subject-driven generation — the source subject is pasted into the output without contextual adaptation.
Stylistic dissonance during inpainting — the inpainted region does not match the surrounding style.
Incomplete execution in compositional editing — only part of the multi-faceted instruction is followed.

The paper attributes TAG-MoE's success on these cases to the Predictive Alignment Regularization, which "effectively decouples and routes conflicting sub-tasks (e.g. local semantic edits versus global style preservation) to specialized experts."

Figure 5 presents qualitative comparisons with specialized baselines across editing and subject-driven tasks. For editing, specialized baselines struggle with "significant structural or geometric changes" — the paper provides two examples: a silver car case where baselines fail to execute the motion of turning around (producing only minor texture changes), and a shelf structure case where baselines fail to synthesize the side view. TAG-MoE handles these 3D-aware edits, which the paper attributes to "the structural diversity and geometric awareness implicitly learned from subject-driven data." For subject-driven tasks, specialized baselines either "lose facial identity/clothing details" (OmniControl) or "fail to render the office context" (UNO) for the human subject case, and generate "incorrect upright poses" for a toy subject that requires a handstand. TAG-MoE maintains "robust identity while adhering to complex motion instructions," attributed to "the high consistency derived from editing alignment data during unified training."

This bidirectional cross-task benefit — editing accuracy from generation data, generation fidelity from editing data — is presented as the qualitative signature of successful interference resolution.

User Study

Figure 7 reports the results of a user study with 65 participants evaluating 50 cases from ICE-Bench across 350 total sets. Participants selected the single best result according to three criteria:

Reference Alignment (consistency with the source image)
Prompt Alignment (faithfulness to the textual instruction)
Overall Preference (overall visual quality)

The aggregated results "reveal a clear and consistent preference for our method, which achieved the highest selection rate across all three evaluation criteria." The specific selection percentages for TAG-MoE versus each baseline are shown in Figure 7 but not quoted in the text. The paper does not report whether the preferences are statistically significant (no p-values or confidence intervals), nor does it describe the participant demographics or whether any randomization or blinding was used in the evaluation protocol.

Ablation Studies and Robustness Checks

Effectiveness of the MoE architecture (Table 4, Figure 6 left): The paper compares the full MoE model to "a dense baseline of an equivalent activated parameter count." The dense model "shows a severe performance drop on ICE-Bench metrics" (exact numbers in Table 4) and exhibits "slower convergence" as shown in Figure 6 (left), where the training loss curve for the dense model remains consistently above the MoE curve throughout training. The paper concludes: "This validates that the sparse architecture is fundamentally more effective at mitigating the severe task interference inherent in the unified task space than a computationally-equivalent dense model." This is an important control — it demonstrates that the MoE architecture itself (even without semantic alignment) provides benefits beyond mere capacity scaling, since both models have equivalent activated parameters per token.

Effect of Predictive Alignment Regularization (Table 4): Removing the alignment loss ℒ_align (the "MoE w/o ℒ_align" variant) produces "substantial degradation across all major metrics" in Table 4. The paper emphasizes: "This finding is key: a sparse MoE architecture alone is not sufficient. ℒ_align is what enables semantically guided routing, which is essential for mitigating task interference." Notably, the MoE w/o ℒ_align variant still outperforms the dense baseline, which is attributed to "the larger effective capacity of the sparse MoE structure, which allows exploration of a richer solution space under the same computational budget." This establishes a clear hierarchy: Dense < Task-Agnostic MoE < Task-Aware MoE (TAG-MoE). The performance gap between the last two is attributed specifically to the semantic alignment mechanism.

Analysis of Expert Specialization (Figure 6, middle and right): This is a qualitative analysis rather than a quantitative ablation, but it provides the most direct evidence for the paper's central claim about task-aware specialization. The analysis is a two-step process described in Section 4.5:

Expert Utilization Rate heatmap (Figure 6, middle): For each MoE layer, the percentage of total image tokens routed to each expert is computed. A utilization of 0% (blue) or 100% (red) indicates no specialization — either the expert is never used or always used, meaning routing is degenerate. Mixed utilization (near white) indicates functional specialization, where experts are selectively activated for different inputs. The analysis focuses on these differentiated layers.
Per-token routing score visualization (Figure 6, right): For the active, differentiated layers, per-token routing scores for each expert are reshaped to the image's spatial dimensions and visualized as heatmaps. High scores (shown in blue) indicate that the corresponding image tokens are strongly routed to that specific expert.

The results reveal two types of specialization:

Task-specific specialization: For "Change Material" versus "Change Color," the model activates "distinct combinations of experts," meaning different tasks route to different expert sets.
Spatially-aware specialization: The token heatmaps for active experts show that "computation is spatially concentrated on the backpack's pixels, precisely the region relevant to the edit." Non-relevant background tokens are correctly routed to other experts or have near-zero activation for these experts.

The paper interprets this as "strong evidence that our model has learned a sophisticated specialization that is both task-specific (using unique expert combinations for different tasks) and spatially-aware (experts learn to process semantically relevant image regions)."

Negative result — ReST^EM revision model (if present): The paper does not appear to include a negative result analogous to the ReST^EM experiment in the reference example. The ablation study focuses exclusively on validating the proposed components (MoE architecture, alignment loss) and analyzing expert behavior. No experiments are reported where an attempted improvement backfired or where sensitivity to hyperparameters is explored.

Missing ablations: Several experiments that would strengthen the paper are absent:

Sensitivity to number of MoE layers: The paper places MoE layers in the final 10 blocks but does not ablate this choice — would 5 layers suffice? Would 15 be better?
Sensitivity to number of experts: The paper uses 4 experts with top-1 routing — how does performance change with 2, 8, or 16 experts? Does the alignment loss become more or less important as expert count changes?
Sensitivity to the alignment loss weight λ_align: No sweep over this hyperparameter is reported, making it impossible to assess whether performance is robust to this choice or whether precise tuning is required.
Alternative alignment mechanisms: The paper does not compare predictive alignment against simpler approaches like directly conditioning the gate on a task embedding vector (concatenating it to the gate input) — this is the most obvious baseline for "task-aware routing" and its absence is a significant gap.
Alternative tag aggregation methods: The paper uses summation to combine tag embeddings into s. No comparison with averaging, concatenation, or attention-based pooling is reported.
Inference-time instruction rewriting ablation: The paper does not ablate the VLM-based instruction rewriting at inference time — what happens if raw instructions are used directly without rewriting? This is relevant because the rewritten instructions are claimed to carry the implicit task semantics.

Critical Assessment

Claim 1: Task-aware routing via predictive alignment resolves task interference and outperforms task-agnostic MoE and dense baselines.

Assessment: The ablation study in Table 4 provides direct support for this claim within the specific experimental setup. The clear hierarchy — Dense < MoE w/o ℒ_align < TAG-MoE — demonstrates that both the MoE architecture and the semantic alignment loss contribute independently to performance. The expert activation analysis in Figure 6 provides qualitative corroboration by showing that the learned specializations are semantically and spatially meaningful.

However, the evidence has important boundary conditions that the paper does not fully explore:

The comparison is strictly within the Qwen-Image architecture family. Whether the same hierarchy would hold for other base models (Stable Diffusion 3, Flux, Playground) is untested. The paper argues Qwen-Image is representative, but the claim of generality is unverified.

The "task-agnostic MoE" baseline may not be fully optimized. The paper uses the same architecture (4 experts, top-1 routing, 10 MoE layers) for both the task-agnostic and task-aware variants. The optimal architecture for task-agnostic MoE might differ — perhaps more experts, different routing, or different layer placement would close some of the gap. The paper does not explore whether the architectural choices were co-optimized with the alignment loss.

The alternative of direct task conditioning is not tested. The most natural baseline for "task-aware routing" is to simply provide the task embedding as an additional input to the gating network (concatenating s to the token representation x before the gate MLP). The paper's predictive alignment approach is more elegant (it forces the gate to encode task information in its output rather than consuming it as input), but without comparing against the simpler alternative, it is impossible to quantify how much the specific predictive alignment mechanism matters versus simply making the gate task-aware by any means. This is the single most important missing baseline. If direct task conditioning achieves comparable performance, the paper's core architectural innovation becomes less significant.

The ablation does not isolate the effect of the tag embedding matrix W_tag from the alignment loss. The tag embeddings are learned parameters — they could serve as a form of task-conditioning signal even without the alignment loss, if used differently. An ablation that replaces learned tag embeddings with fixed (e.g., one-hot) embeddings would clarify whether the semantic structure comes from the alignment loss or from the learned embeddings themselves.

Claim 2: TAG-MoE achieves state-of-the-art overall performance among open-source baselines on unified generation benchmarks.

Assessment: Supported by Table 1 on ICE-Bench, with the qualification that "overall performance" is a multi-dimensional concept. TAG-MoE leads on three of four metrics (Aesthetic, CLIP-cap, vllmqa) and trails on one (CLIP-src). The paper's framing — that CLIP-src is less meaningful because it can be gamed by copy-paste behavior — is reasonable, but it means the "overall" claim depends on a value judgment about which metrics matter more. A different weighting of metrics could produce a different ranking.

Several details weaken the strength of this claim:

Exact numerical values are not quoted in the text for Table 1, making it impossible to assess the magnitude of TAG-MoE's advantage from the paper text alone. A 0.5% lead on CLIP-cap is different from a 5% lead.
The 26-category breakdown in Figure 4 is a radar chart, which makes precise comparison difficult. The paper states TAG-MoE leads in "the vast majority" of categories, but this is a visual judgment rather than a counted statistic. The number of categories where TAG-MoE leads versus trails is not reported, nor are the per-category margins.
ICE-Bench is a single benchmark, and while it covers diverse task types, it represents one specific evaluation philosophy and one specific data distribution. The paper does evaluate on four additional benchmarks (EmuEdit, GEdit, DreamBench++, OmniContext), but these are specialized benchmarks where different subsets of baselines are evaluated, making the "overall" comparison across all five benchmarks difficult to synthesize into a single ranking.
No statistical testing is reported. With 500 test questions on ICE-Bench (implied but not specified for this benchmark) and point estimates without confidence intervals, the observed differences may not be statistically significant. The user study provides preference data but does not report statistical tests.

Claim 3: TAG-MoE matches or exceeds specialized models on their own benchmarks, demonstrating that task interference has been resolved.

Assessment: The strongest evidence for this claim is the vllmqa leadership on both editing benchmarks (Table 2) and the Face-ref leadership on both subject-driven benchmarks (Table 3). These are direct comparisons against task-specific specialists, and winning on the same metrics those specialists were designed to optimize is a meaningful result.

However, several important caveats apply:

TAG-MoE was trained on 11M samples including data from diverse tasks, while the specialized baselines were trained on task-specific data of varying (and unreported) scales. The comparison is between a model trained on massive multi-task data and models trained on narrower data distributions — the data scale and diversity advantage may contribute to TAG-MoE's performance independently of the architectural contribution. A fairer comparison would control for total training data volume.
Not all specialized baselines are evaluated on all benchmarks. EmuEdit cannot be evaluated on GEdit-Bench because it is not open-source. This creates an asymmetric comparison where the strongest editing specialist on one benchmark is missing from the other.
The claim of "matching or exceeding" is metric-dependent. TAG-MoE does not lead on every metric — the paper acknowledges it "does not achieve top-1 performance on every metric" in Table 2. The focus on vllmqa as "the most important indicator" is justified by the paper's argument that static similarity metrics are less reliable, but this argument itself is not experimentally validated within the paper (e.g., no correlation analysis between different metrics and human judgment is reported beyond the user study, which uses different criteria).
The claim of "resolution" of task interference is supported by the bidirectional cross-task benefits observed in qualitative results (editing benefiting from generation data, generation benefiting from editing data). But this evidence is qualitative and example-based — no quantitative metric directly measures "interference" as a separate construct from overall performance. The paper infers interference resolution from the pattern of results rather than measuring it directly.

Claim 4: Expert specialization is both task-specific and spatially-aware, confirming the mechanism works as intended.

Assessment: The expert activation analysis in Figure 6 is the most direct mechanistic evidence in the paper and is genuinely informative. The visualization showing different expert combinations for "Change Material" versus "Change Color," and the spatial concentration on the backpack region, provides compelling qualitative evidence for the claimed specialization pattern.

However, the analysis has significant limitations:

It is qualitative and example-based. The paper shows one or a few examples in Figure 6. There is no quantitative aggregation of specialization patterns across the full test set — what fraction of tasks show distinct expert combinations? What is the average spatial concentration (e.g., measured by entropy of the spatial activation distribution)? Without such aggregation, the claim generalizes from anecdotal examples to the entire model's behavior.
The analysis is performed only on the final trained model. There is no comparison of expert activation patterns for the "MoE w/o ℒ_align" variant, which would directly show whether the alignment loss causes the observed specialization or whether task-agnostic MoE also develops similar (but perhaps less clean) specializations. This is the most important missing analysis for validating the mechanism — if task-agnostic MoE also shows task-specific and spatially-aware expert activation, the alignment loss's contribution is incremental rather than fundamental.
Causality is not established. The paper shows that TAG-MoE performs better AND that its experts show specialized activation patterns. But it does not demonstrate that the specialization causes the performance improvement. The specialization could be a correlate rather than a cause — perhaps the alignment loss improves training dynamics in some other way, and the specialization patterns are a side effect.
The analysis covers only "some" tasks. The full space of 26 ICE-Bench task types is not systematically analyzed. The specific examples shown (Change Material, Change Color) are both editing tasks — similar analyses for subject-driven generation tasks are not presented. This leaves open whether the specialization patterns are consistent across the full task spectrum.

Missing Experiments That Would Have Strengthened the Paper

Direct task-conditioning baseline. The most important missing experiment is a comparison against a simple gate-conditioning approach: concatenate the semantic embedding s (or a learned projection of it) to each token representation x before the gate MLP, and train without the predictive alignment loss. This would isolate whether the specific predictive alignment mechanism matters or whether any form of task awareness produces similar benefits.

Ablation of the number and placement of MoE layers. The choice of 10 MoE layers in the final blocks is justified by reference to general MoE literature but not validated for this specific setting. An ablation over layer count (0, 5, 10, 15, all layers) and placement (early, middle, late) would clarify whether the gains saturate or whether more aggressive MoE-ification would help further.

Ablation of expert count. With 4 experts and top-1 routing, each token sees exactly one of four possible FFNs. Does performance improve with 8 or 16 experts? Does the alignment loss become more or less important with more experts (since the routing space is larger and potentially harder to train)?

Quantitative specialization metrics. Beyond the qualitative visualization in Figure 6, quantitative metrics for expert specialization would strengthen the mechanistic claim. Candidate metrics include: mutual information between task type and expert selection, spatial entropy of per-expert activation maps, expert utilization entropy (are experts used uniformly or does one dominate?), and layer-wise specialization scores. These would enable comparison across configurations and ablation of the alignment loss's effect on specialization patterns.

Scaling analysis. The paper trains on 11M samples. How does performance scale with dataset size? Does the alignment loss provide more benefit at smaller data scales (where interference might be harder to disentangle) or at larger scales (where there is more data for each task type)? A data scaling ablation would inform practical deployment decisions.

Inference-time rewriting ablation. The paper uses VLM-based instruction rewriting at inference time to provide enriched prompts that implicitly carry task semantics. What happens if raw instructions are used directly without rewriting? This is critical because the rewriting step adds latency and requires a separate VLM call — if the model works nearly as well without it, the deployment complexity could be reduced.

Expert activation analysis for the task-agnostic MoE variant. This is perhaps the single most important missing analysis. Showing that the MoE w/o ℒ_align variant produces diffuse, non-specialized, or spatially unfocused activation patterns would provide direct causal evidence that the alignment loss creates the observed specialization. Without this comparison, the mechanistic interpretation remains correlational.

Broader model family evaluation. All results are on Qwen-Image. Replicating the key findings (dense vs. task-agnostic MoE vs. task-aware MoE) on at least one other base architecture (e.g., Flux, Stable Diffusion 3) would establish generality. The paper's claims are framed as architectural innovations rather than Qwen-specific optimizations, but the evidence is entirely single-model-family.

In summary, the experimental analysis provides credible support for the paper's main claims within the specific setting tested, but several of the most important claims (mechanism of alignment, generality of the approach, statistical reliability of the improvements) rest on evidence that is more suggestive than conclusive. The missing comparisons — particularly the direct task-conditioning baseline and the expert activation analysis for the task-agnostic variant — leave room for alternative interpretations of the results that the current experiments cannot rule out. The paper's contributions would be significantly strengthened by quantitative metrics for specialization, broader model family evaluation, and rigorous ablation of the architectural choices (layer count, expert count) that are currently justified by reference to external literature rather than demonstrated through experiment.

6. Limitations and Trade-offs

The Predictive Alignment Mechanism Is Not Compared Against Direct Task Conditioning

The paper's central technical contribution is the Predictive Alignment Regularization (Section 3.3): forcing the aggregated routing signature $\mathbf{g}$ to predict the global semantic embedding $\mathbf{s}$ via a cosine similarity loss. The paper frames this as "a fundamentally different inductive bias from standard task conditioning" and characterizes standard conditioning as merely telling the gate "here's what task we're doing; now route accordingly" versus predictive alignment which says "route in whatever way you want, but your aggregate routing pattern must encode what task we're doing."

However, no experiment compares predictive alignment against the simpler alternative of directly conditioning the gating network on task semantics. The natural baseline — concatenating the semantic embedding $\mathbf{s}$ (or a learned projection) to each token representation $x$ before the gate MLP, without the prediction head $\mathcal{H}_{\text{pred}}$ or the alignment loss $\mathcal{L}_{\text{align}}$ — is never implemented or evaluated. This is the most important missing baseline in the entire paper, because it would directly test whether the specific predictive alignment mechanism provides benefits beyond simply making task information available to the gate.

The consequence is that we cannot determine whether the paper's core architectural innovation matters. If direct task conditioning achieves comparable performance, then the predictive alignment mechanism — the prediction head, the aggregation of routing signatures, the cosine similarity loss, and the gradient flow from $\hat{\mathbf{s}}$ back through all gating networks — is unnecessary complexity. The contribution would reduce to the observation that task awareness helps MoE routing, which is a less significant claim. Conversely, if predictive alignment substantially outperforms direct conditioning, the paper's stronger claim about the mechanism forcing meaningful specialization would be validated.

No evidence in the paper addresses this question. The ablation study (Table 4) compares only three configurations: dense baseline, MoE without $\mathcal{L}_{\text{align}}$ , and full TAG-MoE. The "MoE without $\mathcal{L}_{\text{align}}$ " variant is truly task-agnostic — it has no access to the hierarchical tags at all. A direct task-conditioned MoE (with $\mathbf{s}$ as additional gate input but no predictive alignment loss) would be an intermediate point between these extremes, and its absence means the ablation cannot distinguish "any form of task awareness helps" from "predictive alignment specifically helps."

The paper does not acknowledge this as a limitation and does not suggest it as future work. The mitigation status is: not addressed.

The Expert Activation Analysis Has No Quantitative Comparison Against Task-Agnostic MoE

Section 4.5 and Figure 6 present the expert activation analysis as direct mechanistic evidence for the paper's central claim: that the alignment loss produces task-specific and spatially-aware expert specialization. The analysis shows that different tasks activate distinct expert combinations, and that active experts concentrate computation on precisely the relevant image region (the backpack pixels for a material edit). The paper interprets this as "strong evidence that our model has learned a sophisticated specialization."

However, this analysis is performed only on the full TAG-MoE model — the task-agnostic MoE variant (without $\mathcal{L}_{\text{align}}$ ) is never analyzed in the same way. We cannot see whether task-agnostic MoE also develops similar (perhaps less clean) specialization patterns. This is a critical gap because it means the causal link between the alignment loss and the observed specialization is unsubstantiated. It is entirely possible that the MoE architecture itself — with its sparse top-1 routing and load balancing — naturally develops task-specific and spatially-aware expert activation patterns through standard gradient descent, and that the alignment loss merely sharpens or regularizes specializations that would have emerged anyway. The paper provides no evidence to rule out this alternative interpretation.

The consequence is that the paper's interpretive claim about mechanism is correlational, not causal. The logic presented is: TAG-MoE performs better than task-agnostic MoE; TAG-MoE shows specialized expert activation; therefore, the alignment loss caused the specialization, which caused the performance improvement. But without analyzing the task-agnostic variant's activation patterns, the first causal arrow (alignment loss → specialization) is unexamined, and the second (specialization → performance) is assumed. The performance improvement could arise from better training dynamics induced by the alignment loss gradient, with the specialization patterns as an epiphenomenon rather than the mechanism. Or the task-agnostic MoE could show equally specialized activations that are simply less well-aligned with the specific taxonomy the paper uses to visualize them.

The analysis is also purely qualitative and example-based. The paper shows a few examples in Figure 6 without any quantitative aggregation across the full test set. No metrics are reported for specialization quality: mutual information between task type and expert selection, spatial concentration entropy, expert utilization uniformity, or layer-wise specialization scores. Without such metrics, it is impossible to compare specialization patterns across configurations or to assess how representative the shown examples are.

The paper does not acknowledge this as a limitation. The mitigation status is: not addressed, though the analysis is presented as confirmatory evidence rather than as a rigorous mechanistic study. Adding the same visualization for the "MoE w/o $\mathcal{L}_{\text{align}}$ " variant would directly address this concern.

All Results Are on a Single Model Family With Unreported Scale

The paper builds TAG-MoE exclusively on the Qwen-Image T2I model (Section 4.1) and evaluates it against baselines that, while diverse in their task coverage, are not systematically varied in architecture family. The paper states that it replaces the FFNs of the image stream in the final 10 transformer blocks with MoE layers, but the total parameter count of neither the base Qwen-Image model nor the TAG-MoE variant is reported anywhere in the paper. This makes it impossible to situate the results in the context of model scale — we do not know whether TAG-MoE is a 1B-parameter model, a 10B-parameter model, or larger, and therefore we cannot assess whether the claimed benefits of task-aware MoE are specific to a particular scale regime.

The consequence extends beyond missing metadata. All claims about the effectiveness of predictive alignment, expert specialization, and interference resolution are conditional on the Qwen-Image architecture. We cannot determine whether the same hierarchy (Dense < Task-Agnostic MoE < TAG-MoE) would hold for other diffusion transformer families such as Flux, Stable Diffusion 3, or Playground. The gating network's ability to learn task-semantic routing may depend on properties of the Qwen-Image architecture — its hidden dimension, its attention patterns, its pretraining data distribution — that do not transfer. Similarly, the expert specialization patterns observed in Figure 6 may be specific to how Qwen-Image represents visual features.

The only evidence bearing on generality is the paper's claim that Qwen-Image is "representative of the capabilities of many contemporary LLMs" (a phrasing that actually refers to language models in the reference example, but the analogous claim here would be about image generation models). No such claim is explicitly made, and no cross-architecture validation is performed. The paper does not acknowledge this limitation explicitly. The mitigation status is: not addressed. Replicating the key dense vs. task-agnostic MoE vs. TAG-MoE comparison on at least one additional architecture family would be needed to establish that the method generalizes beyond Qwen-Image.

The Hierarchical Annotation Pipeline's Reliability Is Not Validated

The Hierarchical Task Semantic Annotation scheme (Section 3.2) is the linchpin that provides the supervisory signal for the entire predictive alignment mechanism. The scheme uses Qwen-VL to automatically analyze training triplets (source image, instruction, target image) and output structured tags across Scope, Type, and Preservation dimensions. The paper states that the rule set defining valid tag assignments is "continuously refined to maintain consistency and semantic quality," implying an iterative human-in-the-loop process.

However, no quantitative measure of annotation quality is reported. Key questions are left unanswered: What is the inter-annotator agreement between Qwen-VL and human annotators (if such a comparison was done)? What is the consistency of Qwen-VL's annotations across similar instructions? What fraction of the 11M training samples received annotations that were manually verified? What error rate is expected in the tag assignments, and how sensitive is the predictive alignment loss to annotation noise?

The consequence is that the entire semantic supervision signal is of unknown reliability. If Qwen-VL systematically misclassifies certain task types or produces inconsistent tags for semantically equivalent instructions, the semantic embedding $\mathbf{s}$ will contain noise. The cosine similarity loss $\mathcal{L}_{\text{align}}$ would then train the gating network to align its routing with a noisy target, potentially producing less effective specializations than a cleaner signal would. Conversely, the fact that the method works well despite potentially noisy annotations could be evidence that the alignment loss is robust — but we cannot distinguish "the annotations are clean and the method works" from "the annotations are noisy but the method is robust" without annotation quality metrics.

The paper partially mitigates this concern by describing the annotation pipeline's rule refinement process, which suggests human oversight. But this is a qualitative description, not a quantitative validation. The paper does not report annotation accuracy, does not compare against human-annotated ground truth for even a subset of the data, and does not ablate annotation quality (e.g., by training on cleanly-annotated versus noisily-annotated subsets to measure sensitivity). The mitigation status is: partially addressed through process description but not validated.

This limitation is practically important because any practitioner seeking to replicate TAG-MoE on a new domain or dataset would need to implement a similar annotation pipeline. Without knowing the required annotation quality, they cannot determine whether their own VLM-based or rule-based annotation scheme is sufficient. If annotation quality matters substantially, the cost of human verification could make the approach less scalable than the paper suggests.

Inference-Time Instruction Rewriting Cost Is Unquantified and Not Ablated

The training pipeline relies on the Hierarchical Task Semantic Annotation scheme with explicit structured tags, but at inference time, these tags are unavailable. The paper's solution (Section 3.2) is to use Qwen-VL to rewrite the user's raw instruction into a "more detailed, descriptive prompt" that implicitly carries the same task semantics. This rewritten prompt is then encoded as the text embedding $C$ and fed into the MM-DiT, where the trained gating networks presumably extract the task-relevant semantic information.

However, the paper provides no ablation comparing performance with versus without instruction rewriting. There are no results showing what happens if raw user instructions are used directly at inference time, without the VLM rewriting step. This is a critical gap because the rewriting step has practical costs: it requires an additional VLM inference call per user query, adding latency; it requires a separate model (Qwen-VL) to be available at deployment; and it introduces a potential failure mode where the VLM produces a poor or misleading rewrite that degrades the downstream generation. If TAG-MoE performs nearly as well with raw instructions, the rewriting step is unnecessary complexity and overhead. If rewriting is essential, then the method's inference-time efficiency claims (Section 3.2, suggesting "minimal inference overhead from the TAG-MoE contributions") must account for the VLM call.

The consequence is that the deployment cost of TAG-MoE is unknown. The paper frames the method as adding no inference overhead beyond standard MoE routing, but this is only true if instruction rewriting is treated as external to the method. From a practitioner's perspective, if rewriting is required for the method to achieve its reported performance, it is part of the system cost. The VLM rewriting step could dominate inference latency — Qwen-VL processing an image and instruction and generating a detailed prompt likely takes significantly longer than the diffusion model's per-token routing decisions. Without quantitative latency measurements or an ablation showing rewriting is optional, the paper's efficiency claims are incomplete.

No evidence in the paper addresses this. The mitigation status is: not addressed. The paper describes the rewriting step as a "lightweight pre-processing step" but provides no characterization of its cost or its necessity. An ablation comparing raw-versus-rewritten instruction performance, plus latency measurements, would directly address this limitation.

Hard Tasks That Require Content-Based Reasoning Remain Fundamentally Unsolved

Section 5 (Limitations and Future Work) explicitly identifies a sharp capability boundary:

"Our model relies on pre-processed instructions (the intent) and cannot jointly reason over this intent and the visual content of the source image. This separation restricts tasks requiring integrated semantic and perceptual understanding. For instance, our model fails at content-based reasoning (e.g., solving a math problem in an image) because it understands the editing intent (e.g., scope, type) but not the contextual information in the pixels themselves."

This is a structural limitation of the framework's design: the task semantics are extracted from the instruction alone (via the hierarchical annotation pipeline during training or via instruction rewriting during inference), without a mechanism for the routing to depend on the interaction between the instruction semantics and the actual pixel content. The gating network learns to route based on task type ("this is a local editing task with identity preservation"), but it cannot adjust routing based on how difficult or unusual the specific source image is for that task type.

The consequence is a hard failure boundary: any task whose correct execution depends on understanding the specific content of the source image — rather than just its category-level properties — is outside the method's capability envelope. The paper gives the example of solving a math problem in an image, but other examples would include: editing instructions that require object recognition ("remove the third person from the left"), tasks that depend on spatial relationships ("move the cup to the left of the plate"), or tasks where the appropriate expert routing depends on image complexity rather than task type (a simple "change color" on a complex textured surface versus a flat one). In all these cases, the task tag ("local editing, object removal, background preservation") does not capture the content-dependent difficulty or the specific visual reasoning required.

The paper provides no quantitative characterization of how common such tasks are in the evaluation benchmarks or what the performance drop is on these task types versus tasks that are well-captured by the hierarchical tags. The 26-category ICE-Bench radar chart (Figure 4) shows leadership in "the vast majority" of categories but does not identify which categories might fall into this limitation regime.

The paper acknowledges this limitation explicitly in Section 5 and proposes future work: "an end-to-end system incorporating a multimodal reasoning engine to unify perceptual understanding (content), intent comprehension (command), and conceptual generation (reasoning)." However, this is a speculative direction rather than a concrete mitigation. The current framework provides no pathway to addressing content-dependent routing within the predictive alignment paradigm, since the alignment signal comes entirely from instruction-derived tags.

This limitation matters practically because it defines the boundary of what TAG-MoE can handle: tasks that are semantically well-described by the Scope-Type-Preservation taxonomy but do not require joint reasoning over instruction and pixel content. For many real-world editing scenarios — particularly those involving fine-grained spatial reasoning or object-level understanding — this boundary may exclude a significant fraction of user intents. Practitioners deploying TAG-MoE would need to understand this limitation to set appropriate user expectations and to implement fallback mechanisms (e.g., routing content-based-reasoning tasks to a separate system).

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a diagnostic reframing that redefines the bottleneck in unified image generation models. Before TAG-MoE, the field understood task interference in unified models as a capacity problem — the shared parameter space simply needed to be larger, deeper, or trained on more data to accommodate conflicting objectives. MoE architectures were adopted as a capacity-scaling solution, with the implicit assumption that data-driven routing would naturally discover useful expert specializations. The paper's core reframing is that the primary bottleneck is not capacity but information architecture: standard MoE gating networks operate on local token features alone and have zero access to the global task identity that determines correct routing. This is not a problem that more experts or more data can solve — it is a structural information gap.

The magnitude of this shift is more than incremental but less than a paradigm change. It does not propose a fundamentally new generative architecture (the MM-DiT backbone remains standard) or a new training objective (flow matching is inherited from Qwen-Image). Rather, it redefines what makes MoE routing effective in multi-task settings: not the number of experts or the routing algorithm, but whether the gate has access to the semantic variables that distinguish task requirements. This shifts research attention from scaling MoE capacity (more experts, deeper routing networks) toward injecting structured task information into routing mechanisms — a design axis that was largely overlooked in prior MoE-for-generation work.

The paper's reconciliation of conflicting evidence is implicit but important. The field faced a tension: specialized editing models (InstructPix2Pix, MagicBrush, ICEdit) perform well on their target tasks but cannot generalize to generation; unified models (ACE++, Flux-Kontext, OmniGen2) cover both task families but degrade relative to specialists; and prior MoE-for-generation work (DiT-MoE, Dense2MoE) focused exclusively on single-task text-to-image generation where task interference does not arise. TAG-MoE provides a unified explanation: task-agnostic routing cannot structurally disentangle conflicting objectives, which is why prior unified models hit a "mediocre compromise" ceiling, and why prior MoE work never observed the interference problem (they operated in single-task regimes). The resolution — task-aware routing via semantic supervision — demonstrates that unified models are not inherently limited by interference if the architecture provides mechanisms for task-conditional computation. This converts the field's narrative from "unified models face an unavoidable tradeoff" to "unified models need task-aware routing architectures," which is a more actionable and optimistic framing.

Several research directions become more attractive as a result:

Structured task representations for multi-task generative models become a first-class design concern rather than an afterthought. The paper's hierarchical annotation scheme (Scope, Type, Preservation) is one instantiation, but the principle — decompose task semantics into dimensions that correspond to computational conflicts — is generalizable to video generation, 3D synthesis, audio processing, and any domain where different task instances require different processing strategies.
Predictive alignment as a general mechanism for injecting global information into local routing extends beyond image generation. Any architecture with local routing decisions (mixture-of-experts, dynamic module selection, adaptive computation) that needs to be aware of a global context variable (task identity, language, domain, user intent) could benefit from an auxiliary objective that forces routing patterns to encode that variable, rather than providing it as an additional input.
Improving verifier / routing robustness through semantic supervision becomes a natural priority. The paper shows that verifier over-optimization (in the sense of routing collapse to degenerate patterns) can be prevented by tying routing to semantic structure. This has implications for RLHF-based training of generative models, where reward hacking often manifests as the model exploiting the reward model through degenerate generation patterns — semantic alignment of internal routing could provide a structural regularizer.

Conversely, some directions become less attractive:

Pure capacity scaling of dense unified models faces a stronger burden of proof. The paper's ablation (Table 4, Figure 6 left) shows that a dense model with equivalent activated parameters performs substantially worse and converges more slowly, despite having identical per-token FLOPs. This suggests that simply making dense unified models larger — the default approach in the field — hits diminishing returns due to task interference, not just capacity saturation.
Purely data-driven MoE routing for multi-task settings is shown to be insufficient. The "MoE w/o ℒ_align" variant in Table 4 benefits from capacity but leaves substantial performance on the table. This suggests that future MoE designs for multi-task generative models should invest in task-awareness mechanisms rather than trusting gradient descent to discover useful specializations from data alone.

Follow-Up Research This Work Enables

Direct task-conditioning baseline: ablate predictive alignment against gate-input conditioning. The most urgent missing experiment is a comparison between TAG-MoE's predictive alignment mechanism and the simpler alternative of concatenating the semantic embedding $\mathbf{s}$ (or a learned projection) to each token representation before the gating MLP. This baseline — which makes task information available as gate input without the prediction head $\mathcal{H}_{\text{pred}}$ , the routing signature aggregation, or the cosine similarity loss — would directly test whether the specific predictive alignment mechanism matters or whether any form of task-aware routing produces similar benefits. A strong follow-up would implement this baseline on the same Qwen-Image backbone with the same 11M-sample training set, compare performance across all five benchmarks (ICE-Bench, EmuEdit, GEdit, DreamBench++, OmniContext), and analyze expert activation patterns using the same visualization methodology from Figure 6. If direct conditioning matches TAG-MoE's performance, the paper's core architectural contribution is unnecessary complexity and the contribution reduces to the observation that task awareness helps. If TAG-MoE substantially outperforms direct conditioning, the predictive alignment mechanism is validated as producing qualitatively different (and better) routing behavior. The paper's own expert activation analysis — showing spatially-aware, task-specific specialization — would be strengthened by comparing activation patterns for the direct-conditioning variant: does it produce similarly focused specializations, or diffuse routing that merely correlates with task type?

Cross-architecture replication on Flux or Stable Diffusion 3. All results in the paper are on the Qwen-Image backbone. The paper's claims are framed as architectural innovations ("a novel task-aware sparse MoE framework") rather than Qwen-specific optimizations, but the evidence is entirely single-model-family. A strong follow-up would replicate the core ablation (dense vs. task-agnostic MoE vs. TAG-MoE) on at least one other diffusion transformer architecture — Flux (Black Forest Labs) and Stable Diffusion 3 (Stability AI) are natural candidates since they have publicly available weights and established unified-generation fine-tuning recipes. The key question is whether the same hierarchy (Dense < Task-Agnostic MoE < TAG-MoE) holds across architectures with different hidden dimensions, attention mechanisms, and pretraining distributions. If the gap between task-agnostic and task-aware MoE narrows or disappears on certain architectures, it would reveal that the benefit of predictive alignment depends on architectural properties (e.g., how much task-relevant information is already encoded in token representations at the depth where MoE layers are placed). The paper does not report the Qwen-Image model scale, which is a barrier to controlled replication — any follow-up should report parameter counts and FLOPs for both base and MoE-augmented models.

Quantitative expert specialization metrics with task-agnostic comparison. The paper's expert activation analysis (Figure 6) is qualitative and performed only on the full TAG-MoE model. A rigorous follow-up would define and compute quantitative specialization metrics for both TAG-MoE and the task-agnostic MoE variant (without $\mathcal{L}_{\text{align}}$ ). Candidate metrics include: (1) Mutual information between task type (from the 26 ICE-Bench categories) and expert selection (which of the 4 experts a token is routed to), computed per MoE layer — this measures whether routing decisions carry task-identifying information beyond what chance would produce. (2) Spatial concentration entropy of per-expert activation maps — for each expert, compute the entropy of the spatial distribution of tokens routed to that expert; low entropy indicates spatially focused specialization. (3) Task-pair routing divergence — for pairs of tasks with different semantic profiles (e.g., "change material" vs. "style transfer"), compute the Jensen-Shannon divergence between their per-expert routing distributions; high divergence indicates task-specific routing. (4) Layer-wise specialization progression — compute these metrics per MoE layer (layers 1 through 10) to test whether specialization increases with depth, as the paper's design choice (MoE only in later layers) would predict. Computing these metrics for both TAG-MoE and the task-agnostic variant would directly test whether the alignment loss causes the observed specialization patterns or whether task-agnostic MoE also develops similar patterns through standard training. If the task-agnostic variant shows comparable specialization scores, the paper's mechanistic interpretation is weakened; if TAG-MoE shows substantially higher scores, the causal link is strengthened.

Inference-time instruction rewriting ablation. The paper uses VLM-based instruction rewriting at inference to provide enriched prompts that implicitly carry task semantics, but no ablation compares performance with raw versus rewritten instructions. A necessary follow-up would evaluate TAG-MoE on ICE-Bench (or a subset of its 26 categories) using both raw user instructions and VLM-rewritten instructions, measuring the performance gap on each metric (Aesthetic, CLIP-cap, vllmqa, CLIP-src). This would quantify whether rewriting is essential for the method's reported performance or merely a helpful augmentation. If the gap is small (e.g., <2 percentage points on vllmqa), deployment could skip the VLM call, reducing latency and removing the external dependency on Qwen-VL. If the gap is large, the rewriting step is a de facto requirement, and the paper's claim of "minimal inference overhead from the TAG-MoE contributions" must be qualified by the cost of the VLM inference. Beyond performance, a latency analysis comparing the VLM rewriting step to the diffusion sampling step would inform practical deployment: if rewriting takes 200ms and diffusion takes 2s, the overhead is acceptable; if rewriting takes 2s on its own, the system latency doubles. The ablation should also test robustness to rewriting quality — what happens if the VLM produces a poor or misleading rewrite (e.g., misidentifying the task type)? This would reveal whether the gating network has learned routing patterns that are robust to noisy text inputs or whether it overfits to the clean rewritten prompts.

Difficulty-conditioned or content-dependent routing extension. The paper's own limitation (Section 5) identifies that TAG-MoE "cannot jointly reason over intent and the visual content of the source image" — routing depends on task semantics from the instruction alone, not on how difficult or unusual the specific source image is. A natural extension would augment the semantic embedding $\mathbf{s}$ with a content-dependent component: for example, using the PRM or a lightweight visual encoder to produce a "difficulty embedding" that captures image-specific complexity (clutter, unusual viewpoints, fine-grained textures), and incorporating this into either the semantic target $\mathbf{s}$ (by adding a learned content embedding to the tag summation) or as an additional alignment target (a separate predictive head mapping from $\mathbf{g}$ to the content embedding). The experiment would test whether content-aware routing improves performance on tasks identified as challenging in Figure 4 — specifically, categories where the instruction is semantically simple but the source image is visually complex. The ICE-Bench per-category breakdown provides a natural testbed: identify categories where TAG-MoE's absolute performance is low despite the task type being well-represented in training (suggesting image content, not task semantics, is the bottleneck), and measure whether content-conditioned routing closes the gap. A negative result — content-conditioning provides no benefit — would clarify that TAG-MoE's limitation is more fundamental than content-blind routing and may require architectural changes beyond the gating mechanism.

Scaling behavior of the alignment loss with dataset size and expert count. The paper trains on 11M samples with 4 experts but provides no analysis of how the benefit of $\mathcal{L}_{\text{align}}$ scales with data volume or expert count. A scaling study would vary training set size (e.g., 1M, 3M, 6M, 11M samples) and measure the performance gap between task-agnostic MoE and TAG-MoE at each scale. The key question: does the alignment loss provide more benefit in low-data regimes (where task interference is harder to disentangle from limited examples) or in high-data regimes (where there are enough examples per task type for the alignment signal to be precise)? If the gap narrows at large data scales, the alignment loss is primarily a data-efficiency tool; if it persists or widens, the mechanism provides benefits beyond what data alone can achieve. Similarly, varying expert count (2, 4, 8, 16) with and without $\mathcal{L}_{\text{align}}$ would test whether semantic alignment becomes more or less important as the routing space grows — with more experts, the gate has more freedom to discover task-specific specializations, but the routing problem is also higher-dimensional and potentially harder to train without semantic guidance. These experiments would provide practical guidance for practitioners: at what data scale and expert count is the alignment loss worth the implementation complexity and hyperparameter tuning?

Practical Applications and Downstream Use Cases

Unified image editing APIs with per-request task routing. The most direct application is a single deployment that handles the full spectrum of user image manipulation requests — object removal, style transfer, subject insertion, identity-preserving customization, background replacement, and so on — without maintaining separate models for each task family. Current production systems (e.g., OpenAI's GPT-4o image generation, Adobe Firefly, Canva's AI tools) either use separate specialized models behind a unified API (routing requests to different models based on intent classification) or use a single dense model that exhibits the "mediocre compromise" behavior the paper documents. TAG-MoE's architecture enables a genuinely unified model where the routing to specialized computational pathways happens internally, based on the task semantics embedded in the (potentially rewritten) user instruction. The practical benefit is quantified by the paper's results: on ICE-Bench, TAG-MoE achieves the highest open-source scores on Aesthetic, CLIP-cap, and vllmqa metrics (Table 1), meaning it produces outputs that are more beautiful, more instruction-aligned, and more likely to be judged correct by a VLM evaluator. On specialized benchmarks, it matches or exceeds task-specific specialists on their own metrics — highest vllmqa on both EmuEdit-Bench and GEdit-Bench (Table 2), highest Face-ref on both DreamBench++ and OmniContext (Table 3). For a production system, this means one model can replace multiple specialist models without the quality degradation that typically accompanies unification, reducing operational complexity (one training pipeline, one deployment, one monitoring dashboard) while maintaining or improving output quality.

Data generation pipelines for multi-task visual instruction tuning. Large-scale visual instruction tuning datasets (used to train models like GPT-4o, Gemini, and open-source VLMs) require diverse, high-quality image editing and generation examples — source image, instruction, and target image triplets covering many task types. Generating these triplets with a single unified model rather than orchestrating multiple specialist models simplifies the data pipeline and ensures stylistic consistency across task types. TAG-MoE's demonstrated ability to handle 3D-aware edits and maintain identity preservation (Figure 5) is particularly valuable for this use case: the generated target images need to be geometrically plausible and identity-consistent, since downstream models trained on these triplets will learn from any artifacts. The paper's data construction pipeline (Section 3.4) — which uses GPT-4o for instruction generation and specialist/generalist models for target image generation — could be replaced or augmented by TAG-MoE itself, creating a self-improvement loop where the model generates training data for future iterations. The 11M-sample training scale demonstrates that the approach is viable at production data volumes, and the bidirectional cross-task benefits observed qualitatively (editing accuracy benefits from subject-driven data, generation fidelity benefits from editing data) suggest that the generated data would have the multi-task diversity needed for robust instruction tuning.

On-device or edge deployment of unified generation models. The paper's MoE architecture with top-1 routing provides a capacity-expansion mechanism that keeps per-token inference cost comparable to the dense backbone. Each token activates exactly one of four FFN experts, so the activated parameter count per token is identical to the dense model, while the total parameter count is approximately 4× larger (for the MoE-ified layers). This means TAG-MoE achieves the benefits of a larger model (better handling of diverse tasks, more capacity for specialization) while maintaining the inference latency and memory footprint of a smaller dense model. For edge deployment scenarios — mobile photo editing apps, on-device generative AI features, real-time video editing — where memory and latency constraints are tight, this parameter-efficiency-through-sparsity is directly valuable. The paper does not report memory usage or latency measurements, which would need to be characterized for a specific deployment target, but the architecture's design (identical expert structure to the FFN it replaces, top-1 routing) is inherently deployment-friendly compared to alternatives that activate multiple experts per token or use larger expert networks. A mobile deployment could use a distilled or quantized version of TAG-MoE where the same architectural principle (task-aware sparse experts) applies at a smaller scale.

Interactive creative tools with instruction-based editing and customization. Creative tools (photo editors, design software, AR filters) are increasingly integrating text-instructed image manipulation. A common failure mode in current tools is that users must carefully phrase their instructions to get the desired result — "make the background a beach" might work, but "make it look like a beach" produces inconsistent results depending on whether the model interprets "it" as the background or the subject. TAG-MoE's semantic-aligned routing, which learns to associate instruction semantics with expert selection patterns, could produce more robust instruction understanding, since the VLM rewriting step at inference time explicitly disambiguates the instruction (expanding "make it look like a beach" into a detailed prompt specifying the scope, type, and preservation constraints). The user study results (Figure 7) show that TAG-MoE is consistently preferred by human evaluators on reference alignment, prompt alignment, and overall preference compared to unified baselines — for an interactive creative tool, this translates to fewer user retries, less frustration, and higher-quality outputs on the first attempt. The 26-category leadership on ICE-Bench (Figure 4) suggests the robustness spans a wide range of creative tasks, from subtle local edits to global style transformations.