mlsanityData to 14 Oct 2026

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

arXiv 2412.18164 ↗First posted 24 Dec 2024cs.LG math.OCcited by 19

On the map: Diffusion models → Diffusion and flow policies for robotics

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.

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.

Builds on

54 of 78 references.

Closest ancestors · 8

Same direction, newest first

  1. 2024Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control13 Sep
  2. 2024Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning12 Sep
  3. 2024Reward-Directed Score-Based Diffusion Models via q-Learning7 Sep
  4. 2024Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review18 Jul
  5. 2024Adding Conditional Control to Diffusion Models with Reinforcement Learning17 Jun
  6. 2024Feedback Efficient Online Fine-Tuning of Diffusion Models26 Feb
  7. 2023Diffusion Model Alignment Using Direct Preference Optimization21 Nov
  8. 2023Training Diffusion Models with Reinforcement Learning22 May

Other references · 46

  1. 2020Denoising Diffusion Probabilistic Modelscited by 21,076
  2. 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,863
  3. 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,380
  4. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,122
  5. 2021Diffusion Models Beat GANs on Image Synthesiscited by 8,255
  6. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,058
  7. 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,585
  8. 2022Classifier-Free Diffusion Guidancecited by 5,847
  9. 2023SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesiscited by 4,356
  10. 2019Generative Modeling by Estimating Gradients of the Data Distributioncited by 4,101
  11. 2022DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generationcited by 3,300
  12. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,621
34 more
  1. 2022Scaling Laws for Reward Model Overoptimizationcited by 1,095
  2. 2023ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generationcited by 1,045
  3. 2019Provably Efficient Reinforcement Learning with Linear Function Approximationcited by 520
  4. 2023Aligning Text-to-Image Models using Human Feedbackcited by 422
  5. 2016Tutorial on Variational Autoencoderscited by 409
  6. 2023Emu: Enhancing Image Generation Models Using Photogenic Needles in a Haystackcited by 279
  7. 2018Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulatorcited by 250
  8. 2020Fast Global Convergence of Natural Policy Gradient Methods with Entropy Regularizationcited by 176
  9. 2019Neural Policy Gradient Methods: Global Optimality and Rates of Convergencecited by 159
  10. 2019Global Optimality Guarantees For Policy Gradient Methodscited by 150
  11. 2022An optimal control perspective on diffusion-based generative modelingcited by 147
  12. 2022On the Convergence Rates of Policy Gradient Methodscited by 119
  13. 2020Variational Policy Gradient Method for Reinforcement Learning with General Utilitiescited by 111
  14. 2021Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergencecited by 84
  15. 2023Optimizing DDPM Sampling with Shortcut Fine-Tuningcited by 81
  16. 2022An Improved Analysis of (Variance-Reduced) Policy Gradient and Natural Policy Gradient Methodscited by 68
  17. 2018Derivative-Free Methods for Policy Optimization: Guarantees for Linear Quadratic Systemscited by 67
  18. 2023Stochastic Policy Gradient Methods: Improved Sample Complexity for Fisher-non-degenerate Policiescited by 58
  19. 2020On Linear Convergence of Policy Gradient Methods for Finite MDPscited by 52
  20. 2020Policy Gradient Methods for the Noisy Linear Quadratic Regulator over a Finite Horizoncited by 50
  21. 2024Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalizationcited by 42
  22. 2024Fine-tuning of diffusion models via stochastic control: entropy regularization and beyondcited by 42
  23. 2019LQR through the Lens of First Order Methods: Discrete-time Casecited by 40
  24. 2020Single-Timescale Actor-Critic Provably Finds Globally Optimal Policycited by 34
  25. 2019Understanding the Limitations of Conditional Generative Modelscited by 28
  26. 2023Improved Sample Complexity Analysis of Natural Policy Gradient Algorithm with General Parameterization for Infinite Horizon Discounted Reward Markov Decision Processescited by 25
  27. 2024Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Modelscited by 25
  28. 2021Doubly Robust Off-Policy Actor-Critic: Convergence and Optimalitycited by 21
  29. 2021Exploration-exploitation trade-off for continuous-time episodic reinforcement learning with linear-convex modelscited by 21
  30. 2022Single Time-scale Actor-critic Method to Solve the Linear Quadratic Regulator with Convergence Guaranteescited by 14
  31. 2023A Policy Gradient Framework for Stochastic Optimal Control Problems with Global Convergence Guaranteecited by 10
  32. 2024Solving Time-Continuous Stochastic Optimal Control Problems: Algorithm Design and Convergence Analysis of Actor-Critic Flowcited by 8
  33. 2024On Bellman equations for continuous-time policy evaluation I: discretization and approximationcited by 7
  34. 2023Policy Gradient Converges to the Globally Optimal Policy for Nearly Linear-Quadratic Regulatorscited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 7; 2026: 12.

19 papers.

Sort:

By area:

  1. 2025TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusioncited by 14
  2. 2026Conditional Diffusion Guidance under Hard Constraint: A Stochastic Analysis Approachcited by 11
  3. 2025Composition and Alignment of Diffusion Models using Constrained Learningcited by 10
  4. 2025Understanding Sampler Stochasticity in Training Diffusion Models for RLHFcited by 7
  5. 2026Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Controlcited by 6
  6. 2026Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matchingcited by 5
  7. 2026Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipelinecited by 5
  8. 2025Continuous Q-Score Matching: Diffusion Guided Reinforcement Learning for Continuous-Time Controlcited by 4
  9. 2025Is RL fine-tuning harder than regression? A PDE learning approach for diffusion modelscited by 3
  10. 2025Schr"odinger bridge for generative AI: Soft-constrained formulation and convergence analysiscited by 3
  11. 2026Supervised Guidance Training for Infinite-Dimensional Diffusion Modelscited by 3
  12. 2026Diffusion Controller: Framework, Algorithms and Parameterizationcited by 2
7 more
  1. 2026Inference-Time Attribute Distribution Alignment for Unconditional Diffusioncited by 2
  2. 2025Coffee: Controllable Diffusion Fine-tuningcited by 1
  3. 2026FluxLite: Inference-Time Proposal Control for Discrete Diffusion Modelscited by 1
  4. 2026Continuous-time reinforcement learning: ellipticity enables model-free value function approximationcited by 0
  5. 2026Sobolev Regularized Score Difference Estimation in Diffusion Modelscited by 0
  6. 2026Error-Corrected Inference-Time Scaling for Imperfect Diffusion Modelscited by 0
  7. 2026Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Costcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Fine-tuning Diffusion Policies with Backpropagation Through Diffusion Timesteps
  2. 2024Avoiding mode collapse in diffusion models fine-tuned with reinforcement learning
  3. 2025Inference-Time Alignment Control for Diffusion Models with Reinforcement Learning Guidance
  4. 2026A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models
  5. 2023Towards Controllable Diffusion Models via Reward-Guided Exploration
  6. 2025Score as Action: Fine-Tuning Diffusion Generative Models by Continuous-time Reinforcement Learning
  7. 2026Diffusion Policy through Conditional Proximal Policy Optimization
  8. 2024Diffusion Tuning: Transferring Diffusion Models via Chain of Forgetting