Adding Conditional Control to Diffusion Models with Reinforcement Learning
Yulai Zhao, Masatoshi Uehara, Gabriele Scalia, Sunyuan Kung, Tommaso Biancalani, Sergey Levine, Ehsan Hajiramezanali
TL;DR
This work tackles the problem of adding new conditional controls to pre-trained diffusion models using offline data. It reframes conditioning as a reinforcement-learning task and introduces CTRL, an augmented-diffusion approach that learns a soft-optimal drift to sample from the target distribution p(x|c,y) while penalizing divergence from the pre-trained model. CTRL achieves greater sample efficiency by leveraging conditional independence to simplify offline-data requirements and by training a calibrated reward model p(y|x,c) rather than directly modeling p(x|c,y). The authors provide theoretical connections to classifier guidance, detail an implementable three-step algorithm, analyze potential error sources, and validate performance on both single-task compressibility conditioning and multi-task conditioning, demonstrating notable improvements over baselines.
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.