Generalized Fine-Tuning of Diffusion Models via Stochastic Control and FBSDEs

Abstract

We propose a generalized fine-tuning framework for diffusion models from the perspective of stochastic control. Beyond entropy-regularized formulations, we introduce a general running cost that induces a relative generalized path cost, encompassing both Kullback-Leibler divergence and optimal transport metrics as special cases. This leads to a fully nonlinear Hamilton-Jacobi-Bellman equation. We characterize the value function through a forward-backward stochastic differential equation system and establish existence and uniqueness under standard regularity conditions. The optimal control admits a nonlinear feedback form driven by the backward stochastic differential equation gradient component. Finally, we show that generalized fine-tuning naturally introduces an additional gradient-dependent penalty, providing a unified framework for diffusion fine-tuning under general distributional constraints.