cs.LGMar 7, 2026
SaveDiffusion Controller: Framework, Algorithms and Parameterization
Organizations: Carnegie Mellon University; Intern at Google Research · Google · Yale University
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 -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.
Figures & tables
Table 1: HPS-v2 win rate against the pretrained model on the HPS-v2 test prompt set for all setups (each at its reported checkpoint).
Figure 2: Curves of HPS-v2 win rate against the pretrained model for SFT (left), RWL (middle), and PPO (right).
Figure 3: End-of-training win-rate vs. baselines. Each subplot reports three paired comparisons: (gray-box) DiffCon vs. DiffCon-Naive; (white-box) DiffCon-J vs. LoRA; (white-box) DiffCon-S vs. LoRA. (a)-(c): HPS-v2 win rates for SFT/RWL/PPO with orange error bars showing standard deviation; (d): human-evaluated win rate for PPO.
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
| Block Name | Structure & Sequence |
| CrossAttnDownBlock2D | ResnetBlock2D Transformer2DModel Downsample2D (optional) |
| DownBlock2D | ResnetBlock2D Downsample2D (optional) |
| CrossAttnMidBlock2D | ResnetBlock2D Transformer2DModel ResnetBlock2D |
| CrossAttnUpBlock2D | ResnetBlock2D Transformer2DModel Upsample2D (optional) |
| UpBlock2D | ResnetBlock2D Upsample2D (optional) |
Table 2: Block definitions used in . Common components: ResnetBlock2D : Standard ResNet block with GroupNorm, SiLU, and two Convolutional layers with kernel size and stride . Transformer2DModel : spatial transformer with cross-attention. Downsample2D : a convolutional layer with kernel size and stride . Upsample2D : Nearest Neighbor Resize + conv, stride .
| DiffCon: | DiffCon-Naive: | |
| Backbone | lightweight UNet | same backbone (by design) |
| Side-net input | ||
| Output parameterization | single tensor (no split) | |
| Output channels (last conv) | (extra channel for ) | (one fewer channel) |
| How output is used | controller-form correction ( 27 ), ( 28 ) | direct residual: |
| Intuition | gated, structured control term | ungated additive correction |
Table 4: Side network variants in gray-box finetuning. Both methods share the same side-net backbone, but DiffCon uses a structured output applied in the controller-form correction, while DiffCon-Naive directly adds an unsplit residual.
| Prompt | Pretrained | DiffCon | DiffCon-Naive | LoRA | DiffCon-J | DiffCon-S |
| ”An intricate, biomechanical blue planet by Bruce Pennington” | ||||||
| ”A bluejay is eating spaghetti” | ||||||
| ”A crazy looking fish swimming in Alcatraz in the style of artgerm” | ||||||
| ”An oil painting by Frank Frazetta depicting a humanoid spawn of dragons with scales, big eyes, and a menacing appearance” | ||||||
| ”An oil painting by Rembrandt depicting a muscular cat wielding a weapon with dramatic clouds in the background” |
Table 5: Qualitative comparison of different methods using SFT training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S.
| Prompt | Pretrained | DiffCon | DiffCon-Naive | LoRA | DiffCon-J | DiffCon-S |
| ”A Victorian woman quietly sings by a lake at night surrounded by fireflies, moon, and stars, painted by Vincent van Gogh and Jacques-Louis David” | ||||||
| ”Tetsuo and Kaneda engage in a race through Neo Tokyo in a pencil drawing featuring a scribbled style” | ||||||
| ”A black cat wearing a suit and smoking a cigar” | ||||||
| ”A dog sits in a white car with the door open” | ||||||
| ”A painting of a cityscape with a red cloud shining light over a sea and bridge, by Greg Rutkowski and Thomas Kinkade” |
Table 6: Qualitative comparison of different methods using RWL training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S.
| Prompt | Pretrained | DiffCon | DiffCon-Naive | LoRA | DiffCon-J | DiffCon-S |
| ”A bike parked in front of a doorway” | ||||||
| ”A cliff covered in blue blood” | ||||||
| ”A knight in black armour holds a detailed silver sword in front of a castle” | ||||||
| ”A lizard wearing sunglasses” | ||||||
| ”A creepy eldritch monster in a Swedish forest photographed from a low angle with detailed, realistic features and soft colors, inspired by Lovecraftian horror and created by artist Simon Stålenhag” |
Table 7: Qualitative comparison of different methods using PPO training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S.
Figure 4: Generations with different guidance strengths using SFT with DiffCon. Prompt: “A black cat wearing a suit and smoking a cigar”.
Figure 5: Generations with different guidance strengths using SFT with DiffCon-J. Prompt: “A bluejay is eating spaghetti”.
Figure 6: Generations with different guidance strengths using RWL with DiffCon. Prompt: “A crazy looking fish swimming in Alcatraz in the style of artgerm”.
Figure 7: Generations with different guidance strengths using RWL with DiffCon-J. Prompt: “A cute humanoid cat soldier wears a yellow raincoat, carries a rifle, and ventures through a dense forest, looking back over their shoulder”.
Table 8: SFT evaluation results on HPS-v2 test prompt set at step 1000.
Table 9: RWL evaluation results on HPS-v2 test prompt set at step 2000.
Table 10: PPO evaluation results on HPS-v2 test prompt set at step 2400.
Figure 8: HPS-v2 win rate comparison between our polynomial reward weighting (LoRA, ours) and DPOK-RWL (LoRA) during RWL finetuning. (a) uses teacher sampling where the images for training at each iteration are generated by the pretrained model, and (b) uses online sampling where the images for training are generated by the current diffusion policy.
Figure 9: Learning rate ablation for RWL and PPO on HPS-v2 win rate (comparing 1e-4 vs. 1e-5).
Figure 10: Ablation studies on for PPO on HPS-v2 win rate (five learning curves correspond to 0, 1e-5, 1e-4, 5e-4, 1e-3 ). 1e-4 leads to the best performance.
Figure 11: Ablation studies on for RWL (three HPS-v2 win rate curves in each subfigure correspond to ). leads to the best performance.
Figure 12: RWL under the exponential reward weighting function (Eq. ( 20 )) with on HPS-v2 win rate.
| Label | # trainable params | down_block_types | up_block_types | block_out_channels |
| 1Res | 11,702,717 | DownBlock2D | UpBlock2D | 256 |
| 1Res+1Attn | 13,391,549 | CrossAttnDownBlock2D , DownBlock2D | UpBlock2D , CrossAttnUpBlock2D | (128, 256) |
| 2Attn | 18,519,485 | CrossAttnDownBlock2D , CrossAttnDownBlock2D | CrossAttnUpBlock2D , CrossAttnUpBlock2D | (128, 256) |
| 2Res | 11,810,621 | DownBlock2D , DownBlock2D | UpBlock2D , UpBlock2D | (128, 256) |
Table 11: Side network structure variants used in the ablation study.
Figure 13: Ablation studies on DiffCon block structures evaluated by the HPS-v2 win rate curve. Each plot contains four curves corresponding to 1Res , 1Res+1Attn , 2Attn , and 2Res (main-paper setting).
Figure 14: Ablation studies on the input of DiffCon for RWL, evaluated by the HPS-v2 win rate curve. The plot contains two curves corresponding to input_mean =True or False, i.e., using the mean or as the input to the side network .