Research map / Image generation / Autoregressive and latent generation
Diff-MoE: Diffusion Transformer with Time-Aware and Space-Adaptive 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 Diff-MoE, a Diffusion Transformer whose MLPs are replaced by mixture-of-experts with expert-specific timestep conditioning, global feature recalibration, and low-rank decomposition, trained from scratch on ImageNet 256x256.
Problem
Dense Diffusion Transformers activate all parameters at every timestep and token, and treat all denoising stages uniformly, limiting scalability; prior MoE-for-diffusion work captures only temporal or only spatial specialization.
Approach
Jointly combines spatial token routing with expert-specific timestep AdaLN conditioning so each expert adapts to the denoising stage, plus a globally pooled context embedding that recalibrates expert channels, with low-rank factorization to limit parameter overhead.
Key results
- Diff-MoE-XL/2 reaches FID 2.69 vs 3.42 for DiT-MoE-XL/2 at cfg=1.5 with similar activated parameters (1.5B) on ImageNet 256x256 (400K iterations).
- Diff-MoE-S reduces FID from 53.90 (SiT-LLaMA-S) to 44.27 at cfg=1.0.
- Diff-MoE-S w/o rectified flow gets FID 31.18 vs 37.43 for DTR-S at cfg=1.5 with matched activated params.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace SiT-LLaMA-S/2 with Diff-MoE-S/2 | ImageNet class-conditional generation · ImageNet 256×256, 400K, cfg=1.0 | FID: 53.90 -> 44.27 | Helped |
| replace SiT-LLaMA-B/2 with Diff-MoE-B/2 | ImageNet class-conditional generation · ImageNet 256×256, 400K, cfg=1.0 | FID: 29.53 -> 24.88 | Helped |
| replace SiT-LLaMA-L/2 with Diff-MoE-L/2 | ImageNet class-conditional generation · ImageNet 256×256, 400K, cfg=1.0 | FID: 14.32 -> 13.98 | Helped |
| use Diff-MoE instead of temporal MoE baselines | ImageNet 256×256 generation, small and base scales · ImageNet 256×256, 400K, cfg=1.5 | FID: 37.43 -> 31.18; 33.99 -> 31.18; 16.58 -> 11.97; 16.21 -> 11.97 | Helped |
| use Diff-MoE instead of DiT-MoE | ImageNet 256×256 generation, small, base, and XL scales · ImageNet 256×256, 400K, cfg=1.5 | FID: 17.92 -> 9.27; 7.19 -> 4.00; 3.42 -> 2.69 | Helped |
| add 2D RoPE (ablation) | DiT-S/2 · ImageNet 256×256, 400K, cfg=1.5 | FID: 43.02 -> 39.80 | Helped |
| add GLU (ablation) | DiT-S/2 with RoPE 2D · ImageNet 256×256, 400K, cfg=1.5 | FID: 39.80 -> 38.20 | Helped |
| add depthwise convolution (ablation) | DiT-S/2 with RoPE 2D and GLU · ImageNet 256×256, 400K, cfg=1.5 | FID: 38.20 -> 35.85 | Helped |
| add spatial MoE (ablation) | DiT-S/2 with RoPE 2D, GLU, and depthwise convolution · ImageNet 256×256, 400K, cfg=1.5 | FID: 35.85 -> 32.74 | Helped |
| add expert-specific timestep conditioning (ablation) | MoE Baseline · ImageNet 256×256, 400K, cfg=1.5 | FID: 32.74 -> 28.71 | Helped |
| add global feature recalibration (ablation) | MoE Baseline with expert-specific timestep conditioning · ImageNet 256×256, 400K, cfg=1.5 | FID: 28.71 -> 26.33 | Helped |
| add low-rank decomposition (r=64) (ablation) | MoE Baseline with expert-specific timestep conditioning and global feature recalibration · ImageNet 256×256, 400K, cfg=1.5 | FID: 26.33 -> 27.59 | Hurt |
Limitations
- Large-scale variants such as Diff-MoE-XL were not fully evaluated at longer training (e.g., 7M steps) due to compute constraints.
- Experiments limited to class-conditional ImageNet 256x256 image generation.
- Some comparisons use own variants (e.g., w/o rectified flow) rather than identical training protocols.
Experimental setup
Datasets: ImageNet
Benchmarks: ImageNet 256x256
Metrics: FID IS
Baselines: DiT SiT SiT-LLaMA DTR Switch-DiT DiT-MoE
Builds on: DiT DiT-MoE DTR Switch-DiT SiT-LLaMA GLU
Closest papers by meaning
- Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts (ICML 2025 · similarity 0.90)
- EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing (ICLR 2025 · similarity 0.90)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.88)
- Boosting Generative Image Modeling via Joint Image-Feature Synthesis (NeurIPS 2025 · similarity 0.88)
- LaVin-DiT: Large Vision Diffusion Transformer (CVPR 2025 · similarity 0.87)
- Diff-Plugin: Revitalizing Details for Diffusion-based Low-level Tasks (CVPR 2024 · similarity 0.87)
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.