Research map / Image generation / Fast diffusion sampling
Decouple-Then-Merge: Finetune Diffusion Models as Multi-Task Learning
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.
Decouples diffusion model finetuning by training separate models on non-overlapping timestep ranges (with consistency loss, probabilistic sampling, channel-wise projection) and merges them in parameter space into one model.
Problem
Parameter-shared diffusion denoisers suffer gradient conflicts across timesteps, which can degrade image generation; existing fixes (timestep-wise ensembles, loss reweighting) add inference cost or only partially mitigate conflicts.
Approach
Finetune one pretrained diffusion model separately on N disjoint timestep ranges with consistency loss, probabilistic sampling and channel-wise projection to retain shared knowledge, then merge the N task vectors into a single model via weighted addition so inference cost is unchanged.
Key results
- On LSUN-Church with DDPM, FID drops from 10.69 to 7.27 after merging (−3.42), versus at most about −0.1 for prior loss-reweighting baselines.
- On Stable Diffusion MS-COCO, merged model reduces FID from 13.42 to 13.06 and raises CLIP score from 29.88 to 30.11.
- On CIFAR10 with N=8 and all techniques, FID improves from 4.40 to 3.87 in the ablation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| decouple training into timestep ranges and merge the finetuned models | DDPM · CIFAR10 | FID: 4.42 -> 3.51 | Helped |
| decouple training into timestep ranges and merge the finetuned models | DDPM · LSUN-Church | FID: 10.69 -> 7.27 | Helped |
| decouple training into timestep ranges and merge the finetuned models | Stable Diffusion · MS-COCO | FID: 13.42 -> 13.06 | Helped |
| merge the finetuned models rather than use the ensemble | DeMe · LSUN-Church | FID: 9.57 -> 7.27 | Helped |
| merge the finetuned models rather than use the ensemble | DeMe · MS-COCO | CLIP Score: 29.85 -> 30.11 | Helped |
| divide overall timesteps into 8 non-overlapping ranges (ablation) | DDPM · CIFAR10 | FID: 4.40 -> 4.32 | Helped |
| add probabilistic sampling (ablation) | DDPM with N=8 · CIFAR10 | FID: 0.08 FID reduction | Helped |
| add consistency loss (ablation) | DDPM with N=8 · CIFAR10 | FID: 0.05 FID reduction | Helped |
| apply channel-wise projection with N=1 (ablation) | DDPM · CIFAR10 | FID: 4.40 -> 4.45 | Hurt |
| apply all training techniques during finetuning (ablation) | DDPM with N=8 · CIFAR10 | FID: 0.53 FID reduction | Helped |
Surprising findings
- Merged model outperforms the timestep-wise ensemble on LSUN-Church (FID 7.27 vs 9.57) and on ImageNet FID, despite using a single model.
Limitations
- Ensemble inference requires N models, multiplying storage and memory costs; the merged model requires a grid search over merge weights.
- Gains on some metrics are small (e.g., ImageNet CLIP and PartiPrompts results for merged model are mixed).
Experimental setup
Datasets: LAION-Aesthetics V2 6.5+ CIFAR10 LSUN-Church LSUN-Bedroom
Benchmarks: CIFAR10 LSUN-Church LSUN-Bedroom MS-COCO 30K ImageNet1K PartiPrompts
Metrics: FID CLIP score
Baselines: SNR+1 Truncated SNR Min-SNR-γ P2 weighting ANT-NashMTL ANT-UW Before-finetuning
Builds on: DDPM Min-SNR weighting task arithmetic model merging ANT DDIM
Closest papers by meaning
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.90)
- Improving Training Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architecture (CVPR 2024 · similarity 0.89)
- Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising (NeurIPS 2024 · similarity 0.89)
- Diffusion Model Patching via Mixture-of-Prompts (AAAI 2025 · similarity 0.89)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.89)
- Denoising Task Routing for Diffusion Models (ICLR 2024 · 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.