Diffusion models are powerful generative models that allow for precise control over the characteristics of the generated samples. While these diffusion models trained on large datasets have achieved success, there is often a need to introduce additional controls in downstream fine-tuning processes, treating these powerful models as pre-trained diffusion models. This work presents a novel method based on reinforcement learning (RL) to add additional controls, leveraging an offline dataset comprising inputs and corresponding labels. We formulate this task as an RL problem, with the classifier learned from the offline dataset and the KL divergence against pre-trained models serving as the reward functions. We introduce our method, $\textbf{CTRL}$ ($\textbf{C}$onditioning pre-$\textbf{T}$rained diffusion models with $\textbf{R}$einforcement $\textbf{L}$earning), which produces soft-optimal policies that maximize the abovementioned reward functions. We formally demonstrate that our method enables sampling from the conditional distribution conditioned on additional controls during inference. Our RL-based approach offers several advantages over existing methods. Compared to commonly used classifier-free guidance, our approach improves sample efficiency, and can greatly simplify offline dataset construction by exploiting conditional independence between the inputs and additional controls. Furthermore, unlike classifier guidance, we avoid the need to train classifiers from intermediate states to additional controls.

翻译：扩散模型是强大的生成模型，能够对生成样本的特征进行精确控制。尽管在大规模数据集上训练的扩散模型已取得成功，但在下游微调过程中，通常需要引入额外控制，将这些强大模型视为预训练扩散模型。本研究提出了一种基于强化学习（RL）的新方法，利用包含输入及对应标签的离线数据集来添加额外控制。我们将此任务构建为RL问题，将从离线数据集中学习到的分类器以及与预训练模型之间的KL散度作为奖励函数。我们提出的方法$\textbf{CTRL}$（基于$\textbf{R}$einforcement $\textbf{L}$earning对预-$\textbf{T}$rained扩散模型进行$\textbf{C}$onditioning）能够生成最大化上述奖励函数的软最优策略。我们通过形式化证明表明，该方法能够在推理阶段从以额外控制为条件的条件分布中进行采样。与现有方法相比，我们基于RL的方法具有多重优势：相较于常用的无分类器引导技术，本方法提升了采样效率，并能通过利用输入与额外控制间的条件独立性大幅简化离线数据集构建过程；此外，与分类器引导方法不同，本方法无需训练从中间状态到额外控制的分类器。