Diffusion Model for Data-Driven Black-Box Optimization^††thanks: $\diamond$ denotes equal contribution listed in alphabetical order. $\ddagger$ denotes co-correspondence.^††thanks: This paper is a journal version of Yuan et al. (2023).

Diffusion models successfully incorporate diverse guidance in practical applications. For image generation, 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 externally trained, i.e., classifier guidance (Dhariwal and Nichol, 2021) or implicitly specified by a conditioned diffusion model, i.e., classifier-free guidance (Ho and Salimans, 2022). Classifier-free guidance has become a standard technique in the state-of-the-art text-to-image diffusion models (Rombach et al., 2022; Ramesh et al., 2022; Balaji et al., 2022). Other types of guidance are also explored in Nichol et al. (2021); Graikos et al. (2022); Bansal et al. (2023).

Similar ideas have been explored in sequence modeling problems. In offline reinforcement learning, Decision Diffuser (Ajay et al., 2023) 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 (Li et al., 2022) 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. These appealing empirical performance raises condensed curiosity of their theoretical underpinnings.

The theoretical foundation of diffusion models are gradually developing in the recent two years. A line of work studies diffusion models from a sampling perspective (De Bortoli et al., 2021; Albergo et al., 2023; Block et al., 2020; Lee et al., 2022a; Chen et al., 2022b; Lee et al., 2022b; Chen et al., 2023c, b; Benton et al., 2023; Montanari and Wu, 2023). When assuming access to a score function that can accurately approximate the ground truth score function in $L^{\infty}$ or $L^{2}$ norm, Chen et al. (2022b); Lee et al. (2023a) provide polynomial convergence guarantees of score-based diffusion models. De Bortoli (2022) further studies diffusion models under the manifold hypothesis.

Recently, Chen et al. (2023a); Oko et al. (2023); Mei and Wu (2023); Wibisono et al. (2024) provide an end-to-end analysis of diffusion models, leveraging both sampling and statistical tools. In particular, they develop score estimation and distribution estimation guarantees of diffusion models. Statistical sample complexities are established and explain the success of diffusion models by their minimax optimality and adaptivity to data intrinsic structures. These results largely motivate our theory, whereas, we are the first to consider conditional score matching and statistical analysis of conditional diffusion models.

In this paper, we concentrate on black-box optimization using a pre-collected dataset, while additional first/second-order information cannot be accessed beyond the given dataset. This is sometimes termed the offline black-box or zero-order optimization (Fu and Levine, 2021; Shahriari et al., 2015; Luenberger et al., 2021), and is different from the widely studied setting where active interactions with the objective function are allowed (Peters and Schaal, 2007; Snoek et al., 2014; Hebbal et al., 2019; Snoek et al., 2012). To solve the black-box optimization without active function interactions, Kumar and Levine (2020) propose a model inversion network that learns a stochastic inverse mapping from a reward value to the data instances. Then they generate new instances corresponding to the best possible reward. Such a method demonstrates empirical success, leaving theoretical understandings unexplored. Notably, Krishnamoorthy et al. (2023) also empirically study offline black-box optimization with diffusion model with classifier-free guidance.

Our proposed method and analysis is related the offline bandit/RL theory (Munos and Szepesvári, 2008; Liu et al., 2018; Chen and Jiang, 2019; Fan et al., 2020; Jin et al., 2021; Nguyen-Tang et al., 2021; Brandfonbrener et al., 2021). In particular, our theory extensively deals with distribution shift in the pre-collected dataset by class restricted divergence measures, which are commonly adopted in offline RL (Duan et al., 2020; Xie et al., 2021; Ji et al., 2022). Moreover, our established sub-optimality bound of the generated data consists of an error term that coincides with off-policy linear bandits. However, our analysis goes far beyond the scope of bandit/RL.

We present our main algorithm in Algorithm 1 and also see a graphical illustration in Figure 3. In order to generate novel samples with both high fidelity and high rewards, we propose Reward-Conditioned Generation via Diffusion Models (RCGDM). Specifically, we learn the reward function $\widehat{f}$ from the dataset $\mathcal{D}_{\rm label}$ via regression, where $L(f,\mathcal{D}_{\rm label})$ is the corresponding loss function. In the real-valued reward setting, we specify

where $\rho(f)$ is a regularization term. When $f$ is a linear function of $x$, we are actually doing ridge regression. In the human preference setting, we specify

which is essentially an MLE objective. We then use $\widehat{f}$ to augment the unlabeled data $\mathcal{D}_{\rm 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{\sf N}(0,\nu^{2})$ of a small variance $\nu^{2}$. Here, we add 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 $\widehat{y}=a$, we can generate conditioned samples from the distribution $\widehat{P}(x|\widehat{y}=a)$ by backward diffusion.

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{\mathcal{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})$.

The working flow of conditional diffusion models is nearly identical to that of unconditioned diffusion models reviewed in Chen et al. (2023a). 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,

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 $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 (Song et al., 2020b). A conceptual way is to minimize the following quadratic loss with ${\mathcal{S}}$, a concept class:

Unfortunately, the loss in (4.5) is intractable since $\nabla\log p_{t}(x_{t}|y)$ is unknown. Inspired by Hyvärinen and Dayan (2005) and Vincent (2011), we choose a new objective (4.2) and show their equivalence in the following Proposition. The proof is provided in Appendix A.1.

The proof of Proposition 4.1 can be found in Appendix A.1. Here Equation (4.1) allows an efficient implementation, since $P_{x\widehat{y}}$ can be approximated by the empirical data distribution in $\widetilde{\mathcal{D}}$ and $x^{\prime}$ is easy to sample. Integrating (4.1) over time $t$ leads to a practical conditional score matching object

where $t_{0}>0$ is an early-stopping time to stabilize the training (Song and Ermon, 2020; Vahdat et al., 2021) and $\widehat{\mathbb{E}}$ denotes the empirical distribution.

Constructing a function class ${\mathcal{S}}$ adaptive to data structure is beneficial for learning the conditional score. In the same spirit of Chen et al. (2023a), we propose the score network architecture (see Figure 2(c) for an illustration):

with $V$ being any $D\times d$ matirx with orthonormal columns and $\Phi$ a customizable function class. This design has a linear encoder-decoder structure, catering to the latent subspace structure in data. Also $-\frac{1}{h(t)}x$ is included as a shortcut connection.

Sampling from the model is realized by running a discretized backward process with step size $\eta>0$ described as follows:

initialized with $\widetilde{X}_{t}^{t,\Leftarrow}\sim{\sf N}(0,I_{D})$ and $\overline{W}_{t}$ is a reversed Wiener process. Note that in (4.3), the unknown conditional score $\nabla p_{t}(x|y)$ is substituted by $\widehat{s}(x,y,t)$.

We introduce the performance metric for a generated distribution with respect to a certain target reward value. Let $y^{*}$ be a target reward value and $P$ be a distribution learned by Algorithm 1 (we omit the hat notation on $P$ for simplicity). We define the sub-optimality of $P$ as

which measures the gap between the expected reward of $x\sim P$ and the target value $y^{*}$. In the language of bandit learning, this gap can be viewed as a form of off-policy sub-optimality.

We introduce the following proposition, which decomposes the sub-optimality $\texttt{SubOpt}(P;y^{*})$. Specifically, we show that $\texttt{SubOpt}(P;y^{*}=a)$ comprises of three components: off-policy bandit regret which comes from the estimation error of the reward function, on-support and off-support errors coming from estimating conditional distributions with diffusion.

The proof is provided in Appendix B.3.1. Error $\mathcal{E}_{1}$ corresponds to the reward function estimation error, errors $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$ are induced by the training of diffusion model.

In this section, we analyze the fidelity of the conditional generation process specified by Algorithm 1. Specifically, we assess its ability to capture the underlying representation, and how the data generated by the conditional diffusion model deviates the ground truth representation. Recall that Algorithm 1 has two outputs: the generated distribution $\widehat{P}(\cdot|\widehat{y}=a)$ and the learned representation matrix $V$. We use notations $\widehat{P}_{a}:=\widehat{P}(\cdot|\widehat{y}=a)$ (generated distribution) and $P_{a}:=P(\cdot|\widehat{y}=a)$ (target distribution) for better clarity in the rest of the presentation. We define two metrics

Here, $\angle({V},{A})$ is defined for matrices $V$ and $A$, where $A$ is the matrix encoding the ground truth subspace. Clearly, $\angle({V},{A})$ measures the difference in the column span of $V$ and $A$, which is also known as the subspace angle. Term $\mathbb{E}_{x\sim\widehat{P}_{a}}[\|x_{\perp}\|_{2}]$ is defined as the expected $l_{2}$ distance between $x$ and the subspace.

We impose the following commonly adopted light-tail assumption (Wainwright, 2019) on the latent distribution $z$ for quantifying sample fidelity.

We remark that the realizability condition is for technical convenience, which can be removed in our nonparametric extension in Section 6.2. Assumption 5.2 is also mild and covers arbitrary sub-Gaussian latent data distributions, which include the following Gaussian assumption as a special case.

Now we present sample fidelity guarantees.

The proof is provided in Appendix B.2. In Theorem 5.4, we provide a fidelity guarantee for the learned subspace $V$. Specifically, when the covariance of the underlying representation $z$ is invertible, we prove that the error in subspace learning converges in a parametric rate.

In this subsection, we consider the case when we have access to a noisy reward label for every $x_{i}$ in $\mathcal{D}_{\rm label}$, i.e. $\mathcal{D}_{\rm label}=\{(x_{i}.y_{i})\}_{i\in[n_{2}]}$, where $y_{i}$ is sampled according to (3.1). Recall that in Assumption 3.2, the ground truth reward has the decomposition $f^{*}(x)=g^{*}(x_{\parallel})+h^{*}(x_{\perp})$, where $x_{\parallel}$ is the projection of $x$ on subspace $A$, and $x_{\perp}$ is the part orthogonal to $A$. We consider two setups for $g^{*}(x_{\parallel})$: 1) the parametric setting, where $g^{*}(x_{\parallel})$ is a linear function of $x_{\parallel}$; 2) the nonparametric setting, where $g^{*}(x_{\parallel})$ is a general Hölder continuous function, and we use a deep ReLU neural network for reward estimation.

Our theoretical analysis, in its full generality, extends to using general nonparametric function approximation for both the reward and score functions. Built upon the insights from the linear setting, we provide an analysis of the nonparametric reward and general data distribution setting. We generalize Assumption 6.1 to the following.

Hölder continuity is widely studied in nonparametric statistics literature (Györfi et al., 2002; Tsybakov, 2008). Under Assumption 6.4, we use nonparametric regression for estimating $f^{*}$. Specifically, we specialize (4.1) in Algorithm 1 to

where $\mathcal{F}$ is chosen to be a class of neural networks. Hyperparameters in $\mathcal{F}$ will be chosen properly in Theorem 6.7.

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 6.5 says $P_{z}$ has a light tail and also encodes distributions with compact support. Furthermore, we assume that the curated data $(x,\widehat{y})$ induces Lipschitz conditional scores. Motivated by Chen et al. (2023a), we show that the linear subspace structure in $x$ leads to a similar conditional score decomposition $\nabla\log p_{t}(x|\widehat{y})=s_{\parallel}(x,\widehat{y},t)+s_{\perp}(x,\widehat{y},t)$, where $s_{\parallel}$ is the on-support score and $s_{\perp}$ is the orthogonal score. The formal statement and proof for this decomposition are deferred to Appendix E. We make the following assumption on $s_{\parallel}$.

Lipschitz score is commonly adopted in existing works (Chen et al., 2022b; Lee et al., 2023a). Yet Assumption 6.6 only requires the Lipschitz continuity of the on-support score, which matches the weak regularity conditions in Lee et al. (2023a); Chen et al. (2023a). 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 $\mathcal{F}$ and ${\mathcal{S}}$ for nonparametric regression and conditional score matching. 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

For the conditional score network, we will additionally impose some Lipschitz continuity requirement, i.e., $\|f(x)-f(y)\|_{2}\leq c_{\rm lip}\|x-y\|_{2}$ for some Lipschitz coefficient $c_{\rm lip}$. Now we are ready to bound $\texttt{SubOpt}(\widehat{P}_{a};y^{*}=a)$ in Theorem 6.7 in terms of non-parametric regression error, score matching error and distribution shifts in both regression and score matching.

The proof is provided in Appendix D.2. Here the $\delta(n)$ term accounts for the unbounded domain of $x$, which is negligible when $n$ is large. As a consequency, we prove that under a nonparametric setting, Algorithm 1 enjoys a $\widetilde{\mathcal{O}}\big{(}\texttt{DistroShift}(a)\cdot\big{(}n_{2}^{-\frac{\alpha}{2\alpha+d}}+n_{1}^{-\frac{2}{3(d+6)}}\big{)}\big{)}$ suboptimality when the neural network is properly chosen.