Research map / Architectures and efficiency / Mixture of experts
Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Race-DiT with Expert Race, a global top-K routing over batch, sequence and expert dimensions for MoE diffusion transformers, plus router similarity loss, per-layer regularization and an EMA threshold for inference.
Problem
Standard MoE routing (token-choice, expert-choice) in diffusion transformers allocates a fixed number of experts per token or per expert and ignores timestep and spatial complexity variation, limiting allocation flexibility when scaling DiT with MoE.
Approach
Expert Race performs a single global top-K selection over all token-expert gating scores across batch, timesteps, spatial tokens and experts, with identity gating, an EMA-learned threshold for sample-independent inference, a router similarity loss on expert co-selection patterns, and a per-layer auxiliary regression loss for shallow layers.
Key results
- On ImageNet 256x256, RaceDiT-XL/2-4in32 (0.707B activated, 2.788B total) reaches FID 2.06 with 1.7M x 1024 samples, versus 3.02 for the dense DiT-XL/2* trained under the same budget.
- In ablations on DiT-B/2 with 2-in-8 MoE, Expert Race with identity gating, learnable threshold, per-layer regularization and router similarity loss improves FID from 16.47 (softmax baseline) to 8.03.
- Against DiT-MoE-B/2-8E2A under identical 500K-iteration training, RaceDiT-B/2-2in8 gets FID 8.02 vs 9.06 with fewer activated parameters (0.135B vs 0.286B).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace dense FFNs with Race-DiT 4-in-32 MoE at comparable activated parameter count | DiT-B/2 · ImageNet | FID: 18.03 -> 7.35 | Helped |
| replace dense FFNs with Race-DiT 4-in-32 MoE at comparable activated parameter count | DiT-M/2 · ImageNet | FID: 11.18 -> 5.16 | Helped |
| replace dense FFNs with Race-DiT 4-in-32 MoE at comparable activated parameter count | DiT-L/2 · ImageNet | FID: 7.88 -> 4.04 | Helped |
| replace dense FFNs with Race-DiT 4-in-32 MoE at comparable activated parameter count | DiT-XL/2 · ImageNet | FID: 6.31 -> 3.31 | Helped |
| use Expert Race instead of Token Choice (ablation) | DiT-B/2 2-in-8 MoE · ImageNet | FID: 9.50 -> 8.03 | Helped |
| use Expert Race instead of Expert Choice (ablation) | DiT-B/2 2-in-8 MoE · ImageNet | FID: 10.13 -> 8.03 | Helped |
| replace softmax with identity gating (ablation) | Expert Race · ImageNet, DiT-B/2 2-in-8 MoE | FID: 16.47 -> 13.66 | Helped |
| replace softmax with sigmoid gating (ablation) | Expert Race · ImageNet, DiT-B/2 2-in-8 MoE | FID: 16.47 -> 13.85 | Helped |
| add learnable threshold (ablation) | Expert Race with identity gating · ImageNet, DiT-B/2 2-in-8 MoE | FID: 13.66 -> 11.56 | Helped |
| add per-layer regularization (ablation) | Expert Race with learnable threshold · ImageNet, DiT-B/2 2-in-8 MoE | FID: 11.56 -> 8.95 | Helped |
| add router similarity loss (ablation) | Expert Race with per-layer regularization · ImageNet, DiT-B/2 2-in-8 MoE | FID: 8.95 -> 8.03 | Helped |
| replace conventional balance loss with router similarity loss (ablation) | 4-in-32 MoE · ImageNet | FID: 11.67 -> 10.77 | Helped |
Surprising findings
- Softmax gating hurts Expert Race substantially; identity gating (no normalization) gives the best results, since softmax disrupts cross-token score ordering.
- Per-layer regularization and the router similarity loss each give large gains on top of the learnable threshold, which is a larger improvement than the routing change alone in some ablations.
Limitations
- Evaluation limited to class-conditional ImageNet generation; extension to other diffusion tasks is left to future work.
- Expert Race requires a global top-K over the batch, and training-inference mismatch is handled by an EMA threshold, so inference behavior depends on the learned threshold.
Experimental setup
Datasets: ImageNet
Benchmarks: ImageNet 256x256
Metrics: FID CMMD CLIP Score IS Precision Recall MaxVio Combination Usage
Baselines: Token Choice Expert Choice DiT-MoE-XL/2 DiT-XL/2 SiT-XL MDT ADM-G LDM-8-G MaskDiT BL-Choice BE-Choice LE-Choice
Builds on: DiT DiT-MoE Expert Choice routing fine-grained expert segmentation DeepSeekMoE
Closest papers by meaning
- EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing (ICLR 2025 · similarity 0.91)
- Diff-MoE: Diffusion Transformer with Time-Aware and Space-Adaptive Experts (ICML 2025 · similarity 0.90)
- Rethinking Cross-Layer Information Routing in Diffusion Transformer (NeurIPS 2026 · similarity 0.89)
- EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching (NeurIPS 2024 · similarity 0.88)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.88)
- Towards Precise Scaling Laws for Video Diffusion Transformers (CVPR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.