Research map / Image generation / Fast diffusion sampling
Improving Training Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architecture
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 a multi-stage diffusion training framework that splits the timestep range into intervals via an optimal-denoiser-based clustering and uses a shared encoder with per-stage decoders in the U-Net, improving training and sampling efficiency for DPM-Solver, EDM, and LDM.
Problem
Diffusion models train and sample slowly because a single large network must handle all noise levels with dissimilar gradients and varying capacity needs, while fully separate per-interval networks overfit.
Approach
Partition timesteps into stages using a functional-similarity clustering derived from the optimal denoiser, and use one shared U-Net encoder with stage-specific decoders of tailored widths, combining the benefits of unified and separate architectures.
Key results
- On CIFAR-10, multistage DPM-Solver reaches FID 2.71 vs 2.73 for DPM-Solver with about 56% of training iterations.
- Multistage EDM improves CelebA 32x32 training PFLOPs by 82% reduction (8.28 vs 10.06) while FID goes 1.55 to 1.44.
- Multistage LDM on CelebA 256x256 cuts training PFLOPs to 30% of baseline (12.95 vs 43.31) with FID 8.38 vs 8.29.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply the multistage framework | DPM-Solver · CIFAR-10 | FID: 2.73 -> 2.35 | Helped |
| Apply the multistage framework | EDM · CIFAR-10 | FID: 2.05 -> 1.96 | Helped |
| Apply the multistage framework | DPM-Solver · CIFAR-10 | Total training PFLOPs: 7.94 -> 4.66 (59%) | Helped |
| Apply the multistage framework | EDM · CelebA 32 × 32 | Total training PFLOPs: 10.06 -> 8.28 (82%) | Helped |
| Apply the multistage framework | LDM · CelebA 256 × 256 | Total training PFLOPs: 43.31 -> 12.95 (30.0%) | Helped |
| Use separate U-Nets for each interval instead of a unified U-Net (ablation) | CIFAR-10 U-Net · CIFAR-10 | FID: 2.73 -> 2.87 | Hurt |
| Add early stopping to separate U-Nets (ablation) | CIFAR-10 U-Net · CIFAR-10 | FID: 2.87 -> 2.80 | Helped |
| Use tailored parameters in separate U-Nets (ablation) | CIFAR-10 U-Net · CIFAR-10 | FID: 2.87 -> 2.68 | Helped |
| Add early stopping and tailored parameters to separate U-Nets (ablation) | CIFAR-10 U-Net · CIFAR-10 | FID: 2.87 -> 2.52 | Helped |
| Use a shared encoder and interval-specific decoders instead of a unified U-Net (ablation) | CIFAR-10 U-Net · CIFAR-10 | FID: 2.73 -> 2.35 | Helped |
| Use timestep-based clustering instead of optimal denoiser-based clustering (ablation) | Multistage DPM-Solver · CIFAR-10 | FID: 2.35 -> 3.12 | Hurt |
| Use SNR-based clustering instead of optimal denoiser-based clustering (ablation) | Multistage DPM-Solver · CIFAR-10 | FID: 2.35 -> 2.72 | Hurt |
Surprising findings
- Separate per-interval architectures overfit and perform worse than a unified network with equal parameters, motivating the shared encoder.
Limitations
- Evaluated only on unconditional image generation; conditional and inverse-problem settings left to future work.
- Three-stage configuration chosen empirically; general multi-stage ablation deferred to appendix.
Experimental setup
Datasets: CIFAR-10 CelebA 32x32 CelebA 256x256
Benchmarks: CIFAR-10 CelebA
Metrics: FID NFE GFLOPs total training PFLOPs training iterations
Baselines: DDPM Score SDE PFGM DDIM gDDIM DEIS DPM-Solver EDM LDM
Builds on: EDM DPM-Solver Latent Diffusion Models Mask R-CNN multi-head design expert denoisers
Closest papers by meaning
- Adaptive Non-Uniform Timestep Sampling for Accelerating Diffusion Model Training (CVPR 2025 · similarity 0.91)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.90)
- Multi-Architecture Multi-Expert Diffusion Models (AAAI 2024 · similarity 0.90)
- Early-Bird Diffusion: Investigating and Leveraging Timestep-Aware Early-Bird Tickets in Diffusion Models for Efficient Training (CVPR 2025 · similarity 0.90)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.89)
- Decouple-Then-Merge: Finetune Diffusion Models as Multi-Task Learning (CVPR 2025 · 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.