Research map / Image generation / Autoregressive and latent generation
Scaling Diffusion Transformers Efficiently via $\mu$P
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.
Extends μP theory to diffusion Transformers (U-ViT, DiT, PixArt-α, MMDiT) via Tensor Programs, then empirically validates μTransfer of base learning rate for DiT, PixArt-α, and MMDiT up to 18B with reduced tuning cost.
Problem
Hyperparameter tuning for diffusion Transformers becomes prohibitively expensive at billions of parameters; it is unclear whether μP, derived for vanilla (autoregressive) Transformers, applies to diffusion Transformers with their extra text/timestep conditioning modules and iterative denoising.
Approach
Prove via the Tensor Programs (NEXUS/Tensor Program) framework that forward passes of U-ViT, DiT, PixArt-α and MMDiT are representable, so their μP matches standard μP; then apply μTransfer of base HPs from small proxy models (width 288 for DiT/PixArt, 512 for MMDiT) to large targets.
Key results
- DiT-XL-2-muP with transferred base LR 2^-10 reaches FID-50K comparable to DiT-XL-2 final at 2.4M steps versus 7M steps (2.9x faster convergence).
- PixArt-alpha-muP (0.61B) outperforms PixArt-alpha at every checkpoint; at 30 epochs GenEval 0.26 vs 0.15 and MJHQ FID-30K 29.96 vs 42.71, with tuning costing 5.5% of pretraining FLOPs.
- MMDiT-muP-18B achieves GenEval overall 0.8218 vs 0.8154 for expert-tuned MMDiT-18B, with tuning at about 3% of the human-tuning FLOPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Transfer base learning rate 2^{-10} under μP to DiT-XL-2 | DiT-XL-2 · ImageNet | convergence steps: 7M steps -> 2.4M steps; 2.9× faster convergence | Helped |
| Transfer base learning rate 2^{-10} and scale PixArt-α to 0.61B | PixArt-α · SAM/SA-1B, epoch 30, MJHQ-30K | GenEval: 0.15 -> 0.26 | Helped |
| Transfer base learning rate 2^{-10} and scale PixArt-α to 0.61B | PixArt-α · SAM/SA-1B, epoch 30, MJHQ-30K | FID-30K: 42.71 -> 29.96 | Helped |
| Transfer base learning rate 2^{-10} and scale PixArt-α to 0.61B | PixArt-α · SAM/SA-1B, epoch 30, MJHQ-30K | CLIP Score: 26.25 -> 27.13 | Helped |
| Transfer base learning rate 2^{-10} and scale PixArt-α to 0.61B | PixArt-α · SAM/SA-1B, epoch 30, MS-COCO-30K | FID-30K: 37.61 -> 25.84 | Helped |
| Transfer base learning rate 2^{-10} and scale PixArt-α to 0.61B | PixArt-α · SAM/SA-1B, epoch 30, MS-COCO-30K | CLIP Score: 28.91 -> 29.58 | Helped |
| Transfer base hyperparameters from a 0.18B proxy to an 18B model | MMDiT · MMDiT text-to-image, GenEval | Overall: 0.8154 -> 0.8218 | Helped |
| Transfer base hyperparameters from a 0.18B proxy to an 18B model | MMDiT · MMDiT text-to-image, GenEval | Counting: 81.88 -> 79.69 | Hurt |
| Transfer base hyperparameters from a 0.18B proxy to an 18B model | MMDiT · MMDiT text-to-image, GenEval | Color Attribution: 68.75 -> 68.5 | Hurt |
| Search the PixArt-α-μP proxy learning rate; choose 2^{-10} (ablation) | PixArt-α-μP proxy · SAM proxy task, 0.04B | GenEval: 0.078 -> 0.083 | Helped |
| Search the PixArt-α-μP proxy learning rate; choose 2^{-10} (ablation) | PixArt-α-μP proxy · SAM proxy task, 0.04B | FID-30K (MS-COCO): 49.24 -> 47.46 | Helped |
| Search the PixArt-α-μP proxy learning rate; choose 2^{-10} (ablation) | PixArt-α-μP proxy · SAM proxy task, 0.04B | FID-30K (MJHQ): 46.31 -> 47.71 | Hurt |
Surprising findings
- Optimal base learning rate for MMDiT-muP proxies (~2.5e-4) is not near the maximum stable value, unlike DiT and PixArt where muP favored near-maximal LR.
- PixArt-alpha baseline degrades after 20 epochs (overfitting), whereas the muP variant keeps improving.
Limitations
- The optimal proxy task size balancing tuning cost and target performance was not identified due to limited compute.
- Some experiments are on internal proprietary data (820M image-text pairs) and an internal human-evaluated benchmark, limiting reproducibility.
- Only base learning rate was thoroughly searched for DiT and PixArt; MMDiT search covered four base HPs with 80 trials.
Experimental setup
Datasets: ImageNet SA-1B (SAM) SAM-LLaVA captions internal 820M image-text pairs MJHQ-30K MS-COCO-30K
Benchmarks: ImageNet 256x256 FID-50K MJHQ-30K MS-COCO-30K GenEval internal 150-prompt alignment benchmark
Metrics: FID sFID Inception Score precision recall CLIP Score GenEval overall training loss human alignment score
Baselines: DiT-XL-2 PixArt-alpha MMDiT-18B (expert-tuned) AuraFlow v0.1
Builds on: muP (Maximal Update Parametrization) Tensor Programs DiT PixArt-alpha MMDiT U-ViT
Closest papers by meaning
- Towards Precise Scaling Laws for Video Diffusion Transformers (CVPR 2025 · similarity 0.92)
- SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer (ICML 2025 · similarity 0.90)
- PTQ4DiT: Post-training Quantization for Diffusion Transformers (NeurIPS 2024 · similarity 0.90)
- EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing (ICLR 2025 · similarity 0.90)
- Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget (CVPR 2025 · similarity 0.90)
- Plug-and-Play Diffusion Distillation (CVPR 2024 · similarity 0.89)
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.