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

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.

Paper

Similar papers

Peer review

Reviewer uwg26/10 · confidence 2/52023-07-04

Summary

The paper addresses conditional generation with reward-conditioned diffusion models. They propose to learn a reward function from a small subset of labeled data. The paper aims to answer an intriguing research question: "How can we reliably estimate the reward-conditioned distribution through diffusions and balance the trade-off between the reward signal and generating quality?" Additionally, the paper claims that the reward-conditioned diffusion model implicitly learns the latent subspace representation of x.

Strengths

Strength: 1) The paper is well written and the problem they try to tackle is interesting. 2) The paper provides an insightful theory for conditional distribution learned through a reward-based diffusion model

Weaknesses

1) The paper lacks a comparison to other similar models, such as classifier-guided diffusion models. 2) The assumptions made for the theoretical work appear to be overly simplified.

Questions

Some suggestions: 1) It would be good to add what each color represents in Figure 2. 2) I think the main paper should include at least a summary of related work so that reader can have a better picture of where this work stands when compared to the existing works. 3) In Figure 1, in the block for the diffusion model, X_0 is represented as the noise, I think it is better to keep the same notation as in other literature to avoid confusion. Questions 1) Regarding Figure 6, it is mentioned that those are picked examples, does it mean the presented results are "cherry-picked" and not generalizable? also, I think here y represents the colorfulness level ( as the ground truth reward model favors colorful images), there seems the context is also changed, not only the color, do you have any insight on that? 2) The model seems very similar to the classifier-guided diffusion model, but it never compared to that in the experiments, any reason for that? 3) I am wondering if it is possible to train the reward function together with the diffusion model instead of pertaining to the reward network. 4) The paper assumes generated data is a linear transformation of the latent z (assumption 3.1), I am wondering if this assumption is too strong and does not hold in general. 5) Assumption 3.2 assume a linear reward, however, the ground truth reward in the experiment used the imageNet with a randomly initiated final prediction layer. I am not sure how to interpret this part. 6) After reading the paper, I am still confused about the assumption that the model implicitly learns the latent subspace representation of x. Some suggestions: 1) It would be beneficial to add a legend in Figure 2 that explains the meaning of each color. It is recommended to include a summary of related work in the main paper to provide readers with a better understanding of where this work stands in comparison to existing literature. 2) In Figure 1, the block representing the diffusion model uses X_0 to denote noise. It may be clearer to use consistent notation with other literature to avoid confusion. Questions: 1) Regarding Figure 6, it is mentioned that the examples shown are picked intentionally. Does this imply that the presented results are cherry-picked and may not be generalizable? Additionally, it appears that y represents the colorfulness level (as the ground truth reward model favors colorful images), but there seems to be a change in context as well. Do you have any insights on this? 2) The model appears to be quite similar to the classifier-guided diffusion model, yet there is no comparison with it in the experiments. Is there a specific reason for this omission? 3) I'm curious if it is possible to train the reward function concurrently with the diffusion model instead of relying on a separate reward network. 4) The paper assumes that the generated data is a linear transformation of the latent z (assumption 3.1). However, I'm wondering if this assumption is overly strong and may not hold in general. 5) Assumption 3.2 assumes a linear reward, but the ground truth reward used in the experiments involved ImageNet with a randomly initiated final prediction layer. I'm unsure how to interpret this aspect. 6) Even after reading the paper, I'm still unclear about the assumption that the model implicitly learns the latent subspace representation of x.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The paper did not discuss any limitations. The paper has no potential negative societal impact.

Reviewer mN9F6/10 · confidence 3/52023-07-06

Summary

In this work, authors explore the problem of reward-directed generation using conditional diffusion models in a semi-supervised learning setup. More specifically, they consider a dataset which has a small subset of it labeled with rewards and the majority of it unlabeled. Using the small labeled subset, they first train a reward approximator with regression, then use the learned reward approximator to label the unlabeled portion of the dataset with pseudo labels. Finally, they train the conditional diffusion model using the samples and their pseudo labels. Authors present theoretical explanations for conditioned diffusion models and reward improvement guarantees for reward-conditioned generations. Finally, experimental results are provided that demonstrated the findings of the theory.

Strengths

Overall Strengths: 1. This paper is very well-organized and well-written and I was able to follow along fairly easily. So, I highly commend the authors for the great job they have done. In particular, all the sections leading up to the main theory are organized well and give a nice background and trajectory to where the theory lies. 2. The high-level idea of theoretically studying how reward-conditioned distribution of diffusion models change and how to balance reward to trade-off sample quality and reward maximization is well-worth pursuing. 3. Theory is well-driven and thoroughly discussed.

Weaknesses

Overall Weaknesses: 1. One of the main questions that this paper aims to answer is “How to balance the reward signal and distribution-shift effect to have high-reward and high-quality samples?” which is what I was most excited to learn about. However, I'm not confident that I truly got an answer to this question. 2. The experimental setup is limited and the computed metrics do not thoroughly cover the theoretical findings of the papers. 3. Although figure 2 looks very nice, I’m not sure if it has any added value in the main body of the paper.

Questions

1. How will the results generalize to settings where data doesn’t have a latent linear representation? 2. How will the results of the current model (conditional diffusion model) compare to using an unconditional diffusion model with classifier guidance to achieve high rewards? 3. For text-to-image generation, it is said that higher reward refers to more vividly colored images. Do you have a numerical metric for this?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

In addition to the points provided in the weaknesses section, I believe a deeper effort on the experiments section would be of tremendous value for this work. For me, the experimental results fall short in supporting the theory of the paper and I’m not sure if I’m convinced the full message is conveyed.

Reviewer ZNG86/10 · confidence 2/52023-07-06

Summary

This paper presents an approach to generation using diffusion models augmented with a reward function. It does so by setting up a semi-supervised learning setup, where the reward function is learned from a small set of data. The reward is then used to learn a reward conditioned score function, which is subsequently used to generate data conditioned on requested reward. Assuming a linear subspace, the analysis approaches this setting through the lens of linear bandits, and characterizes the error or suboptimality in terms of regret to the target or requested reward, off-distribution error, and on-distribution error. Experiments evaluate the above interplay, and how requesting higher rewards leads to distribution shift in the generated samples. Further experiments show how pre-trained models can be adapted to generate reward-conditioned samples.

Strengths

* Presents a practical way to add subjective rewards to generate data beyond concrete prompts like text * Analyzes this reward conditioned generation setting and identifies the interplay between rewards and distribution shift. * Experimentally verifies the claims made in the analysis.

Weaknesses

* The paper does not discuss a practical manner in to identify when the generative model starts deviating from the training distribution. * Comparing the diffusion model's error to a linear bandit setting is interesting, but the exact connection between the two seems a bit murky. The training process does not seem to take advantage of any bandit learning algorithms. * While it shows that the model can be adjusted to arbitrary rewards, it does not showcase a practical use case. * The technical novelty is unclear. ============================= ### Post-Rebuttal * The author's rebuttal and other reviews have made the technical contribution of the paper abundantly clear. * Additionally, the connection to _offline_ bandits is more evident and does add value. * The additional RL experiment grounds the claims made in the paper more concretely than the previous experiments. As such, the majority of my concerns have been allayed with the author response.

Questions

* Section 2.1 gives the scale of $\sigma$ but does not mention any such constraint about $f^*$. Section 4.2 constructs a random reward, but this setting also makes no mention of the scale of the reward. How does the scale of the reward affect learning? * If someone were to try and reproduce these results, how would they go about it? Perhaps an experiment that specifically tries to maximize some specific property that is hard to communicate through a text prompt could be shown here, to communicate the effectiveness of this approach, as well as to assist reproducibility. * The labeled data is based on the CIFAR-10 dataset. Is the unlabeled dataset from a corresponding distribution? Does it have similar resolution, size, and kinds of pictures? * The paper states `To optimally choose a target value, we must trade off between the two counteractive effects.` Are there any practical methods to do so? * Reinforcement learning faces the problem of reward design. The analysis done in this paper could be useful for this problem of reward design. It would be useful to deepen the discussion on reward design and reference some related work [1, 2] ### References [1] Booth, S., Knox, W.B., Shah, J., Niekum, S., Stone, P. and Allievi, A., 2023, June. The perils of trial-and-error reward design: misdesign through overfitting and invalid task specifications. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 37, No. 5, pp. 5920-5929). [2] Knox, W.B., Allievi, A., Banzhaf, H., Schmitt, F. and Stone, P., 2023. Reward (mis) design for autonomous driving. Artificial Intelligence, 316, p.103829.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

while the proposed approach opens up avenues to communicate and generate data using feedback other than text prompts, it does not sufficiently address problems that might arise from such freeform feedback.

Reviewer umuF6/10 · confidence 3/52023-07-09

Summary

The paper addresses the problem of conditional generation with diffusion models in a self-supervised setting, where the conditional generation is guided by a learned regressor on the small labeled subset. This is referred to as reward-directed conditional diffusion. Assuming the inputs have a latent linear representation, it is shown that reward-conditioned diffusion models implicitly learn this latent representation. Further, assuming a linear reward model is used, it is shown that reward-conditional generation can be viewed as off-policy bandit learning in latent feature space. The theory is also extended to non-nonparametric reward and score functions. Experiments on synthetic and text-to-image data support the theory.

Strengths

Steering generate models (especially diffusion models) towards generating samples with desirable properties is a topic of wide interest. The semi-supervised setting is a particularly important special case of that, which comes up in many real-world applications. As the authors note, there is not much theoretical work in this space yet, making this a valuable contribution. I found the theoretical analysis to be quite insightful. I particularly like the analysis regarding the trade-off between distribution shift and reward maximization. The simulation experiment seems to align well with the theory. I also appreciate that the theoretical analysis extends to more general reward and score functions.

Weaknesses

I found some parts of the paper a bit difficult to read. First, the motivation and derivation of the score network architecture is unclear without reading the reference ([8]). For example, it is not obvious that this functional form follows from the linearity assumption of x = Az. Second, a lot of notation is not introduced, e.g. $k$ in Eq. 2.3, $x_\parallel$ and $x_\perp$ in Assumption 3.2. This makes some equations difficult to understand. Third, it does not become clear until halfway through the paper why the pseudo-labeller is called a reward function. Perhaps a forward reference would help to clarify this. I am also not sure what the practical implications are. I could have (qualitatively) predicted the results for text-to-image generation without the theory, simply because you are asking the model to extrapolate beyond the training data. The quantitative results also only make qualitative statements. Apart from that, I found the experimental setup of using a random reward model a bit strange. Minor comments: - Given the similarity of Fig. 2c with Fig 2 in [8], citing [8] is probably warranted ("Figure adapted from [8]")

Questions

1. It is not clear to me why Algorithm 1 (especially Eqs 2.2 and 2.3) is necessary. Why couldn't you use either classifier-free or classifier-based guidance? Why does Algorithm 1 offer "theoretical cleanness"? 2. I do not understand why Assumption 3.2 is required. To me, it feels more like a user choice to penalize off-support data in the reward as this should also be reflected in the unconditional score (recall that $\nabla_x \log p(x|y) = \nabla_x \log p(y|x) + \nabla_x \log p(x)$). Could you clarify this? 3. In the experiments, what is the reward distribution of the training data? At which reward values are you asking the model to extrapolate? Minor comments: - In the problem setup section, why is $1 > \sigma$ necessary?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors have clearly laid out the assumptions for the theoretical analysis. I do not see any major limitations within this scope.

Reviewer ZNG82023-08-14

Thank you for the Response

I thank the authors for the detailed response. The confusion about the connection to linear bandits likely arose from this sentence in the paper on line 67: ``` In the case of a linear reward model, we show that the regret mimics the off-policy regret of linear bandits with full knowledge of the subspace feature. ``` I agree that in the offline setting, typical bandit algorithm approaches are not applicable. Additionally, my main concern about technical novelty has been allayed. The experiment in the RL domain has especially helped to ground the idea in a concrete problem. I have raised my score accordingly.

Authorsrebuttal2023-08-15

Thank you for your prompt feedback on our rebuttal. The sentence you pointed out does need more clarity to avoid confusion. We will modify it based on the facts in **A2** to ``` Given a target reward value, we analyze the statistical error of reward-directed generation, measured by the difference between the target value and the average reward of the generated population. In the case of a linear reward model, we show that this error includes the suboptimality gap of linear off-policy bandits with full knowledge of the subspace feature, if taking the target to be the maximum possible. ``` In addition, thanks for your especially positive comment on our new experiment. Your positive rating flipped from the negative is very encouraging and reassuring to us :)

Reviewer mN9F2023-08-14

I appreciate the authors' thorough response. After a careful re-evaluation of the work, considering the input from other reviewers and taking into account the authors' rebuttal efforsts, I have decided to increase my rating. I particularly find the newly added experiment on all baselines helpful.

Authorsrebuttal2023-08-15

Author's Follow-up

Dear Reviewer, We want to follow up to gently check in your opinions on our rebuttal posted one week ago. A brief summary of our rebuttal to your concerns is: - We compared our method to guidance methods in detail in 2. Comparison of [our rebuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE). It's evident from the comparison that Alg.1 and its analysis captures the essence of both classifier guidance and classifier free guidance methods. - "The assumptions appear to be overly simplified" is a misunderstanding. In **Q2&A2**, we gave a fine-grind breakdown of what assumptions we made and what results we include in our paper. By the way, we added a [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf) on solving RL tasks via Alg.1 to showcase its versatility. Please let us know if you have further suggestions. Thank you very much!

Reviewer umuF2023-08-15

Thank you very much for the comprehensive reply. I really appreciate the additional explanations and experiments. The rebuttal confirms my initially positive take on the paper. Thus, I will maintain my score. In addition, I am slightly raising my confidence score given the additional clarity.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC