Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement
Authors: Hui Yuan
Kaixuan Huang
Chengzhuo Ni
Minshuo Chen
Mengdi Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: Theory
Diffusion Model
Reward Optimization
Low-dimensional Data
Distribution estimation

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 dataset consists of majorly unlabeled data and a small set of data with noisy reward labels. Our approach leverages a learned reward function on the smaller data set as a pseudolabeler to label the unlabelled data. After pseudo-labelling, a conditional diffusion model (CDM) is trained on the data and samples are generated by setting a target value $a$ as the condition in CDM. From a theoretical standpoint, we show that this directed generator can effectively learn and sample from the reward-conditioned data distribution: 1. our model is capable of recovering the data's latent subspace representation. 2. the model generates samples moving closer to the user-specified target. The improvement in rewards of samples is influenced by a 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.

Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement
Hui Yuan
Kaixuan Huang
Chengzhuo
Minshuo Chen
Mengdi Wang"
Abstract
We explore the methodology and theory of reward-directed generation via condi- tional diffusion models. Directed generation aims to generate samples with desired properties as measured by a reward function, which has broad applications generative AI, reinforcement learning, and computational biology. We consider the common learning scenario where the dataset consists of majorly unlabeled data and small set of data with noisy reward labels. Our approach leverages a learned reward function on the smaller data set as a pseudolabeler to label the unlabelled data. After pscudo-labelling. conditional diffusion model (CDM) is trained the data and samples are generated by setting a target value a as the condition in CDM. effectively
From theoretical standpoint, show that this directed generator learn and sample from the reward-conditioned data distribution: 1. our model is capable of recovering the data's latent subspace representation. 2. the model generates samples moving closer to the user-specified target. The improvement in rewards of samples is influenced by a interplay between the strength of the reward signal, the distribution shift, and the cost of off-support extrapolation. We provide empirical results validate our theory and highlight the relationship between the strength of extrapolation and the quality of generated samples. Our code is //g available https: //g /gi thub. com/eafffaeeae/cc Kaffalj dhmah2/RCGDD
Introduction
Controlling the behavior of generative models towards desired properties is a major problem for deploying deep learning models for real-world usage. As large and powerful pre-trained generative models achieve steady improvements over the years, one increasingly important question is how to adopt generative models to fit the needs of specific domain and ensure the generation results satisfying certain constraints (e.g., safety, fairness, physical constraints) without sabotaging the power of the original pre- pre-trained model [35, 27, 54, 41].
In this paper, we focus on directing the generation of diffusion models [19, 43], family of score- matching generative models that have demonstrated the state-of-the-art performances in various domains, such as image generation [39, 38, 4] and audio generation, with fascinating potentials in broader domains, including text modeling [3, 27], reinforcement learning [21, 1, 36, 28) and protein structure modeling [26]. Diffusion models are trained to predict a clean version of the noised input, and generate data by sequentially removing noises and trying find cleaner version of the input. The denoising network (a.k.a. score network) s(z, t) approximates the score function v log Pe (x) [45, 46). and controls the behavior of diffusion models. People can incorporate any control information y as an additional input to s(x, y, t) during the training and inference [38, 54].
Department of Electrical and Computer Engineering, Princeton University. Authors" emails are: (huiyuan, kaixuanh, cn10, mc0750, mengdiv)@princetts edu
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Here we abstract various control goals as a scalar reward y, measuring how well the generated instance satisfies our desired properties. In this way, the directed generation problem becomes finding plausible instances with higher rewards and can be tackled via reward-conditoon subtlety of this problem lies in that the two goals potentially conflict with each other: diffusion models are learned to generate instances similar to the training distribution, while maximizing the rewards of the generation drives the model to deviate from the training distribution. In other words, the model needs to "interpolate" and "extrapolate" at the same time. A higher value of y provides a stronger signal that guides the diffusion model towards higher rewards, while the increasing distribution shift may hurt the generated samples' quality. In the sequel, we provide theoretical guarantees for the reward-condie conditioed diffusion models, aiming to answer the following question:
diffusion models. The
How to provably estimate the reward signal and distribution-shij effect, and ensure reward improvement in generated samples?
distribution via diffussio?? How to balance the
Our Approach. To answer both questions, we consider a semi-supervised learning setting, where we are given small dataset Diabel with annotated rewards and a massive unlabeled dataset uullabel. We estimate the reward function using Diabel and then use the estimator for pseudo-labeling on Dunlabel- Then we train a reward-conditione diffusion model using the pseudo-labeled data. Our approach is illustrated in Figure 1. In real-world applications, there are other ways to incorporate the knowledge from the massive dataset D Punlabel. e.g., finetuning from a pre-trained model [35, 54]. We focus on the pseudo-labeline on data size and distribution shift. The intuition behind and the message are applicable other semi-supervised approaches; see experiments in Section 5.2.
approach, as it provides a cleaner formulation and exposes the error dependency
From a theoretical standpoint, we consider data point x having a latent linear representation. Specifi- cally, we assume x = Az for some matrix A with orthonormal columns and being a latent variable. The latent variable often has smaller dimension, reflecting the fact that practical data sets often exhibit intrinsic tow-dimensional structures [13, 48, 37J. The representation matrix A should be learned to promote sample efficiency and generation quality [8]. Our theoretical analysis reveals an intricate interplay between reward guidance, distribution shift, and implicit representation learning; see Figure for illustration.
Contributions. Our results are summarized as follows.
(1) We show that the reward-conditionss diffusion model implicitly learns the latent subspace representation of x. Consequently, the model provably generates high-fideliyy data that stay close to the subspace (Theorem 4.5).
(2) Given a target reward value, we analyze the statistical error of reward-directed generation, mea- sured 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, the other two components of this error correspond to the distribution shift in score matching and the cost of off-support extrapolation (Theorem 4.6).
(3) We further extend our theory to nonparametric reward and distribution configurations where reward prediction and score matching are approximated by general function class, which covers the wildly adopted ReLU Neural Networks in real-world implementation (Section 4.3 and Appendix F).
(4) We provide numerical experiments under both synthesized setting and more realistic settings such text-to-image generation (stable diffusion) and reinforcement learning (decision-diffuser) to support our theory (Section and Appendix 1).
To our best knowledge, our results present the first statistical theory for conditioned diffusion models generation.
and provably reward improvement guarantees for
2 Related Work
Guided Diffusions. For image generations, guiding the backward diffusion process towards higher log probabilities predicted by classifier (which can be viewed as the reward signal) leads to improved sample quality, where the classifier can either be separated trained, i.e., classifier-guided [1]] or implicitly specified by the conditioned diffusion models, i.e., classifier-free [18]. Classifier- free guidance has become a standard technique in the ofate-of-the-art text-to-image diffusion models [39, 38, 4]. Other types of guidance are also explored [33, 14]. Similar idcas have been explored in
prediccte

---

Target Reward
D Dunlabel
Conditional a(Leature, label, r)
D Dacw
Lobel
Modid
High-cewrdd Generited Sumplo
Diabel
Reward--oonditionel Diffasion
Step 1: Reward Learning
Step 2: Pscudo Labeling
Step 3: Conditional Dilfusion Model Training
Step 4: Guided Generation
Figure I: 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, train reward-conditioned diffusion model using the pscudo-labeled data. Using the reward-conditioned diffusion model, we are able to generate high-reward samples.
Sherkcuk
-A"
A
00
Subspace
roward
2
R
R
R
81.00 8.13
Samples
[b) Distribution
(c) Score Network
Figure 2: Illustrations of distribution shifts in samples, reward, and encoder-decoder score network. When performing ceward-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) (Adapted from [8]) the score network for reward-directed conditioned diffusion adopts Encoder-Decodee structure.
sequence modelling problems. In offline reinforcement learning, Decision Diffuser [1] is diffusion model trained on offline trajectories and can be conditioned to generate new trajectories with high return,, satisfying certain safety constraints,
composing skills. For discrete generations, Diffusion LM [27] manages to train diffusion models on discrete text space with an additional embedding layer and rounding step. The authors further show that gradients of any classifier can incorporated to control and guide the text generation.
Theory of Diffusion Models A line of work studies diffusion models from sampling perspective. When assuming access to score function that can accurately approximate the ground truth score function in - L norm, [9, 25] provide polynomial convergence guarantees of score-based diffusion models. "Convergence of denoising diffusion models under the manifold hypothesis" by Valentin De Bortoli further studies diffusion models under the manifold hypothesis. Recently, [8] and [34] provide an end-to-end analysis of diffusion models. In particular, they develop score estimation and distribution estimation guarantees using the estimated score function. These results largely motivate our theory, whereas, we are the first to consider conditional score matching and statistical analysis of conditional diffusion models.
Connection to Offline Bandit/RL Our off-policy regret analysis generated samples related offline bandit/RL theory [30, 29, 6, 12, 22, 32, 5]. In particular, our theory extensively deals with distribution shift the offline data set by class restricted divergence measures, which are commonly adopted in offline RL. Moreover, our regret bound of generated samples consists of an error term that coincides with off off-policy linear bandits. However, our analysis goes far beyond the scope of bandit/RI.
3 Reward-Directed Generation via Conditional Diffusion Models
In this section, we develop conditioned diffusion model-based method to generate high-fidelity samples with desired properties. In real-world applications such image/text generation and protein design, one often has access to abundant unlabeled data, but relatively limited number of labeled data. This motivates to consider semi-supervised learning setting.

---

Notation: ay denotes ground truth joint distribution of z and its label y. Ps is the marginal of x. Any piece of data in Diabel follows Pay and any data in Danlahel follows Pz. P is used denote a distribution and p denotes its corresponding density. P(x y a) and P(x, y a) are the conditionals of Pay Similarly, we also use notation Pag. P(x 9 = a) for the joint and conditional of (x. 77, where ī is predicted by the learnt reward model. Also, denote generated distribution using diffusion by P (density p) followed by the same argument in parentheses as the true distribution approximates, e.g. P(x y = a) is generated as an approximation of P(xr y = =a.
3.1 Problem Setup
Suppose we are given an unlabeled data set Dunlabel = رزز and labeled data set Diabel {( (2 where it is often the case that > n2. Assume without loss of generality that Diabel and uulabel are independent. In both datasets, suppose z is sampled from an unknown population distribution P.. In our subsequent analysis, we focus on the case where Pa is supported on latent subspace, meaning that the raw data r admits a low-dimensional representation (see Assumption 4.1). We model y as a noisy measurement of reward function determined by given by
y = /" (x) + for e~ N(O, o²) with 1>a> 0.
user can specify a target reward value, i.e., y = a. Then the objective of directed generation is to sample from the conditional distribution P(æly = a). Given f", Pa or the low-dimensional support of Pz are unknown, we need to learn these unknowns explicitly and implicitly through reward-conditioned diffusion.
3.2 Meta Algorithm
Algorithm 1 Reeaadd Generation via Diffusion Model (RCGDM)
1: Input: Datasets Dunlahel, Diabel, target reward value a, early-stopping time In. noise level v. (Note: in the following psuedo-code, or (3r) is the Gaussian density and v is the step size of discrete backward SDE, see $3.3 for elaborations on conditional diffusion)
2: Reward Learning: Estimate the reward function by
f E argmin FEF
Σ
{((x)))):
(3.1)
(EL.
where é is a loss and F a function class.
3: Pseudo labeling: Use the learned function f to evaluate unlabeled data Dunlabel and augment it with pseudo labeles: D = {(z,, ky) = f(z,) + E, 3 for Ej iid N(O, x2).
4: Conditional score matching: Minimize over s E S (S constructed as 3.8) on data set D via
S" B Be
@ E argmin SES
| log ()(((())))) s(x 0, 11 dt. (3.2)
5: Conditioned generation: Use the estimated score s(,, -) to sample from the backward SDE:
1 Vv.= + *(80/= kog a. T' kn) dt + ddW for E [Ky,, (k + 1)7), k km
T
dit=
[_]
(3.3)
6: Return: Generated population P(-1ğ a). learned subspace representation V contained in s.
In order to generate novel samples with both high fidelity and high rewards, we propose Reward- Conditioned Generation via Diffusion Models (RCGDM); see Algorithm 1 for details. By using the labeled data Diabel, we approximately estimate the reward function f* by regression, then we obtain an estimated reward function f. We then use f to augment the unlabeled data D Dunlabel with "pseudo labeling" and additive noise, i.e., D = {{2,,, 23 = f(x;) + } with Ej N(O, 1²) of small variance v² Here, added noise &, merely for technical reasons in the proof. We denote the

---

joint distribution of (z. () as Pog. Next, we train conditional diffusion model using the augmented dataset D. If specify a target value of the reward, for example letting y a, we can generate conditioned samples from the distribution P(x|ū = a) by backward diffusion.
Alternative approaches. In Line 4, Algorithm trains the conditional diffusion model via conditional score matching. This approach is suitable when we have access to the unlabeled dataset and need to train a brand-new diffusion model from scratch. Empirically, in order to realize conditional generation, we can utilize a pre-trained diffusion model and incorporate with reward signals to be conditioned on. Existing methods falling in to this category include classifier- based guidance [11], fine-tuning [54], and distillatiionn [47]. For theoretical cleanness, focus on analysing Algorithm I as it shares the same core essence with other alternative methods, which is approximating of the conditional score (oll)))
3.3 Training of Conditional Diffusion Model
In this section, we provide details about the training and sampling of conditioned diffusion in Algorithm I (Line 4: conditional score matching and Line 5: conditional generation). In Algorithm 1, conditional diffusion model learned with D {(T} Øj = f(z;) + - where (x, 77 ~ Pag- For simplicity, till the end of this section we use y instead of 7 to denote the condition variable. The diffusion model is to approximate the conditional probability P(x (7).
Conditional Score Matching. The working flow of conditional diffusion models is nearly identical to that of unconditioned diffusion models reviewed in Appendix A. A major difference is we learn a conditional score V log Pt (aly) instead of the unconditional one. Here Pt denotes the marginal density function at time of the following forward O-U process,
dxp- = (() + Vg(l)dW, with X5 Po(xly) and E (0,2),
1
(3.4)
where similarly T is terminal time, (We) is a Wiener process, and the initial distribution Po(zly) is induced by the (z, D) -pair distribution Pag. Note here the noise is only added on but not on Throughout the paper, we consider g(t) for simplicity. We denote by (x the distribution of X and let pt (xrll) be its density and P. (x y) be the corresponding joint, shorthanded as P,. A key step to estimate the unknown V log (x through denoising score matching [46]. conceptual way is to minimize the following quadratic loss with S, concept class.
argmin #ES
TE E... [[V (o s(x,, y, 1) dt,
(3.5)
Unfortunately, the loss in (3.5) is intractable since log Pr (g) is unknown. Inspired by Hyvärinen and Dayan [20] and Vincent [52], we choose a new objective (3.2) and show their equivalence in the following Proposition. The proof is provided in Appendix C.I.
Proposition 3.1 (Score Matching Objective for Implementation). For any t and score estimator s, there exists a constant C2 independent of such that
E (...))>P [[|V log Pe (arlv) - s(x,, V, 1]]]]
(E ((a N(a(t),, [[[ a' log (z'(x)) s( (,, C..
(3.6)
where V log (x) z'=0(t)z where (2'(x)) is the density of (o(()),,((t)T))) with 1(0)
a(t) exp(-t/2) and ((t) = exx(
Equation (3.6) allows an efficient implementation, since Pag can be approximated by the empirical data distribution in D and x' is easy to sample. Integrating (3.6) over time t leads to practical conditional score matching object
1' B.p.n) B. Es Eor [[| log (x'(x)) - s(x', y. t]1] dt,
argmin #ES
(3.7)
where ln > is an early-stopping time to stabilize the training [44, 50) and È. denotes the empirical distribution.
5

---

Constructing a function class S adaptive to data structure is beneficial for learning the conditional score. In the same spirit of [8], we propose the network architecture (see Figure 2(c) for an illustration):
1 (V- VVV x, =): VE Pxd + E V. R x [4o, T7 Rd S = (sN = h((t)"
(3.8)
with V being any D d matirx with orthonormal columns and + customizable function class. This design has linear encoder-decoder structure, catering for the latent subspace structure in data. Also Moj" J is includes shortcut connection.
Conditioned Generation. Sampling from the model is realized by running discretized backward process with step size 1) > 0 described as follows:
1 - + 3(XMn= kn y. T kn) dt + dw, for t [KT], (k + 1)n]. (3.3 revisited)
dx; x
2
initialized with X ~ N(O, ID) and W. is a reversed Wiener process. Note that in (3.3), the unknown conditional score (2.(2))) is substituted by 3(r,
4
Main Theory
In this section, we analyze the conditional generation process specified by Algorithm We will focus on the scenario where samples admit low-dimensional subspace representation, stated as the following assumption.
Assumption 4.1. Data sampling distribution Pa is supported low-dimensional linear subspace, ic., I = Az for an unknown Dxd with orthonormal columns and E IR is a latent variable.
Note that our setup covers the full- dimennionn setting as special case when d = D. Yet the case of d < D is much more interesting, as practical datasets are rich in intrinsic geometric structures [13, 37, 48]. Furthermore, the representation matrix A may encode critical constraints on the generated data. For example, in protein design, the generated samples need to be similar to natural proteins and abide rules of biology, otherwise they easily fail to stay stable, leading to substantial reward decay. In those applications, off-support data may be risky and suffer from a large degradation of rewards, which we model using function / as follows.
Assumption 4.2 The ground truth reward (x)x g (mi) h²" (x1),, where g g0(x) ) = (0*)T>)
where 6* = A for some 8 E Rd and [0*|| ||2 = || 1| = and h" (x1) is non-decreasing in terms
of || laa with h" (0) O.
Assumption 4.2 adopts simple linear reward model for ease of presentation. In this case, we estimate 8 by ridge regression, and (3.1) in Algorithm becomes argming Σen(0 x; - y1)² + A| for positive coefficient X. Later in Section 3.3 and Appendix F, we extend our results beyond linear models to deep ReLU networks.
(8T
4.1 Conditional DM Learns Subspace Representation
Recall that Algorithm has two outputs: generated population P--10 a) and learned representation matrix V. Use notation P., = P(-lğ = a) (generated distribution) and P. = P(-1ğ = a) (target distribution) for better clarity in result presentation. To assess the quality of subspace learning, utilize two metrics defined as
<(V = I - AATIF and E E [|]]]
(4.1)
<VV A) is defined for matrices V. A, where A is the matrix encoding the ground truth subspace. Clearly, ((V, A) measures the difference in the column span of V and A, which is also known as the subspace angle. E E. [||=ا|]] || is defined the expected distance between and the true subspace. Theorem 4.5 provides guarantees on this two metrics under following assumptions, proof and Interpretation of Theorem 4.5 are deferred to Appendix D.2.
To ease the presentation, we consider a Gaussian design v, i.e. the latent is Gaussian as stated in Assumption 4.4. Since our guarantee on <V, A) holds under milder assumption than Gaussian, we also list the Assumption 4.3.

---

Assumption 4.3. The latent variable follows distribution Ps with density Ps. such that there exists
constants B, C1 C2 verifying (z) < (2T) ( exp (-C2ll=11//2) whenever || > B. And ground truth score is realizable: V log Pr (z ?7 E S.
Assumption 4.4 Further assume
N(0, Σ) with its covariance matrix Σ satisfying Amin la z
Σ 3 La for 0 < Amin Amax - 1.
Theorem 4.5 (Subspace Fidelity of Generated Data). Under Assumption 4.1, if Assumption 4.3 holds with cola a 3 Ex [22] J. then with high probability on data,
N(S, 1/mi)D
A(V, =
(4.2)
with ((S, 1 1/mi) being the log covering number of function class S as in (3.8). When Assumption
4.4 holds, ((S,1/n)) = ((d - Dd) log( (D)))) and thus <(V, A) 0 Further under Assumption 4.2, holds that
Dd
D
Auma
2ー Pa - [لاا o (ViDD (VooD+VZ(V.A
a² ||3+
+
(4.3)
where 3 is groundtruth parameter of linear model.
4.2 Provable Reward Improvement via Conditional Generation
Let y" be a target reward value and P be generated distribution. Define the suboptimality of P SubOpt(P: = y² Ee-ppff(x))]
which measures the gap between the expected reward of x and the target value y". In the language of bandit learning, this gap also be viewed derive guarantees for SubOpt(Pa:
form of off-policy regret. Given
target value y" a, we denotes the generated distribution. In Theorem 4.6, we show SubOpt(P., (Pa: = a) comprises of three components: off- off-policy bandit regret which comes from the estimation error of f. on-support and off-support errors coming from approximating conditional distributions with diffusion.
a). recall P. = P(-ğ a)
Theorem 4.6 (Off-policy Regret of Generated Samples). Suppose Assumption 4.1, 4.2 and 4.4 hold. We choose À = 1, to ((Dr + D2d)/n)) and I/VD. With high probability, running Algorithm with a target reward value a gives rise to
BubOpt(Pa:i' =a)
V IIPa) Ep.)) o o
(Valog
+ Ep. (21)) EEp 19" (=1))
Ep. 1h (#1))
122
E2:on-support diffusiio error
Ex support diffissio emror
Es :nff-poolyy handit regret
(4.4)
where Σx:= = (X x AT) where X is the stack matrix of Diabel and Σp, = Ep, [ax_]]
Implications and Discussions: (1) Equation (4.4) decomposes the suboptimality gap into two
separate parts of error: error from reward learning (E1) and error coming from diffusion (Ez and Ex).
(2) It is also worth mentioning that Eg and Eg depend on to and that by taking to (DD&+ + DDd)/n.)) 1/6 one gets good trade-off and small Eg Ez.
(3) & suggests diffusion model essentially doing representation learning, reducing to smaller latent dimension d. It can be seen from Tr Ep.) "p,,) O
a² + +d) when m2 Ω )
(3) If we ignore the diffusion errors when ni is large enough, the suboptimatliy gap resembles the suboptimatliy of off off-policy bandit learning in feature subspace [22, Section 3.2), [32, 5]. It shows
7

---

the major source of error occurs when moving towards the target distributions and the error behaves similar bandit.
(4) On-support diffusion error entangles with distribution shift in complicated ways. We show
8 (DistroShift(a) ( + DDd) 1/6 -1/6 .)).
where (istroShift(a) quantifies the distribution shift depending different reward values. In the special case of the latent covariance matrix is known, we can quantify the distribution shift as DistroShift(a) O(a V d). We observe an interesting phase shift. When d, the training data have a sufficient coverage with respect to the generated distribution Pa- Therefore, the on-support diffusion error has lenient linear dependence
However, when d, the data coverage is very poor and Eg becomes quadratic in a, which quickly amplifies.
(5) When generated samples deviate away from the latent space, the reward may substantially degrade as determined by the nature of h.
To the authors' best knowledge, this is a first theoretical attempt to understand reward improvement of conditional diffusion. These results imply potential connection between diffusion theory and off-policy bandit learning, which is interesting for more future research. See proof's in Appendix D.3.
4.3 Extension to Nonparametrie Function Class
Our theoretical analysis, in its full generality, extends to using general nonparametrie function approximation for both the reward and score functions. To keep our paper succinct, we refer to Appendix F and Theorem E4 for details of our nonparametric theory for reward-conditioned generation. Informally, the regret of generated samples is bounded by
SubOpt(Pa:g' = = = (000=00(tstann Diistroshif(()) 1122
)) + Ep, [h" (x (x1)
2o+a
with high probability. Additionally, the nonparamtric generators is able to estimate the representation
matrix A up to error of A)V Ö(n オंं Here the score is assumed to be Lipschitz continuous and a the smoothness parameter the reward function, and DistroShift(a) is class-restricted distribution shift measure. Our results on nonparametric function approximation covers the use of deep ReLU networks as special cases.
5 Numerical Experiments
5.1 Simulation
We first perform the numerical simulation of Algorithm following the setup in Assumption 4. 1, 4.2 and 4.4. We choose d = 16, D = 64, gxx) == 51|m/|33 and generate 3* by uniformly sampling from the unit sphere. The latent variable 2 generated from N(O, la), which is then used construct Az with some randomly generated orthonormal matrix A. We use the 1-dimensional version of UNet [40] to approximate the score function. More details are deferred to Appendix I.
Figure 3 shows the average reward of the generated samples under different target reward values. We also plot the distribution shift and off-support deviation in terms of the 2-norm distance from the support. For small target reward values, the generation average reward almost scales linearly with the target value, which is consistent with the theory the distribution shift remains small for these target values. The generation reward begins to decrease we further increase the target reward value, and the reason two fold. Firstly, the off-support deviation of the generated samples becomes large in this case, which prevents the generation reward from further going up. Secondly, the distribution shift increases rapidly as we further increase the target value, making the theoretical guarantee no longer valid. In Figure 4, we show the distribution of the rewards in the generated samples. As we increase the target reward values, the generation rewards become less concentrated and are shifted to the left of the target value, which is also due to the distribution shift and off-support deviation.

---

0.25
0.50
Swweed
0.2
0.20
2.5 Target Roward
Toget
Targot Volae
Figure 3: Quality of generated samples 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.
Begget wase
Terpet Seaard Ware
Terget Smaurd Wiiee
Thrpet Seward
ST Eze
20
200
2.5
Figure Shifting reward distribution of the generated population.
5.2 Directed Text-to-Image Generation
Next, we empirically verify our theory through directed text-to-image generation. Instead training diffusion model from scratch, Stable Diffusion operates on the latent space its Variational Auto-Encoder and can incorporate text conditions. We show that by training a reward model we can further guide the Stable Diffusion model to generate images of desired properties.
Stable Diffusion v1.5 [39], pre-trained LAION dataset [42].
Ground-truth Reward Model. We start from ImageNet [10] pre-trained ResNet- 18 [17] model and replace the final prediction layer with a randomly initialized linear layer of scalar outputs. Then this model as the ground trouhd-truth reward model. To investigate the meaning of this randomly- generated reward model, we generate random samples and manually inspect the images with high rewards and low rewards. The ground-truth reward model seems to favor colorful and vivid natural scenes against monochrome and dull images; see Appendix for sample images.
Labelled Dataset. We use the ground-truth reward model to compute scalar output for each instance in the CIFAR-10 [24] training dataset and perturb the output by adding a Gaussian noise from (0, 0.01). We use the images and the corresponding outputs the training dataset.
Reward-network Training. To avoid adding additional input to the diffusion model and tuning the new parameters, we introduce a new network and approximate ()(()r)) by (x), of). For simplicity, set o² as tunable hyperparameter, We share network parameters for different noise levels so our Ho has no additional input of t. We train 44 by minimizing the expected KI. divergence between Pe (y(lx) and N(Ho(x),
Ily no (xx)11
E,E. (p.(a) (( ) (x..))
Constant.
20
Equivalently, train the reward model μą to predict the noisy reward y from the noisy inputs æt Also, notice that the minimizers of the objective do not depend on the choice of o².
Reward-network-basse Directed Diffusion. To perform reward-directed conditional diffusion, observe that V_ log (z(lv) = Va (x) V. log pe (y)z), and (x(l)rr) exp 2n²
Therefore,
V logp((y/>))) = = 11/2 V
u (z)
our implementation, we compute the gradient by pack-propagation through μ4 and incorporate this gradient guidance into each denoising step of the DDIM sampler [43] following [11 (equation (14)). We see that 1/o² corresponds to the weights of the gradient with respect to unconditioned score. In the sequel, we refer to 1/o² the "guidance level", and y the "target value".

---

Quantitative Results. We vary 1/0² in {25, 50, 100, 200, 400} and in (1, 2, 4, 16). For each combination, we generate 100 images with the text prompt "A nice photo" and calculate the mean and the standard variation of the predicted rewards and the ground-truth rewards. The results are plotted in Figure 5. From the plot, we see similar effects of increasing the target value y at different guidance levels 1 1 A larger target value puts more weight on the guidance signals V (z)) which successfully drives the generated images towards higher predicted rewards, but suffers more from the distribution-shii effects between the training distribution and the reward-conditiones distribution, which renders larger gaps between the predicted rewards and the ground-truth rewards. To optimally choose a target value, must trade off between the two counteractive effects.
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 the expected ground-truth reward for undirected generations.
Figure 5:
Qualitative Results. To qualitatively test the effects of the reward conditioning, we generate a set of images with increasing target values under different text prompts and investigate the visual properties of the produced images. We isolate the effect reward conditioning by fixing all the randomness during the generation processes, so the generated images have similar semantic layouts. After syper-parameter tuning, we find that setting 1/o" = 100 and E {2, 4, 6, 8, 10} achieves good results across different text prompts and random seeds. We pick out typical examples and summarized the results in Figure 6, which demonstrates that as we increase the target value, the generated images become more colorful at the expense of degradations of the image qualities.
improves target increases
Quality degrades overly extrapolated
Figure 6:
The effects of the reward-directed diffusion. Increasing the target value directs the images to be more colorful and vivid at the cost of degradation of the image qualities. Leftmost: without reward conditioning. Second-to-Last target value y = 2, 4, 6, 8, 10. The guidance level 1/o² is fixed to 100. The text prompts are "A cat with glass of water.", ", "An astronaut the horseback".
6 Conclusion
In the paper, we study the problem of generating high-reward and high-quality samples using reward-directed conditional diffusion models, focusing on the semi-supervised setting where massive unlabeled data and limited labeled data are given. We provide theoretical results for subspace recovery and reward improvement, demonstrating the trade-off between the strength of the reward target and the distribution shift. Numerical results support theory well.
10

---

Acknowledgments and Disclosure of Funding
Mengdi Wang acknowledges the support by NSF grants CPS-2312093, DMS-1953686, IIS-2107304, CMMI1653435, ONR grant 1006977 and C3.AI.
References
[1] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023. URL https: httpp://oppeerrrreee openrevien....e... net//orun?Id=sPPfooo0 wwt.idddpPPPPoPeeeeeaeee
[2] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 13(3):31--326, 1982.
[3] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17 17811999,,,99,, 2021.
[4] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, et al. ediffi: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01322 2022.
[5] David Brandfonbrener, William Whitney, Rajesh Ranganath, and Joan Bruna. Offline contextual bandits with overparameterized models. In International Conference on Machine Learning. pages 1059-1058. PMLR, 2021.
[6] Jinglin Chen and Nan Jiang. cnoommtiiontheorette considerations in batch reinforcement learning. In International Conference on Machine Learning, pages 1042-1051. PMLR, 2019,
[7] Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Ad- vances in Neural Information Processing Systems, 11-22134-22145. 2020.
[8] Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estima- tion and distribution recovery of diffusion models on low-dimensional data. arXiv preprint arXiv:2302.07194 2023.
[9] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint irXiv:2209.1121: 2022.
[10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and patlern recognition, pages 248-255. Icce, 2009.
[11] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[12] Jianqing Fan, Zhaoran Wang, Yuchen Xie, and Zhuoran Yang. A theoretical analysis of deep q-learning. In Learning for Dynamics and Control, pages 486-489. PMLR, 2020.
[13] Sixue Gong, Vishnu Naresh Boddeti, and Anil K. Jain. On the intrinsic dimensionality image representations. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 3987-3996, 2019.
[14] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusson models plug-and- priors. arXiv preprint arXiv arXiv:2206.09012. 2022.
[15] László Györfi, Michael Köhler, Adam Krzyžak, and Harro Walk. A distribution-free theory of nonparametric regression, volume 1. Springer, 2002.
[16] Ulrich Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, pages 1188-1205, 1986.

---

[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[18] Jonathan Ho and Tim Salimans. arXiv:2207.1223 12598, 2022.
Classifier-free diffusion guidance.
arXiv preprini
[19] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 6851, 2020.
[20] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4). 2005.
[21] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[22] Ying Jin, Zhuoran Yang. and Zhaoran Wang. pessimism provably efficient for offline ri? In International Conference Machine Learning, pages 5084 5096. PMLR, 2021.
[23] Durk Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018.
[24] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
(25] Holden Lee, Jianfeng Lu,, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory, pages 946-985. PMLR, 2023.
[26] Jin Sub Lee, Jisun Kim, and Philip M. Kim. Proteinsgm: Score-based generative modeling for de novo protein design. bioRxiv, 2023. doi: 10.1 13210001222200.13449999666 URL https: / //www. biorxiv. wwweiiioorr///ooooo (ow.tenn/eaaal//////////////////
499967.
[27] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing Systems, 35:4328 4343, 2022.
[28] Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdif- fuser: Diffusion models as adaptive self-evolving planners. arXiv preprint a7X11:2302.01877, 2023.
[29] Qiang Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou Breaking the curse horizon: Infinite-horizon off-policy estimation. Advances in Neural Information Processing Systems, 31, 2018.
[30] Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5), 2008
[31] Ryumei Nakada and Masaaki Imaizumi. Adaptive approximation and generalization of deep neural network with intrinsic dimensionality. The Journal of Machine Learning Research, 21 (1):7018-7055, 2020.
[32] Thanh Nguyen-Tang, Sunil Gupta, A Tuan Nguyen, and Svetha Venkatesh. Offline neural con- textual bandits: Pessimism, optimization and generalization. arXiv preprint arXiv: /3807, 2021.
[33] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint 12. 2021.
[34] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In ICI. 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023, URL https: /ttp:///eppeeeeeee/////eee wttt:/oorumm110000611111111.
12

---

[35] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida. Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730-27744, 2022.
[36] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam De- vlin. Imitating human behaviour with diffusion models. In The Eleventh International Con- ference on Learning Representations, 2023, URL https://openreview..ne net/ forum?id= Pv1GPQZRrC8.
[37] Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning, arXiv preprint arXiv:2104.08894 2021.
[38] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
[39] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[40] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCA. 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
[41] Timo Schick, Sahana Udupa, and Hinrich Schütze. Self-diagnosis and self-debiasing: A proposal for reducing corpus-based bias in nlp. Computing Research Repository, arXiv:2103.00453, 2021. URL httt::ttta/iiiii org/abs/ absp/1as./2.../ 00453.
[42] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et Laion- 5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402 2022.
[43] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021. URL
[44] Yang Song and Stefano Ermon. Improved techniques for training score- generative models. Advances in neural information processing systems, 33:12438- 2020.
[45] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pages 584, PMLR, 2020.
[46] Yang Song, Jascha Sohl-Dickstein Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
[47] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01466) 2023.
[48] Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, 199055500:2319922222 2000.
[49] Alexandre B. Tsybakov. Introduction to Nonparametric Estimation. Springer Publishing
Company, Incorporated, 1st edition, 2008. ISBN 0387790519.
[50] Arash Vahdat, Karsten Kreis, and Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:11287-11300 2021.
13

---

[51] Roman Vershynin. High-dimennsooaal probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
[52] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674, 2011.
[53] Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge university press, 2019.
[54] Lymin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543 2023.
14

---

Index of Supplementary Materials
SA
Preliminaries on Diffusion Models
B
Limitations and Open Problems
SC
Omitted Proof in Section 3
SC.1
Proof of Proposition 3.1
SD
Omitted Proofs in Section 4
$D.1
Parametric Conditional Score Matching Error
$D.2
Proof of Theorem 4.5
$D.3
Proof of Theorem 4.6
SB
Supporting Lemmas and Proofs for $D
$F
Extensions to Nonparametric Function Class
SG
Omitted Proofs in $F
SH
Proof for $D.1
SI
Additional Experimental Results
A Preliminaries
A.1 Diffusion Models
We first provide a brief review of diffusion models and its training/sampling procedure. We consider diffusion in continuous time [23, 46], where diffusion is described as forward and backward SDEs.
Forward SDE and Score Matching. In the forward process, noise is added to original data progressively as an Ornstein-Ulhenbeel process for instance:
dX, 2 (t) Vo(tjdW, for g(t) 0,
(A.1)
where initial Xo Paata and (Wr) is standard Wiener process, and g(t) is nondecreasing weighting function. practice, the forward process (A.1) terminates at sufficiently large T such that the corrupted X7 is close to the standard Gaussian N(O, /p). To enable data generation in future, the score logpe( (r(-) at f the key to learn, here pt denotes the marginal density X,. We often use estimated score function 8(., trained by minimizing score matching loss.
Backward SDE for Generation. Diffusion models generate samples through backward SDE (A.2) reversing the time in (A.1) [2, 16], i.e..
dxx
5 g - + g(T t) Po dt + Vg(T t)dW,.
(A.2)
where W, a reversed Wiener process. In practice, the backward process initialized with N(0, TD) and the unknown conditional score V log Pc(-) is replaced by estimated counterpart 3(-
B Limitations
We do not outstanding limitations in our analysis. The linear subspace assumption initiates the study of conditional diffusion models on low-dimensional data. We expect to stimulate more sophisticated analyses under general assumptions such as manifold data.
C Omitted Proof in Section 3
C.1 Proof of Proposition 3.1
Proof. For any 2 0, it hold that Vz. log (x y) = Vz, log y) since the gradient is taken
w.r.t. 2, only. Then plugging this equation and expanding the norm square on the LHS gives
E( [[/V,, log Pr y) s(x,, y, [| (x.,,)~ [[] y, 1)
2(V log (z v), s(1, t)) +
15

---

Then it suffices to prove ((Faa)- [(Vx, log Pe y), s(12, y, t))] ((x.)))P.; No00:
s(x', , t)
Using integration by parts to rewrite the inner product we have
(x <V ( v). s(1, y, 1))] = / pe( log Pe v), s vs d)d
/ (V 1r, v), s( U. t)) dmy
/ p((z,,)) v) div(s( y, t))dd,dy,
where denote by (x((2)) the density of N(a(t)=, with a(t) = exp(-t/2) and h(t) =
1 - exp(-tt (t), then
-f pe v) div(s((aaa)))) d,tdddd ((r.)) / dr x) div(s(d', , d)
= (B(r (V ((, y, dz)
(a.w)) Baa (V |x), s(a', y, (),)
D Omitted Proofs in Section
Additional Notations: We follow the notations in the main paper along with some additional ones. Use P.. (z) to denote the low-dimensional distribution on z corrupted by diffusion noise. Formally, p! (z) f oi (2' (z)d with (+(+)2 being the density of ((a(t)2, h(t)La). P. (z f(Az) = a) the corresponding conditional distribution on f(Az) = a at to, with shorthand as Pb (a). Also give Ps(= f(Az) a)) shorthand as PLD((), In theorems, O hides constant factors and higher order terms in n and nz and, Ô further hides logarithmic terms and can also hide factors in d.
D.I Parametric Conditional Score Matching Error
Theorems presented in Section 4 are established upon the conditional score estimation error, which has been studied in [8] for general distributions, but in Lemma D.I we provide a new one specific our setting where the true score is linear in input (xx,, 7) due to the Gaussian design. Despite the linearity of score in Gaussian case, we emphasize matching score in (3.2) is not simply linear regression S consists of encoder-decoder structure for estimating matrix A to reduce dimension (see ŞH for S construction and more proof details).
In the following lemma, first present general result for the case where the true score is within S, which is constructed as a parametric function class. Then the score matching error is bounded in terms of A/SS 1/n1), the log covering number of S, recall n; is the size of Dunlabel- Instantiating this general result, we derive score matching error for Gaussian case by upper bounding NSS, 1 1/n1) in this special case.
Lemma D.1. Under Assumption 4.1, if log | v) E S, where
S { (v.0(z, h(t) niee (V (VV x,y,t) z): VerRxe WE
Rd [to: T Bd
(3.8) revisited)
with & parametric. Then for 2 O, with probability 1-8, 8, the square score matching error is bounded A(S.1/n, (d n1
by cd.
0
i.e.,
1
Lo
B [[[ (ogg)())))))) 3(2,, y, y, 2) Eir
(D.1)
T to
16

---

recall P. comes from Pag by noising z at t in the forward process. Under Assumption 4.4 and given f(x) d% æ and ÿ = f(x) + E. N(O, v²), the score function (r ()) to approximate is linear in and y. When approximated by S with v linear, 1/m1) = O((d² Dd) log( Ddn1)).
Proof. Proof is in SH.
To provide fidelity and reward guarantees of Pa: the generated distribution of æ given condition ÿ = a, we will need the following lemma. It provides a subspace recovery guarantee between (score matching output) and A(ground truth), as well as a distance measure between distributions Pa and Pa. given score matching error édiff-
Note Pa and Pa are over z, which admits an underlying low-dimensional structure = A2. Thus we measure distance between Pa and Pa by defining
Definition D.2. (P.) drv (PLD( (= [(Az) a), (UTVT) )
with notations:
.rv(:.) is the TV distance between two distribution.
f&P denotes push-forward measure, i.e., for any measurable Q (SAP) = P()-1()))
(VT)+ Pa pushes generated Pa forward the low dimensional subspace using learned subspace matrix V. U an orthonormal matrix dimension d.
to ( f(Az) = a) is close to (A # Pa. with to taking account for the early stopping in backward process.
We note that there is a distribution shift between the training and the generated data, which has profound impact on the generative performance. We quantify the influence of distribution shift by the following class restricted divergence measure.
Definition D.3. Distribution shift between two arbitrary distributions Pi and P2 restricted under function class L is defined
T(P, P2; C) supiec Ee Bet/ with arbitrary two distributions PI, Pa-
Definition D.3 is well perceived bandit and RL literature [30, 29, 6, 12].
Lemma D.4. Given the square score matching error (D.1) upper bounded by e and when P. satisfying Assumption 4.3 with cola I EsnP. xzz T] it guarantees on for Pa and Z(V, A) ||VVT AAT 1/3 that
(TD - VV jx N(O,A), A < dOOD,
(D.2)
<(V, = δ ( to
Ed:11
(D.3)
co
In addition,
TV(P.)=G (1 T(P(x. = a), Pasj; S)
cails)
(D.4)
Co
with S T-
1 S Exgla V Px (v) s(x,, w. ( 1d s TV(P.) and T(P(a,
a), Pag: 5) are defined in Definition D.2 and D.3.
Proof. Proof is in SE.2
17

---

D.2 Proof of Theorem 4.5
Proof. Proof of <(V, By Lemma [8], we have
<V, A) O to
Cass)
co
when the latent satisfying Assumption 4.3 and cola I Ezup, [=z [= with high probability that
Therefore, by (D.1), we have
1 co
N(S.
<(V δ
71
When Assumption 4.4 holds, plugging in Amin and MndN(65 ((S,1/ni) = ((d + Dd) log log(Ddn))). gives
1
A) ō
Amnin
71
where Ô hides logarithmic terms. When D > d², which is often the case in practical applications. we have
1
Dd + D2d
A(V,) =
Amin
n1
Proof of E E P. [||a1llla
[12]. By the definition of æ1 that x1 = -ID - AAT)a
E E |||x [|2] E. [II AA )x1aa1121 /B IlI(ID )
Score matching returns V approximation of A, then
[1 - AA >zll2 & li - VVT)zlla || (VV - AAT )zl(2. AAT)zll2.
E R P. [II[DD AA AAT)=1 |= < 2E P -V
)= 2 [II[VV -A )=1))
where by (D.2) in Lemma D.4 we have
(ID - VV 1 Dz
A ct
for some constant e 2 0. Thus
E EapP, [1[(Tn- (((ID-VT]])) ) Tr(A)= s doD.
(D.5)
On the other hand,
II(VV AA )ar S IIVV - AA 131111133 S V AA AAT 11
where IVVT AAT has upper bound as in (D.3) and E E. A. [|| (2]] is bounded in Lemma E.3 by
[|]:11]]] = -o OdoD +(DD+M(a) (1 (1+
e P.
with ((a) = 0
1 + d
|89
Therefore, to combine things together, have
E 22 p |||a-|1l1
$ V 2E an [1l(1 VVT) + 2E P (VVV AA AAT)=11
< <d' D 2(VV(V,
E BB ||
= (o(v/
(Vto to VZ(V, V
O hides multiplicative constant and /Z(V,A)toD, V ((V which are terms with higher power of n than the leading term.
18

---

Remark of Theorem 4.5, 1. Guarantee (4.2) applies to general distributions with light tail as assumed Assumption 4.3.
Guarantee (4.3) guarantees high fidelity of generated data terms of staying in the subspace when have access to large unlabeled dataset.
Guarantee (4.3) shows that P, [||=] [2] scales up when to goes up, which aligns with the dynamic backward process that samples are concentrating to the learned subspace as to goes
Taking lo - p ||| i III has the decay in O(n ). However, taking - is not ideal for the sake of high reward of z, take the best trade-off of to later in Theorem 4.6.
D.3 Proof of Theorem 4.6
Proof of Theorem 4.6 and that of some results in "Implications and Discussions" following the theorem in main paper are provided in this section. This section breaks down into three parts: Suboptimality Decomposition, Bounding & Relating to Offline Bandits, Bounding Eg and the Distribution Shift in Diffusion.
D.3.1 Subopt(Payy* a) Decomposition
Proof. Recall notations P. P(-lū = a) (generated distribution) and P. := P(-5 a) (target distribution) and f (x) = g" (z1) h- (z1). B [x)( can be decomposed into 3 terms:
E, (x)) P. P. | (x) BB F. [f (x)] - Ez (x)]
EkwP P. [f(x)] - ExnPs || f(x)) - (x)
E (f)] E. [f
> ERap [xx]] - Earra [[(x] - /² (x)
Ex 19" (211)] - E P. lg" [2])]]
E BaNP, [h (x1)),
Es
where E.,P, [f(x)] = Bangla) and we use 211- (xx) g" (z) when ~ Pa. Therefore
SubOpt(Pas a) B Bew P. [f)(
EypP
( 0-) T
+ EawP, [0" (x11) Raw,, P. 10" (21)))
zm [1'((1))
ε comes from regression: preddicion/geneeraizztioo error onto Pa, which is independent from any error of distribution estimation that occurs in diffuuionn. Eg and &3 do not measure regression-predictts f. thus they are independent from the prediction error in J for pseudo-labeling. E2 measures the disparity between Pa and Pa on the subspace support and &g measures the off-subspace component in generated P...
D.3.2 Bounding & Relating to Offline Bandits
For all x, Dlabel, V1 (xx) + = g(x;) + eg. Thus, trained Diabel the prediction model , is essentially approximating g. By estimating 0* with ridge regression on Diabels we have f(xx) = 0 with
= (XT x + X1) T (XO- + m) where X =(21,-.. Di, Fna) and (E1, Ess ... En
(D.6)
19

---

Lemma D.5. Under Assumption 4.1 and 4.2 and given ~ N(O, define Va = XI, Σx:= = Vi and Σr. Ee rx' the covariance matrix (uncentered) of Pa. and take 1, then with high probability
& VIEE )pp)
O (Vdlogn
(D.7)
V02
Proof. Proof is in SE.3.
Lemma D.6. Under Assumption 4.1, 4.2 and 4.4, when X 1, Pa has shift from the empirical marginal of x in dataset by
Tr Ep.) 50 a²
||3 + d)
(D.8)
when n2 {m
}).
Amtn
Proof. Proof is in $E.4.
D.3.3 Bounding & and the Distribution Shift in Diffusion Lemma D.7. Under Assumption 4.1, 4.2 and 4.4, when 4g = =(Dee + DDd)/nn) 1/6
T(P(x, x = Pagi
Did2 + D2d à
Eg =
Amin
Proof. Proof is in $E.5.
Note that T(P(a, a), P2 depends and measures the distribution shift between the desired distribution Pvx, a) and the data distribution Pag. To understand this distribution's dependency on a, what follows we give ((P(x, T = a), Pag; S) shorthand as DistroShift and give it upper bound in one special case of the problem.
Distribution Shift In the special case of covariance Σ of = is known and || V O (IAAA - -V l12), we showcase a bound on the distribution shift in Eg, as promised in the discussion following Theorem 4.6. We have
Ep 9=n [1(z, v] Ep. [0(x, y: 8)]
(istroShift?())
where ((z, 9: = T-to S E V (z')x()')zz) - $(x', vy Bdd. By Proposition 3.1, it suffices to bound
Ep 6=a 1 V logpt 8(x, v, 13
Epr
DistroShift'(a)
Ep, U
V log (r. 5(x v. 1) 133d
We expand the difference
V (or, v) 8(x., y, 111 by
| log v) ((x.,,t)ll -
[I(A - V)B x + レー 2 || B. V) V2Tx2 x013
(t)
n [|A - V -
+ I| v
|| V1 +2))
12
where we recall B, is defined in (H.2) and in the last inequality, we use (a + b)2 < 2n² + 2b2. In the case of covariance matrix Σ is known, i.e., B, is known, we also consider matrix V directly matches A without rotation. Then by [8, Lemma 3 and 17], we have | 3 =O (IIAA VV [2:
20

---

O ( To this end, we only need to find Ep [|]=1]] Since we consider on-support which can be represented find the conditional distribution of z = a. Fortunately, we know (z. y) jointly Gaussian, with mean 0 and covariance
Az, we have ||o|| = ||=||2- Thus, we only need to
Σ E
3
E3 22
Consequently, the conditional distribution of 2 is still Gaussian, with mean _3 Σ3 v²) and covariance E ΣB8 Σ/(3 3B + v²) Hence, we have
E E =| E3 + 22)² (ia² - Σ3 -")3"" + TMEE= =
We integrate over t for the numerator in DistroShift(a) to obtain Eaggee U | log y)
(ar, y, (2de o ((a² Ep. v:)) Note the cancellation between the numerator and denominator, we conclude
DistroShift(o)- o (
As d is natural upper bound of vd and viewing cp as constant, we have DistroShift(a) O (a V d) as desired.
E
Supporting lemmas and proofs
E.1 Supporting lemmas
Lemma E.1. The estimated subspace satisfies
|VU A O (då /Z(,A)
(E.1)
for some orthogonal matrix U E Rdd
Proof. Proof is in $E.6.
Lemma E.2. Suppose P1 and P2 are two distributions over Id and is function defined R4 then E,,,, [m[z)] Ex [m (z) can be bounded in terms of drv P2), specifically when Pi and P2 are Gaussians and m 2)
Eg ||| 1 O (BawPa drv
(E.2)
When P1 and Pz are Gaussians and m²(2) = ||
Be [|2]12] Es e = V
1||=1111 E
v 1l.=111)
drv (P1, Pb)
(F.3)
Proof. Proof is in $E.7.
Lemma E.3. We compute R. pm (a) Rx. [1) F, [||-|12), B B
[|]||]]
in this Lemma.
(1 + v2 za² + trace Σ3 3 53 33 + u²) -1
3 32
E [
IT
(E.4)
Let M(a) = Exppan (a) |1| which has an upper bound M(a) = 0 ( +d).
E By~ (4) || < M(a) + tod.
(E.5)
の~ ||| s o (cloD +
-1 - (V.))) (1)))())))
(E.6)
Proof. Proof is in SE.8.
21

---

E.2 Proof of Lemma D.4
Proof. The first two assertions (D.2) and (D.3) are consequences of [8, Theorem item and 3]. To show (D.4), we first have the conditional score matching eтror under distribution shift being T(P a): Pag: Cass
where T( a), Pag: 8) accounts for the distribution shift in the parametric (Lemma D.7). Then we apply 18, Theorem 3, item 2] to conclude
TV(P.) = (v
T(P(z, M = Payi S)
Edifs
00
The proof is complete.
E.3 Proof of Lemma D.5
Proof. Given
E = Ep,
x T00 -0) s Ep, læll,, 1: -100 - -1llva
then things to prove
E |sclll, (Iaace(VV'Eppp
(E.7)
|10 0llv, S 0 (ooVVlay Vdllgg n2
(E.8)
where the second inequality is to proven with high probability war. the randomness in Diubel. For
(E.7), Ep, Il. S B.p x -v BBp. (V
trace(V xx" T).
For (E.8), what's new to prove compared to classic bandit derivation is its d dependency instead of D, due to the linear subspace structure in r. From the closed form solution of 0, we have
-V-1x AV
(E.9)
Therefore,
|10 Blvs S ||x vllvy- -110 11-
(E.10)
where 1|10"||v-'
vx11011a
2 S V and
|| nli = 17 X (x x AT)) x
= = (XX + (La2) n7-
Let Z = (Zi, "., Znz) s.l. Az; = Ii, then it holds that X = ZAT, and XX
ZAT AZ T = ZZ thus
|| n/ 1ー- = = XXT (XXT + +1m)))
= =n 77. (77 + +Nnn))
-1) 2 2 Ala)
11 T1ZTlles / 2
With probability 1 8, [ d /dlog (2m²) == L Vi E [n2]. Then Theorem I "Improved algorithms for linear stochastic bandits" (by Yasin Abbasi-Yadkony David Pal, and Csaba Szepesvari) gives rise to
1 -+ / (og((2/)) + d log(1 n12 1.2 //2)((
with probability - 5/2. Combine things together and plugging in À = 1, I2 = d + dlog (2m2),
we have with high probability
|100 -00o =o (Valos (12 (navlog(na))
(V dlog n2 + - _dlo((logns)
1
- =o(Vdlogg
22

---

E,4 Proof of Lemma D.6
Proof. Recall the definition of Σx and Ep, that
Σx
1
X
A
Ip:
n2
112
Ep Eza Pa [a:
where X are stack matrix of data supported on A and P. is also supported on A. is the subspace
encoded by matrix A. The following lemma shows it is equivalent to measure x 'Ep.)) on A subspace.
Lemma E.4. For any P.S.D. matrices M1, 23 Rddd and RDdd such that A A La we have
Tr ((Vn AΣ, A ) AEZA Tr (AATa E1) Eg
The lemma above allows abuse notations ΣDx and Ep, in the following way while keeping the same traco(Eū 1E Ep,)) value:
1
Σx
_T Z +
la.
122
712
Ep, (t.Pp.)
[2z
where T (21, 2/r 2n2) s.t. Az; 2; and recall notataion (a) P. (x (x1f(d))).
Given 2 N(,, Σ), as proof artifact, let f(x) θ + N(O, 2²) where we will let in the end, then let 3 = AT E Rd, (z, ((Az) has joint distribution
(z N ( 14
Σ
3 E3++
Σ3
(E.11)
Then we have the conditional distribution f(Az) = a following
Ps (z f(Az) = =a) = |++ Σ3 I + +1²)
(a- (a-3TTw).T)
(E.12)
with I' : -ー 23
3 33 1²) T. E.
When μ = 0, we compute trace E-11E 'Ep,)
Σ 233T Σ
(race(Sx12P) = trace
za²
trace (r-
[1313 + u²)
3 TEΣ 202
5-1 Σ m
+ trace (Σ~1Z) trace
trace
(11813 + v2
|| + n2
x1/2331 31/281/88-1111
trace
a²
[1613 + v2
ΣΣΣ ΣΣ 1 llpp |/
2 + trace
-
1 + v/2
By the Lemma in [7], holds that
| - -all2 50
(
(E.13)
23

---

Therefore,
trace (Ex -Ep))
1Eppp
1 Vimin
a² +0(a(
dd
1
||
Vimin
Then, what left is to bound On one hand,
|| = |10 2 11001AAAT - -110 10-00"1 lAAATT by triangle inequality.
||0 | = ||
(E.14) 14)
On the other hand,
100 0" (AAAT O (10 0" 1|ss,
|00 6* llv.
vnz
d dlog(n2)
0
122
with high probability. Thus when n2 22
11 1
|| 2 || 1s
Therefore
trace( ( Ep.)) So
( 1+ Vaminn a + ((1 1t 1
(011
a²
+ +d). (E.15)
||8+|x
11+11
when n2 Qmmax
1 })-
Amiin
E.5 Proof of lemma D.7
Proof. Recall the definition of g(x) that
g(z) = = "AAAA x,
note that g(x) = 8r T x when æ is supported on A. Thus,
Eg [g(x)] - E E [o(=)]
EawP, 10~ T xx] - -E 10 AA x]]
< Eonp, 10 T x] - Bz, 10 VV x]
+ E En 100
VV
x] - E
P 10= T
AAT. I
where
E E, 10 (VVT- AAT) a1
61
S E P. [[[[(VV--A)) x|ll
$ VVV - AA -ATTIE ||r
VEWP [1<:12]]
Use notation PLD (LD(a) = P(= f(Az) = a), PLD(()) = Pto f(Az) = a)
Enw,, 100 x]
B B. P. 10 T VVVxx) 20]
[E 10 Az] - B
10 VU2)
(e)pLo(e)10 T Az] - E V 10 T VU2)
EENPLP 10
Az] TAz]
Ea Az]
24

---

here
(( 10 Az] + ((to)) 10~ Aud (LD) 10 Az]
$ (1 - a(to)) Ez (1 A2]
1 (ap [||=
Then what is left to bound is
E (2)
10 Az] E IUT
J 10 10 VU=]
E BL 10- T VU2) - E SUT
10 (oTTVu2) + Bam VU2l E P (() 10TTAA Aa]
Then for term Eg, by Lemma E.2, we get
z Bep (a) E T |I|=11a||
= o (TV(PP) (g VIe a 11=
where use E 2 VT
P. [||=||2] R P. [||]||]]
For 64, we have
44 (a) (VU -A)=1
= ato) EZnP(a) 10+ (VU -))²]
- VVU A|lF (a.
Therefore, by combining things together, we have
& Se1 + e2 + E3 + 64
VI - AA F BBM + (IIVU Al + 10) V((a)
+ +o (TV(P))) (VM(a) + lad + B 11.x1111111)))
By Lemma D.4 and Lemma E.I, we have
TV(P.)=ā T(P(s, i= a), Pagji S)
caiss
Amin
vlo
I AAT ||
Ediss) (
Vduin
IVU A/I=
-AA ||p) -AATTl
And by Lemma E.3
P P. [|]| = o OcoD ((
Therefore. leading term in Eg is
Eg = = + to) +Mo))
By plugging in score matching error eaesy = θ t/
+ Dad
we have
TV(P.)= = (1
T(P(x, 1 a), Pagi: S)
Ddd2 D2d t
Auin
vio
25

---

When to ((Dd2 DDdd)/n1) 1/6 it admits the best trade off in &g and Eg bounded by
(DiP + D2d à
T(P(x, T = a), Pag;S)
&2 =
a
Amin
E.6 Proof' of Lemma E.1
Proof. From Lemma 17 in [8], we have
| - V AAlA=====- = O - IF).
Then it suffices to bound
IVU AA1 - UU - V A1
where
IIU AlE 2d trace (U V V A + V
IIU V ATl?? d + trace (A VV A) trace V A+ VV)).
Thus
|IVV A IU V AT d trace (A VVV [IVTA|=
trace (AA (VVT AAT))),
which is because trace VV A) is calcualted
trace VV A) trace (AATVV
trace (AAT AAT) ) + trace (AA (VVT (AAT(VVT AAT)) + trace (AAT (VVT- AAT)).
Then will bound (traee (AAT (VVT AAT)) by IIVV AATIEE. F|F
trace (AA - AAT S trace (AAT) trace (IVVT AAT)
s d trace (VVV AA
- s d 22 VVT AATI
Thus, IIVU AllF (18 (ZVVV((
E.7
Proof of Lemma E.2
Proof. When P1 and Pa Gaussian, m(z) ||=1|
B py [m())] E2-rz(m(=)]]
f m(=) (pI(2)
p2())) -P(())))
M 1| (P)(2) p2((2) dz 1 -/alssan =1=1/20 (2)dz
/
1|=1|>>0 / |1=12r2(=)d=
(2l1gp(22)de + / Mallzz 1121@pp(22ddz
< (Pdry((P.P) +
||
Since Pi and P2 are Gaussains, S =|||>> |1:1/pp (z) and Sitlla> || l2 (z)dz are bounded by some constant C₁ when R² C2 max (EZWA, [|| Ry [|| [8].
as suggested by Lemma in
26

---

Therefore,
Eg-Pi |1 - Rz-r3 ||| + C2 max (E,, (a
drv + 2C1
< P3 ||= + C2 ||| + - 111==1]) dIv P2) + 20I.
Then
E P ||
| O
+ Ea || dTv(PI,PP))
since d P2) decays with ni Similarly, when m=)) = ||=1|2
Bg [m()]] Ez [m(2)]
f m(z) (pi(z) - p2(=))
IellaR 1|=112 (p((2) p2(()) dz / /Illssnn 11=122((2)
/ M M 11)12p2(2)d2
< Rdrrv (Pi, + Vheles 11=1B0 +
|| 111)p2(z)dde,
where Jj 12 and S d2 are bounded by some constant C₁ when
R 2 C2 max (E [||=||3] Ex [[]= as suggested by Lemma 16 [8].
Therefore,
B P. [||=||2] Eg 11122112211 S
C2 max Ep, 11 Ep, ||==1|1 dTv P2) 2C
<
V CG2Ez-p111=: + C2Be-py11l133 [
drv P2) + 20
((V (V 11==1|] + ||=
drv Pa)
E.8 Proof of Lemma E.3
Proof. Recall from (E.12) that
PLD (a) = P2(2 f(Az) - a) (u(a),
with p(a)= :- =E (8 EB 1 v2)
=- Σ} ( x2) ITE.
Bewpro((a [||:16] ulaa7 p(aa + trace(T)
3 3 a² + -traceE--
+ x2) (3
-3TE)
(13113 + v2)
=: M(a).
3 TT 523
M(a)= =0
za² + (race(E)
(113133
a²
=0
)
d
1|31 3
27

---

and by Lemma D.6
|| - 2 |1
Thus (L) (a) [||:|0] (M) o ( ( + +d).
Thus after adding diffusion noise at lo, we have for ail) = e/2/ and h(t) e
P/ (a) [|| = (a
||
EsNPLD(a)) aa (1o) 11 + d h(to)]
(1o) E Ppo (a) || + d h((o)
=e E PLo |||اا1|||
(1 e (0) d.
Thus p
(e) [|| S M(a) + lod.
By orthogonal decomposition have
E E P
|| I Ba P. [[[TD VV >>>) + EzNP, [[VV x112]
= E [1([-- vv )=1[ + EzrP. [ V x x112]
vrT)=11 |=|[]] is bounded by (D.5) and the distribution of U V x. which is (T Pa. is close to P (a) up to (P.),, which is defined in Definition D.2. Then by Lemma E.2, we have
where E. [lI(TD VV
B
V x0 = (Bmpge [1[:12] (1+
P
Thus E P. P. |
O (ctoD - (M(a) + tod) (1+TV(P.))
E.9
Proof of Lemma E.4
Proof. Firstly, one verify the following two equations by direct calculation:
(ATD + AE 1
X (TD A(N -
+ -1)+ i (Ta - (Na + -)))))))
Then we have
AE,A AΣ₂A
(In -l--Alla-- + -E))-E,A AΣ2AT
(AΣZAT -A(Aa + En -EEAAT)
Therefore,
Tr ((VD + -
=Ir
{ -AN:A--((M++2))) -((N+++
|
=Tr
1 (Ee (Ala -M)+,,,
=Ir
(Ia (Al + 1aaa-Aaa+
- En
=Tr
(Ala Σn) -1)-
which has finished the proof.
28

---

F Theory in Nonparametric Setting
Built upon the insights from Section 4, we provide analysis to the nonparametric reward and general data sampling setting. We generalize Assumption 4.2 to the following.
Assumption F.1. The ground truth reward f* is decomposed as
(x(x) g ( h" (xr
where 9" (x is o-Hölder continuous for 2 1 and h (ar 1) is nondecreasing in terms of || llz with h (0) 0. Moreover, 9" has a bounded Hölder norm, i.c., 1g" S 1.
Hölder continuity is widely studied in nonparametric statistics literature [15, 49). h" here penalizes off-support extrapolation.
Under Assumption F.I, we use nonparametric regression for estimating f". Specifically, we specialize (3.1) in Algorithm to
fo E argmin foEF 2n
(ffo(x:) y:)
where F NN(L, M, K,))
x) is chosen to be a class of neural networks. Hyperparameters F will be chosen properly Theorem F.4.
Our theory also considers generic sampling distributions on a. Since æ lies in a low-dimensional subspace, this translates sampling distribution assumption latent variable
Assumption F.2 The latent variable follows distribution Pa with density Ps. such that there exists
constants B, C1, C2 verifying P=(2) I (2+)-(4+1)/2C, exp (-Cal|=1B/2) whenever ||ı||| > B. And cola E Ex~P, [.
Assumption R.2 says Pz has a light tail, which is standard in bigh-dimensional statistics [51, 53). Assumption F.2 also encodes distributions with compact support. Furthermore, assume that the curated data (x, !) induces Lipschitz conditional scores. Motivated by Chen et al. [8], we show that the linear subspace structure in leads a similar conditional score decomposition V log Pr (x] = (x, 7, t) + 1 (I, 1), where lI is the on-support score and s is the orthogonal score. The decomposition for conditional score is (H.I), which applies to both parametric and non-parametric cases. The following assumption is imposed on $||.
Assumption F3. The on-support conditional score function 81 (z., F7, () is Lipschitz with respect to for any & E 7), i.e., there exists a constant Clip. such that for any T and x', T. it holds
||>|| (x. T, () ll (2', ".,.7'1)) l)l < Cuip || 4'112 + Ctp 7 - 7/ 1/12.
Lipschitz score commonly adopted in existing works [9, 25]. Yet Assumption F.3 only requires the Lipschitz continuity of the on-support score, which matches the weak regularity conditions in Lee et al. [25], Chen et al. [8]. We then choose the score network architecture similar to that the linear reward setting, except replace m by nonlinear network. Recall the linear encoder and decoder estimate the representation matrix A.
We consider feedforward networks with ReLU activation functions as concept classes F and S for nonparametric regression and conditional score matching. Generalization different network architectures poses no real difficulty. Given an input neural networks compute
(NN(T) = Wro(... (Wi + bi)...) + br,
(F.I)
where W; and i; are weight matrices and intercepts, respectively. We then define a class of neural networks
NN(L, {f: ff in the form of (F.1) with L layers and width bounded by M,
sup [f < K, maxx||b;lle |W & x for i = 1..... L, and (1 + +1bbllo) - $J}.
29

---

For the conditional score network, we will additionally impose some Lipschitz continuity requirement, i... [l. f(y)ll2 S Cipl|æ yll2 for some Lipschitz coefficient Clip-
Recall the distribution shift defined in Definition D.3 that
T(P,, P2; C) suPtec Rz- (x)))/ [x))]
for arbitrary two distributions Pt, P2 and function class L. Similar to the parametric case, use notation Pa P-- a) and P. := PP a). Then we can bound SubOpt (a;; y" = a) in Theorem F.4 in terms of non-parametric regression error, score matching error and distribution shifts in both regression and score matching.
Theorem E4. Suppose Assumption F.I, E2 and F3 hold. Let (()) = dlogg l Properly chosen log
F and S, with high probability, running Algorithm with a target reward value and stopping at
2-25(m
4+4
2- 262mm
ded
to = 11 + Dn1
a+6
gives rise to <(, A) S Ō
1
11
+ Dn1
and
SubOpt(P Paiu² aa
2-final 2n4d
T(P(x = Ps;FF. δ f2
+ D/n2
E
T(P(x, = a)) Pags S)
σ ((1
28(m1)
di4
||g0||0 10x + M(a)
Dn1
C0
E2
+ E
pnP, P. (x1)),
where M(a) == E2-P(a) || and
8 - T 6-(T-10
/"' Ex, V log y) s, :
F -
Ea penalizes the component in Pa that is off the truth subspace. The function classes F and S are chosen F NN(L,, Mf, Jy, K k5) with
f (log n2), My O
n12
21 (log n2 nz)
Jj JJ= =0=
n12 0:20
(log 122) 4/221
K3 1, Kig 0 =V[iogn
and S NN(L,, Ms, Ka K.)) with
- +=3 L, Olognn + dd, M, = o ( (2//2 (logn1))11 J. = 0 (1/3
(
246
(logn1) d/2+1 (40gm1)4/2+1
0 = (dlog(dnz)) Ko 0 ( (Vdlog(m)d)
Moreover, S also Lipschitz with respect to (x. y) and the Lipschitz coefficient Clip O (10dCrpp)
The proof is provided in Appendix G.2. Here we correct a typo on n 3(4+5) main paper Section 4.3. Quantities [x a), F) and T(P(x, 7 a), Pag: S) depend on characterizing the distribution shift. The 8(n) terms account for the unbounded domain of z, which is negligible when is large. In the main paper, omit ((z) in the regret bound.
Remark.
30

---

G Omitted Proofs in Section F
G.1
Conditional Score Decomposition and Score Matching Error
Lemma G.I. Under Assumption 4.1, E.2 and F.3, with high probability
1
[|((,t)
v log (o) s casss(1))
T to J
2-26(o1)
) for 8(n1) #+8
with (m1) 0
to
111
Dn
dog log Ri. log m
Proof. [8, Theorem I] is easily adapted here to prove Lemma G.I with the input dimension d + and the Lipschitzness Assumption F3. Network size of is implied by [8, Theorem 1] with कमक accounting for the additional dimension of reward ÿ and then the score matching error e= n, follows.
G.2 Proof Theorem E.4
Additional Notations: Similar as before, PLD to denote the low-dimensional distribution on corrupted by diffusion noise. Formally, pt LD (z) = (><<((=))= (=)dz with (x(+)) being the density of ((a(()z,h(()Ta)) (= f(Az) f(Az) = a at lo. with shorthand as (a). Also give (=(= f(Az) = a) shorthand as (LD
a) the corresponding conditional distribution
G.2.1 S = a) Decomposition
By the same argument in SD.3.1, we have
SubDpt(Pa:v' = a) s Pawpa [M())) - [(m]]
RawPa [g" (x1)) E Ban,, (g" (x11)
E (x1))
G.2.2 ε1 Nonparamtric Regression Induced Error
Nonparametric Regression Error of f Since Pa has light tail due to Assumption F.2, by union bound and [8. Lemma 16], we have
6/241
IP33 2, with ||al|22 R for = 11...,712) S n2 C2r(d/2+11)
Rd 2 10-ppxx
where C1, O2 are constants and r(.) is the Gamma function. Choosing O /dlog d log )
ensures P(3 x; with ||>;|l2 2 > R for 1, 1...., doog log 22, we have 1...., n denoting 5(n2) log
... 122) 5. On the event E f||za112 2 S R for all
2(u digel)
1f - 71 =
n2
2420
by [31, Theorem 7] with a new covering number of S, when ny is sufficiently large. The corresponding network architecture follows from Theorem 2 in "Nonparametric Regression on www-Dimensional Manifolds using Deep ReLU Networks Function Approximation and Statistical Recovery".
We remark that linear subspace is special case of low Minkowski dimension. Moreover, 5(12) is asymptotically negligible and accounts for the truncation radius R of I; 's (see also [8, Theorem 2 and
3)). The covering number of S is δ (dd/2, # (log n2 4/2 + Dd) as appear [8, Proof of Theorem
31

---

2]. Therefore
EanPa Pa [1(( - EeNPA F(a)] 2
T(P(o = a), F): -11. -1i
stna
a)) P2;F) Ö 122
2ard
+ D/nn2)
G.2.3 Eg: Diffusion Induced On-support Error
Suppose L2 score matching error is Cdffs (11), i...
lo s (x, s (df)((n);
T to 3o
We revoke Definition D.2 measuring the distance between Pa Pa that (V(Pa) drv (P% F(A2)=a)) vU =
Lemma D.4 applies nonparametrie setting, so we have
(ID VV ) N(O.A), A ctulD, <(V δ tu Ed (
(G.1)
(G.2)
co
In addition,
T(P(,, a), Pag: S)
TV(
ears (m1)
(G.3)
co
Eg will be bounded by
Eg Ex |9 (x)] E P [x)]]]
I E 1g (AA
x:] [AATx)]]
Enp [g (VVT a)} + E E loV - (AATx=))
where for E Bani, (VV 2) - (AA x)] we have
P. 1 (VV - 9" (AAT 2)]{ $ E E. V x- AA AA IVVT AA 1 2. [|
(G.4)
For the other term Ie x)] E B (VV x)) will bound it with (P..).
1g' x)] E E. [g (VV x))|
x E (en (a) Ig (A=)] - E - a 2 (V=(VV]] + Be (Az)] - E E2nR (a)lg) (Az)]
Since any =~ (a. can be represented by a(lg)= + No(o)u, where = P(a), N(O, Ia), then
Eo (Az)]
= No Ah(t)Au))
<
< E, + 11 +
32

---

thus
E Faos(1) (Az)] (Az)) - E Pio (a) 0 (Az)] to Bewp(a)ll + d,
where we further use 1 a(46) = 1 t0/2 10/2, h(6o) $1 1.
As for (a-Pen (Az)] - B 2M 1 P. 19¹ (V 2]
we have
E E (a) 0 (Az)] E S~(UTTT) p. 10 (VUz)]
E Eew (a) 00 (VUz)] E B 10 (VU+)] + (Enee, (Az)] EewP (a) (VU=))
where
R Ean (a) (0 E 20 VT) P. (g' (VV(vV))) STV 11,"
E (o (V $ $ - 31A-vUUl
(a
and
Combining things up, we have
&g <|IVV AA || P, [1|= 1|4 VUlFF (a)
to (a) [1|= + TV (P.) 119 bl
Similar to parametric case, Let (a) == (a) ||| 1 then
Ex (a) || S M(aa + Lod,
expect for in nonparametric case, we can not compute (aa out as it is not Gaussian. But still, with higher-order terms in ni hided, we have
& = (Pa( 110 + (oM(a))
= T(P(2, 7 a), Pagi
caff((nn) 1ox + +oM((a)
co
H Parametric Conditional Score Estimation: Proof of Lemma D.1
Proof. We first derive a decomposition of the conditional score function similar to [8]. We have
/ p((z,ulz)p
pc(x,
/ (((el=)p(ula))e
C f exp
(- 2h(t)
---)) (-2nle) ||0 (11) A113) exp (-23 oz (0-=-u)² u)²) p=(2)dz
1
Cexp
(-2hte) II(D AAT AAT) )x11
(- (2h(t)
f exp
1 2h(t)
1 (0T - v)²) pa (z)
|A x a(t)&l1
exp
where equality (i) follows from the fact AAT x 1 (ID AAT x and C' the normalizing constant of Gaussian densities. Taking logarithm and then derivative with respect to x on pe(x, v), we obtain
Va z (x.
a(t) Af =xp
(-zA(os |IAT a(()=113) exp (-1 0/ (OT= =)²) Pz (=)dd
2b(0 ||ATx - =(1)=111 exp (- 1 (OT - u)² (=(2) h(()) x.
h(t))
f exp
33

---

Note that the first term in the right-hand side above only depends on A x and y. Therefore, we can compactly write V, log po (z, y) as
V log h(t)
Au(A" x,y.t) h(t) 2,
(H.1)
where mapping u represents
a(t) exp
2h(() IA a(1)2113
exp (-0; (0T= - u)) p= (=(d)d
exp
2h(1) || x - (t)= exp
(0T 2 u)² P= (z)dz
We observe that (H.1) motivates our choice of the neural network architecture S in (3.8). In particular, t attempts to estimate u and matrix V attempts to estimate A.
In the Gaussian design case (Assumption 4.4), we instantiate P=(2) to the Gaussian density (2 (E)) ( 4/2 exp (- -- 2). Some algebra on the Gaussian integral gives rise
Vx log Po = a(t) ABB14 (x. v) h(t)
1
1 h(t)
(TD - AA
AAT
h(t)
a(t)
h(t) V²
AB;
(at))A x
y0
I,
(H.2)
h(t)
h(t))
where we have denoted
h(t) (e(x,y) = a(t)A x + va yo and B = a 2 (t)
h(t)
va 00 + h(t))=
Score Estimation Error Recall that we estimate the conditional score function via minimizing the denoising score matching loss in Proposition 3.1. To ease the presentation, we denote
1 [x, yi T to
Ep Var log (x (xx) s(,',,,
as the loss function for pair clean data (x, y) and conditional score function 8. Further, we denote the population loss
C(s) = Er. (1(
whose empirical counterpart is denoted as [(s) N Emt f(x;, s1:
To bound the score estimation error, we begin with an oracle inequality. Denote (trune(s) as truncated loss function defined as
Ctrune (s) E[e(x; ys () R. S R)),
where > 0 is a truncation radius chosen as O( / dlog d + log K + log ). Here is uniform upper bound of s(x, y. 1) [|]r]|| < R, yl - R} for s E S, i.c., sup.es [| y. { I R. |y| < R) K. To this end, we have
( = (()) [(8) + ((3)
-(3) [(3) inf ((s)
SES
@ (3) 2(3
$ L(S) trune + Ctrune (8)
trrnne (s)
sup Chrnne Eerunc + sup L(s) Ltrune
(A)
(B)
where equality (1) holds since S contains the ground truth score function. We bound term (A) by PAC-learning concentration argument. Using the same argument in [8, Theorem 2, term (A)], we have
sup gtrune (x, s) = O ( 10(T - 1o) (K² + ((2+12)) (6
s ES
34

---

Applying the standard metric entropy and symmetrization technique, we can show
K + 12 Lo to)
log 2 2m
(A)= R(S)
where Pi the empirical Rademacher complexity of S. Unfamiliar readers can refer to Theorem 3.3 in "Foundations of Machine Learning", second edition for details. The remaining step is to bound the Rademacher complexity by Dudley's entropy integral. Indeed, have
K² Vn
4E
12
(SS)
= inf
VN(S, 112)dc.
n1
emphasize that the log covering number considers x, y in the truncated region. Taking e gives rise
= (A)= ((12--8)
R to)
NV(S, 1/m1) log
(A)
Here K is instance dependent and majorly depends on d. In the Gaussian design case, we can verify
that is (()d). To this end, we deduce (A) = Ō 1 to 2 ((S.1/m)))
In practice, is often
much smaller than D (see for example [37], where ImageNet has intrinsic dimension no than 43 in contrast to image resolution of 224 x 224 x 3). In this way, we can upper bound d2 by D, yet d2 is often tighter upper bound.
For term (B), we invoke the same upper bound in [8, Theorem 2, term (B)] to obtain
(B) (B)= = O InlooTT ( 10)
1
which is negligible compared to (A). Therefore, summing up (A) and (B), we deduce
N(S, D) log
Eass = to
71
Gaussian Design We only need to find the covering number under the Gaussian design case. Using (H.2), we can construct covering from coverings on matrices V and Σ Σ-1. Suppose Vi, V/2 are two matrices with V1 V2 l I πу for some 7 > 0. Meanwhile, let Σ Eng be two covariance
matrices with | Egg 1|2 < 1ps. Then we bound
sup ||8 (s. y, $12,25 (a. y, t)
V (V1 x.y, Vi (Vz x, 1)112
sup |R
h(tt)
Vi (V₂ v. F, Vinh (V₂2 x. v. 1) + |1 (V₂ x. x.y,t) V2 (V, x. y. 1)]l
1 h(t))
(2Rmy + 2v 2v~2Rmm))
where for bounding we invoke the identity [[(T + A) (T+ B)- < B 1|2- Further taking supremum over & [fo, T] leads to
sup R.lllsR
||,,,
18. (8, &VG (s,
(2Bny + 2v -2
to
35

---

for any t E [tg. T). Therefore, the inequality above suggests that coverings on V and Σ-1 form covering on S. The covering numbers of V and Σ- can be directly obtained by a volume ratio argument; we have
2vd '/V
2vd Amin Mps
WV, (1a) < Do 1
and N ,E,, 112) df log 1+
Thus, the log covering number of S is
MSS, n, (12) =
11 122122) + +N(E= t02 1122)
dD todmin?
- (Dd + d²) log
1
where we have plugged y² = 1/D into the last inequality. Setting r 1/ni and substituting into 2 fiff yield the desired result.
We remark that the analysis here does not try to optimize the error bounds, but aims to provide a provable guarantee for conditional score estimation using finite samples. We foresee that sharper analysis Bernstein-type concentration may result in better dependence on n. Nonetheless, the optimal dependence should not beat 1/nt-rate.
I Additional Experimental Results
I.1
Simulation
We generate the latent sample = from standard normal distribution = N(O, la) and set A for randomly generated orthonormal matrix Dxd The dimensions are be d = 16, D 64. The reward function is set to be f(x) = (6*) 21 + 51| 122, where 0* is defined by A,8 We generate 8 by uniformly sampling from the unit sphere.
When estimating 0, mentation from https: / //github. com/lucidin muncirainss/ deiiin//deoiiiinnn... cennoiiinnddiffsiiii onssion-pp on-pytorch, where we modified the class embedding so accepts continuous input. The predictor trained using 8192 samples and the score function is trained using 65536 samples. When training the score function, we choose Adam as the optimizer with learning rate 10 10-5 We train the score function for 10 epochs, each epoch doing a full iteration over the whole training dataset with batch size 32.
1.0. The score matching network is based on the UNet imple-
For evaluation, the statistics is computed using 2048 samples generated from the diffusion model. The curve in the figures is computed by averaging over 5 runs.
1.2 Directed Text-to-Image Generation
Samples of high rewards and low rewards from the ground-truth reward model. Section 5.2, the ground-trudd -tutth reward model is built by replacing the final prediction layer of the ImageNet pre- trained ResNet-18 18 model with randomly initialized linear layer of scalar outputs. To investigate the meaning of this randomly-generated generated reward model, we generate images using Stable Diffusion and filter images with rewards 2 0.4 (positive samples) and rewards < -0.4 (negative samples) and pick two typical images for each; see Figure 7, We note that in real-world use cases, the ground-truth rewards are often measured and annotated by human labors according to the demands.
Training Details. our implementation, as the Stable Diffusion model operates on the latent space of its VAE, we build 3-layer ConvNet with residual connections and batch normalizations on top of the VAE latent space. We train the network using Adam optimizer with learning rate 0.001 for 100 epochs.
I.3 Decision-Diffuser [1]
We replicate the results in Decision Diffuser [1] under (Med-Exxert, Hopper) setting. In Decision Diffuser, the RI. trajectory and the final reward are jointly modeled with conditioned diffusion model. The policy is given by first performing reward-directed conditional diffusion to sample
36

---

(a) A positive sample
(b) A positive sample
(c) A negative sample
(d) negative sample
Figure 7: Random samples with high rewards and low rewards.
trajectory of high reward and then extracting action sequences from the trajectory using trained inverse dynamics model. We plot the mean and standard deviation the total returns (averaged across 10 independent episodes) v.s. the target rewards in Figure 8. The theoretical maximum of the reward 400. Therefore, trajectories with rewards greater than 400 are never seen during training. We observe that when increase the target reward beyond 320, the actual total reward decreases. According to our theory, as we increase the reward guidance signal, the condition effect becomes stronger but the distribution-shift effect also becomes stronger.
120
100
Decision-Diffuser
COL
IOL
TT
MOReL Diffuser
40
20-
86.0
160.0
240.0
320.0
400.0
0
800.0
1000.0
1200.0
1400.0
1600.0
3200.0
64000.0 0
Figure 8: Total Reward v.s. Target Reward for Decision Diffuser (Med-Exppl Expe-E,,et,, Hopper) setting.
37


Reviews:
[{"review_id": "IAMIjWgacA", "reviewer": "Reviewer_umuF", "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.", "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\"?\n2. 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?\n3. In the experiments, what is the reward distribution of the training data? At which reward values are you asking the model to extrapolate?\n\nMinor comments:\n- In the problem setup section, why is $1 > \\sigma$ necessary?", "limitations": "The authors have clearly laid out the assumptions for the theoretical analysis. I do not see any major limitations within this scope.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\nI 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.\n", "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.\n\nI 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.\n\nApart from that, I found the experimental setup of using a random reward model a bit strange.\n\nMinor comments:\n- Given the similarity of Fig. 2c with Fig 2 in [8], citing [8] is probably warranted (\"Figure adapted from [8]\") ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful and insightful review.  We’ve revised our paper to clarify the notations and added more explanations around the score network according to your suggestions. \n\n>**Q1**. What are the practical implications. I could have (qualitatively) predicted the results for text-to-image generation without the theory.\n\n**A1**. Our theory is not limited to understanding image generation, but also covers applying diffusion for RL and control tasks (see [1,2] and our [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf). You are absolutely right that there’s an extrapolation from training data happening, but there are more questions that we answered by our theorems.\n\nImplications of theory: \n- Fidelity of generated samples: conditional score matching network with architecture (2.8) accurately recovers any latent subspace structure in data.\n- Trade-off between high target reward $a$ and large distribution-shift and a comprehensive exposition of the distribution-shift associated with $a$. \n- Low-dimensionality dependency: error coming from reward learning only depends on the subspace dimension $d$, spot a relation to offline bandit/RL.\n- Nonparametric results beyond linear settings\n\nPlease refer to the \"impact and novelty\" section in [our rebuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE)  for a more detailed exposition on this question.\n\n>**Q2**. the experimental setup of using a random reward model a bit strange.\n\n**A2**: Our paper focuses on general reward-guided diffusion rather than the specific application of images, thus we use a random function for reward for generality. In particular, in the image example, we constructed the reward by using a random linear projection layer on top of an ImageNet pre-trained ResNet model. It can reflect real-world applications because (1) in practice rewards may be unknown and require expensive data collection processes; (2) the pre-trained ResNet models are believed to extract useful representations for images that are semantically meaningful, thus it’s very likely a real-world reward is an easy (MLP) function of  these representations. In addition, we provided a [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf) on RL where the reward is predefined in the environment but has to be learned through jointly modeling the trajectories and the final rewards using a conditional diffusion model.\n\n>**Q3**. 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\"?\n\n**A3**: Alg. 1 does offer cleanness to our analysis. To clarify, classifier and classifier-free guidance can be formulated by the exact same backward sampling as (2.3), by plugging in a similar but different score from what is learned by (2.2). Comparison [link] discusses the relation between Alg.1 and guidance methods. Alg.1 captures the essence of both classifier guidance and classifier free guidance, i.e. to estimate the conditional score $\\nabla \\log p_t (x_t | y)$, but with an ease of putting aside other less-essential components, such as the multiple classifiers in classifier guidance and the mix of conditioned and unconditioned scores in classifier-free guidance.\n\n>**Q4**. 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.\n\n**A4**. Assumption 3.2 configures the ground truth reward by its $g$ and $h$ components. Linearity in $g$ is required for bounding the term $\\mathcal{E}_1$ in Theorem 3.6. As for the $h$ component, we make it explicit otherwise the reward could be ill-defined. We agree that $h$ is for penalizing off-support data in the reward,  and it can be problem dependent. Our result (Theorem 3.6) still holds with any choice of $h$, depending on problem nature or user choice. \n\nYes, the off-support portion of x will be reflected in its score $\\nabla_x \\log p(x|y)$. If the conditional diffusion model learns the score $\\nabla_x \\log p(x|y)$ well, then generated data form it should have little off-support component, which is what we proved in Theorem 3.5 by upper bounding $x^{\\perp}$ showing $x$ has high fidelity to the support. The ground-truth reward in Assumption 3.2 models a combination of data utility and fidelity.\n\n>**Q5**. In the experiments, what is the reward distribution of the training data? At which reward values are you asking the model to extrapolate?\n\n**A5**: In our synthetic experiment (Section 4.1), the unconditioned reward distribution follows the standard Gaussian distribution $N(0, 1)$. When we set target reward value $\\ge 3$, this falls outside the $3\\sigma$ region of the Gaussian distribution, and we view any reward level beyond 3 as extrapolation. The visualization of the reward distribution can be found in Figure 2 (Left) in our supplementary pdf file.\n\nIn our image generation experiment (Section 4.2), the mean and the standard deviation of the reward of the training data are mean -0.9113 std = 0.4130\nWhen we set target reward value $0.4$ (outside $3\\sigma$), we can view it as asking the model to extrapolate. The visualization of the reward distribution can be found in Figure 2 (Right) in our supplementary pdf file.\n\nIn our new RL experiment, the mean and the standard deviation of the reward are\n mean =294.8 std = 91.2. (Note: This is an offline dataset generated by human experts, so it is higher than all reported offline RL algorithms.)  The visualization of the reward distribution can be found in Figure 1 (Right) in our supplementary pdf file.\n \n>**Q6**. In the problem setup section, why is 1>sigma necessary?\n\n**A6**: We let $\\sigma<1$ for simplicity. It is only a scaling constant. If $\\sigma>1$, Theorem 3.6 would still hold and the term $\\mathcal{E}_1$ will linearly scale up with $\\sigma$.\n"}, {"review_id": "ZpeZcr5wHD", "reviewer": "Reviewer_ZNG8", "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.\nAssuming 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.\n\nExperiments 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.", "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?\n* 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.\n* 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?\n* 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?\n* 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]\n\n### References\n[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).\n\n[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.", "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Presents a practical way to add subjective rewards to generate data beyond concrete prompts like text\n* Analyzes this reward conditioned generation setting and identifies the interplay between rewards and distribution shift.\n* 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.\n* 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.\n* While it shows that the model can be adjusted to arbitrary rewards, it does not showcase a practical use case.\n* The technical novelty is unclear.\n\n=============================\n### Post-Rebuttal\n* The author's rebuttal and other reviews have made the technical contribution of the paper abundantly clear.\n* Additionally, the connection to _offline_ bandits is more evident and does add value.\n* The additional RL experiment grounds the claims made in the paper more concretely than the previous experiments.\n\nAs such, the majority of my concerns have been allayed with the author response.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1**. The paper does not discuss a practical manner in identifying when the generative model starts deviating from the training distribution.\n\n**A1**. Our focus is theory and implications are listed in “impact and novelty” in [our rebuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE). We agree that to identify a proper target so that generated samples are deviating from training distribution towards higher reward but not deviating too much is the key of success for empirical applications. In practice, the guidance level is viewed as a tuning parameter, and optimal choice is achieved by hyperparameter tuning. In theory, designing an explicit way for identifying the optimal target is difficult, as the interplay between the guidance level and the distribution shift is data dependent and complicated (see our discussion following Theorem 3.6 (Line 214)). There is unlikely to be a general rule of thumb or a one-fits-all way to identify.\n\n>**Q2**. 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 bandit algorithms.\n\n**A2**. Thank you for bringing up this confusion. We agree that the theoretical connection between CDM and bandit is intriguing and we are the first to observe this potential connection.  A detailed discussion:\n\nOur problem is similar to *offline* bandits in two ways: (1) both problems deal with offline data of the form {(x,y)}; (2) the goal is to find new x with improved reward value y. Given its offline nature, it does not benefit from any online bandit exploration techniques such as UCB or Thompson sampling. \n\nNote that our theory is not about linear bandits, and we never established any form of equivalence. The statistical error of condition diffusion consists of several parts (Theorem 3.6): $\\mathcal{E}_1$ (due to reward function estimation) and $\\mathcal{E}_2, \\mathcal{E}_3$ (due to conditional score matching)\nAmong these terms, only a single term $\\mathcal{E}_1$ resembles the suboptimality gap in off-policy bandit/RL (if we pick target value to be the max value). Our full analysis of reward-directed DM is much more complex beyond $\\mathcal{E}_1$. In addition, our analysis is not limited to linear reward, but allows general nonparametric reward (Section 3.3). \n\nFurther, in our [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf), we tested CDM on an RL problem and obtained improved reward performance close to the best known off-policy RL benchmark. This observation is also consistent with our theory.\n\n>**Q3**. While it shows that the model can be adjusted to arbitrary rewards, it does not showcase a practical use case.\n\n**A3**: Our primary focus is to establish thoery for CDM, especially for reward maximization. Many papers have already showcased the practical usage of CDM for generating high-reward samples in various contexts, not limited to image generation [3] but also control and RL [1,2]. We were motivated by those empirical successes and decided to go deep in theory. To this end, we formulate the practical scenario of semi-supervised learning and provided Alg. 1 as a meta algorithm (see \"2. Comparison\" in [our rebuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE) for its close relation to guidance methods).  Our results are not limited to a specific use case, instead they provide insights on diverse applications.\n- [1] https://arxiv.org/abs/2211.15657 (2022)\n- [2] https://arxiv.org/abs/2302.01877 (2023).\n- [3] https://arxiv.org/abs/2305.13301 (2023).\n\n>**Q4**. The technical novelty is unclear.\n\n**A4**. The theory of (conditional) diffusion models is widely open. There are very limited results providing statistical theories. To the best of our knowledge, we are the first to give theory for conditional diffusion models and connect them to bandits. In our theory, we provide novel analysis on conditional score estimation from both a parametric and nonparametric point of view. Moreover, in Theorem 3.6, we developed an oracle-type decomposition of the suboptimality gap into three terms, indicating the trade-off on the guidance level $a$ and distribution shift. These results are all derived by using novel analysis and new techniques.\n\n>**Q5**. Section 2.1 gives the scale of sigma but does not mention any such constraint about $f^*$ . \n\n**A5**. In section 2.1, constraint on the scale of $f^*(x)$ is imposed by the assumption on $g*$ (on-support component of $f^*$) and the distribution of $z$ (latent of $x$). In the non-parametric section, $f^*(x)$ has a similar constraint. The scale of reward reflects on term $\\mathcal{E}_1$ in Theorem 3.6, which grows linearly with the scare of $f^*$. In practice, It is common to normalize the reward by preprocessing data before training so that the scale does not influence learning.\nIn both of our experiments, the scales of the reward function fall into normal ranges. Please refer to Figure 2 (Left, Right) for the training distributions of the rewards. In the new RL experiment (see general response above), the rewards are normalized to line in [0,1].\n\n>**Q6**. Reproducibility of experiments.\n\n**A6**. Our main results are Theorem 3.5, 3.6 and their extensions to nonparametric setting in Section 3.3, with proofs fully fleshed out in appendix and easily reproducible. While the experiment is only for illustration, we will release code and demos in the final paper. Any user can supply a reward model oracle and directly use our code to generate high-reward data. \n\n \n>**Q7**. 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?\n\n**A7**.The unlabeled dataset is LAION-5B. LAION-5B dataset has various image resolutions and sizes and covers various scenes. Both datasets correspond to natural images.\n\n\n\n"}, {"review_id": "P5xSLf9l8N", "reviewer": "Reviewer_mN9F", "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.\n", "questions": "1. How will the results generalize to settings where data doesn’t have a latent linear representation?\n2. 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?\n3. 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?  \n", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall Strengths:\n1. 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. \n2. 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.\n3. Theory is well-driven and thoroughly discussed.\n", "weaknesses": "Overall Weaknesses:\n1. 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.\n2. The experimental setup is limited and the computed metrics do not thoroughly cover the theoretical findings of the papers.\n3. Although figure 2 looks very nice, I’m not sure if it has any added value in the main body of the paper. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments!\n\n>**Q1**. “How to balance the reward signal and distribution-shift effect to have high-reward and high-quality samples?” Experimental setup is limited\n\n**A1**. Our paper provides the first theory for conditional diffusion and use of reward-conditioned diffusion to generate better samples. There is a rich body of empirical works using CDM for reward maximization in various contexts [1].  \n\nIn general, designing an explicit formula for the optimal target is difficult, as the interplay between the guidance level and the distribution shift is data dependent and complicated. There is unlikely a one-fits-all solution, which is why we need theory research. Theorem 3.6 is the first rigorous result characterizing this complicated tradeoff under very general assumptions. In practice, one can use a simple doubling trick (ie binary search) to tune the target value since it is only 1-dim, at a cost of only a log factor in time. \n\nOur experiments are for the purpose of illustration and complements the theory. For rebuttal, we add a [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf) on RL and showcased the use of CDM with varying target values to maximize value. We welcome any further suggestions on the experiment part. But we want to keep our focus on deep learning theory, more extensive experiments would better fit an application paper. \n\n[1] https://arxiv.org/abs/2211.15657  \n \n\n>**Q2**. Although figure 2 looks very nice, I’m not sure if it has any added value in the main body of the paper.\n\n**A2**. Thank you for the comments. Figure 2 is closely related to our method and results, but we realize that it lacks annotations and might have confused you. We have updated Figure 2 with clear annotations and explanations ([a preview is here](https://openreview.net/attachment?id=npWKushsuE&name=pdf)).\n\nSubfigure (a) and (b) illustrate the distribution shift when extrapolating reward prediction and the distribution shift in extrapolating diffusion respectively, these two are the key components in the error decomposition shown in Theorem 3.6. Subfigure (a) corresponds to term $\\mathcal{E}_1$ and Subfigure (b) depicts the on-support situation corresponding to $\\mathcal{E}_2$, $\\mathcal{E}_3$ occurs in the space perpendicular to the subspace in (b) and is not explicitly drawn out . We updated the text annotations wrapped around figure 2 pointing its connection to theory. (c) illustrates the architecture of score matching network we analyzed.\n\n>**Q3**. How will the results generalize to settings where data doesn’t have a latent linear representation?\n\n**A3**. Our theory applies to general data distributions. We show that the conditional diffusion model naturally adapts to data distribution (high-dimensional data simply corresponds to the case where $A = I_D$, and when data lies in an unknown subspace, it corresponds to $A$ with a smaller column dimension). \n\nOur results also apply to nonlinear data such as Fourier series or data from kernel spaces with a known basis transformation. For example, when data $x = \\phi^{-1}(z)$ for a known mapping $\\phi$ and latent variable $z$, we consider linear conditional diffusion on $z = \\phi(x)$ and our analysis immediately applies with an additional $\\phi^{-1}$ transform.\n\n>**Q4**. 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?\n\n**A4**. Our model can be viewed as a special case of classifier-free guidance. It is simpler than classifier guidance and doesn’t need multiple classifier $c_t$’s. Please see \"Q2.Comparison\" in [our renuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE) for detailed discussion.\n\n>**Q5**. 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?\n\n**A5**. In our text-to-image generation experiment, a higher reward only \"favors\" more colorful images, that is, we observe a high correlation of large reward and colorfulness. As our target is to maximize the reward (which is not exactly the same as colorfulness), we only need to track the reward values of the generated samples. Therefore, the numerical metrics can be evaluated by running samples through our ground-truth reward model, which only *correlates* with more vividly colored images.\n"}, {"review_id": "WEObpQW3ty", "reviewer": "Reviewer_uwg2", "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.", "questions": "Some suggestions:\n1) It would be good to add what each color represents in Figure 2. \n2) 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. \n3) 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. \n\nQuestions\n1) 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?\n2) The model seems very similar to the classifier-guided diffusion model, but it never compared to that in the experiments, any reason for that?\n3) I am wondering if it is possible to train the reward function together with the diffusion model instead of pertaining to the reward network. \n4)  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. \n5) 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. \n6)  After reading the paper, I am still confused about the assumption that the model implicitly learns the latent subspace representation of x.\n\nSome suggestions:\n1) It would be beneficial to add a legend in Figure 2 that explains the meaning of each color.\nIt 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.\n2) 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.\n\n\nQuestions:\n\n1) 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?\n2) 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?\n3) 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.\n4) 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.\n5) 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.\n6) Even after reading the paper, I'm still unclear about the assumption that the model implicitly learns the latent subspace representation of x.\n\n", "limitations": "The paper did not discuss any limitations. \nThe paper has no potential negative societal impact. ", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Strength:\n1)  The paper is well written and the problem they try to tackle is interesting.\n2) The paper provides an insightful theory for conditional distribution learned through a reward-based diffusion model\n", "weaknesses": "1) The paper lacks a comparison to other similar models, such as classifier-guided diffusion models.\n\n2) The assumptions made for the theoretical work appear to be overly simplified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1**. The paper lacks a comparison to other similar models, such as classifier-guided diffusion models.\n\n**A1**. Alg 1 is not an alternative to classifier-guided diffusion. Instead, it is a simplification of it and also generalized to continuous reward and semi-supervised learning. Please refer to \"2.Comparison\" in [our rebuttal](https://openreview.net/forum?id=58HwnnEdtF&noteId=npWKushsuE) for a detailed comparison between Alg.1 and guidance methods. Rather than proposing an algorithm that outperforms the existing ones like classifier(-free)-guided diffusion, the purpose of Alg.1 is to give a formal mathematical statement of the conditional diffusion procedure with rigorous mathematical guarantees. \n\n>**Q2**. The assumptions appear to be overly simplified.\n\n**A2**. There is a misunderstanding. Our assumptions are in fact mild and general. We allow practical encoder-decoder architecture for score matching, data with latent representation, nonlinear reward, noparametric reward learning, nonparametric distribution with mild regularity conditions. \n\nNext we discuss two of our assumptions: (1) low-dimensional linear representation that x = Az. (2) parametric/nonparametric form of the ground-truth reward and regularities of data distribution. \n\nFor (1) , our theorems hold in either case where data lies in an unknown subspace spanned by columns of $A$ with a smaller column dimension, or the case $A = I_D$ meaning that data is full-dimensional. Thus, our theory adapts to data distributions with an arbitrary intrinsic dimension. Our analysis also applies to nonlinear data through kernel transformation. For example,  when data $x = \\phi^{-1}(z)$ for a known mapping $\\phi$ (for example if we have kernel or feature map) and latent variable $z$, we consider linear conditional diffusion on $z = \\phi(x)$ and our results immediately apply with a simple transform.\n\nFor (2), in section 3 we present in detail the theorems for a parametric setting of reward and data distribution; and then in Section 3.3 and Appendix G an extension to nonparametric setting that allows nonlinear reward and general data distributions.\nIn the parametric setting, reward is assumed to be linear and x is assumed to be Gaussian. As the first theory for conditional diffusion, this parametric configuration is a go-to setting to study since it builds up the theoretical foundation and gives most insights for other advanced settings such as logistic, kernel, neural network (NTK) models. The nonparametric setting further gives generality of our theorems so that it covers the practical scenario where deep ReLU networks are adopted to approximate the reward and score.\n\n>**Q3**. It would be beneficial to add a legend in Figure 2 that explains the meaning of each color and include a summary of related work in the main paper\n\n**A3**. Figure 2 is updated with legend ([preview is here](https://openreview.net/attachment?id=npWKushsuE&name=pdf)). Thanks for your suggestion. Due to the space limit, we had to defer related work to Appendix A. We will try to move it back in the final paper. \n\n>**Q4**. 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.\n\n**A4**. The block in Figure 1 represents the backward process of diffusion models (note the left arrow in the superscript), which is not to be confused with the often-seen forward process. Therefore, the starting point of the backward process is pure noise. This notation is also adopted by some other papers in diffusion literature such as [1].\n\n[1]  Chen et al. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data.  arXiv:2302.07194, 2023.\n\n>**Q5**. Figure 6, are examples cherry picked? why does context changes with colorfulness?\n\n**A5**. The examples are typical and we didn’t cherry pick. The reward was randomly constructed for generality. We will release the model checkpoints and our code upon acceptance.  During rebuttal, we also conducted a [new experiment](https://openreview.net/attachment?id=npWKushsuE&name=pdf) on RL task.\n\n>**Q6**. Is it possible to train the reward function with the diffusion model instead of relying on a separate reward network.\n\n**A6**. In practice it’s possible. In theory, we use a separate reward trained on labeled data that is independent from unlabeled data, for ease of theoretical analysis and better separating the error from reward training and diffusion training in our guarantees.\n\n>**Q7**. 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.\n\n**A7**. Our main paper only presented results with linear reward in detail due to page limit, but our theory extends to nonparametric nonlinear rewards and general distributions (section 3.3). Our synthetic experiment in section 4.1 directly tested the linear reward, and our imagenet experiment tested with nonlinear reward which can be interpreted using the neural tangent kernel theory. During rebuttal, we also provided a new RL experiment where y is the value of policy and depends on x via complicated nonlinear relations. \n\n>**Q8**. I'm still unclear about the assumption that the model implicitly learns the latent subspace representation of x.\n\n**A8**. There is a misunderstanding. We never assume “The model implicitly learns the latent subspace representation of x”. Instead, we proved it as one of our main results, i.e.,  (3.2) in Theorem 3.5 where the subspace angle between the ground truth $A$ and the learned one $V$ is upper bounded by $\\frac{1}{\\sqrt{n_1}}$, proving that diffusion model learns the latent subspace representation of $x$. The implication is that the DM is able to learn the true natural space of data and thus generate high-fidelity samples.\n"}]

Author rebuttal:
We would like to thank all the reviewers for valuable comments!

**Q1 Impact and novelty of theory**

**A1:** Conditional diffusion models (CDM) have emerged as a powerful generative model with diverse applications from image generation to control and RL[1, 2, 3].  In sharp contrast to abundant empirical successes, theoretical understanding of diffusion is limited, let alone theory of conditional diffusion. Our paper established the *first finite-sample statistical theory for training CDM* and *the first provable efficiency guarantee for CDM applied to optimizing reward*. We assume a practical encoder-decoder architecture for score matching and provide theoretical results that apply to both linear and nonparametric nonlinear reward (Theorem 3.5, 3.6), general distribution with mild regularity assumptions (Section 3.3), which covers the use of ReLU networks for training.

We summarize implications of our theory:

*a. Fidelity of generated samples:* Theorem 3.5 proves that conditional score matching network with architecture (2.8) accurately recovers any latent subspace structure in data.

*b. Trade-off between high target reward $a$ and large distribution-shift:* Theorem 3.6 shows that the average reward of generated samples is lower bounded by $a - Error(a)$. $Error(a)$ further decomposes into distribution shift penalties $\mathcal{E}_1$ (due to reward function estimation) and $\mathcal{E}_2$ (due to conditional score matching). It provides a comprehensive exposition of the interplay between reward signal and distribution-dependent factors. 

*c.Low-dimensionality dependency:* $\mathcal{E}_1$ only depends on the subspace dimension $d$, free of the ambient dimension $D$. Moreover,  $\mathcal{E}_1$ turns out to be the distribution shift term commonly seen in offline bandit/RL, whereas they require prior knowledge of the latent space but our conditional diffusion automatically learns the latent space.

*d. Beyond linear settings:* We extend to nonparametric settings in Section 3.3, allowing general distributions and nonlinear rewards.

**Q2 Comparison between Alg.1, diffusion with classifier guidance or classifier free guidance**

**A2:** Alg. 1 is not an alternative to classifier(-free) guidance, but it is a simplification and a generalization of them for theoretical purposes. We will add the following extensive remark about this point in our final paper:

>Our work draws motivation from the guidance-based diffusion classifier which was originally designed for discrete labels-conditioned generation. We generalize it with provable guarantees to continuous reward by formulating Alg.1. Our analysis provides theoretical justification for reward-conditioned diffusion in control and RL [2,3], not limited to classification. 
Notably, we provide the first theoretical guarantee for general conditional diffusion models. We show that CDM with score network as in (2.8) (FIgure 2(c)) guarantees generation fidelity in Theorem 3.5 for general light-tailed distribution with mild assumptions on score regularity.
 
>*Alg.1 and its analysis captures the essence of both classifier guidance and classifier free guidance methods.*  From a mathematical point of view, classifier guidance and classifier-free guidance share the same foundational objective, which is also the objective used in Alg. 1 — to estimate the conditional score $\nabla \log p_t (x_t | y)$. Accordingly, the foundational theoretical question is the same: understanding score matching and finite-sample distribution approximation, see Lemma E.1 in Appendix. Thus, one can view Alg. 1 as a meta algorithm and a simplification of these two guidance diffusion methods for elegance of mathematical analysis.  We provide more detailed discussion below:
>By the Bayes’ rule, it holds that
$$\nabla \log p_t\left(x_t \mid y\right)=\nabla \log p_t\left(x_t\right)+\nabla \log c_t\left(y \mid x_t\right). \qquad (\star)$$ Recall that Alg.1 directly learns the conditional score $\nabla \log p_t\left(x_t \mid y\right)$ on the LHS of ($\star$). The score network is trained on $(x, \hat{f}(x))$ pseudo-labelled by reward prediction $ \hat{f}$. Next we discuss its relation to classifier guidance and classifier-free guidance:
> + Classifier guidance [1] focuses on the RHS of ($\star$) and estimates $\nabla \log p_t\left(x_t\right)$ and $\nabla \log c_t\left(y \mid x_t\right)$ separately. Here $c_t$ is trained to take in noisy input $x_t$ and the time index $t$ to predict the label $y$. In contrast, Alg. 1 avoids the hassle of training a classifier/reward model with noisy inputs and additional input dimension $t$, and also offers theoretical cleanness.
> + Classifier-free guidance learns the same conditional score $\nabla \log p_t\left(x_t \mid y\right)$ as Alg.1 does, and uses a linear combination of conditioned and unconditioned scores for inference: $$\widetilde{s}_\theta(x, y, t)=(1+\eta) \widehat{s}(x, y, t)-\eta \widehat{s}(x, \emptyset, t).$$ When $\eta = 1$, Alg.1 is equivalent to classifier free guidance.

>We are not aware of any existing theory for classifier guidance and classifier-free guidance. Given that Alg.1 resembles yet simplifies both of them and keeps the essential spirit, we believe one can extrapolate from our theory to better understand these two methods.

**3 New Experiments on RL**

**A3:** To verify our theory beyond synthetic experiments and text-to-image (sections 4.1, 4.2), we conducted a new experiment using condition diffusion for offline RL, following [2]. See the pdf attached for details. 
Further, we were able to match best known benchmarks of offline RL. Empirical observations match the our theory. 

[1] "Classifier-free diffusion guidance." arXiv preprint arXiv:2207.12598 (2022).
[2] "Is conditional generative modeling all you need for decision-making?." arXiv preprint arXiv:2211.15657 (2022).(ICLR 2023 Oral)
[3]. "Adaptdiffuser: Diffusion models as adaptive self-evolving planners." arXiv preprint arXiv:2302.01877 (2023).

Meta-review: The authors provided a theoretical justification of the reward-conditioned diffusion model. The main theoretical results are two-fold: 

* 1) diffusion model can recover the linear subspace;
* 2) reward-conditioned diffusion model can adapt to the conditioned reward;

The results are technically sound and all of the reviewers leaned towards acceptance for this submission. I do hope the authors can include the discussion on the following two questions in the final version if possible:

* 1) what will happen if we mis-specify the latent dimension in our parameterization, say the latent dimension is $d$ but we specify it as $d'$ with $d' > d$ or $d' < d$?
* 2) why should we focus on the if we can if the expected reward on the conditional distribution match the target $y^*$, instead of directly bound some distances between the ground truth conditional distribution and learned conditional distribution?
Decision: Accept (poster)