# Direct Distributional Optimization for Provable Alignment of Diffusion Models

arXiv: https://arxiv.org/abs/2502.02954 · HTML: https://mlsanity.com/paper/2502.02954/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 33 of its 52 references that are in the map.
- First posted on arXiv 5 Feb 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- LLM alignment / RLHF / preference optimisation / [Diffusion and multimodal model alignment](https://mlsanity.com/alignment/diffusion-and-multimodal-model-alignment/)

## Builds on (most cited first, 30 of 33)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Directly Fine-Tuning Diffusion Models on Differentiable Rewards](https://mlsanity.com/paper/2309.17400/) (2023)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions](https://mlsanity.com/paper/2211.01916/) (2022)
- [Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control](https://mlsanity.com/paper/2402.15194/) (2024)
- [Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding](https://mlsanity.com/paper/2408.08252/) (2024)
- [Aligning Diffusion Models by Optimizing Human Utility](https://mlsanity.com/paper/2404.04465/) (2024)
- [Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review](https://mlsanity.com/paper/2407.13734/) (2024)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks](https://mlsanity.com/paper/1905.07769/) (2019)
- [Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond](https://mlsanity.com/paper/2403.06279/) (2024)
- [Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images](https://mlsanity.com/paper/2308.02062/) (2023)
- [Diffusion Model-based Probabilistic Downscaling for 180-year East Asian Climate Reconstruction](https://mlsanity.com/paper/2402.06646/) (2024)
- [Implicit Diffusion: Efficient Optimization through Stochastic Sampling](https://mlsanity.com/paper/2402.05468/) (2024)
- [Particle Dual Averaging: Optimization of Mean Field Neural Networks with Global Convergence Rate Analysis](https://mlsanity.com/paper/2012.15477/) (2020)
- [Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction](https://mlsanity.com/paper/2306.07221/) (2023)
- [Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning](https://mlsanity.com/paper/2409.08400/) (2024)
- [Entropic Fictitious Play for Mean Field Optimization Problem](https://mlsanity.com/paper/2202.05841/) (2022)
- [Improved Particle Approximation Error for Mean Field Neural Networks](https://mlsanity.com/paper/2405.15767/) (2024)
- [Polyak-L ojasiewicz inequality on the space of measures and convergence of mean-field birth-death processes](https://mlsanity.com/paper/2206.02774/) (2022)

## Built on by (most cited first, 5 of 5)

- [Diffusion Models with Double Guidance: Generate with aggregated datasets](https://mlsanity.com/paper/2505.13213/) (2025)
- [Reg-DPO: SFT-Regularized Direct Preference Optimization with GT-Pair for Improving Video Generation](https://mlsanity.com/paper/2511.01450/) (2025)
- [Regret Analysis of Guided Diffusion for Black-Box Optimization over Structured Inputs](https://mlsanity.com/paper/2605.10385/) (2026)
- [Steepest Guidance: A Practical and Principled Approach to Inference-Time Alignment of Flow and Diffusion-based Models](https://mlsanity.com/paper/2609.39091/) (2026)
- [SaTPO: Text-to-Image Safety Alignment via Transition Preference Optimization](https://mlsanity.com/paper/2610.19563/) (2026)

## Abstract (arXiv)

We introduce a novel alignment method for diffusion models from distribution optimization perspectives while providing rigorous convergence guarantees. We first formulate the problem as a generic regularized loss minimization over probability distributions and directly optimize the distribution using the Dual Averaging method. Next, we enable sampling from the learned distribution by approximating its score function via Doob's h-transform technique. The proposed framework is supported by rigorous convergence guarantees and an end-to-end bound on the sampling error, which imply that when the original distribution's score is known accurately, the complexity of sampling from shifted distributions is independent of isoperimetric conditions. This framework is broadly applicable to general distribution optimization problems, including alignment tasks in Reinforcement Learning with Human Feedback (RLHF), Direct Preference Optimization (DPO), and Kahneman-Tversky Optimization (KTO). We empirically validate its performance on synthetic and image datasets using the DPO objective.
