# Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence

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

## Key facts

- Cited by 19 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 54 of its 78 references that are in the map.
- First posted on arXiv 24 Dec 2024.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation](https://mlsanity.com/paper/2208.12242/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Aligning Text-to-Image Models using Human Feedback](https://mlsanity.com/paper/2302.12192/) (2023)
- [Tutorial on Variational Autoencoders](https://mlsanity.com/paper/1606.05908/) (2016)
- [Emu: Enhancing Image Generation Models Using Photogenic Needles in a Haystack](https://mlsanity.com/paper/2309.15807/) (2023)
- [Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulator](https://mlsanity.com/paper/1801.05039/) (2018)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Fast Global Convergence of Natural Policy Gradient Methods with Entropy Regularization](https://mlsanity.com/paper/2007.06558/) (2020)
- [Neural Policy Gradient Methods: Global Optimality and Rates of Convergence](https://mlsanity.com/paper/1909.01150/) (2019)
- [Global Optimality Guarantees For Policy Gradient Methods](https://mlsanity.com/paper/1906.01786/) (2019)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [On the Convergence Rates of Policy Gradient Methods](https://mlsanity.com/paper/2201.07443/) (2022)
- [Variational Policy Gradient Method for Reinforcement Learning with General Utilities](https://mlsanity.com/paper/2007.02151/) (2020)
- [Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review](https://mlsanity.com/paper/2407.13734/) (2024)
- [Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergence](https://mlsanity.com/paper/2105.11066/) (2021)

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

- [TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://mlsanity.com/paper/2509.25171/) (2025)
- [Conditional Diffusion Guidance under Hard Constraint: A Stochastic Analysis Approach](https://mlsanity.com/paper/2602.05533/) (2026)
- [Composition and Alignment of Diffusion Models using Constrained Learning](https://mlsanity.com/paper/2508.19104/) (2025)
- [Understanding Sampler Stochasticity in Training Diffusion Models for RLHF](https://mlsanity.com/paper/2510.10767/) (2025)
- [Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Control](https://mlsanity.com/paper/2604.08580/) (2026)
- [Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching](https://mlsanity.com/paper/2601.06514/) (2026)
- [Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline](https://mlsanity.com/paper/2605.06583/) (2026)
- [Continuous Q-Score Matching: Diffusion Guided Reinforcement Learning for Continuous-Time Control](https://mlsanity.com/paper/2510.17122/) (2025)
- [Is RL fine-tuning harder than regression? A PDE learning approach for diffusion models](https://mlsanity.com/paper/2509.02528/) (2025)
- [Schr"odinger bridge for generative AI: Soft-constrained formulation and convergence analysis](https://mlsanity.com/paper/2510.11829/) (2025)
- [Supervised Guidance Training for Infinite-Dimensional Diffusion Models](https://mlsanity.com/paper/2601.20756/) (2026)
- [Diffusion Controller: Framework, Algorithms and Parameterization](https://mlsanity.com/paper/2603.06981/) (2026)
- [Inference-Time Attribute Distribution Alignment for Unconditional Diffusion](https://mlsanity.com/paper/2605.07456/) (2026)
- [Coffee: Controllable Diffusion Fine-tuning](https://mlsanity.com/paper/2511.14113/) (2025)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [Continuous-time reinforcement learning: ellipticity enables model-free value function approximation](https://mlsanity.com/paper/2602.06930/) (2026)
- [Sobolev Regularized Score Difference Estimation in Diffusion Models](https://mlsanity.com/paper/2608.18237/) (2026)
- [Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models](https://mlsanity.com/paper/2610.01933/) (2026)
- [Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Cost](https://mlsanity.com/paper/2610.20035/) (2026)

## Abstract (arXiv)

Diffusion models have emerged as powerful tools for generative modeling, demonstrating exceptional capability in capturing target data distributions from large datasets. However, fine-tuning these massive models for specific downstream tasks, constraints, and human preferences remains a critical challenge. While recent advances have leveraged reinforcement learning algorithms to tackle this problem, much of the progress has been empirical, with limited theoretical understanding. To bridge this gap, we propose a stochastic control framework for fine-tuning diffusion models. Building on denoising diffusion probabilistic models as the pre-trained reference dynamics, our approach integrates linear dynamics control with Kullback-Leibler regularization. We establish the well-posedness and regularity of the stochastic control problem and develop a policy iteration algorithm (PI-FT) for numerical solution. We show that PI-FT achieves global convergence at a linear rate. Unlike existing work that assumes regularities throughout training, we prove that the control and value sequences generated by the algorithm maintain the regularity. Additionally, we explore extensions of our framework to parametric settings and continuous-time formulations, and demonstrate the practical effectiveness of the proposed PI-FT algorithm through numerical experiments. Our code is available at https://github.com/yinbinhan/fine-tuning-of-diffusion-models.