English

Diffusion Controller: Framework, Algorithms and Parameterization

Machine Learning 2026-03-10 v1 Artificial Intelligence

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 ff-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.

Keywords

Cite

@article{arxiv.2603.06981,
  title  = {Diffusion Controller: Framework, Algorithms and Parameterization},
  author = {Tong Yang and Moonkyung Ryu and Chih-Wei Hsu and Guy Tennenholtz and Yuejie Chi and Craig Boutilier and Bo Dai},
  journal= {arXiv preprint arXiv:2603.06981},
  year   = {2026}
}