SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hui Yuan
Chengzhuo Ni
Kaixuan Huang
Minshuo Chen
Mengdi Wang *
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.
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 score-matching generative models that have demonstrated the state-of-the-art performances in various domains, such as image generation [39, 38, 4] and audio generation, with fascinating potentials in broader domains, including text modeling [3, 27], reinforcement learning [21, 1, 36, 28] and protein structure modeling [26]. Diffusion models are trained to predict a clean version of the noised input, and generate data by sequentially removing noises and trying to find a cleaner version of the input. The denoising network (a.k.a. score network)  $s(x,t)$  approximates the score function  $\nabla \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].
*Department of Electrical and Computer Engineering, Princeton University. Authors' emails are: {huiyuan, kaixuanh, cn10, mc0750, mengdiw}@princeton.edu
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  $\mathcal{D}_{\mathrm{label}}$  with annotated rewards and a massive unlabeled dataset  $\mathcal{D}_{\mathrm{unlabel}}$ . We estimate the reward function using  $\mathcal{D}_{\mathrm{label}}$  and then use the estimator for pseudo-labeling on  $\mathcal{D}_{\mathrm{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  $\mathcal{D}_{\mathrm{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 = A z$  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.
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
Figure 1: Overview of reward-directed generation via conditional diffusion model. We estimate the reward function from the labeled dataset. Then we compute the estimated reward for each instance of the unlabeled dataset. Finally, we train a reward-conditioned diffusion model using the pseudo-labeled data. Using the reward-conditioned diffusion model, we are able to generate high-reward samples.
(a) Distribution Shift in Generated Samples
(b) Distribution Shift in Reward
(c) Score Network
Figure 2: Illustrations of distribution shifts in samples, reward, and encoder-decoder score network. When performing reward-directed conditional diffusion, (a) the distribution of the generated data shifts, but still stays close to the feasible data support; (b) the distribution of rewards for the next generation shifts and the mean reward improves. (c) (Adapted from [8]) the score network for reward-directed conditioned diffusion adopts an Encoder-Decoder structure.
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^{\infty}$  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  $\mathcal{D}_{\mathrm{label}}$  follows  $P_{xy}$  and any data in  $\mathcal{D}_{\mathrm{unlabel}}$  follows  $P_x$ .  $P$  is used to denote a distribution and  $p$  denotes its corresponding density.  $P(x \mid y = a)$  and  $P(x, y = a)$  are the conditionals of  $P_{xy}$ . Similarly, we also use notation  $P_{x\widehat{y}}$ ,  $P(x \mid \widehat{y} = a)$  for the joint and conditional of  $(x, \widehat{y})$ , where  $\widehat{y}$  is predicted by the learnt reward model. Also, denote a generated distribution using diffusion by  $\widehat{P}$  (density  $\widehat{p}$ ) followed by the same argument in parentheses as the true distribution it approximates, e.g.  $\widehat{P}(x \mid y = a)$  is generated as an approximation of  $P(x \mid y = a)$ .
Suppose we are given an unlabeled data set  $\mathcal{D}_{\mathrm{unlabel}} = \{x_j\}_{j=1}^{n_1}$  and a labeled data set  $\mathcal{D}_{\mathrm{label}} = \{(x_i, y_i)\}_{i=1}^{n_2}$ , where it is often the case that  $n_1 \gg n_2$ . Assume without loss of generality that  $\mathcal{D}_{\mathrm{label}}$  and  $\mathcal{D}_{\mathrm{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
$$
y = f ^ {*} (x) + \epsilon \quad \text {f o r} \quad \epsilon \sim \mathsf {N} (0, \sigma^ {2}) \quad \text {w i t h} \quad 1 > \sigma > 0.
$$
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  $\mathcal{D}_{\mathrm{unlabel}}$ ,  $\mathcal{D}_{\mathrm{label}}$ , target reward value  $a$ , early-stopping time  $t_0$ , noise level  $\nu$ . (Note: in the following psuedo-code,  $\phi_t(x)$  is the Gaussian density and  $\eta$  is the step size of discrete backward SDE, see §3.3 for elaborations on conditional diffusion)
2: Reward Learning: Estimate the reward function by
$$
\widehat {f} \in \underset {f \in \mathcal {F}} {\operatorname {a r g m i n}} \sum_ {\left(x _ {i}, y _ {i}\right) \in \mathcal {D} _ {\text {l a b e l}}} \ell \left(f \left(x _ {i}\right), y _ {i}\right), \tag {3.1}
$$
where  $\ell$  is a loss and  $\mathcal{F}$  is a function class.
3: Pseudo labeling: Use the learned function  $\widehat{f}$  to evaluate unlabeled data  $\mathcal{D}_{\mathrm{unlabel}}$  and augment it with pseudo labels:  $\widetilde{\mathcal{D}} = \{(x_j,\widehat{y}_j) = \widehat{f} (x_j) + \xi_j\}_{j = 1}^{n_1}$  for  $\xi_{j}\stackrel {\mathrm{i.i.d.}}{\sim}\mathsf{N}(0,\nu^{2})$
4: Conditional score matching: Minimize over  $s \in S$  ( $S$  constructed as 3.8) on data set  $\widetilde{D}$  via
$$
\widehat {s} \in \underset {s \in \mathcal {S}} {\operatorname {a r g m i n}} \int_ {t _ {0}} ^ {T} \widehat {\mathbb {E}} _ {(x, \widehat {y}) \in \tilde {\mathcal {D}}} \mathbb {E} _ {x ^ {\prime} \sim \mathrm {N} (\alpha (t) x, h (t) I _ {D})} \left[ \| \nabla_ {x ^ {\prime}} \log \phi_ {t} \left(x ^ {\prime} | x\right) - s \left(x ^ {\prime}, \widehat {y}, t\right) \| _ {2} ^ {2} \right] d t. \tag {3.2}
$$
5: Conditioned generation: Use the estimated score  $\widehat{s} (\cdot ,a,\cdot)$  to sample from the backward SDE:
$$
\mathrm {d} \widetilde {X} _ {t} ^ {t, \Leftarrow} = \left[ \frac {1}{2} \widetilde {X} _ {k \eta} ^ {y, \Leftarrow} + \widehat {s} \left(\widetilde {X} _ {k \eta} ^ {y, \Leftarrow}, a, T - k \eta\right) \right] \mathrm {d} t + \mathrm {d} \bar {W} _ {t} \quad \text {f o r} t \in [ k \eta , (k + 1) \eta ], k \in \lfloor \frac {T}{\eta} \rfloor . \tag {3.3}
$$
6: Return: Generated population  $\widehat{P} (\cdot |\widehat{y} = a)$ , learned subspace representation  $V$  contained in  $\widehat{s}$ .
In order to generate novel samples with both high fidelity and high rewards, we propose Reward-Conditioned Generation via Diffusion Models (RCGDM); see Algorithm 1 for details. By using the labeled data  $\mathcal{D}_{\mathrm{label}}$ , we approximately estimate the reward function  $f^{*}$  by regression, then we obtain an estimated reward function  $\widehat{f}$ . We then use  $\widehat{f}$  to augment the unlabeled data  $\mathcal{D}_{\mathrm{unlabel}}$  with "pseudo labeling" and additive noise, i.e.,  $\widetilde{\mathcal{D}} = \{(x_{j},\widehat{y}_{j} = \widehat{f} (x_{j}) + \xi_{j})\}_{j = 1}^{n_{1}}$  with  $\xi_{j}\sim \mathsf{N}(0,\nu^{2})$  of a small variance  $\nu^2$ . Here, we added noise  $\xi_{j}$  merely for technical reasons in the proof. We denote the
joint distribution of  $(x, \widehat{y})$  as  $P_{x\widehat{y}}$ . Next, we train a conditional diffusion model using the augmented dataset  $\widetilde{D}$ . If we specify a target value of the reward, for example letting  $y = a$ , we can generate conditioned samples from the distribution  $\widehat{P}(x|\widehat{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  $\nabla \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  $\widetilde{D} = \{(x_j,\widehat{y}_j = \widehat{f} (x_j) + \xi_j)\}_{j = 1}^{n_1}$ , where  $(x,\widehat{y})\sim P_{x\widehat{y}}$ . For simplicity, till the end of this section we use  $y$  instead of  $\widehat{y}$  to denote the condition variable. The diffusion model is to approximate the conditional probability  $P(x\mid \widehat{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  $\nabla \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,
$$
\mathrm {d} X _ {t} ^ {y} = - \frac {1}{2} g (t) X _ {t} ^ {y} \mathrm {d} t + \sqrt {g (t)} \mathrm {d} W _ {t} \quad \text {w i t h} \quad X _ {0} ^ {y} \sim P _ {0} (x | y) \text {a n d} t \in (0, T ], \tag {3.4}
$$
where similarly  $T$  is a terminal time,  $(W_{t})_{t\geq 0}$  is a Wiener process, and the initial distribution  $P_0(x|y)$  is induced by the  $(x,\widehat{y})$ -pair distribution  $\bar{P}_{x\widehat{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_{t}^{y}$  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  $\nabla \log p_t(x_t|y)$  through denoising score matching [46]. A conceptual way is to minimize the following quadratic loss with  $\mathcal{S}$ , a concept class.
$$
\underset {s \in S} {\operatorname {a r g m i n}} \int_ {0} ^ {T} \mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} \left[ \| \nabla \log p _ {t} (x _ {t} | y) - s (x _ {t}, y, t) \| _ {2} ^ {2} \right] d t, \tag {3.5}
$$
Unfortunately, the loss in (3.5) is intractable since  $\nabla \log p_t(x_t|y)$  is unknown. Inspired by Hyvarinen 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
$$
\begin{array}{l} \mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} \left[ \| \nabla \log p _ {t} (x _ {t} | y) - s (x _ {t}, y, t) \| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {(x, y) \sim P _ {x \bar {y}}} \mathbb {E} _ {x ^ {\prime} \sim \mathrm {N} (\alpha (t) x, h (t) I _ {D})} \left[ \| \nabla_ {x ^ {\prime}} \log \phi_ {t} \left(x ^ {\prime} | x\right) - s \left(x ^ {\prime}, y, t\right) \| _ {2} ^ {2} \right] + C _ {t}, \tag {3.6} \\ \end{array}
$$
where  $\nabla_{x'}\log \phi_t(x'|x) = -\frac{x' - \alpha(t)x}{h(t)}$ , where  $\phi_t(x'|x)$  is the density of  $\mathsf{N}(\alpha (t)x,h(t)I_D)$  with  $\alpha (t) = \exp (-t / 2)$  and  $h(t) = 1 - \exp (-t)$ .
Equation (3.6) allows an efficient implementation, since  $P_{x\widehat{y}}$  can be approximated by the empirical data distribution in  $\widetilde{\mathcal{D}}$  and  $x'$  is easy to sample. Integrating (3.6) over time  $t$  leads to a practical conditional score matching object
$$
\underset {s \in \mathcal {S}} {\operatorname {a r g m i n}} \int_ {t _ {0}} ^ {T} \widehat {\mathbb {E}} _ {(x, y) \sim P _ {x \widehat {y}}} \mathbb {E} _ {x ^ {\prime} \sim \mathrm {N} (\alpha (t) x, h (t) I _ {D})} \left[ \| \nabla_ {x ^ {\prime}} \log \phi_ {t} \left(x ^ {\prime} | x\right) - s \left(x ^ {\prime}, y, t\right) \| _ {2} ^ {2} \right] d t, \tag {3.7}
$$
where  $t_0 > 0$  is an early-stopping time to stabilize the training [44, 50] and  $\widehat{\mathbb{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):
$$
\mathcal {S} = \left\{\mathbf {s} _ {V, \psi} (x, y, t) = \frac {1}{h (t)} \left(V \cdot \psi \left(V ^ {\top} x, y, t\right) - x\right): V \in \mathbb {R} ^ {D \times d}, \psi \in \Psi : \mathbb {R} ^ {d + 1} \times \left[ t _ {0}, T \right]\rightarrow \mathbb {R} ^ {d} \right\}, \tag {3.8}
$$
with  $V$  being any  $D \times d$  matrix with orthonormal columns and  $\Phi$  a customizable function class. This design has a linear encoder-decoder structure, catering for the latent subspace structure in data. Also  $-\frac{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  $\eta > 0$  described as follows:
$$
\mathrm {d} \widetilde {X} _ {t} ^ {t, \Leftarrow} = \left[ \frac {1}{2} \widetilde {X} _ {k \eta} ^ {y, \Leftarrow} + \widehat {s} \left(\widetilde {X} _ {k \eta} ^ {y, \Leftarrow}, y, T - k \eta\right) \right] \mathrm {d} t + \mathrm {d} \overline {{W}} _ {t} \quad \text {f o r} \quad t \in [ k \eta , (k + 1) \eta ]. \tag {3.3 revisited}
$$
initialized with  $\widetilde{X}_t^{t,\Leftarrow} \sim \mathsf{N}(0, I_D)$  and  $\overline{W}_t$  is a reversed Wiener process. Note that in (3.3), the unknown conditional score  $\nabla p_t(x|y)$  is substituted by  $\widehat{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 = A z$  for an unknown  $A \in \mathbb{R}^{D \times d}$  with orthonormal columns and  $z \in \mathbb{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.