Research map / Image generation / Fast diffusion sampling
Diffusion Model Patching via Mixture-of-Prompts
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 Diffusion Model Patching (DMP), which freezes a converged diffusion model and trains a small pool of timestep-gated prompts added to the input of each block on the original pre-training data, improving FID.
Problem
Pre-trained diffusion models that have already converged are assumed to gain little from further in-domain training; the paper targets improving them with minimal extra parameters and no external data.
Approach
Insert learnable prompts directly added (not prepended) to each DiT block's input and use a timestep- and depth-conditioned softmax gating network to form a mixture-of-prompts at every denoising step, with load and importance balancing losses, while keeping the backbone frozen.
Key results
- DMP improves converged DiT-L/2 FID on FFHQ 256x256 from 6.26 to 5.61 (-0.65, about 10.38%) after 50K further iterations with 1.43% extra parameters, while fine-tuning and naive prompt tuning degrade it.
- On DiT-B/2 FFHQ, DMP reaches FID 5.87 vs 6.27 baseline, while full fine-tuning gives 6.57 and prompt tuning 6.81.
- On ImageNet DiT-XL/2, DMP reduces FID from 2.29 to 2.25 with 20K iterations, while fine-tuning worsens it to 5.04.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add dynamically gated learnable prompts | DiT-B/2 · FFHQ 256×256 | FID: 6.27 -> 5.87 | Helped |
| add dynamically gated learnable prompts | DiT-B/2 · MS-COCO 256×256 | FID: 7.33 -> 7.12 | Helped |
| add dynamically gated learnable prompts | Stable Diffusion v1.5 · Laion5B 512×512 | FID: 47.18 -> 35.44 | Helped |
| add dynamically gated learnable prompts | DiT-L/2 · FFHQ 256×256 | FID: 6.26 -> 5.61 | Helped |
| add dynamically gated learnable prompts | DiT-XL/2 · ImageNet 256×256 | FID: 2.29 -> 2.25 | Helped |
| use timestep-only gating instead of timestep-plus-class-condition gating (ablation) | DMP on DiT-XL/2 · ImageNet 256×256 | FID: 2.25 -> 2.25 | No clear effect |
| include the text condition embedding in gating (ablation) | DMP on DiT-B/2 · MS-COCO | FID: 7.30 -> 7.12 | Helped |
| add prompts across more DiT blocks (ablation) | DMP on DiT-B/2 · FFHQ | FID | Helped |
| replace linear gating with attention gating (ablation) | DMP on DiT-B/2 · FFHQ | FID: 5.87 vs 6.41 | Hurt |
| use hard top-192 prompt selection instead of soft weighted prompt selection (ablation) | DMP · FFHQ | FID | Hurt |
| use the same prompt selection at every block instead of block-specific selection (ablation) | DMP · FFHQ | FID | Hurt |
| prepend prompts instead of adding them element-wise to image tokens (ablation) | DMP on DiT-B/2 · FFHQ | FID: 5.87 vs 6.79 | Hurt |
Surprising findings
- Naive fine-tuning and prompt tuning on the original pre-training data degrade converged diffusion models, whereas DMP improves them.
- LoRA on the pre-training data also worsened FID (7.11 vs 6.27 baseline on DiT-B/2 FFHQ).
Limitations
- Fixed number of prompts; prompt-adding strategy limits scalability to other configurations.
- Gains are reported on in-domain further training only; cross-dataset and personalization settings are not evaluated.
- Structural bias observed in prompt activations (temporal and spatial), with effects not fully characterized.
Experimental setup
Datasets: FFHQ MS-COCO Laion5B ImageNet
Benchmarks: FFHQ MS-COCO Laion5B ImageNet
Metrics: FID precision recall
Baselines: full fine-tuning naive prompt tuning LoRA
Builds on: prompt tuning DiT Mixture-of-Experts load balancing
Closest papers by meaning
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.89)
- DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents (ICML 2024 · similarity 0.89)
- Decouple-Then-Merge: Finetune Diffusion Models as Multi-Task Learning (CVPR 2025 · similarity 0.89)
- Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising (NeurIPS 2024 · similarity 0.89)
- Learning Diffusion Models with Flexible Representation Guidance (NeurIPS 2025 · similarity 0.89)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 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.