Diffusion Controller: Framework, Algorithms and Parameterization

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.

Modern AI image generators can create high-quality pictures from text, but it is still difficult to reliably guide them toward what people actually want—such as more appealing, better-aligned, or safer images—without hurting image quality. This paper introduces Diffusion Controller, a new way to understand and improve this kind of guidance. Instead of treating many existing fine-tuning methods as separate tricks, the paper shows that they can be viewed through one common idea: gently steering a pretrained image generator while keeping it close to what it already knows. This view leads to new training methods for improving image generators using either example images or feedback-based rewards. It also suggests a practical design: keep the large pretrained model fixed and add a small “controller” network that learns how to adjust the generation process. This is useful when the full model cannot be changed directly. In experiments with Stable Diffusion v1.4, the approach improves preference-alignment scores and often gives a better quality–efficiency trade-off than several baseline methods, including LoRA in some settings.