Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

AI summary

Overview

Research area: Computer vision / generative modeling — specifically Mixture-of-Experts (MoE) architectures for Diffusion Transformers (DiTs) used in image generation.

Technical level: Intermediate. The paper assumes familiarity with diffusion models, transformer architectures, and the basic MoE routing paradigm, but its core argument is accessible: image tokens behave differently from text tokens, and standard MoE routers do not account for that.

Scope: The paper diagnoses why MoE — a major success in large language models — yields only marginal gains in Diffusion Transformers, and proposes ProMoE, a router design with explicit guidance signals that restores expert specialization.

What This Paper Is About

Mixture-of-Experts layers let models grow their total parameter count while keeping per-token compute roughly constant, by sending each token to only a few specialized sub-networks ("experts"). This works spectacularly in language models but has produced disappointing results in Diffusion Transformers, where MoE models often match or barely beat dense baselines of equal activated size. This paper argues the failure is not incidental: image tokens are spatially redundant and come in two functionally distinct flavors (conditional and unconditional, due to classifier-free guidance), so a router that treats all tokens uniformly cannot induce meaningful expert specialization. ProMoE fixes this by routing tokens according to their functional role first, then according to their semantic content, and by training the routing with a contrastive objective.

Key Contributions

1. A diagnosis of the language-vision MoE gap. The authors quantify that LLM tokens form tight, well-separated semantic clusters (inter/intra-class distance ratio of 19.283), while visual tokens are diffuse (ratio of 0.748), and identify functional heterogeneity from classifier-free guidance as a second obstacle.

2. A two-step router. Step one, conditional routing, hard-partitions image tokens into unconditional and conditional sets, dispatching them to dedicated unconditional experts or routed experts respectively. Step two, prototypical routing, assigns conditional tokens to experts via cosine similarity against a set of learnable per-expert prototypes in latent space.

3. A routing contrastive loss (RCL). This loss pulls each prototype toward the centroid of the tokens it was assigned and pushes it away from other experts' centroids, explicitly enforcing intra-expert coherence and inter-expert diversity without needing manual labels — unlike the classification-based guidance they also test.

4. Extensive validation across training objectives and scales. ProMoE is evaluated at four sizes (33M–675M activated parameters) under both Rectified Flow and DDPM objectives on ImageNet, plus a text-to-image experiment on GenEval, consistently outperforming dense DiTs and existing MoE methods.

Main Findings

• Semantic guidance is the missing ingredient. Adding either explicit (class-label classification) or implicit (k-means clustering) routing guidance to a base-size MoE dramatically improves results: FID at CFG=1.5 drops from 9.02 (dense) and 8.22 (DiffMoE) to 5.91 (classification-guided) and 6.24 (k-means-guided), with the load-balancing loss disabled in both guided cases.

• Each component contributes incrementally. Ablating on the base model at CFG=1.5: dense DiT gives FID 9.02; adding prototypical routing gives 7.92; adding RCL gives 6.75; adding conditional routing gives 6.39 (with IS rising from 131.13 to 154.21).

• Identity beats sigmoid and softmax for routing scores. Since prototypical routing computes similarities in latent space, the activation choice matters: identity (24.44 FID) outperforms sigmoid (25.49) and softmax (25.74), which the authors attribute to stable top-K selection.

• ProMoE beats dense models of larger activated size. ProMoE-L-Flow (458M activated) reaches FID 11.61/2.79 at CFG 1.0/1.5, versus 13.38/3.23 for Dense-DiT-XL-Flow (675M activated). At XL scale, ProMoE reaches FID 2.59 and IS 265.62.

• Parameter efficiency against MoE SOTAs. ProMoE-L-Flow with 1.063B total parameters outperforms DiffMoE-L-Flow with 16 experts and 1.846B total parameters (FID 2.79 vs 3.30 at CFG=1.5), while also using fewer GFLOPs.

• Gains are larger under Rectified Flow. Without CFG, ProMoE-L-Flow cuts FID by 24.8% and raises IS by 19.7% relative to the dense counterpart; ProMoE-XL-Flow cuts FID by 29.4% — a bigger relative improvement than under DDPM.

• Convergence is faster. Training loss curves show ProMoE converging lower and quicker than both dense and MoE baselines at the L scale, and this persists at XL even after 1.2M steps.

• Text-to-image generalizes too. On GenEval after 400K steps with a 3B-activated model, ProMoE scores 0.463 overall versus 0.390 for dense and 0.417 for Token-Choice MoE, with the largest gains on Position (0.212 vs 0.095 dense) and Colors (0.675 vs 0.611).

• Expert utilization confirms real specialization. Measured across two disjoint 200-class subsets, DiT-MoE produces nearly identical token-per-expert distributions across subsets (poor differentiation), whereas ProMoE shows distinct utilization patterns across subsets and varied frequencies within each, without starving or overloading any expert.

• Conditional routing is method-agnostic. Removing it from the k-means-based routing variant degrades FID from 6.24 to 8.75, suggesting the functional-role split helps other routing schemes, not just ProMoE's own.

Methodology in Plain English

The authors start by asking why MoE, which works so well in language models, has faltered in image generation. They trace it to two properties of image tokens. First, image patches overlap and correlate heavily, so they carry redundant information that gives the router little signal to distinguish them. Second, classifier-free guidance — the standard technique for improving diffusion sample quality — means the model processes two functionally different kinds of input at once: image tokens under a real condition (a class label, a text prompt) and the same image tokens under an empty condition. A conventional router has no way to know these mean different things.

ProMoE's router handles this in two explicit stages. The first stage is a hard rule, not a learned one: if a token came from the null-conditioned branch, it goes to a dedicated unconditional expert; otherwise it heads to the second stage. This forces the unconditional and conditional experts to stop competing for the same job.

The second stage assigns conditional tokens by measuring cosine similarity between each token's hidden representation and a set of learnable "prototype" vectors, one per expert. Each prototype acts as the centroid of an expert's semantic territory. Tokens go to the top-K most similar prototypes. The authors use the raw similarity score as the routing weight (identity activation) rather than softmax, because softmax is sensitive to sequence length and blurs the ranking.

To train this, they add a contrastive loss: within each batch, the prototype is pulled toward the average of the tokens routed to it (keeping an expert's content coherent) and pushed away from the averages of other experts' assigned tokens (keeping experts distinct). This gives the router explicit gradient signal about what each expert should own. The push-away term also turns out to double as a semantic load balancer, reducing reliance on the traditional load-balancing loss.

To keep compute fixed, the experts' hidden dimensions are halved so that activating one shared expert plus one routed/unconditional expert matches the FLOPs of a dense model of the same activated size.

Why This Matters

Impact on research. The paper reframes the DiT-MoE problem from "make the router smarter" to "recognize that vision tokens have different structure than language tokens." That shift suggests many MoE techniques borrowed wholesale from LLMs need adaptation before they transfer to vision. The finding that explicit semantic guidance is essential — and that a contrastive loss can supply it without labels — gives the field a concrete, principled lever rather than more routing heuristics. The demonstrated parameter efficiency (beating models with 1.7x more total parameters) also matters for anyone training large generative models on a fixed budget.

Real-world applications:

• Text-to-image and image generation services, where fewer activated parameters per forward pass translates directly into lower inference cost at the same or better quality.
• Advertising, e-commerce, and design tooling that generate product imagery or marketing assets at scale, where FID and IS improvements mean fewer rejected outputs per batch.
• Synthetic training data generation for downstream vision models, where higher-fidelity conditional samples improve the utility of the generated datasets.
• Video and multi-modal generation pipelines, which share the DiT backbone and the same classifier-free guidance structure, and where compute costs are far higher than for images.

Industry relevance. The work comes from a collaboration including Alibaba's Tongyi Lab, and the efficiency argument is squarely industrial: matching dense-model quality while activating fewer parameters and using fewer GFLOPs than competing MoE methods reduces serving costs for large generative models. The code is publicly released.

Future Directions

• Video diffusion. The authors cite video generation work but only evaluate on images and one text-to-image benchmark. Video tokens are even more spatially and temporally redundant, so the routing guidance argument should apply, but has not been tested.

• Adaptive expert and prototype counts. The prototype design assumes a fixed number of experts per layer. Whether the number of prototypes should vary across layers or scale with data diversity is unexplored, and the k-means baseline's sensitivity to cluster count suggests this is not trivial.

• Generalizing conditional routing beyond classifier-free guidance. The split relies on diffusion models' two-branch inference scheme. Whether the principle extends to other functional heterogeneity (timestep regimes, modality-specific tokens in multi-modal DiTs) is an open question the paper does not address.

• Better evaluation than FID and IS. The authors explicitly flag this as a limitation — FID50K and IS may not capture fine-grained perceptual quality or semantic faithfulness, which matters especially for the text-to-image results where the reported GenEval gains are modest in absolute terms.

Target Audience

Researchers and engineers working on efficient generative model architectures, particularly those scaling diffusion transformers or adapting MoE designs to non-language modalities. It is most useful to readers who already understand transformer routing and diffusion training objectives, and who want a concrete, well-ablated account of why standard MoE transfers poorly to vision and what design changes fix it. Practitioners deploying large image generation models under compute constraints will also find the parameter-efficiency comparisons directly actionable.

Authors’ abstract

Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fundamental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert specialization in vision MoE. To this end, we present ProMoE, an MoE framework featuring a two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and refine the assignments of conditional image tokens through prototypical routing with learnable prototypes based on semantic content. Moreover, the similarity-based expert allocation in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and we validate that such guidance is crucial for vision MoE. Building on this, we propose a routing contrastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on ImageNet benchmark demonstrate that ProMoE surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code is available at https://github.com/ali-vilab/ProMoE.