Are we really tilting? The mechanics of reward guidance in flow and diffusion models

Sanjit Dandapanthula; Nicholas M. Boffi

cs.LGJun 1, 2026

Are we really tilting? The mechanics of reward guidance in flow and diffusion models

Authors: Sanjit Dandapanthula, Nicholas M. Boffi

Organizations: Carnegie Mellon University

Abstract

Reward guidance algorithms steer a learned generative process toward the reward-tilted measure at inference time. While empirically powerful, these methods are prone to reward hacking: the guided model over-optimizes the reward at the cost of fidelity to the learned distribution. Prior work has attributed this to the complexity of neural reward functions or implicit biases in diffusion training, but its fundamental origins remain poorly understood. We show that reward hacking arises from an approximation made in most practical implementations of reward-guided diffusion -- finite-particle plug-in estimation of the Doob h-function -- even in the simplest non-trivial settings of Gaussian and Gaussian mixture targets with quadratic rewards. In closed form, we isolate two distinct failure modes of the plug-in estimator: it leads to reward hacking within each mode and it cannot select high-reward modes. We propose a closed-form reward damping schedule that corrects the within-mode bias with no additional compute, and clarify the role of best-of-n sampling in compensating for the mode selection failure. Experiments on Gaussian mixture targets, a 2D checkerboard, and FLUX.1 text-to-image generation confirm that our theoretical insights carry over to practical settings.

Explore similar work

Apr 19, 2026cs.LG

Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models

Reward-based fine-tuning steers a pretrained diffusion or flow-based generative model toward higher-reward samples while remaining close to the pretrained model. Although existing methods are derived from different perspectives, we show that many can be written under a common framework, which we call reward score matching (RSM). Under this view, alignment becomes score matching against a value-guided target, and the main differences across methods reduce to the construction of the value-guidance estimator and the effective optimization strength across timesteps. This unification clarifies the bias-variance-compute tradeoffs of existing designs, and distinguishes core optimization components from auxiliary mechanisms that add complexity without clear benefit. Guided by this perspective, we develop simpler, more efficient redesigns across representative differentiable and black-box reward alignment tasks. Overall, RSM turns a seemingly fragmented collection of reward-based fine-tuning methods into a smaller, more interpretable, and more actionable design space. Code is available at https://github.com/jaylee2000/rsm

Jeongjae Lee, Jinho Chang, Jeongsol Kim +1

Score-Based Diffusion Model Training-Free Reward Alignment Reinforcement Fine-Tuning Generative Flow Networks

Ask PDF

Jun 16, 2026cs.LG

NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment

We introduce the Noise-Tilted Reverse Kernel (NTRK), a reward-guided diffusion sampler that injects reward gradients through the noise term, leaving the pretrained reverse kernel unchanged and requiring only a single sample per step. Reward-guided sampling at inference time has greatly expanded the versatility of pretrained diffusion models. Yet existing methods face a trade-off. Gradient-based guidance shifts the reverse mean, steering generation but pushing intermediate states outside the region that the model was trained on and degrading quality. Search-based methods preserve quality but gain no gradient signal. No prior method achieves both. NTRK resolves this by keeping the reverse mean fixed and biasing the noise term toward high reward. This is enabled by a whitening operator, the central mechanism behind NTRK, which converts reward gradients into noise-compatible perturbations without losing their guiding signal. Across various reward alignment tasks, NTRK outperforms recent state-of-the-art baselines without losing sample quality. Remarkably, on aesthetic generation, NTRK surpasses the reward of the best baseline at 500 NFEs using only 25 NFEs, a 20 times reduction in compute.

Jisung Hwang, Yunhong Min, Jaihoon Kim +2

Diffusion Sampling Reward Gradients

Ask PDF

May 12, 2026cs.LG

The tractability landscape of diffusion alignment: regularization, rewards, and computational primitives

Inference-time reward alignment asks how to turn a pre-trained diffusion model with base law

p

into a sampler that favors a reward

r

while remaining close to

p

. Since there is no canonical distributional distance for this closeness constraint, different choices lead to different "reward-aligned" laws and, just as importantly, different algorithmic problems. We develop a primitive-based approach to reward alignment: rather than assuming arbitrary reward-aligned laws can be sampled, we ask which simple algorithmic primitives suffice to implement alignment for non-trivial reward classes. If closeness is measured in KL distance, the target law is

q(x) \propto p(x) \exp(λ^{-1}r(x))

. For this setting, we show that linear exponential tilts of the form

q(x)\propto p(x)\exp(\langle θ, x \rangle)

-- which according to recent work [MRR26] can be efficiently sampled from -- are a sufficient primitive for aligning to a very broad class of convex low-dimensional rewards. If closeness is measured in Wasserstein distance, the corresponding primitive is a proximal transport oracle: given

x

, solve

\mbox{argmax}_y \{r(y)- λc(x,y)\}

. This oracle can be efficiently implemented for concave or low-dimensional Lipschitz rewards

r(x)=f(Ax)

. Together, these results illustrate that the choice of distribution distance for alignment affects the computational primitive and the tractable reward class.

Ankur Moitra, Andrej Risteski, Dhruv Rohatgi

Regularization Techniques Computable

Ask PDF