Efficient Controllable Diffusion via Optimal Classifier Guidance

TL;DR

This paper tackles controllable diffusion without resorting to heavyweight RL by reframing the problem as KL-regularized reward maximization and solving it via supervised learning. The authors introduce SLCD, a data-aggregation, DAgger-inspired method that iteratively learns a reward distribution \hat{R} to guide diffusion through a gradient $\mathbf{f}^n(\mathbf{x}_t,t) = \nabla_{\mathbf{x}_t}\ln \mathbb{E}_{r\sim \hat{R}^n(\cdot|\mathbf{x}_t,t)} e^{\eta r}$. They prove a no-regret based convergence guarantee to the KL-optimal target, with a bound $\mathrm{KL}(\hat{p}_T \| p_T) \le \epsilon_T + \frac{1}{2}T\|g\|_{\infty}^2 L^2 \gamma_N$, under standard assumptions. Empirically, SLCD improves reward while preserving near-baseline inference time across continuous image diffusion and discrete sequence diffusion, and shows favorable reward-FID trade-offs compared to SVDD variants. This yields a scalable, principled, and efficient route for controllable diffusion applicable to diverse modalities, with code available at the referenced repository.

Abstract

The controllable generation of diffusion models aims to steer the model to generate samples that optimize some given objective functions. It is desirable for a variety of applications including image generation, molecule generation, and DNA/sequence generation. Reinforcement Learning (RL) based fine-tuning of the base model is a popular approach but it can overfit the reward function while requiring significant resources. We frame controllable generation as a problem of finding a distribution that optimizes a KL-regularized objective function. We present SLCD -- Supervised Learning based Controllable Diffusion, which iteratively generates online data and trains a small classifier to guide the generation of the diffusion model. Similar to the standard classifier-guided diffusion, SLCD's key computation primitive is classification and does not involve any complex concepts from RL or control. Via a reduction to no-regret online learning analysis, we show that under KL divergence, the output from SLCD provably converges to the optimal solution of the KL-regularized objective. Further, we empirically demonstrate that SLCD can generate high quality samples with nearly the same inference time as the base model in both image generation with continuous diffusion and biological sequence generation with discrete diffusion. Our code is available at https://github.com/Owen-Oertell/slcd

Figures (9)

• Figure 1: An overview of the main experimental results. Top: Qualitative examples for continuous diffusion image tasks (image compression and aesthetic maximization). Relaxing the KL constraint at test time (larger $\eta$) consistently increases the score. Bottom left: SLCD stays closer to the initial image distribution (lower FID score) for the same reward. Bottom right: SLCD is likewise effective at controlling discrete diffusion models.
• Figure 2: Covariate shift (left) and data collection in our approach (right). The left figure illustrates covariate shift. In the offline naive approach, classifier will be trained under the green samples. However in inference time, the classifier will be applied at the red samples -- samples generated by using the classifier itself as guidance. The difference in green samples (training distribution) and red samples (testing distribution) is the covariate shift. Our approach (right) mitigates this by iteratively augmenting the training set with samples drawn from guided diffusion. We rollin with the classifier-guided diffusion to get to $\mathbf{x}_t$. We then rollout with the prior's score function to get to $x_T$ and query a reward $r(\mathbf{x}_T)$. The triple $(t,\mathbf{x}_t, r(\mathbf{x}_T))$ will be used to refine the classifier.
• Figure 3: Reward vs. number of iterations of SLCD. The reward increases as the restart‐state distribution becomes richer, and only a few iterations are needed to reach good performance.
• Figure 4: Images generated by SLCD with varying $\eta$ values and their rewards . The first column shows results from the base model, SD1.5, which corresponds to our method with $\eta = 0$. As $\eta$ increases, the KL penalty is relaxed, allowing the generated images to be more strongly optimized for the reward function, and consequently, they diverge further from the base model’s original distribution.
• Figure 5: Distribution of rewards for RNA sequences (5'UTR) across different methods. SLCD achieves higher median rewards and better overall distribution compared to baseline approaches.

Theorems & Definitions (9)

• Remark 1: Modeling the one-dimensional distribution as a classifier
• Remark 2: Comparison to SVDD uehara2024feedback
• Theorem 7
• Lemma 8
• Lemma 9
• Lemma 10
• Lemma 11
• Lemma 12: Section 5 of anderson1982reverse
• Lemma 13: Lemma 6 of chen2023improved