Research map / Reinforcement learning / RL from feedback
Direct Distributional Optimization for Provable Alignment of Diffusion Models
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 aligning diffusion models by directly optimizing a regularized distributional loss with Dual Averaging and sampling the aligned distribution via Doob's h-transform, with convergence guarantees free of isoperimetric conditions.
Problem
Aligning pretrained diffusion models (RLHF, DPO, KTO) requires distribution optimization of a regularized loss over densities that are inaccessible and highly multimodal, where Langevin-based mean-field methods converge slowly due to LSI dependence.
Approach
Combine the Dual Averaging method for distribution optimization (iteratively fitting potentials f_k to functional derivatives via samples from the reference model) with Doob's h-transform to construct a score for the aligned distribution, yielding convergence bounds without isoperimetric (LSI) assumptions.
Key results
- On the 2D GMM, the proposed method reaches true DPO loss 0.328 vs 0.340 for Diffusion-DPO (50 iter) and 0.343 (200 iter).
- Metric loss on GMM is 2.022 for the proposed method vs 3.011 and 3.303 for Diffusion-DPO variants.
- Convergence guarantees are given for convex (O(1/K)) and nonconvex (dual variable gradient converges) objectives without LSI assumptions.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use direct DPO optimization with Dual Averaging instead of Diffusion-DPO | score-based diffusion model · Gaussian Mixture Model (GMM) | True DPO loss: 0.340 -> 0.328 | Helped |
| use direct DPO optimization with Dual Averaging instead of Diffusion-DPO | score-based diffusion model · Gaussian Mixture Model (GMM) | Metric loss: 3.011 -> 2.022 | Helped |
Surprising findings
- Diffusion-DPO optimizing an upper bound failed to fully control the true DPO loss, with the 200-iteration regularized run ending higher than the 50-iteration run.
Limitations
- Doob h-transform sampling requires Monte Carlo conditional expectations with O(L^2) time complexity and large memory.
- Experiments are small-scale (2D GMM and low-resolution images); no large-scale text-to-image alignment.
- Convergence analysis relies on bounded and TV-Lipschitz functional derivatives and approximation error assumptions.
Experimental setup
Datasets: Gaussian Mixture Model (synthetic) Diffusion Models Course butterfly images Medical MNIST Head CT
Benchmarks: Gaussian Mixture Model DPO alignment butterfly color alignment Head CT tilt correction
Metrics: true DPO loss upper bound loss metric loss (mean distance to target) optimization time GPU memory
Baselines: Diffusion-DPO
Builds on: Dual Averaging particle dual averaging mean-field Langevin dynamics Doob's h-transform Diffusion-DPO
Closest papers by meaning
- Inference-Time Alignment of Diffusion Models with Direct Noise Optimization (ICML 2025 · similarity 0.92)
- Test-time Alignment of Diffusion Models without Reward Over-optimization (ICLR 2025 · similarity 0.90)
- D-Fusion: Direct Preference Optimization for Aligning Diffusion Models with Visually Consistent Samples (ICML 2025 · similarity 0.89)
- Forward KL Regularized Preference Optimization for Aligning Diffusion Policies (AAAI 2025 · similarity 0.88)
- Composition and Alignment of Diffusion Models using Constrained Learning (NeurIPS 2025 · similarity 0.88)
- Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards (CVPR 2025 · 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.