PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yujie Wei, Shiwei Zhang, Hangjie Yuan, Yujin Han, Zhekai Chen, Jiayu Wang, Difan Zou, Xihui Liu, Yingya Zhang, Yu Liu, Hongming Shan

arXiv:2510.24711v2cs.CV

TL;DR

DiT MoE methods have shown limited gains because visual tokens are spatially redundant and functionally heterogeneous, unlike semantically dense language tokens. ProMoE uses conditional and prototypical routing with explicit semantic guidance and routing contrastive learning, and it outperforms dense and state-of-the-art MoE methods across Rectified Flow and DDPM settings. Its guidance strategies have scope and robustness trade-offs, while generative-model risks remain.

  • Problem

    Existing MoE applications to Diffusion Transformers show limited gains because visual-token redundancy and functional heterogeneity hinder expert specialization.

  • Method

    ProMoE uses conditional routing for functional roles, prototypical routing for semantic content, and routing contrastive loss for explicit guidance.

  • Results

    ProMoE outperforms dense models and state-of-the-art MoE methods across both Rectified Flow and DDPM objectives, including comparisons with larger models.

  • Takeaways & Limitations

    Explicit semantic guidance promotes intra-expert coherence and inter-expert diversity, supporting effective and parameter-efficient scaling of DiTs with MoE.

  • Takeaways & Limitations

    Classification guidance requires costly manual annotations, while clustering guidance is limited to top-1 assignment and is sensitive to cluster count and initialization.

Abstract

from arXiv · show

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.

1 INTRODUCTION

DiT models incur high computation from dense activation, while visual-token redundancy and functional heterogeneity impede expert specialization compared with language models. ProMoE addresses this with guided two-step routing and achieves broad gains across diffusion objectives and model settings.

  • 1 INTRODUCTION: Dense DiTs activate all parameters regardless of input or task, creating substantial computational overhead that limits further scalability.MoE is introduced as a capacity-scaling paradigm that sparsely activates specialized subnetworks.
  • 1 INTRODUCTION: Visual tokens’ spatial redundancy and functional heterogeneity make expert specialization harder than for semantically dense, separable language tokens.These properties can lead experts to learn homogeneous features and help explain limited gains from prior DiT-MoE approaches.
  • 1 INTRODUCTION: ProMoE combines conditional routing by functional role with prototypical routing by semantic content to promote expert specialization in DiTs.Conditional routing separates unconditional and conditional image tokens; prototypical routing refines conditional assignments using learnable prototypes.
  • 1 INTRODUCTION: Routing contrastive loss uses semantic similarity to promote intra-expert coherence and inter-expert diversity without manual labels.The loss strengthens prototypical routing by grouping semantically similar tokens and separating dissimilar token distributions.
  • 1 INTRODUCTION: ProMoE outperforms dense models and state-of-the-art MoE methods across diverse settings, including both Flow Matching and DDPM paradigms.The framework uses fewer activated parameters than dense counterparts while surpassing methods with 1.7× more total parameters.

2 RELATED WORK

Diffusion models have shifted from U-Net architectures toward scalable Transformer backbones, while MoE selectively activates subnetworks to expand capacity efficiently. Related work applies MoE to diffusion models, motivated by successful scaling in language models.

  • 2 RELATED WORK: Recent diffusion models have shifted from U-Net toward Transformer architectures for improved scalability and generative quality.Earlier diffusion work primarily used U-Net models trained with DDPM objectives.
  • 2 RELATED WORK: MoE expands model capacity while minimizing computation by sparsely activating subnetworks for distinct inputs.This paradigm inspired applications to diffusion models after strong results in LLMs.

3 PRELIMINARIES

Diffusion models reverse a noising process with a denoising network conditioned on inputs such as class labels or text prompts. MoE routes tokens to selected experts and combines their outputs, using token-choice or expert-choice allocation.

  • 3 PRELIMINARIES: Diffusion models learn data distributions by reversing a forward noising process, with x_t = α_t x_0 + σ_tϵ.The timestep controls decreasing signal and increasing noise, while the denoiser predicts a target conditioned on c.
  • 3 PRELIMINARIES: The denoising target is Gaussian noise ϵ for DDPM models or the vector field (ϵ − x_0) for Rectified Flow models.The denoising network is trained at each timestep using the chosen objective.
  • 3 PRELIMINARIES: A standard MoE router maps token representations to token–expert affinity scores, activates the top-K experts, and combines their outputs through gating.The input has batch, token-length, and hidden-dimension axes, while the expert dimension indexes available subnetworks.
  • 3 PRELIMINARIES: Token-choice routing lets each token select its top-K experts, whereas expert-choice routing lets each expert select a fixed number of top-K tokens.These are the two common routing paradigms described for MoE layers.

4 PROMOE

ProMoE uses explicit routing guidance in a two-step router to separate visual tokens by functional role and semantic content, enhancing expert specialization. Routing contrastive learning further promotes coherent within-expert assignments and diverse expert functions, while experiments evaluate semantic guidance, scalability, and model configurations.

  • Two-Step Router: ProMoE first partitions visual tokens into conditional and unconditional sets, then routes conditional tokens among experts using learnable semantic prototypes.Conditional routing uses input conditioning, while prototypical routing allocates tokens through cosine similarity to expert-associated prototypes.
  • Two-Step Router: Shared experts process all tokens, while routed or unconditional experts selectively process tokens according to their conditional or unconditional type.The MoE output sums shared-expert outputs with a token-type-dependent selective output.
  • Semantic Routing Guidance: Semantic routing guidance improves visual MoE assignment through classification-based supervision or clustering-based similarity, addressing spatial redundancy that otherwise produces homogeneous expert features.The classification router uses ImageNet superclasses, whereas k-means assigns semantically similar tokens to the same expert.
  • Semantic Routing Guidance: Classification guidance requires costly manual annotations and limits token-level flexibility, while clustering guidance is restricted to top-1 assignment and is sensitive to cluster count and initialization.These limitations motivate learnable prototypical routing with routing contrastive learning.
  • Routing Contrastive Learning: Routing Contrastive Loss pulls prototypes toward assigned-token centroids and pushes them away from other clusters, enforcing intra-expert coherence and inter-expert diversity.The loss uses positive tokens assigned to each prototype and negative tokens assigned to other prototypes; its push-away operation also acts as semantic load balancing.
  • Experiments: ProMoE is evaluated against dense and MoE baselines under Rectified Flow and DDPM objectives, with model sizes aligned to established DiT benchmarks.The experimental setup includes equivalent activated parameters for MoE comparisons and configurations spanning S, B, L, and XL sizes.

5 EXPERIMENT

Experiments evaluate ProMoE against dense and MoE DiT baselines across objectives, scales, guidance settings, efficiency measures, and routing ablations. ProMoE consistently improves performance and scalability while using fewer activated or total parameters.

  • 5.2 MAIN RESULTS: ProMoE consistently surpasses dense counterparts at equivalent activated parameters across model sizes, training objectives, and CFG settings.Under Rectified Flow, FID reductions range from 19.8% to 29.4% in the reported comparisons.
  • 5.2 MAIN RESULTS: ProMoE outperforms MoE baselines across both DDPM and Rectified Flow, including a smaller model surpassing a 16-expert DiffMoE with 1.846B parameters.ProMoE-L-Flow has 1.063B parameters, compared with 1.846B for the larger DiffMoE model.
  • 5.2 MAIN RESULTS: ProMoE achieves lower inference time and fewer GFLOPs than DiffMoE while maintaining comparable GFLOPs to other MoE baselines.The computational-cost comparison evaluates training time, inference time, and FLOPs under identical hardware settings.
  • 5.3 ANALYSIS: At the L scale, ProMoE reduces training loss and converges faster than dense and MoE baselines, with the convergence advantage persisting at XL.The XL comparison extends training to 1.2M steps.
  • 5.4 ABLATION STUDIES: Adding routing contrastive learning reduces FID by 10.6% and increases IS by 9.8% over prototypical routing alone.The ablation also finds that prototypical routing improves performance over DiT-MoE-B-Flow and DiffMoE-B-Flow.
  • 5.4 ABLATION STUDIES: Identity activation performs best for prototypical-routing scores, while sigmoid ranks second and softmax performs worst.The identity mapping is therefore adopted in the final method.
  • 5.4 ABLATION STUDIES: Removing conditional routing significantly degrades K-Means-based routing performance, supporting its method-agnostic value.The conditional-routing component is evaluated within an alternative K-Means-based routing scheme.

6 CONCLUSION

The paper concludes that ProMoE addresses visual-token challenges in DiT MoE through explicit functional and semantic routing guidance. Its experiments show improved performance over dense DiT and existing MoE methods with fewer activated or total parameters.

  • 6 CONCLUSION: ProMoE uses conditional and prototypical routing plus routing contrastive learning to promote intra-expert coherence and inter-expert diversity.Conditional routing partitions tokens by functional role, while prototypical routing refines assignments using semantic content.
  • 6 CONCLUSION: On ImageNet, ProMoE outperforms dense DiT and state-of-the-art MoE methods despite using fewer activated or total parameters.The conclusion presents this as a solution for scaling DiT with MoE.

ETHICS STATEMENT

The ethics statement acknowledges generative-model risks, including the potential creation of fake data, and notes the use of forgery detection and ethical guidelines.

  • ETHICS STATEMENT: ProMoE inherits generative-model risks such as potentially creating fake data, motivating robust image-forgery detection.The authors also state that ethical guidelines were followed in all experiments.

A EXPERIMENTAL SETUP

Experiments compare ProMoE with dense and MoE baselines under matched or comparable settings across DDPM and Rectified Flow. Evaluation uses standard ImageNet and text-to-image protocols alongside computational-cost measurements.

  • A EXPERIMENTAL SETUP: Baselines include Dense-DiT, DiT-MoE, EC-DiT, and DiffMoE, evaluated with equivalent activated parameters and comparable total parameters.The models are trained under both DDPM and Rectified Flow objectives and ProMoE is scaled across S, B, L, and XL sizes.
  • A EXPERIMENTAL SETUP: Computational cost is reported through training time, inference time, and FLOPs, with inference measured at CFG 1.5, batch size 128, and 250 sampling steps.All computational comparisons use four NVIDIA A800 GPUs under identical hardware settings.
  • A EXPERIMENTAL SETUP: Text-to-image experiments use an MM-DiT architecture trained on a 2M-image LAION-5B subset with matched 12B total and 3B activated parameter budgets.Token-Choice MoE and ProMoE each use five experts under this comparison.

B IMPLEMENTATION ALGORITHMS

Additional analyses show that ProMoE separates expert assignments more effectively and consistently outperforms dense and MoE baselines across reported ImageNet comparisons.

  • Token analysis: Additional t-SNE analyses compare token embeddings across layers and diffusion timesteps for DiT-XL/2 and Llama-3 8B.The embeddings are clustered into 10 groups using k-means for comparison.
  • Expert allocation: ProMoE produces more structured expert token assignments than DiT-MoE in the t-SNE visualization analysis.The comparison visualizes intermediate-layer assignments, with each color representing one expert.
  • Generation results: ProMoE generates high-quality images across simple and challenging categories, including text-to-image samples.Additional generation examples are provided for ProMoE-XL-Flow and the text-to-image task.
  • Additional comparisons: Across both training objectives and CFG settings, ProMoE consistently outperforms the dense model and existing MoE state-of-the-art methods.The comparisons include DDPM and Rectified Flow evaluations on ImageNet.

C.5 MORE COMPARISON RESULTS WITH EXTENDED TRAINING STEPS

Extended-training experiments show that ProMoE improves performance and parameter efficiency relative to dense and MoE baselines across model sizes, guidance scales, and routing configurations.

  • Extended training comparisons: ProMoE-L at 500K steps surpasses Dense-DiT-L at 1M steps on FID, while ProMoE-XL at 500K steps surpasses Dense-DiT-XL at 1M steps on FID and IS.With longer training, ProMoE-L at 1M steps also outperforms Dense-DiT-L at 2M steps and Dense-DiT-XL at 1M steps.
  • Extended training comparisons: At CFG=4.0, ProMoE reduces FID by ∼21% over dense models and 10% over DiffMoE for L, and by ∼30% and 16% for XL, respectively.ProMoE also achieves the best IS for XL and retains improvements at lower CFG values.
  • Activated experts: Increasing activated routed experts from 1 to 3 improves performance while keeping total parameter count unchanged.The experiment increases activated parameters and evaluates the change on ImageNet under Rectified Flow.
  • Efficiency: ProMoE’s performance gains are primarily attributed to methodological design rather than increased computational overhead.The computational analysis compares training time, inference time, and GFLOPs with dense and MoE baselines.
  • Routing baselines: ProMoE consistently outperforms GMM-based routing across FID and IS, CFG scales, and top-1 or top-3 activation settings.The comparison is conducted under Rectified Flow on the large model.

D MORE RESULTS ON SCALING BEHAVIOR

Scaling experiments show that ProMoE benefits from larger model sizes and more experts, while K-means-based routing is sensitive to expert count and does not scale reliably.

  • Scaling model size: ProMoE-L-Flow surpasses Dense-XL-Flow on FID and Inception Score at CFG=1.0 and 1.5 despite using fewer activated parameters.ProMoE consistently outperforms its dense counterparts across model-size scaling experiments.
  • Scaling expert count: ProMoE performance improves as the number of experts increases, demonstrating scalability across expert counts.The experiment reports Inception Score at CFG=1.0 and 1.5 and FID at CFG=1.5.
  • Scaling expert count: K-means-based routing shows fluctuating performance with increasing expert counts and no clear gain despite increased parameter count.This contrasts with ProMoE’s scaling behavior and indicates sensitivity to the selected expert number.

E.1 ABLATION ON LOAD-BALANCING LOSS

Ablations indicate that routing contrastive learning provides semantic load balancing, while conventional load balancing adds little benefit and can slightly degrade performance.

  • Load-balancing loss: Adding a conventional load-balancing loss on top of ProMoE slightly degrades performance.The ablation attributes this to conventional balancing regularizing token counts without directly using assignment quality or semantics.
  • Routing contrastive learning: The routing contrastive loss’s push-away term acts as a semantic load-balancing regularizer and is more effective than traditional load-balancing loss.Its final contribution is weighted by λ_RCL in the ProMoE training objective.
  • RCL weight: Increasing the RCL loss weight from 1 to 10 yields only marginal gains, so the default weight is set to 1 except for ProMoE-B-DDPM.ProMoE-B-DDPM uses a weight of 10 based on the ablation.
  • Routing algorithm: The ProMoE layer first separates unconditional and conditional image tokens, then routes conditional tokens through prototype-based top-k expert assignment.The algorithm also applies shared-expert processing and computes the RCL auxiliary loss.
Loading 2510.24711v2…