Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement

TL;DR

Reward-Directed Conditional Diffusion extends diffusion models with reward conditioning in a semi-supervised setting, demonstrating that a reward-conditioned diffusion framework can implicitly learn the data's latent subspace and provide provable guarantees for sampling from the reward-conditioned distribution and for reward improvement. The authors develop a meta-algorithm that uses pseudo-labeling of unlabeled data, conditional score matching with a subspace-encoder architecture, and theoretical bounds that decompose suboptimality into offline-bandit-like regret and diffusion-induced errors. They extend the theory to nonparametric function classes and validate the approach with simulations and directed text-to-image generation, highlighting trade-offs between reward strength and distribution shift. By connecting diffusion-based generation with offline RL concepts and providing subspace-recovery results, this work offers a principled path to producing high-reward samples while maintaining fidelity to the data manifold. The framework has potential implications for safer and domain-tailored generative systems with verifiable reward improvements.

Abstract

We explore the methodology and theory of reward-directed generation via conditional diffusion models. Directed generation aims to generate samples with desired properties as measured by a reward function, which has broad applications in generative AI, reinforcement learning, and computational biology. We consider the common learning scenario where the data set consists of unlabeled data along with a smaller set of data with noisy reward labels. Our approach leverages a learned reward function on the smaller data set as a pseudolabeler. From a theoretical standpoint, we show that this directed generator can effectively learn and sample from the reward-conditioned data distribution. Additionally, our model is capable of recovering the latent subspace representation of data. Moreover, we establish that the model generates a new population that moves closer to a user-specified target reward value, where the optimality gap aligns with the off-policy bandit regret in the feature subspace. The improvement in rewards obtained is influenced by the interplay between the strength of the reward signal, the distribution shift, and the cost of off-support extrapolation. We provide empirical results to validate our theory and highlight the relationship between the strength of extrapolation and the quality of generated samples.

Figures (7)

• Figure 1: Overview of reward-directed generation via conditional diffusion model. We estimate the reward function from the labeled dataset. Then we compute the estimated reward for each instance of the unlabeled dataset. Finally, we train a reward-conditioned diffusion model using the pseudo-labeled data. Using the reward-conditioned diffusion model, we are able to generate high-reward samples.
• Figure 2: Illustrations of distribution shifts in samples, reward, and encoder-decoder score network. When performing reward-directed conditional diffusion, (a) the distribution of the generated data shifts, but still stays close to the feasible data support; (b) the distribution of rewards for the next generation shifts and the mean reward improves. (c). The score network for reward-directed conditioned diffusion adopts an Encoder-Decoder structure.
• Figure 3: Quality of generated samples as target reward value increases. Left: Average reward of the generation; Middle: Distribution shift; Right: Off-support deviation. The errorbar is computed by $2$ times the standard deviation over $5$ runs.
• Figure 4: Shifting reward distribution of the generated population.
• Figure 5: The predicted rewards and the ground-truth rewards of the generated images. At each guidance level, increasing the target $y$ successfully directs the generation towards higher predicted rewards, but also increases the error induced by the distribution shift. The reported baseline is the expected ground-truth reward for undirected generations.

Theorems & Definitions (37)

• Proposition 3.1: Score Matching Objective for Implementation
• Theorem 4.5: Subspace Fidelity of Generated Data
• Theorem 4.6: Off-policy Regret of Generated Samples
• proof
• Lemma C.1
• proof
• Definition C.2
• Definition C.3
• Lemma C.4
• proof