Research map / Image generation / Fast diffusion sampling
One More Step: A Versatile Plug-and-Play Module for Rectifying Diffusion Schedule Flaws and Enhancing Low-Frequency Controls
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 One More Step (OMS), a plug-and-play auxiliary UNet trained to map pure Gaussian noise to the terminal noise distribution expected by pre-trained diffusion models, fixing brightness bias without changing the base model's weights.
Problem
Pre-trained variance-preserving diffusion models with non-zero terminal SNR generate images biased toward mean brightness because training terminal latents leak low-frequency data information while inference starts from pure Gaussian noise.
Approach
Insert an extra inference step before the standard sampler: a separately trained text-conditioned v-prediction network maps pure Gaussian noise to the data-adulterated terminal latent the pre-trained model expects. The base model's parameters and schedule stay frozen, and the module can be shared across models in the same latent space; its prompt and CFG also give low-frequency control over bright
Key results
- On SD1.5, OMS raised Precision from 0.60 to 0.64 and reduced Wasserstein distance of mean pixel values from 22.47 to 7.84.
- On SDXL, OMS raised ImageReward from 0.8246 to 0.8730 and PickScore from 22.51 to 22.52, with Wasserstein distance falling from 11.08 to 7.25.
- Even a 3.7M-parameter OMS module still improved the brightness distribution, indicating low sensitivity to module size.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | FID: 12.52 -> 14.74; 14.10 -> 15.72; 13.14 -> 13.29 | Hurt |
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | CLIP Score: 0.2641 -> 0.2645; 0.2624 -> 0.2628; 0.2669 -> 0.2679 | Helped |
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | ImageReward: 0.1991 -> 0.2289; 0.4501 -> 0.4565; 0.8246 -> 0.8730 | Helped |
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | Precision: 0.60 -> 0.64; 0.58 -> 0.61; 0.64 -> 0.65 | Helped |
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | Recall: 0.55 -> 0.46; 0.55 -> 0.48; 0.52 -> 0.49 | Hurt |
| Add One More Step (OMS) | SD1.5, SD2.1, and SDXL · MS COCO zero-shot | Wasserstein: 22.47 -> 7.84; 21.63 -> 7.70; 11.08 -> 7.25 | Helped |
| Vary OMS module parameter count (three network sizes) (ablation) | One More Step · Image generation | image quality | No clear effect |
| Change the OMS text encoder (ablation) | One More Step · Image generation | image quality | No clear effect |
| Use a modified prompt in the OMS stage | One More Step · Image generation | brightness control | Helped |
| Use a modified prompt in the OMS stage | One More Step · Image generation | object color control | Helped |
| Apply classifier-free guidance in the OMS module | One More Step · Image generation | low-frequency image characteristics | Helped |
Surprising findings
- FID worsened after OMS (e.g., SD1.5 12.52 to 14.74), which the authors attribute to FID's weak correlation with aesthetics.
- The OMS module needs very few parameters; the authors suggest the low-SNR early denoising steps mainly set low-frequency traits.
Limitations
- FID got worse on all three models after OMS, so the gains rely on other metrics.
- A separate OMS module was needed for SDXL because its VAE latent domain differs from SD1.5, SD2.1 and LCM.
Experimental setup
Datasets: LAION 2B MS COCO
Benchmarks: MS COCO 10k zero-shot captions
Metrics: FID CLIP Score ImageReward PickScore Precision Recall Density Coverage Wasserstein distance of mean pixel values
Baselines: SD1.5 SD2.1 SDXL LCM
Builds on: DDPM DDIM Offset noise Zero terminal SNR Classifier-free guidance
Closest papers by meaning
- Ambient Diffusion Omni: Training Good Models with Bad Data (NeurIPS 2025 · similarity 0.89)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.89)
- Diffusion Model is Effectively Its Own Teacher (CVPR 2025 · similarity 0.89)
- Multi-Step Denoising Scheduled Sampling: Towards Alleviating Exposure Bias for Diffusion Models (AAAI 2024 · similarity 0.89)
- Your Diffusion Model is Secretly a Noise Classifier and Benefits from Contrastive Training (NeurIPS 2024 · similarity 0.89)
- Simple Distillation for One-Step Diffusion Models (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.