Research map / Image generation / Fast diffusion sampling
Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
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.
Formulates fine-tuning of DDPM-based diffusion models as discrete-time stochastic control with linear dynamics and KL regularization, and proposes a policy iteration algorithm (PI-FT) with linear convergence and proven regularity of value and control sequences.
Problem
Fine-tuning pre-trained diffusion models to maximize a known terminal reward while staying close to the pre-trained model lacks theoretical understanding, especially convergence guarantees.
Approach
Treats the fine-tuned score as a control in a discrete-time DDPM-based linear-dynamics control problem with KL penalty equal to a weighted squared distance between control and pre-trained score; the KL induces strong concavity of the one-step Bellman problem, enabling a policy iteration algorithm with provable regularity maintained across iterations and linear convergence.
Key results
- PI-FT converges globally to the optimal control at a linear rate (1-lambda)^M for M inner iterations.
- Value function and control sequences generated by the algorithm maintain Lipschitz and gradient-Lipschitz regularity uniformly over iterations.
- Optimal control is unique under a sufficiently large beta_t, with beta_t chosen so that 1 - (sigma_t^2/beta_t) L_{1,t+1} >= lambda > 0.
Limitations
- Assumes the reward function is known and smooth (Lipschitz and gradient Lipschitz).
- Requires beta_t sufficiently large, trading off closeness to pre-trained model against convergence speed.
- Analysis is in discrete time with Gaussian dynamics; general reward learning left for future work.
Experimental setup
Metrics: convergence rate lipschitz constants
Builds on: DDPM policy iteration policy gradient methods for LQR
Closest papers by meaning
- Score as Action: Fine Tuning Diffusion Generative Models by Continuous-time Reinforcement Learning (ICML 2025 · similarity 0.89)
- Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control (ICLR 2025 · similarity 0.89)
- Score Regularized Policy Optimization through Diffusion Behavior (ICLR 2024 · similarity 0.89)
- Diffusion Tuning: Transferring Diffusion Models via Chain of Forgetting (NeurIPS 2024 · similarity 0.89)
- Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching (NeurIPS 2026 · similarity 0.89)
- Stochastic Optimal Control for Diffusion Bridges in Function Spaces (NeurIPS 2024 · similarity 0.89)
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.
