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

2024-12-24

arXiv.org

AI Summary

This work addresses the challenge of efficient fine-tuning of diffusion models under downstream tasks, constraints, and human preferences. We propose the first theoretical framework grounded in stochastic optimal control: using the denoising diffusion process as the reference dynamics, we integrate a linear control structure with KL regularization. We establish the first rigorous stochastic control formulation of diffusion fine-tuning, proving its well-posedness, Hölder regularity of the optimal solution, and global linear convergence of the associated algorithm—achieving, for the first time, guaranteed regularity of the iterate sequence without prior regularity assumptions. Based on this, we design Policy Iteration Fine-Tuning (PI-FT), extending it to parametric control and continuous-time settings. Our framework provides the first fine-tuning paradigm for controllable generation endowed with provable optimality guarantees.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• Fine-tuning diffusion models for specific tasks and preferences
• Lack of theoretical understanding in current tuning methods
• Developing a stochastic control framework with proven convergence

Innovation

• Stochastic control framework for fine-tuning diffusion models
• Linear dynamics control with KL regularization
• Policy iteration algorithm with global convergence