Abstract

KL regularized reward optimization in the policy space is the mainstay of aligning generative models, specifically LLMs, to rewards. In the LLM case, the likelihood of producing a rollout is available that enables gradient-based RL algorithms. However, Diffusion models do not provide the marginal likelihood of the final denoised sample they generate. Some existing techniques impose KL regularization on the entire trajectory which suffer from quality issues. In this talk, we first show a general equivalence between the reward-reshaped, KL-regularized RL framework and a general class of acceptance-rejection sampling-based fine-tuning methods leading to a method called GRAFT. We further exploit properties to motivate aligning policy on rewards only on higher noise scales when the number of samples is limited in RL training leading to a method called PGRAFT. PGRAFT effectively exploits a bias variance tradeoff in fine tuning diffusion models. We also introduce a variant for Flow models that corrects only the initial noise distribution to align with the rewards.

This is based on joint work with Dheeraj Nagaraj, Nitish Kannen, Gautham Govind Anil and Shaan Ul Haque: arXiv:2510.02692 (accepted at ICLR 2026).