Title: Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement
Venue: NeurIPS 2023
Accepted: yes
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

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. Our code is available at https://github.com/Kaffaljidhmah2/RCGDM .

1 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 a specific domain and to ensure the generation results satisfying certain constraints (e.g., safety, fairness, physical constraints) without sabotaging the power of the original pre-trained model [35, 27, 54, 41].

In this paper, we focus on directing the generation of diffusion models [19, 43], a family of scorematching 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 to find a cleaner version of the input. The denoising network (a.k.a. score network) s ( x, t ) approximates the score function ∇ log p t ( 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].

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-conditioned diffusion models. The 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-conditioned diffusion models, aiming to answer the following question:

How to provably estimate the reward-conditioned distribution via diffusion? How to balance the reward signal and distribution-shift effect, and ensure reward improvement in generated samples?

Our Approach. To answer both questions, we consider a semi-supervised learning setting, where we are given a small dataset D label with annotated rewards and a massive unlabeled dataset D unlabel . We estimate the reward function using D label and then use the estimator for pseudo-labeling on D unlabel . Then we train a reward-conditioned 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 unlabel , e.g., finetuning from a pre-trained model [35, 54]. We focus on the pseudo-labeling approach, as it provides a cleaner formulation and exposes the error dependency on data size and distribution shift. The intuition behind and the message are applicable to other semi-supervised approaches; see experiments in Section 5.2.

From a theoretical standpoint, we consider data point x having a latent linear representation. Specifically, we assume x = Az for some matrix A with orthonormal columns and z being a latent variable. The latent variable often has a smaller dimension, reflecting the fact that practical data sets often exhibit intrinsic low-dimensional structures [13, 48, 37]. 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 2 for illustration.

Contributions. Our results are summarized as follows.

(1) We show that the reward-conditioned diffusion model implicitly learns the latent subspace representation of x . Consequently, the model provably generates high-fidelity 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, measured by the difference between the target value and the average reward of the generated population. In the case of a linear reward model, we show that this error includes the suboptimality gap of linear off-policy bandits with full knowledge of the subspace feature, if taking the target to be the maximum possible. In addition, 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 as text-to-image generation (stable diffusion) and reinforcement learning (decision-diffuser) to support our theory (Section 5 and Appendix I).

To our best knowledge, our results present the first statistical theory for conditioned diffusion models and provably reward improvement guarantees for reward-directed generation.

2 Related Work

Guided Diffusions. For image generations, guiding the backward diffusion process towards higher log probabilities predicted by a 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 [11] or implicitly specified by the conditioned diffusion models, i.e., classifier-free [18]. Classifier-free guidance has become a standard technique in the state-of-the-art text-to-image diffusion models [39, 38, 4]. Other types of guidance are also explored [33, 14]. Similar ideas have been explored in

<latexi sh

1_b

64="

2

WI+

THp

D

of

3

>A

CF

c

V

Mw

v

K

0

g

9SE

UOY

r

dy

Q

/

Z

Gkz

j

5

m

N

u

n

8

X

<latexi sh

1_b

64="Q/3d

N

q

Tcj

2

Izm8

>A

CFn

VDLS

U

v

M

gE

y

RXw

B

rG

f0

oZ

9

Hu5

W

p

Y

O

7

K

P

<latexi sh

1_b

64="

3B5NYLkA9KG

Mz

HZq

X

E

>

C

c

VD

S

FJ

rPU

7d

R8

w

O

ygn

p

T

mQ

f

+

/

o

u

j

2v

I

0

W

<latexi sh

1_b

64="

EDkjMGS

TmW

7A

Y

8ZQ

I

>

CFn

c

V

L

N

J3

rPU

d

ug

w

R

0

XB

K9

p

z

/

2

y

f+

O

5

o

v

H

q

<latexi sh

1_b

64="Cn3o

P7T

cJ

A

wI

rLRmk

>

F

VD

S

N

U

d

ug

0WXBj

K9

p

z

5Q

/

+

Yf8

M

Z

2

E

O

H

G

v

y

q

<latexi sh

1_b

64="

3fE

795In

TjYpF/

UQCo

>A

L

c

VB

SgN

OyJ

u

G

m

k

0v

Hz

W

2

wM

Z

qr

+

X

8

R

D

K

P

d

<latexi sh

1_b

64="TEkPqA2Q

Rfvo

KmBwzI

8

>

CGX

c

VDLS

N

FJ3U

d

u

5

7

O

0

/

H

r

+

y

njp

W9

Y

g

M

Z

<latexi sh

1_b

64="

/T

qVO

XYEA

pU

PZ

F2

>

C

n

c

DLS

N

J3

v

y

5KIj

W

u

G

0M

m

B9z

H7g

f

8

r

R

+

Q

<latexi sh

1_b

64="

m

NSvz

D5nZM

yR

G

C

E

>A

c

V

L

FJ3U

2

q

g

Xw

B

r

f0IQ

o

k

/

+

u

K

H

7

dW

p

Y

8

O

9

T

P

j

D

unlabel

D

label

…

Reward Model

Label

Pseudo Label

˜

D

X

0

X

T

X

t

X

t

′

Conditional Score

Reward-Conditioned Di

ff

usion

s

(feature

,

label

, t

)

Target Reward

D

new

Step 1: Reward Learning

Step 2: Pseudo Labeling

Step 3: Conditional Di

ff

Step 4: Guided Generation

Subspace

(a) Distribution Shift in Generated Samples

(b) Distribution Shift in Reward

(c) Score Network

reward

sequence modelling problems. In offline reinforcement learning, Decision Diffuser [1] is a diffusion model trained on offline trajectories and can be conditioned to generate new trajectories with high returns, satisfying certain safety constraints, or composing skills. For discrete generations, Diffusion LM [27] manages to train diffusion models on discrete text space with an additional embedding layer and a rounding step. The authors further show that gradients of any classifier can be incorporated to control and guide the text generation.

Theory of Diffusion Models A line of work studies diffusion models from a sampling perspective. When assuming access to a score function that can accurately approximate the ground truth score function in L ∞ or L 2 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 of generated samples is related to offline bandit/RL theory [30, 29, 6, 12, 22, 32, 5]. In particular, our theory extensively deals with distribution shift in 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-policy linear bandits. However, our analysis goes far beyond the scope of bandit/RL.

3 Reward-Directed Generation via Conditional Diffusion Models

In this section, we develop a conditioned diffusion model-based method to generate high-fidelity samples with desired properties. In real-world applications such as image/text generation and protein design, one often has access to abundant unlabeled data, but relatively limited number of labeled data. This motivates us to consider a semi-supervised learning setting.

Notation : P xy denotes ground truth joint distribution of x and its label y , P x is the marginal of x . Any piece of data in D label follows P xy and any data in D unlabel follows P x . P is used to denote a distribution and p denotes its corresponding density. P ( x | y = a ) and P ( x, y = a ) are the conditionals of P xy Similarly, we also use notation P x ̂ y , P ( x | ̂ y = a ) for the joint and conditional of ( x, ̂ y ) , where ̂ y is predicted by the learnt reward model. Also, denote a generated distribution using diffusion by ̂ P (density ̂ p ) followed by the same argument in parentheses as the true distribution it approximates, e.g. ̂ P ( x | y = a ) is generated as an approximation of P ( x | y = a ) .

3.1 Problem Setup

Suppose we are given an unlabeled data set D unlabel = { x j } n 1 j =1 and a labeled data set D label = { ( x i , y i ) } n 2 i =1 , where it is often the case that n 1 ≫ n 2 . Assume without loss of generality that D label and D unlabel are independent. In both datasets, suppose x is sampled from an unknown population distribution P x . In our subsequent analysis, we focus on the case where P x is supported on a latent subspace, meaning that the raw data x admits a low-dimensional representation (see Assumption 4.1). We model y as a noisy measurement of a reward function determined by x , given by

A 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 ( x | y = a ) . Given f ∗ , P x or the low-dimensional support of P x are unknown, we need to learn these unknowns explicitly and implicitly through reward-conditioned diffusion.

3.2 Meta Algorithm

Algorithm 1 Reward-Conditioned Generation via Diffusion Model (RCGDM)

1: Input : Datasets D unlabel , D label , target reward value a , early-stopping time t 0 , noise level ν . (Note: in the following psuedo-code, ϕ t ( x ) is the Gaussian density and η is the step size of discrete backward SDE, see §3.3 for elaborations on conditional diffusion)

2: Reward Learning : Estimate the reward function by

where ℓ is a loss and F is a function class.

4: Conditional score matching : Minimize over s ∈ S ( S constructed as 3.8) on data set ˜ D via

3: Pseudo labeling : Use the learned function ̂ f to evaluate unlabeled data D unlabel and augment it with pseudo labeles: ˜ D = { ( x j , ̂ y j ) = ̂ f ( x j ) + ξ j } n 1 j =1 for ξ j i.i.d. ∼ N (0 , ν 2 ) .

5: Conditioned generation : Use the estimated score ̂ s ( · , a, · ) to sample from the backward SDE:

6: Return : Generated population ̂ P ( ·| ̂ y = a ) , learned subspace representation V contained in ̂ s .

In order to generate novel samples with both high fidelity and high rewards, we propose RewardConditioned Generation via Diffusion Models (RCGDM); see Algorithm 1 for details. By using the labeled data D label , 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 unlabel with 'pseudo labeling" and additive noise, i.e., ˜ D = { ( x j , ̂ y j = ̂ f ( x j ) + ξ j ) } n 1 j =1 with ξ j ∼ N (0 , ν 2 ) of a small variance ν 2 . Here, we added noise ξ j merely for technical reasons in the proof. We denote the

joint distribution of ( x, ̂ y ) as P x ̂ y . Next, we train a conditional diffusion model using the augmented dataset ˜ D . If we specify a target value of the reward, for example letting y = a , we can generate conditioned samples from the distribution ̂ P ( x | ̂ y = a ) by backward diffusion.

Alternative approaches. In Line 4, Algorithm 1 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 self-distillation [47]. For theoretical cleanness, we focus on analysing Algorithm 1 as it shares the same core essence with other alternative methods, which is approximating of the conditional score ∇ log p t ( x t | y ) .

3.3 Training of Conditional Diffusion Model

In this section, we provide details about the training and sampling of conditioned diffusion in Algorithm 1 (Line 4: conditional score matching and Line 5: conditional generation). In Algorithm 1, conditional diffusion model is learned with ˜ D = { ( x j , ̂ y j = ̂ f ( x j ) + ξ j ) } n 1 j =1 , where ( x, ̂ y ) ∼ P x ̂ y . For simplicity, till the end of this section we use y instead of ̂ y to denote the condition variable. The diffusion model is to approximate the conditional probability P ( x | ̂ y ) .

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 ∇ log p t ( x | y ) instead of the unconditional one. Here p t denotes the marginal density function at time t of the following forward O-U process,

where similarly T is a terminal time, ( W t ) t ≥ 0 is a Wiener process, and the initial distribution P 0 ( x | y ) is induced by the ( x, ̂ y ) -pair distribution P x ̂ y . Note here the noise is only added on x but not on y . Throughout the paper, we consider g ( t ) = 1 for simplicity. We denote by P t ( x t | y ) the distribution of X y t and let p t ( x t | y ) be its density and P t ( x t , y ) be the corresponding joint, shorthanded as P t . A key step is to estimate the unknown ∇ log p t ( x t | y ) through denoising score matching [46]. A conceptual way is to minimize the following quadratic loss with S , a concept class.

Unfortunately, the loss in (3.5) is intractable since ∇ log p t ( x t | y ) 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.1.

Proposition 3.1 ( Score Matching Objective for Implementation ) . For any t > 0 and score estimator s , there exists a constant C t independent of s such that

where ∇ x ′ log ϕ t ( x ′ | x ) = -x ′ -α ( t ) x h ( t ) , where ϕ t ( x ′ | x ) is the density of N ( α ( t ) x, h ( t ) I D ) with α ( t ) = exp( -t/ 2) and h ( t ) = 1 -exp( -t ) .

Equation (3.6) allows an efficient implementation, since P x ̂ y can be approximated by the empirical data distribution in ˜ D and x ′ is easy to sample. Integrating (3.6) over time t leads to a practical conditional score matching object

where t 0 > 0 is an early-stopping time to stabilize the training [44, 50] and ̂ E denotes the empirical distribution.

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 score network architecture (see Figure 2(c) for an illustration):

with V being any D × d matirx with orthonormal columns and Φ a customizable function class. This design has a linear encoder-decoder structure, catering for the latent subspace structure in data. Also -1 h ( t ) x is includes as a shortcut connection.

Conditioned Generation. Sampling from the model is realized by running a discretized backward process with step size η > 0 described as follows:

initialized with ˜ X t, ⇐ t ∼ N (0 , I D ) and W t is a reversed Wiener process. Note that in (3.3), the unknown conditional score ∇ p t ( x | y ) is substituted by ̂ s ( x, y, t ) .

4 Main Theory

In this section, we analyze the conditional generation process specified by Algorithm 1. We will focus on the scenario where samples x admit a low-dimensional subspace representation, stated as the following assumption.

Assumption 4.1 . Data sampling distribution P x is supported on a low-dimensional linear subspace, i.e., x = Az for an unknown A ∈ R D × d with orthonormal columns and z ∈ R d is a latent variable.

Note that our setup covers the full-dimensional setting as a 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 a function h as follows.

Assumption 4.2 . The ground truth reward f ∗ ( x ) = g ∗ ( x ∥ ) -h ∗ ( x ⊥ ) , where g ∗ ( x ∥ ) = ( θ ∗ ) ⊤ x ∥ where θ ∗ = Aβ ∗ for some β ∗ ∈ R d and ∥ θ ∗ ∥ 2 = ∥ β ∗ ∥ 2 = 1 and h ∗ ( x ⊥ ) is non-decreasing in terms of ∥ x ⊥ ∥ 2 with h ∗ (0) = 0 .

Assumption 4.2 adopts a simple linear reward model for ease of presentation. In this case, we estimate θ ∗ by ridge regression, and (3.1) in Algorithm 1 becomes ̂ θ = argmin θ ∑ n 2 i =1 ( θ ⊤ x i -y i ) 2 + λ ∥ θ ∥ 2 2 for a positive coefficient λ . Later in Section 3.3 and Appendix F, we extend our results beyond linear models to deep ReLU networks.

4.1 Conditional DM Learns Subspace Representation

Recall that Algorithm 1 has two outputs: generated population ̂ P ( ·| ̂ y = a ) and learned representation matrix V . Use notation ̂ P a := ̂ P ( ·| ̂ y = a ) (generated distribution) and P a := P ( ·| ̂ y = a ) (target distribution) for better clarity in result presentation. To assess the quality of subspace learning, we utilize two metrics defined as

∠ ( V, 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 x ∼ ̂ P a [ ∥ x ⊥ ∥ ] is defined as the expected l 2 distance between x 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 on x , i.e. the latent z 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 z follows distribution P z with density p z , such that there exists constants B,C 1 , C 2 verifying p z ( z ) ≤ (2 π ) -( d +1) / 2 C 1 exp ( -C 2 ∥ z ∥ 2 2 / 2 ) whenever ∥ z ∥ 2 > B . And ground truth score is realizable: ∇ log p t ( x | ̂ y ) ∈ S .

Theorem 4.5 ( Subspace Fidelity of Generated Data ) . Under Assumption 4.1, if Assumption 4.3 holds with c 0 I d ⪯ E z ∼ P z [ zz ⊤ ] , then with high probability on data,

Assumption 4.4 . Further assume z ∼ N (0 , Σ) with its covariance matrix Σ satisfying λ min I d ⪯ Σ ⪯ λ max I d for 0 < λ min ≤ λ max ≤ 1 .

with N ( S , 1 /n 1 ) being the log covering number of function class S as in (3.8). When Assumption 4.4 holds, N ( S , 1 /n 1 ) = O (( d 2 + Dd ) log( Ddn 1 )) and thus ∠ ( V, A ) = ˜ O ( 1 λ min √ ( Dd 2 + D 2 d ) n 1 ) . Further under Assumption 4.2, it holds that

where β ∗ is groundtruth parameter of linear model.

4.2 Provable Reward Improvement via Conditional Generation

Let y ∗ be a target reward value and P be a generated distribution. Define the suboptimality of P as

which measures the gap between the expected reward of x ∼ P and the target value y ∗ . In the language of bandit learning, this gap can also be viewed as a form of off-policy regret . Given a target value y ∗ = a , we want to derive guarantees for SubOpt ( ̂ P a ; y ∗ = a ) , recall ̂ P a := ̂ P ( ·| ̂ y = a ) denotes the generated distribution. In Theorem 4.6, we show SubOpt ( ̂ P a ; y ∗ = a ) comprises of three components: 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.

Theorem 4.6 ( Off-policy Regret of Generated Samples ) . Suppose Assumption 4.1, 4.2 and 4.4 hold. We choose λ = 1 , t 0 = ( ( Dd 2 + D 2 d ) /n 1 ) 1 / 6 and ν = 1 / √ D . With high probability, running Algorithm 1 with a target reward value a gives rise to

where ̂ Σ λ := 1 n 2 ( X ⊤ X + λI ) where X is the stack matrix of D label and Σ P a = E P a [ xx ⊤ ] .

Implications and Discussions: (1) Equation (4.4) decomposes the suboptimality gap into two separate parts of error: error from reward learning ( E 1 ) and error coming from diffusion ( E 2 and E 3 ).

E 1 suggests diffusion model is essentially doing representation learning, reducing D to smaller latent dimension d . It can be seen from Tr( ̂ Σ -1 λ Σ p q ) ≤ O ( a 2 ∥ β ∗ ∥ Σ + d ) when n 2 = Ω( 1 λ min ) .

It is also worth mentioning that E 2 and E 3 depend on t 0 and that by taking t 0 = ( ( Dd 2 + D 2 d ) /n 1 ) 1 / 6 one gets a good trade-off and small E 2 + E 3 .

(3) If we ignore the diffusion errors when n 1 is large enough, the suboptimatliy gap resembles the suboptimatliy of off-policy bandit learning in feature subspace [22, Section 3.2], [32, 5]. It shows

the major source of error occurs when moving towards the target distributions and the error behaves similar to bandit.

(4) On-support diffusion error entangles with distribution shift in complicated ways. We show

where DistroShift ( a ) quantifies the distribution shift depending on 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 ∨ d ) . We observe an interesting phase shift. When a < d , the training data have a sufficient coverage with respect to the generated distribution ̂ P a . Therefore, the on-support diffusion error has a lenient linear dependence on a . However, when a > d , the data coverage is very poor and E 2 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 a potential connection between diffusion theory and off-policy bandit learning, which is interesting for more future research. See proofs in Appendix D.3.

4.3 Extension to Nonparametric Function Class

Our theoretical analysis, in its full generality, extends to using general nonparametric function approximation for both the reward and score functions. To keep our paper succinct, we refer to Appendix F and Theorem F.4 for details of our nonparametric theory for reward-conditioned generation. Informally, the regret of generated samples is bounded by

with high probability. Additionally, the nonparamtric generators is able to estimate the representation matrix A up to an error of ∠ ( V, A ) = ˜ O ( n -2 d +6 1 ) . Here the score is assumed to be Lipschitz continuous and α is the smoothness parameter of the reward function, and DistroShift ( a ) is a 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 1 following the setup in Assumption 4.1, 4.2 and 4.4. We choose d = 16 , D = 64 , g ⋆ ( x ) := 5 ∥ x ∥ 2 2 , and generate β ⋆ by uniformly sampling from the unit sphere. The latent variable z is generated from N ( 0 , I d ) , which is then used to construct x = Az with some randomly generated orthonormal matrix A . We use the 1 -dimensional version of the 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 as the distribution shift remains small for these target values. The generation reward begins to decrease as we further increase the target reward value, and the reason is 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.

Samples

2.8

0.55

3.0

2.5

0.50

Generated

2.5

Shift

2.2

Deviation

0.45

2.0

Distribution

2.0

0.40

1.5

1.8

Off-support

0.35

Reward

1.0

0.30

1.5

Average

0.5

1.2

0.25

0.0

0.20

1.0

手

0.0

1.0

2.0

3.0

4.0

5.0

0.0

1.0

2.0

3.0

4.0

5.0

0.0

1.0

2.0

3.0

4.0

5.0

Target Reward Value

Target Reward Value

Target Reward Value

5.0

Target Reward Value =1.0

5.0

Target Reward Value =2.0

5.0

Target Reward Value =3.0

5.0

Target Reward Value = 4.0

4.0

4.0

4.0

4.0

Density

3.0

Density

3.0

Density

3.0

Density

3.0

2.0

2.0

2.0

2.0

1.0

1.0

1.0

1.0

0.0

-1.0

0.0

1.0

2.0

3.0

4.0

0.0

-1.0

0.0

1.0

2.0

3.0

4.0

0.0

-1.0

0.0

1.0

2.0

3.0

4.0

0.0

-1.0

0.0

1.0

2.0

3.0

4.0

Reward

Reward

Reward

Reward

5.2 Directed Text-to-Image Generation

Next, we empirically verify our theory through directed text-to-image generation. Instead of training a diffusion model from scratch, we use Stable Diffusion v1.5 [39], pre-trained on LAION dataset [42]. Stable Diffusion operates on the latent space of 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.

Ground-truth Reward Model. We start from an ImageNet [10] pre-trained ResNet-18 [17] model and replace the final prediction layer with a randomly initialized linear layer of scalar outputs. Then we use this model as the ground-truth reward model. To investigate the meaning of this randomlygenerated 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 I for sample images.

Labelled Dataset. We use the ground-truth reward model to compute a scalar output for each instance in the CIFAR-10 [24] training dataset and perturb the output by adding a Gaussian noise from N (0 , 0 . 01) . We use the images and the corresponding outputs as 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 p t ( y | x t ) by N ( µ θ ( x t ) , σ 2 ) . For simplicity, we set σ 2 as a tunable hyperparameter. We share network parameters for different noise levels t , so our µ θ has no additional input of t . We train µ θ by minimizing the expected KL divergence between p t ( y | x t ) and N ( µ θ ( x t ) , σ 2 ) :

Equivalently, we train the reward model µ θ to predict the noisy reward y from the noisy inputs x t . Also, notice that the minimizers of the objective do not depend on the choice of σ 2 .

Reward-network-based Directed Diffusion. To perform reward-directed conditional diffusion, observe that ∇ x log p t ( x | y ) = ∇ x log p t ( x ) + ∇ x log p t ( y | x ) , and p t ( y | x ) ∝ exp ( -∥ y -µ θ ( x ) ∥ 2 2 2 σ 2 ) . Therefore,

In our implementation, we compute the gradient by back-propagation through µ θ and incorporate this gradient guidance into each denoising step of the DDIM sampler [43] following [11] (equation (14)). We see that 1 /σ 2 corresponds to the weights of the gradient with respect to unconditioned score. In the sequel, we refer to 1 /σ 2 as the 'guidance level', and y as the 'target value'.

Quantitative Results. We vary 1 /σ 2 in { 25 , 50 , 100 , 200 , 400 } and y in { 1 , 2 , 4 , 8 , 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 /σ 2 . Alarger target value puts more weight on the guidance signals ∇ x µ θ ( x ) , which successfully drives the generated images towards higher predicted rewards, but suffers more from the distribution-shift effects between the training distribution and the reward-conditioned distribution, which renders larger gaps between the predicted rewards and the ground-truth rewards. To optimally choose a target value, we must trade off between the two counteractive effects.

25

50

100

200

400

0

2

4

6

8

Reward

ground truth

prediction

baseline

target=1

target=2

target=4

target=8

target=16

target=1

target=2

target=4

target=8

target=16

Guidance Level

Qualitative Results. To qualitatively test the effects of the reward conditioning, we generate a set of images with increasing target values y under different text prompts and investigate the visual properties of the produced images. We isolate the effect of reward conditioning by fixing all the randomness during the generation processes, so the generated images have similar semantic layouts. After hyper-parameter tuning, we find that setting 1 /σ 2 = 100 and y ∈ { 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.

Reward improves astarget value increases

Qualitydegradeswhenoverlyextrapolated

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 our theory well.

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

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://openreview. net/forum?id=sP1fo2K9DFG .

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

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:17981-17993, 2021.

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.01324 , 2022.

David Brandfonbrener, William Whitney, Rajesh Ranganath, and Joan Bruna. Offline contextual bandits with overparameterized models. In International Conference on Machine Learning , pages 1049-1058. PMLR, 2021.

Jinglin Chen and Nan Jiang. Information-theoretic considerations in batch reinforcement learning. In International Conference on Machine Learning , pages 1042-1051. PMLR, 2019.

Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances in Neural Information Processing Systems , 33:22134-22145, 2020.

Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. arXiv preprint arXiv:2302.07194 , 2023.

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 arXiv:2209.11215 , 2022.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

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.

Sixue Gong, Vishnu Naresh Boddeti, and Anil K Jain. On the intrinsic dimensionality of image representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 3987-3996, 2019.

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. arXiv preprint arXiv:2206.09012 , 2022.

László Györfi, Michael Köhler, Adam Krzy˙ zak, and Harro Walk. A distribution-free theory of nonparametric regression , volume 1. Springer, 2002.

Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability , pages 1188-1205, 1986.

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.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , 2022.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems , 31, 2018.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

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.

Jin Sub Lee, Jisun Kim, and Philip M. Kim. Proteinsgm: Score-based generative modeling for de novo protein design. bioRxiv , 2023. doi: 10.1101/2022.07.13.499967. URL https: //www.biorxiv.org/content/early/2023/02/04/2022.07.13.499967 .

Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. Advances in Neural Information Processing Systems , 35:4328-4343, 2022.

Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners. arXiv preprint arXiv:2302.01877 , 2023.

Qiang Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou. Breaking the curse of horizon: Infinite-horizon off-policy estimation. Advances in Neural Information Processing Systems , 31, 2018.

Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research , 9(5), 2008.

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.

Thanh Nguyen-Tang, Sunil Gupta, A Tuan Nguyen, and Svetha Venkatesh. Offline neural contextual bandits: Pessimism, optimization and generalization. arXiv preprint arXiv:2111.13807 , 2021.

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 arXiv:2112.10741 , 2021.

Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models , 2023. URL https://openreview.net/forum?id=6961CeTSFA .

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.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= Pv1GPQzRrC8 .

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.

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.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

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 http://arxiv.org/abs/2103.00453 .

Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402 , 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=St1giarCHLP .

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems , 33:12438-12448, 2020.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence , pages 574-584. PMLR, 2020.

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.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023.

Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science , 290(5500):2319-2323, 2000.

Alexandre B. Tsybakov. Introduction to Nonparametric Estimation . Springer Publishing Company, Incorporated, 1st edition, 2008. ISBN 0387790519.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems , 34:11287-11302, 2021.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543 , 2023.

Index of Supplementary Materials

§A Preliminaries on Diffusion Models

§B Limitations and Open Problems

§C Omitted Proof in Section 3

-§C.1 Proof of Proposition 3.1

§D 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

§E Supporting Lemmas and Proofs for §D

§F Extensions to Nonparametric Function Class

§G Omitted Proofs in §F

§H Proof for §D.1

§I 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-Ulhenbeck process for instance:

where initial X 0 ∼ P data and ( W t ) t ≥ 0 is a standard Wiener process, and g ( t ) is a nondecreasing weighting function. In practice, the forward process (A.1) terminates at a sufficiently large T > 0 such that the corrupted X T is close to the standard Gaussian N ( 0 , I D ) . To enable data generation in future, the score ∇ log p t ( · ) at t is the key to learn, here p t denotes the marginal density of X t . We often use an estimated score function ̂ s ( · , t ) trained by minimizing a score matching loss.

Backward SDE for Generation. Diffusion models generate samples through a backward SDE (A.2) reversing the time in (A.1) [2, 16], i.e.,

where W t is a reversed Wiener process. In practice, the backward process is initialized with N (0 , I D ) and the unknown conditional score ∇ log p t ( · ) is replaced by an estimated counterpart ̂ s ( · , t ) .

B Limitations

We do not see 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 t ≥ 0 , it hold that ∇ x t log p t ( x t | y ) = ∇ x t log p t ( x t , y ) since the gradient is taken w.r.t. x t only. Then plugging in this equation and expanding the norm square on the LHS gives

Then it suffices to prove

Using integration by parts to rewrite the inner product we have

where denote by ϕ t ( x ′ | x ) the density of N ( α ( t ) x, h ( t ) I D ) with α ( t ) = exp( -t/ 2) and h ( t ) = 1 -exp( -t ) , then

D Omitted Proofs in Section 4

Additional Notations: We follow the notations in the main paper along with some additional ones. Use P LD t ( z ) to denote the low-dimensional distribution on z corrupted by diffusion noise. Formally, p LD t ( z ) = ∫ ϕ t ( z ′ | z ) p z ( z )d z with ϕ t ( ·| z ) being the density of N ( α ( t ) z, h ( t ) I d ) . P LD t 0 ( z | ̂ f ( Az ) = a ) the corresponding conditional distribution on ̂ f ( Az ) = a at t 0 , with shorthand as P LD t 0 ( a ) . Also give P z ( z | ̂ f ( Az ) = a ) a shorthand as P LD ( a ) . In our theorems, O hides constant factors and higher order terms in n -1 1 and n -1 2 and , ˜ O further hides logarithmic terms and can also hide factors in d .

D.1 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.1 we provide a new one specific to our setting where the true score is linear in input ( x t , ̂ y ) 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 as S consists of an encoder-decoder structure for estimating matrix A to reduce dimension (see §H for S construction and more proof details).

In the following lemma, we first present a 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 N ( S , 1 /n 1 ) , the log covering number of S , recall n 1 is the size of D unlabel . Instantiating this general result, we derive score matching error for Gaussian case by upper bounding N ( S , 1 /n 1 ) in this special case.

Lemma D.1. Under Assumption 4.1, if ∇ log p t ( x | y ) ∈ S , where

with parametric. Then for , with probability , the square score matching error is bounded

recall P t comes from P x ̂ y by noising x at t in the forward process. Under Assumption 4.4 and given ̂ f ( x ) = ̂ θ ⊤ x and ̂ y = ̂ f ( x ) + ξ, ξ ∼ N (0 , ν 2 ) , the score function ∇ log p t ( x | ̂ y ) to approximate is linear in x and ̂ y . When approximated by S with Ψ linear, N ( S , 1 /n 1 ) = O (( d 2 + Dd ) log( Ddn 1 )) .

Proof. Proof is in § H.

To provide fidelity and reward guarantees of ̂ P a : the generated distribution of x given condition ̂ y = a , we will need the following lemma. It provides a subspace recovery guarantee between V (score matching output) and A (ground truth), as well as a distance measure between distributions P a and ̂ P a , given score matching error ϵ diff .

Note P a and ̂ P a are over x , which admits an underlying low-dimensional structure x = Az . Thus we measure distance between P a and ̂ P a by defining

d TV ( · , · ) is the TV distance between two distribution.

f ♯ P denotes a push-forward measure, i.e., for any measurable Ω , ( f ♯ P )(Ω) = P ( f -1 (Ω))

( V ⊤ ) # ̂ P a pushes generated ̂ P a forward to the low dimensional subspace using learned subspace matrix V . U is an orthonormal matrix of dimension d .

P LD t 0 ( z | ̂ f ( Az ) = a ) is close to ( A ⊤ ) # P a , with t 0 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 a 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 P 1 and P 2 restricted under function class L is defined as

Definition D.3 is well perceived in bandit and RL literature [30, 29, 6, 12].

Lemma D.4. Given the square score matching error (D.1) upper bounded by ϵ 2 diff , and when P z satisfying Assumption 4.3 with c 0 I d ⪯ E z ∼ P z [ zz ⊤ ] , it guarantees on for x ∼ ̂ P a and ∠ ( V, A ) := ∥ V V ⊤ -AA ⊤ ∥ 2 F that

In addition,

with ¯ S = { 1 T -t 0 ∫ T t 0 E x t | x ∥∇ log p t ( x t | y ) -s ( x t , y, t ) ∥ 2 2 d t : s ∈ S } . TV ( ̂ P a ) and T ( P ( x, ̂ y = a ) , P x ̂ y ; ¯ S ) are defined in Definition D.2 and D.3.

Proof. Proof is in § E.2.

D.2 Proof of Theorem 4.5

Proof. Proof of ∠ ( V, A ) . By Lemma 3 of [8], we have

when the latent z satisfying Assumption 4.3 and c 0 I d ⪯ E z ∼ P z [ zz ⊤ ] . Therefore, by (D.1), we have with high probability that

When Assumption 4.4 holds, plugging in c 0 = λ min and N ( S , 1 /n 1 ) = O (( d 2 + Dd ) log( Ddn 1 )) , it gives

where ˜ O hides logarithmic terms. When D > d 2 , which is often the case in practical applications, we have

Proof of E x ∼ ̂ P a [ ∥ x ⊥ ∥ 2 ] . By the definition of x ⊥ that x ⊥ = ( I D -AA ⊤ ) x ,

Score matching returns V as an approximation of A , then

where by (D.2) in Lemma D.4 we have for some constant c ≥ 0 . Thus

On the other hand,

where ∥ V V ⊤ -AA ⊤ ∥ 2 F has an upper bound as in (D.3) and E x ∼ ̂ P a [ ∥ x ∥ 2 2 ] ) is bounded in Lemma E.3 by

with M ( a ) = O ( a 2 ∥ β ∗ ∥ Σ + d ) .

Therefore, to combine things together, we have

O hides multiplicative constant and √ ∠ ( V, A ) t 0 D , √ ∠ ( V, A ) M ( a ) TV ( ̂ P a ) , which are terms with higher power of n -1 1 than the leading term.

Remark of Theorem 4.5. 1. Guarantee (4.2) applies to general distributions with light tail as assumed in Assumption 4.3.

Guarantee (4.3) guarantees high fidelity of generated data in terms of staying in the subspace when we have access to a large unlabeled dataset.

Guarantee (4.3) shows that E x ∼ ̂ P a [ ∥ x ⊥ ∥ 2 ] scales up when t 0 goes up, which aligns with the dynamic in backward process that samples are concentrating to the learned subspace as t 0 goes to 0 . Taking t 0 → 0 , E x ∼ ̂ P a [ ∥ x ⊥ ∥ ] has the decay in O ( n -1 4 1 ) . However, taking t 0 → 0 is not ideal for the sake of high reward of x , we take the best trade-off of t 0 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 E 1 Relating to Offline Bandits , Bounding E 2 and the Distribution Shift in Diffusion .

D.3.1 SubOpt ( ̂ P a ; y ∗ = a ) Decomposition

Proof. Recall notations ̂ P a := ̂ P ( ·| ̂ y = a ) (generated distribution) and P a := P ( ·| ̂ y = a ) (target distribution) and f ∗ ( x ) = g ∗ ( x ∥ ) -h ∗ ( x ⊥ ) . E x ∼ ̂ P a [ f ⋆ ( x )] can be decomposed into 3 terms:

E 1 comes from regression: prediction/generalization error onto P a , which is independent from any error of distribution estimation that occurs in diffusion. E 2 and E 3 do not measure regression-predicted ̂ f , thus they are independent from the prediction error in ̂ f for pseudo-labeling. E 2 measures the disparity between ̂ P a and P a on the subspace support and E 3 measures the off-subspace component in generated ̂ P a .

D.3.2 Bounding E 1 Relating to Offline Bandits

For all x i ∈ D label , y i = f ∗ ( x i ) + ϵ i = g ( x i ) + ϵ i . Thus, trained on D label the prediction model ̂ f is essentially approximating g . By estimating θ ∗ with ridge regression on D label , we have ̂ f ( x ) = ̂ θ ⊤ x with

Lemma D.5. Under Assumption 4.1 and 4.2 and given ϵ i ∼ N ( 0 , σ 2 ) , define V λ := X ⊤ X + λI , ̂ Σ λ := 1 n 2 V λ and Σ P a := E x ∼ P a xx ⊤ the covariance matrix (uncentered) of P a , and take λ = 1 , then with high probability

Proof. Proof is in § E.3.

Lemma D.6. Under Assumption 4.1, 4.2 and 4.4, when λ = 1 , P a has a shift from the empirical marginal of x in dataset by

when n 2 = Ω(max { 1 λ min , d ∥ β ∗ ∥ 2 Σ } ) .

Proof. Proof is in § E.4.

D.3.3 Bounding E 2 and the Distribution Shift in Diffusion

Lemma D.7. Under Assumption 4.1, 4.2 and 4.4, when t 0 = ( ( Dd 2 + D 2 d ) /n 1 ) 1 / 6

Proof. Proof is in § E.5.

Note that T ( P ( x, ̂ y = a ) , P x ̂ y ; ¯ S ) depends on a and measures the distribution shift between the desired distribution P ( x, ̂ y = a ) and the data distribution P x ̂ y . To understand this distribution's dependency on a , it what follows we give T ( P ( x, ̂ y = a ) , P x ̂ y ; ¯ S ) a shorthand as DistroShift 2 ( a ) and give it an upper bound in one special case of the problem.

Distribution Shift In the special case of covariance Σ of z is known and ∥ A -V ∥ 2 2 = O ( ∥ AA ⊤ -V V ⊤ ∥ 2 F ) , we showcase a bound on the distribution shift in E 2 , as promised in the discussion following Theorem 4.6. We have

where ℓ ( x, y ; ̂ s ) = 1 T -t 0 ∫ T t 0 E x ′ | x ∥∇ x ′ log ϕ t ( x ′ | x ) -̂ s ( x ′ , y, t ) ∥ 2 2 d t . By Proposition 3.1, it suffices to bound

We expand the difference ∥∇ log p t ( x, y ) -̂ s ( x, y, t ) ∥ 2 2 by

where we recall B t is defined in (H.2) and in the last inequality, we use ( a + b ) 2 ≤ 2 a 2 +2 b 2 . In the case of covariance matrix Σ is known, i.e., B t is known, we also consider matrix V directly matches A without rotation. Then by [8, Lemma 3 and 17], we have ∥ A -V ∥ 2 2 = O ( ∥ AA ⊤ -V V ⊤ ∥ 2 F ) =

O ( t 0 /c 0 E P x ̂ y [ ℓ ( x, y ; ̂ s )] ) . To this end, we only need to find E P x | ̂ y = a [ ∥ x ∥ 2 2 ] . Since we consider on-support x , which can be represented as x = Az , we have ∥ x ∥ 2 = ∥ z ∥ 2 . Thus, we only need to find the conditional distribution of z | ̂ y = a . Fortunately, we know ( z, ̂ y ) is jointly Gaussian, with mean 0 and covariance

Consequently, the conditional distribution of z | ̂ y = a is still Gaussian, with mean Σ ̂ βa/ ( ̂ β ⊤ Σ ̂ β + ν 2 ) and covariance Σ -Σ ̂ β ̂ β ⊤ Σ / ( ̂ β ⊤ Σ ̂ β + ν 2 ) . Hence, we have

We integrate over t for the numerator in DistroShift ( a ) to obtain E P x | ̂ y = a [ ∫ T t 0 ∥∇ log p t ( x, y ) -̂ s ( x, y, t ) ∥ 2 2 d t = O ( ( a 2 ∨ d ) 1 c 0 E P x ̂ y [ ℓ ( x, y ; ̂ s )] ) . Note the cancellation between the numerator and denominator, we conclude

As d is a natural upper bound of √ d and viewing c 0 as a constant, we have DistroShift ( a ) = O ( a ∨ d ) as desired.

E Supporting lemmas and proofs

E.1 Supporting lemmas

Lemma E.1. The estimated subspace V satisfies

for some orthogonal matrix U ∈ R d × d .

Proof. Proof is in

Lemma E.2. Suppose P 1 and P 2 are two distributions over R d and m is a function defined on R d , then | E x ∼ P 1 [ m ( z )] -E z ∼ P 2 [ m ( z )] | can be bounded in terms of d TV ( P 1 , P 2 ) , specifically when P 1 and P 2 are Gaussians and m ( z ) = ∥ z ∥ 2 2 :

When P 1 and P 2 are Gaussians and m ( z ) = ∥ z ∥ 2 :

Proof. Proof is in § E.7.

LemmaE.3. Wecompute E z ∼ P LD ( a ) [ ∥ z ∥ 2 2 ] , E z ∼ P t LD 0 ( a ) [ ∥ z ∥ 2 2 ] , E x ∼ ̂ P a [ ∥ x ∥ 2 2 ] , E z ∼ ( U ⊤ V ⊤ ) # ̂ P a [ ∥ z ∥ 2 2 ] in this Lemma.

Let M ( a ) := E z ∼ P LD ( a ) [ ∥ z ∥ 2 2 ] , which has an upper bound M ( a ) = O ( a 2 ∥ β ∗ ∥ Σ + d ) .

Proof. Proof is in § E.8.

E.2 Proof of Lemma D.4

Proof. The first two assertions (D.2) and (D.3) are consequences of [8, Theorem 3, item 1 and 3]. To show (D.4), we first have the conditional score matching error under distribution shift being

where T ( P ( x, ̂ y = a ) , P x ̂ y ; ¯ S ) accounts for the distribution shift as in the parametric case (Lemma D.7). Then we apply [8, Theorem 3, item 2] to conclude

The proof is complete.

E.3 Proof of Lemma D.5

Proof. Given

then things to prove are

where the second inequality is to be proven with high probability w.r.t the randomness in D label . For (E.7), E ̂ P a ∥ x ∥ V -1 λ ≤ √ E ̂ P a x ⊤ V -1 λ x = √ E ̂ P a trace( V -1 λ xx ⊤ ) = √ trace( V -1 λ E ̂ P a xx ⊤ ) .

For (E.8), what's new to prove compared to a classic bandit derivation is its d dependency instead of D , due to the linear subspace structure in x . From the closed form solution of ̂ θ , we have

Therefore,

where λ ∥ θ ∗ ∥ V -1 λ ≤ √ λ ∥ θ ∗ ∥ 2 ≤ √ λ and

Let Z ⊤ = ( z 1 , · · · , z i , · · · , z n 2 ) s.t. Az i = x i , then it holds that X = ZA ⊤ , and XX ⊤ = ZA ⊤ AZ ⊤ = ZZ ⊤ , thus

With probability 1 -δ , ∥ z i ∥ 2 ≤ d + √ d log ( 2 n 2 δ ) := L 2 , ∀ i ∈ [ n 2 ] . Then Theorem 1 in 'Improved algorithms for linear stochastic bandits' (by Yasin Abbasi-Yadkori, David Pal, and Csaba Szepesvari) gives rise to

with probability 1 -δ/ 2 . Combine things together and plugging in λ = 1 , L 2 = d + √ d log ( 2 n 2 δ ) , we have with high probability

E.4 Proof of Lemma D.6

Proof. Recall the definition of ̂ Σ λ and Σ P a that

where X are stack matrix of data supported on A and P a is also supported on A , A is the subspace encoded by matrix A . The following lemma shows it is equivalent to measure trace( ̂ Σ -1 λ Σ P a ) on A subspace.

Lemma E.4. For any P.S.D. matrices Σ 1 , Σ 2 ∈ R d × d and A ∈ R D × d such that A ⊤ A = I d , we have

The lemma above allows us to abuse notations ̂ Σ λ and Σ P a in the following way while keeping the same trace( ̂ Σ -1 λ Σ P a ) value:

where Z ⊤ = ( z 1 , · · · , z i , · · · , z n 2 ) s.t. Az i = x i and recall notataion P LD ( a ) = P z ( z | ̂ f ( Az ) ) .

Given z ∼ N ( µ, Σ) , as a proof artifact, let ̂ f ( x ) = ̂ θ ⊤ x + ξ, ξ ∼ N (0 , ν 2 ) where we will let ν → 0 in the end, then let ̂ β = A ⊤ ̂ θ ∈ R d , ( z, ̂ f ( Az )) has a joint distribution

Then we have the conditional distribution z | ̂ f ( Az ) = a following

When µ = 0 , we compute trace( ̂ Σ -1 λ Σ P a ) as

By the Lemma 3 in [7], it holds that

Therefore,

Then, what left is to bound ∥ ̂ β ∥ Σ = ∥ ̂ θ ∥ A Σ A ⊤ ≥ ∥ θ ∗ ∥ A Σ A ⊤ -∥ ̂ θ -θ ∗ ∥ A Σ A ⊤ by triangle inequality. On one hand, ∥ θ ∗ ∥ A Σ A ⊤ = ∥ β ∗ ∥ Σ . (E.14)

with high probability. Thus when n 2 = Ω( d ∥ β ∗ ∥ 2 Σ

Therefore

E.5 Proof of lemma D.7

Proof. Recall the definition of g ( x ) that

note that g ( x ) = θ ∗⊤ x when x is supported on A . Thus,

where

) .

On the other hand,

here

Then what is left to bound is

Then for term e 3 , by Lemma E.2, we get

where we use E z ∼ ( U ⊤ V ⊤ ) # ̂ P a [ ∥ z ∥ 2 ] ≤ x ∼ ̂ P a [ ∥ x ∥ 2 ] For e 4 , we have

Therefore, by combining things together, we have

By Lemma D.4 and Lemma E.1, we have

And by Lemma E.3

Therefore. leading term in E 2 is

By plugging in score matching error ϵ 2 diff = ˜ O ( 1 t 0 √ Dd 2 + D 2 d n 1 ) , we have

When t 0 = ( ( Dd 2 + D 2 d ) /n 1 ) 1 / 6 , it admits the best trade off in E 2 and E 2 is bounded by

E.6 Proof of Lemma E.1

Proof. From Lemma 17 in [8], we have

Then it suffices to bound where

Thus

which is because trace ( A ⊤ V V ⊤ A ) is calcualted as

Then we will bound ∣ ∣ trace ( AA ⊤ ( V V ⊤ -AA ⊤ ) )∣ ∣ by ∥ V V ⊤ -AA ⊤ ∥ F ,

∥ -∥ O

E.7 Proof of Lemma E.2

Since P 1 and P 2 are Gaussains, ∫ ∥ z ∥ 2 >R ∥ z ∥ 2 2 p 1 ( z )d z and ∫ ∥ z ∥ 2 >R ∥ z ∥ 2 2 p 2 ( z )d z are bounded by some constant C 1 when R 2 ≥ C 2 max { E z ∼ P 1 [ ∥ z ∥ 2 2 ] , E z ∼ P 2 [ ∥ z ∥ 2 2 ] } as suggested by Lemma 16 in [8].

Therefore,

Then

since d TV ( P 1 , P 2 ) decays with n 1 .

Similarly, when m ( z ) = z

where ∫ ∥ z ∥ 2 >R ∥ z ∥ 2 2 p 1 ( z )d z and ∫ ∥ z ∥ 2 >R ∥ z ∥ 2 2 p 2 ( z )d z are bounded by some constant C 1 when R 2 ≥ C 2 max { E z ∼ P 1 [ ∥ z ∥ 2 2 ] , E z ∼ P 2 [ ∥ z ∥ 2 2 ] } as suggested by Lemma 16 in [8]. Therefore,

E.8 Proof of Lemma E.3

Proof. Recall from (E.12) that

and by Lemma D.6

Thus E z ∼ P LD ( a ) [ ∥ z ∥ 2 2 ] = M ( a ) , M ( a ) = O ( a 2 ∥ β ∗ ∥ Σ + d ) .

Thus after adding diffusion noise at t 0 , we have for α ( t ) = e -t/ 2 and h ( t ) = 1 -e -t :

Thus E z ∼ P t LD 0 ( a ) [ ∥ z ∥ 2 2 ] ≤ M ( a ) + t 0 d .

By orthogonal decomposition we have

where E x ∼ ̂ P a [ ∥ ( I D -V V ⊤ ) x ∥ 2 2 ] is bounded by (D.5) and the distribution of U ⊤ V ⊤ x , which is ( U ⊤ V ⊤ ) # ̂ P a , is close to P LD t 0 ( a ) up to TV ( ̂ P a ) , which is defined in Definition D.2. Then by Lemma E.2, we have

E.9 Proof of Lemma E.4

Proof. Firstly, one can verify the following two equations by direct calculation:

Then we have

Therefore,

which has finished the proof.

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

where g ∗ ( x ∥ ) is α -Hölder continuous for α ≥ 1 and h ∗ ( x ⊥ ) is nondecreasing in terms of ∥ x ⊥ ∥ 2 with h ∗ (0) = 0 . Moreover, g ∗ has a bounded Hölder norm, i.e., ∥ g ∗ ∥ H α ≤ 1 .

Hölder continuity is widely studied in nonparametric statistics literature [15, 49]. h ∗ here penalizes off-support extrapolation.

Under Assumption F.1, we use nonparametric regression for estimating f ∗ . Specifically, we specialize (3.1) in Algorithm 1 to

where F = NN( L, M, J, K, κ ) is chosen to be a class of neural networks. Hyperparameters in F will be chosen properly in Theorem F.4.

Our theory also considers generic sampling distributions on x . Since x lies in a low-dimensional subspace, this translates to a sampling distribution assumption on latent variable z .

Assumption F.2 . The latent variable z follows distribution P z with density p z , such that there exists constants B,C 1 , C 2 verifying p z ( z ) ≤ (2 π ) -( d +1) / 2 C 1 exp ( -C 2 ∥ z ∥ 2 2 / 2 ) whenever ∥ z ∥ 2 > B . And c 0 I d ⪯ E z ∼ P z [ zz ⊤ ] .

Assumption F.2 says P z has a light tail, which is standard in high-dimensional statistics [51, 53]. Assumption F.2 also encodes distributions with a compact support. Furthermore, we assume that the curated data ( x, ̂ y ) induces Lipschitz conditional scores. Motivated by Chen et al. [8], we show that the linear subspace structure in x leads to a similar conditional score decomposition ∇ log p t ( x | ̂ y ) = s ∥ ( x, ̂ y, t ) + s ⊥ ( x, ̂ y, t ) , where s ∥ is the on-support score and s ⊥ is the orthogonal score. The decomposition for conditional score is as (H.1), which applies to both parametric and non-parametric cases. The following assumption is imposed on s ∥ .

Assumption F.3 . The on-support conditional score function s ∥ ( x, ̂ y, t ) is Lipschitz with respect to x, ̂ y for any t ∈ (0 , T ] , i.e., there exists a constant C lip , such that for any x, ̂ y and x ′ , ̂ y ′ , it holds

Lipschitz score is 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 in the linear reward setting, except we replace m by a 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 to different network architectures poses no real difficulty. Given an input x , neural networks compute

where W i and b i are weight matrices and intercepts, respectively. We then define a class of neural networks as

NN( L, M, J, K, κ ) = { f : f in the form of (F.1) with L layers and width bounded by M,

For the conditional score network, we will additionally impose some Lipschitz continuity requirement, i.e., ∥ f ( x ) -f ( y ) ∥ 2 ≤ c lip ∥ x -y ∥ 2 for some Lipschitz coefficient c lip .

Recall the distribution shift defined in Definition D.3 that

for arbitrary two distributions P 1 , P 2 and function class L . Similar to the parametric case, use notation ̂ P a := ̂ P ( ·| ̂ y = a ) and P a := P ( ·| ̂ y = a ) . Then we can bound SubOpt ( ̂ P 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 F.4. Suppose Assumption 4.1, F.1, F.2 and F.3 hold. Let δ ( n ) = d log log n log n . Properly chosen F and S , with high probability, running Algorithm 1 with a target reward value a and stopping at 1

where M ( a ) := E z ∼ P ( a ) [ ∥ z ∥ 2 2 ] and

E 3 penalizes the component in ̂ P a that is off the truth subspace. The function classes F and S are chosen as F = NN( L f , M f , J f , K f , κ f ) with

and S = NN( L s , M s , J s , K s , κ s ) with

Moreover, S is also Lipschitz with respect to ( x, y ) and the Lipschitz coefficient is c lip = O (10 dC lip ) .

Remark. The proof is provided in Appendix G.2. Here we correct a typo on n -1 3( d +5) 1 in main paper Section 4.3. Quantities T ( P ( x | ̂ y = a ) , P x ; ¯ F ) and T ( P ( x, ̂ y = a ) , P x ̂ y ; ¯ S ) depend on a characterizing the distribution shift. The δ ( n ) terms account for the unbounded domain of x , which is negligible when n is large. In the main paper, we omit δ ( n ) in the regret bound.

G Omitted Proofs in Section F

G.1 Conditional Score Decomposition and Score Matching Error

Lemma G.1. Under Assumption 4.1, F.2 and F.3, with high probability

Proof. [8, Theorem 1] is easily adapted here to prove Lemma G.1 with the input dimension d +1 and the Lipschitzness in Assumption F.3. Network size of S is implied by [8, Theorem 1] with ϵ = n -1 d +6 1 accounting for the additional dimension of reward ̂ y and then the score matching error follows.

G.2 Proof of Theorem F.4

Additional Notations: Similar as before, use P LD t ( z ) to denote the low-dimensional distribution on z corrupted by diffusion noise. Formally, p LD t ( z ) = ∫ ϕ t ( z ′ | z ) p z ( z )d z with ϕ t ( ·| z ) being the density of N ( α ( t ) z, h ( t ) I d ) . P LD t 0 ( z | ̂ f ( Az ) = a ) the corresponding conditional distribution on ̂ f ( Az ) = a at t 0 , with shorthand as P LD t 0 ( a ) . Also give P z ( z | ̂ f ( Az ) = a ) a shorthand as P LD ( a ) .

G.2.1 SubOpt ( ̂ P a ; y ∗ = a ) Decomposition

By the same argument as in § D.3.1, we have

G.2.2 E 1 : Nonparamtric Regression Induced Error

Nonparametric Regression Error of ̂ f Since P z has a light tail due to Assumption F.2, by union bound and [8, Lemma 16], we have

where C 1 , C 2 are constants and Γ( · ) is the Gamma function. Choosing R = O ( √ d log d +log n δ ) ensures P ( ∃ x i with ∥ x i ∥ 2 > R for i = 1 , . . . , n 2 ) < δ . On the event E = {∥ x i ∥ 2 ≤ R for all i = 1 , . . . , n 2 } , denoting δ ( n 2 ) = d log log n 2 log n 2 , we have

by [31, Theorem 7] with a new covering number of S , when n 2 is sufficiently large. The corresponding network architecture follows from Theorem 2 in 'Nonparametric Regression on Low-Dimensional Manifolds using Deep ReLU Networks : Function Approximation and Statistical Recovery'.

We remark that linear subspace is a special case of low Minkowski dimension. Moreover, δ ( n 2 ) is asymptotically negligible and accounts for the truncation radius R of x i 's (see also [8, Theorem 2 and 3]). The covering number of S is ˜ O ( d d/ 2 n -d α (log n 2 ) d/ 2 + Dd ) as appear in [8, Proof of Theorem

2]. Therefore

G.2.3 E 2 : Diffusion Induced On-support Error

Suppose L 2 score matching error is ϵ 2 diff ( n 1 ) , i.e.

We revoke Definition D.2 measuring the distance between ̂ P a to P a that

Lemma D.4 applies to nonparametric setting, so we have

In addition,

E 2 will be bounded by

where for ∣ ∣ ∣ E x ∼ ̂ P a [ g ∗ ( V V ⊤ x ) -g ∗ ( AA ⊤ x )] ∣ ∣ ∣ , we have

For the other term ∣ ∣ ∣ E x ∼ P a [ g ∗ ( AA ⊤ x )] -E x ∼ ̂ P a [ g ∗ ( V V ⊤ x )] ∣ ∣ ∣ , we will bound it with TV ( ̂ P a ) .

Since any z ∼ P t 0 ( a ) can be represented by α ( t 0 ) z + √ h ( t 0 ) u , where z ∼ P ( a ) , u ∼ N (0 , I d ) , then

thus

where we further use 1 -α ( t 0 ) = 1 -e -t 0 / 2 ≤ t 0 / 2 , h ( t 0 ) ≤ 1 .

where and

Combining things up, we have

Similar to parametric case, Let M ( a ) := E z ∼ P ( a ) [ ∥ z ∥ 2 2 ] , then expect for in nonparametric case, we can not compute M ( a ) out as it is not Gaussian. But still, with higher-order terms in n -1 1 hided, we have

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

where equality ( i ) follows from the fact AA ⊤ x ⊥ ( I D -AA ⊤ ) x and C is the normalizing constant of Gaussian densities. Taking logarithm and then derivative with respect to x on p t ( x, y ) , we obtain

Note that the first term in the right-hand side above only depends on A ⊤ x and y . Therefore, we can compactly write ∇ x log p t ( x, y ) as

where mapping u represents

We observe that (H.1) motivates our choice of the neural network architecture S in (3.8). In particular, ψ attempts to estimate u and matrix V attempts to estimate A .

In the Gaussian design case (Assumption 4.4), we instantiate p z ( z ) to the Gaussian density (2 π | Σ | ) -d/ 2 exp ( -1 2 z ⊤ Σ -1 z ) . Some algebra on the Gaussian integral gives rise to

where we have denoted

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

as the loss function for a pair of clean data ( x, y ) and a conditional score function s . Further, we denote the population loss as

whose empirical counterpart is denoted as ̂ L ( s ) = 1 n 1 ∑ n 1 i =1 ℓ ( x i , y i ; s ) .

To bound the score estimation error, we begin with an oracle inequality. Denote L trunc ( s ) as a truncated loss function defined as

where R > 0 is a truncation radius chosen as O ( √ d log d +log K +log n 1 δ ) . Here K is a uniform upper bound of s ( x, y, t ) 1 {∥ x ∥ 2 ≤ R, | y | ≤ R } for s ∈ S , i.e., sup s ∈S ∥ s ( x, y, t ) 1 {∥ x ∥ 2 ≤ R, | y | ≤ R }∥ 2 ≤ K . To this end, we have where equality ( i ) holds since S contains the ground truth score function. We bound term ( A ) by a PAC-learning concentration argument. Using the same argument in [8, Theorem 2, term ( A ) ], we have

Applying the standard metric entropy and symmetrization technique, we can show

where ̂ R is 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, we have

We emphasize that the log covering number considers x, y in the truncated region. Taking ϵ = 1 n 1 gives rise to

Here K is instance dependent and majorly depends on d . In the Gaussian design case, we can verify that K is O ( √ d ) . To this end, we deduce ( A ) = ˜ O ( 1 t 0 √ d 2 N ( S , 1 /n 1 ) log 1 δ n 1 ) . In practice, d is often much smaller than D (see for example [37], where ImageNet has intrinsic dimension no more than 43 in contrast to image resolution of 224 × 224 × 3 ). In this way, we can upper bound d 2 by D , yet d 2 is often a tighter upper bound.

For term ( B ) , we invoke the same upper bound in [8, Theorem 2, term ( B ) ] to obtain

which is negligible compared to ( A ) . Therefore, summing up ( A ) and ( B ) , we deduce

Gaussian Design We only need to find the covering number under the Gaussian design case. Using (H.2), we can construct a covering from coverings on matrices V and Σ -1 . Suppose V 1 , V 2 are two matrices with ∥ V 1 -V 2 ∥ 2 ≤ η V for some η > 0 . Meanwhile, let Σ -1 1 , Σ -1 2 be two covariance matrices with ∥ Σ -1 1 -Σ -1 2 ∥ 2 ≤ η Σ . Then we bound

where for bounding ( ♠ ) , we invoke the identity ∥ ( I + A ) -1 -( I + B ) -1 ∥ 2 ≤ ∥ B -A ∥ 2 . Further taking supremum over t ∈ [ t 0 , T ] leads to

for any t ∈ [ t 0 , T ] . Therefore, the inequality above suggests that coverings on V and Σ -1 form a covering on S . The covering numbers of V and Σ -1 can be directly obtained by a volume ratio argument; we have

Thus, the log covering number of S is

where we have plugged ν 2 = 1 /D into the last inequality. Setting η = 1 /n 1 and substituting into ϵ 2 diff 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 via Bernstein-type concentration may result in a better dependence on n 1 . Nonetheless, the optimal dependence should not beat a 1 /n 1 -rate.

I Additional Experimental Results

I.1 Simulation

We generate the latent sample z from standard normal distribution z ∼ N ( 0 , I d ) and set x = Az for a randomly generated orthonormal matrix A ∈ R D × d . The dimensions are set to be d = 16 , D = 64 . The reward function is set to be f ( x ) = ( θ ⋆ ) ⊤ x ∥ + 5 ∥ x ⊥ ∥ 2 2 , where θ ⋆ is defined by Aβ ⋆ . We generate β ⋆ by uniformly sampling from the unit sphere.

When estimating ̂ θ , we set λ = 1 . 0 . The score matching network is based on the UNet implementation from https://github.com/lucidrains/denoising-diffusion-pytorch , where we modified the class embedding so it accepts continuous input. The predictor is 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 8 × 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 .

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.

I.2 Directed Text-to-Image Generation

Samples of high rewards and low rewards from the ground-truth reward model. In Section 5.2, the ground-truth reward model is built by replacing the final prediction layer of the ImageNet pretrained ResNet-18 model with a randomly initialized linear layer of scalar outputs. To investigate the meaning of this randomly-generated reward model, we generate images using Stable Diffusion and filter out images with rewards ≥ 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. In our implementation, as the Stable Diffusion model operates on the latent space of its V AE, we build a 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-Expert, Hopper) setting. In Decision Diffuser, the RL trajectory and the final reward are jointly modeled with a conditioned diffusion model. The policy is given by first performing reward-directed conditional diffusion to sample a

Total Reward

(a) A positive sample

(b) A positive sample

(c) A negative sample

(d) A negative sample

NM

恭

米

派

trajectory of high reward and then extracting action sequences from the trajectory using a trained inverse dynamics model. We plot the mean and standard deviation of the total returns (averaged across 10 independent episodes) v.s. the target rewards in Figure 8. The theoretical maximum of the reward is 400. Therefore, trajectories with rewards greater than 400 are never seen during training. We observe that when we 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.

80.0

160.0

240.0

320.0

360.0

400.0

600.0

800.0

1000.0

1200.0

1400.0

1600.0

3200.0

6400.0

0

20

40

60

80

100

120

Decision-Diffuser

BC

CQL

IQL

DT

TT

MOReL

Diffuser

Target Reward

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nStrength:\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\nWEAKNESSES:\n1) 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.\n\nQUESTIONS:\nSome 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\nLIMITATIONS:\nThe paper did not discuss any limitations. \nThe paper has no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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\nSTRENGTHS:\nOverall 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\nWEAKNESSES:\nOverall 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nIn 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nwhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nSteering 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\nWEAKNESSES:\nI 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]\")\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors have clearly laid out the assumptions for the theoretical analysis. I do not see any major limitations within this scope.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}