用随机控制理论优化扩散模型微调,理论保证收敛与稳定性。
Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
- 基于噪声扩散过程构建控制框架,结合KL正则化实现精准微调。
- 算法全局线性收敛,且训练中始终保持解的光滑性。
- 适合关注生成模型理论保障的研究者或工业级微调需求者。
扩散模型已成为强大的生成建模工具,能从大规模数据中捕捉目标分布。然而,针对特定下游任务、约束及人类偏好对这些大模型进行微调仍是关键挑战。尽管近期研究采用强化学习方法解决该问题,但多数进展为经验性,缺乏理论支撑。为此,本文提出一种用于微调扩散模型的随机控制框架。在去噪扩散概率模型作为预训练参考动态的基础上,方法融合线性动力学控制与Kullback-Leibler正则化。我们建立了随机控制问题的适定性与正则性,并开发了策略迭代算法(PI-FT)求解。证明了PI-FT以线性速率实现全局收敛。不同于以往假设全程保持正则性的方法,我们证明该算法生成的控制与价值序列均维持正则性。此外,还拓展了框架至参数化设置和连续时间形式,并通过数值实验验证了所提PI-FT算法的实际有效性。代码已公开于https://github.com/yinbinhan/fine-tuning-of-diffusion-models。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。