无需预训练模型,用强化学习直接优化生成样本的奖励。
Reward-Directed Score-Based Diffusion Models via q-Learning
- 将扩散模型训练转化为连续时间强化学习问题,直接优化生成样本的奖励。
- 通过比率估计获取未知得分函数的噪声观测,实现无需预训练的高效学习。
- 适用于图像生成等高维任务,且可扩展至条件生成与概率流ODE场景。
我们提出一种新的强化学习(RL)框架,用于训练连续时间得分扩散模型,以生成最大化奖励函数的样本,同时保持生成分布接近未知的目标数据分布。与现有方法不同,该方法不依赖于噪声扰动数据分布的预训练得分函数,也不试图学习这些得分函数。相反,我们将问题建模为熵正则化的连续时间强化学习,并证明最优随机策略具有已知协方差矩阵的高斯分布。基于此,我们参数化高斯策略的均值,并开发了一种类似Q-learning的演员-评论家算法来求解该强化学习问题。算法设计中的关键步骤是通过比率估计器从未知得分函数中获取噪声观测。我们的框架还可用于纯得分匹配以及微调预训练模型。数值实验表明,在多个生成任务(包括高维图像生成)上,该方法性能优于两种最先进的基于微调预训练模型的强化学习方法。最后,我们讨论了该强化学习框架在扩散模型的概率流ODE实现及条件扩散模型中的扩展应用。
原文摘要 · Abstract (English)
We propose a new reinforcement learning (RL) formulation for training continuous-time score-based diffusion models for generative AI to generate samples that maximize reward functions while keeping the generated distributions close to the unknown target data distributions. Different from most existing studies, ours does not involve any pretrained model for the unknown score functions of the noise-perturbed data distributions, nor does it attempt to learn the score functions. Instead, we formulate the problem as entropy-regularized continuous-time RL and show that the optimal stochastic policy has a Gaussian distribution with a known covariance matrix. Based on this result, we parameterize the mean of Gaussian policies and develop an actor--critic type (little) q-learning algorithm to solve the RL problem. A key ingredient in our algorithm design is to obtain noisy observations from the unknown score function via a ratio estimator. Our formulation can also be adapted to solve pure score-matching and fine-tuning pretrained models. Numerically, we show the effectiveness of our approach by comparing its performance with two state-of-the-art RL methods that fine-tune pretrained models on several generative tasks including high-dimensional image generations. Finally, we discuss extensions of our RL formulation to probability flow ODE implementation of diffusion models and to conditional diffusion models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。