Research map / Image generation / Fast diffusion sampling
Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
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 Direct Noise Optimization (DNO), an inference-time, tuning-free method that optimizes the injected SDE noise of a pretrained diffusion model to maximize a reward, with a probability regularizer against out-of-distribution reward hacking and hybrid gradient estimators for non-differentiable rewards.
Problem
Aligning pretrained diffusion models with continuous, possibly non-differentiable reward functions without fine-tuning, while avoiding out-of-distribution reward hacking in which optimized noise yields samples outside the pretrained model's support.
Approach
Optimize all injected noise (initial noise plus per-step SDE noise of DDIM with eta=1) directly at inference time, regularize the noise toward high-probability Gaussian regions using concentration-inequality-based log-probabilities over random permutations, and estimate gradients of non-differentiable rewards via a hybrid scheme that perturbs in latent space.
Key results
- On SD v1.5, DNO with 5 min optimization reaches aesthetic 8.587 and PickScore 25.13, versus AlignProp 8.940 aesthetic (fine-tuned, ~12h) and SD v1.5 baseline 5.367 aesthetic.
- With 15 diffusion steps and a 1-minute budget, DNO achieves aesthetic 7.496 and PickScore 24.82.
- DNO runs in under 15GB memory on one A800 GPU, while tuning-based methods typically need 4-8 advanced GPUs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase DNO generation budget from 1 minute to 5 minutes | DNO · Human-aligned rewards | Aesthetic Score: 5.754 -> 8.587 | Helped |
| Increase DNO generation budget from 1 minute to 5 minutes | DNO · Human-aligned rewards | HPS-v2 score: 0.285 -> 0.324 | Helped |
| Increase DNO generation budget from 1 minute to 5 minutes | DNO · Human-aligned rewards | PickScore: 21.25 -> 25.13 | Helped |
| Add probability regularization to noise optimization | DNO · Brightness and darkness rewards with contradictory prompts | CLIP Score and ITM score | Helped |
| Add probability regularization to noise optimization | DNO · Human-aligned reward optimization; other rewards used as test metrics | HPS score and PickScore | Helped |
| Optimize both injected SDE noise and initial noise rather than only initial ODE noise (ablation) | DNO · Toy reward optimization and aesthetic reward | Reward optimization speed | Helped |
| Use Hybrid-2 rather than ZO-SGD or Hybrid-1 gradient estimation (ablation) | DNO for non-differentiable rewards · JPEG Compressibility and Aesthetic Score | Reward score | Helped |
| Use Hybrid-2 rather than ZO-SGD or Hybrid-1 gradient estimation (ablation) | DNO for non-differentiable rewards · JPEG Compressibility and Aesthetic Score | Optimization speed | Helped |
| Use Hybrid-1 gradient estimation for Aesthetic Score instead of Hybrid-2 (ablation) | DNO for non-differentiable rewards · Aesthetic Score | Reward score | Hurt |
| Use k=1 for the probability regularizer instead of k=2 or k=10 (ablation) | Probability-regularized DNO · Darkness reward with white-animal prompts | Reward and consistency score | Hurt |
| Use k=100 for the probability regularizer instead of k=2 or k=10 (ablation) | Probability-regularized DNO · Darkness reward with white-animal prompts | Reward and consistency score | Hurt |
| Increase the diffusion-step count from T=50 to T=15 (ablation) | DNO · 1-minute generation budget | Aesthetic Score: 5.754 -> 7.496 | Helped |
Surprising findings
- Naive noise optimization often drives samples out of distribution, e.g., a 'black duck' prompt turns into a white duck under brightness optimization.
- Optimizing SDE noise (DDIM eta=1) was empirically faster than optimizing only the ODE initial noise.
Limitations
- Inference-time optimization substantially increases generation time versus direct sampling (minutes per sample).
- Probability regularization has a trade-off between convergence speed and OOD avoidance, requiring tuning of gamma.
- Hybrid-1 performs poorly on image-only rewards such as aesthetic score.
Experimental setup
Benchmarks: aesthetic score HPS v2 PickScore JPEG compressibility
Metrics: aesthetic score HPS PickScore CLIP score ITM score
Baselines: LGD SPIN DDPO AlignProp Best-of-N sampling ZO-SGD
Builds on: DDIM LGD DDPO AlignProp ZO-SGD
Closest papers by meaning
- Test-time Alignment of Diffusion Models without Reward Over-optimization (ICLR 2025 · similarity 0.92)
- Direct Distributional Optimization for Provable Alignment of Diffusion Models (ICLR 2025 · similarity 0.92)
- Training-Free Diffusion Model Alignment with Sampling Demons (ICLR 2025 · similarity 0.91)
- Align Your Steps: Optimizing Sampling Schedules in Diffusion Models (ICML 2024 · similarity 0.91)
- Composition and Alignment of Diffusion Models using Constrained Learning (NeurIPS 2025 · similarity 0.90)
- Hierarchical Variational Policies for Reward-Guided Diffusion (NeurIPS 2026 · similarity 0.90)
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.