Diffusion Controller: Framework, Algorithms and Parameterization

TL;DR

This paper introduces Diffusion Controller (DiffCon), a unified control-theoretic framework for diffusion models that improves fine-tuning and control by reweighting transition kernels within a stochastic control setting, leading to better alignment and efficiency.

Contribution

It presents a novel control-theoretic formulation of diffusion processes as LS-MDPs, deriving practical RL-based algorithms for diffusion fine-tuning and a model decomposition for effective gray-box adaptation.

Findings

Consistent improvements in preference-alignment win rates.

Enhanced quality-efficiency trade-offs over baselines.

Effective gray-box adaptation with a lightweight control correction.

Abstract

Controllable diffusion generation often relies on various heuristics that are seemingly disconnected without a unified understanding. We bridge this gap with Diffusion Controller (DiffCon), a unified control-theoretic view that casts reverse diffusion sampling as state-only stochastic control within (generalized) linearly-solvable Markov Decision Processes (LS-MDPs). Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an $f$-divergence cost. From the resulting optimality conditions, we derive practical reinforcement learning methods for diffusion fine-tuning: (i) f-divergence-regularized policy-gradient updates, including a PPO-style rule, and (ii) a regularizer-determined reward-weighted regression objective with a minimizer-preservation guarantee under the Kullback-Leibler (KL) divergence. The LS-MDP framework further implies a principled model form: the optimal score decomposes into a fixed pretrained baseline plus a lightweight control correction, motivating a side-network parameterization conditioned on exposed intermediate denoising outputs, enabling effective gray-box adaptation with a frozen backbone. Experiments on Stable Diffusion v1.4 across supervised and reward-driven finetuning show consistent gains in preference-alignment win rates and improved quality-efficiency trade-offs versus gray-box baselines and even the parameter-efficient white-box adapter LoRA.