Adding Conditional Control to Diffusion Models with Reinforcement Learning

TL;DR

This paper introduces CTRL, a reinforcement learning-based method to add controllability to pre-trained diffusion models, improving efficiency and simplifying dataset requirements for conditional sample generation.

Contribution

The paper proposes a novel RL-based approach to incorporate controls into diffusion models, eliminating the need for classifier training and enhancing sample efficiency.

Findings

CTRL enables sampling from conditional distributions during inference.

The RL approach improves sample efficiency over classifier-free guidance.

It simplifies dataset construction by leveraging conditional independence.

Abstract

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 such controls using an offline dataset comprising inputs and 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. Our method, $\textbf{CTRL}$ ($\textbf{C}$onditioning pre-$\textbf{T}$rained diffusion models with $\textbf{R}$einforcement $\textbf{L}$earning), produces soft-optimal policies that maximize the abovementioned reward functions. We formally demonstrate that our method enables sampling from the conditional distribution with additional controls during inference. Our RL-based approach offers several advantages over existing methods. Compared to classifier-free guidance, it improves sample efficiency and can greatly simplify dataset construction by leveraging conditional independence between the inputs and additional controls. Additionally, unlike classifier guidance, it eliminates the need to train classifiers from intermediate states to additional controls. The code is available at https://github.com/zhaoyl18/CTRL.