Conditional Diffusion Models are Minimax-Optimal and Manifold-Adaptive for Conditional Distribution Estimation

Point 1: simulation and practical guidance.

Reply: Thank you for your valuable suggestions. We have included a simulation in our paper, which can be found in Appendix A of the supplementary material. Please see our global comment for more details.

Point 2: Technical difficulties.

Reply: One main technical challenge in extending the analysis from the unconditional case to the conditional case is developing new analytical approaches to bound the approximation error of the conditional score function using neural networks due to the inclusion of the extra covariate variable $X\in\mathcal M_X$. In particular, as we assume the support of $Y|X=x$ to be a general, unknown nonlinear submanifold $\mathcal M_{Y|X=x}$ dependent on the covariate value $X=x$, the task of handling unknown supports for the conditional distributions requires us to deal with infinitely many unknown and nonlinear submanifolds $\mathcal M_{Y|X=x}$ with $x\in \mathcal{M}_X$. This poses a significant technical challenge that we tackle in our work. In particular, we need to partition the joint space $\mathcal{M}$ of $(X,Y)$ into small pieces with varying resolution levels in $X$ and $Y$, tailored to times $t$: the resolution level in $Y$ should decrease with increasing $t$, reflecting the diminished importance of finer details in $Y$ due to noise injection over time; while the resolution level in $X$ should increase with $t$, emphasizing the growing significance of the global dependence on $X$ as finer details in $Y$ are smoothed out.

Point 3: What is $p_{data}$ on Line 122?

Reply: Apologies for the confusion; there was a typo. The term $p_{\rm data}$ should refer to the target distribution, which is correctly denoted as $\mu^*$.

Point 4: In remark 1, can authors discuss the case when $\alpha_X>1$? What does this mean to the minimax rate?

Reply: The rate is strictly decreasing in terms of $\alpha_X$, so allowing for a larger $\alpha_X$ can lead to a faster convergence rate. We expect that a similar phenomenon to occur in the minimax rate as well.

Point 5: In Remark 2, can the authors clarify why $D_Y$ need to satisfy this condition to have the first term dominate the second one?

Reply: Since the two terms are $n^{-\frac{1}{2+\frac{D_X}{\alpha_X}}}\quad \text{and}\quad n^{-\frac{1+\frac{1}{\alpha_Y}}{2+\frac{D_X}{\alpha_X}+\frac{D_Y}{\alpha_Y}}},$ the transition boundary is $\frac{\frac{1}{\alpha_Y}}{\frac{D_Y}{\alpha_Y}}=\frac{1}{D_Y}=\frac{1}{2+\frac{D_X}{\alpha_X}}$. Thus the dominant term is $n^{-\frac{1+\frac{1}{\alpha_Y}}{2+\frac{D_X}{\alpha_X}+\frac{D_Y}{\alpha_Y}}}$ if $D_Y \geq 2+\frac{D_X}{\alpha_X}$; and the dominant is $n^{-\frac{1}{2+\frac{D_X}{\alpha_X}}}$ if $D_Y\leq 2+\frac{D_X}{\alpha_X}$. The underlying reason for this transition is that a larger $D_Y$ increase the complexity of estimating the conditional density of $Y$, making the task of conditional distribution estimation (distribution regression) in Wasserstein distance more challenging than nonparametric mean regression.

Other comments.

Thank you for pointing out these important issues. We will address them in the revised version of our main text.

Point 1: Practical experiments.

Thank you for your valuable suggestions. We have included a simulation in our paper, which can be found in Appendix A of the updated supplementary material. Please see our global comments for more details.

Point 2: The discussion on "reverse predictor" approach.

Given that the training objective we are considering differs from that of the unconditional diffusion model only in the input of the score network—by the addition of covariates—it should not increase much on the computational efficiency. However, the classifier-guidance method requires extra effort to learn the external classifier $x|y_t$. Furthermore, from a theoretical perspective, our goal is not to learn the score function of the joint density of $X$ and $Y$, but rather the conditional density of $Y$ given $X$. Therefore, we do not need to impose any smoothness assumptions to the marginal density of $X$, but only need to assume that $Y$ smoothly depends on $X$. However, for learning the classifier $p(x|y_t)=p(x)p(y_t|x)/p(y_t)$ in the classifier-guidance method, we might need to assume that the marginal density $p(x)$ of $X$ is also smooth, an assumption that can be avoided in the classifier-free method we are considering.

Point 3: Constraints to NN and implication of Lemma 2.

Reply: Thank you for raising this important point. This architecture is designed so that the estimation error and approximation can be well balanced. If we are choose a NN larger than the specified size, we will achieve a smaller approximation error, i.e., the term $\min_{S\in \mathcal S_i} \mathbb{E} \int_{t_{i-1}}^{t_i} \int_{\mathbb{R}^{D_Y}} || S(w,x,t)-\nabla \log p_t(\cdot\,|\,x)(w)||^2 p_t(\cdot\,|\,x)(w) dwdt$ in Lemma 2, but will also lead to a larger estimation error (or excess risk), i.e., the term $\frac{R_iH_i \log \big(R_iH_i\|W_i\|_{\infty}(B_i \vee 1) n\big) (\log n)^2}{n}$ in Lemma 2. Consequently, this results in a larger total generalization error when these terms are summed. Similarly, selecting a NN smaller than the specified size will lead to a smaller excess risk but a larger approximation error due to the reduced capacity of the NN. This trade-off also culminates in a larger total generalization error. In summary, if these constraints are not satisfied, we may end up with an estimator that exhibits a larger error in terms of the expected Wasserstein distance compared with Theorem 2.

Point 4: Assumption on data space and smoothness.

Sorry for the confusion, we only need the space $\mathcal M_X$ to have a upper Minkowski dimension $d$; it is not necessary to impose any smoothness conditions on $\mathcal M_X$ (e.g., being a smooth manifold). The assumption of $\mathcal M_X=[-1,1]^{D_X}$ and $\mathcal M_Y=[-1,1]^{D_Y}$ in Theorem 1 is only for technical simplicity. The similar technique may also apply to, for example, unit ball, or more generally, compact set with smooth boundary. We assume that we know $\alpha_X$ and $\alpha_Y$ to set the architecture of NN. Theoretically, it may also possible to construct estimators that adaptive to smoothness level using, for example, the Lepski method.

Point 5: Relate to latent diffusion models.

In Theorem 2, we consider the diffusion model that operates directly within the ambient space. Compared with the latent diffusion, it is perhaps less intuitive why this model can effectively adapt to the manifold structure of the data, which serves as a key motivation for the current research described in Theorem 2. For the latent diffusion model, the analysis may be divided into two part: one is the analysis for the encoder-decoder learning, and the second one is the analysis of the score matching in the encoded space, which may be treated as the Euclidean case we considered in Theorem 1.

Point 6: KDE estimator

The KDE is demonstrated to achieve the same rate as our equation (10) only within the Euclidean space, where both $X$ and $Y$ lack any manifold structure. While our Theorem 2 allows both $X$ and $Y$ to have low-dimensional structures. Therefore, compared to KDE, the conditional diffusion model can adapt to the underlying manifold structures, leading to better performance in such scenarios.

Point 7: transition phases.

For the transition phase in terms of smoothness parameters, since the transition boundary is $\frac{\frac{1}{\alpha_Y}}{\frac{D_Y}{\alpha_Y}}=\frac{1}{D_Y}=\frac{1}{2+\frac{D_X}{\alpha_X}}$, the transition boundary is independent of $\alpha_Y$. In terms of $\alpha_X$, the dominant term is $n^{-\frac{1+\frac{1}{\alpha_Y}}{2+\frac{D_X}{\alpha_X}+\frac{D_Y}{\alpha_Y}}}$ if $\alpha_X\geq \frac{D_X}{D_Y-2}$ with $D_Y\geq 2$; and the dominant is $n^{-\frac{1}{2+\frac{D_X}{\alpha_X}}}$ if $\alpha_X\leq \frac{D_X}{D_Y-2}$ or $D_Y\leq 2$. The reason is that when $\alpha_X$ is small, the bottleneck will be estimating the global dependence of $Y$ on $X$, so the dominant term will be the minimax rate of nonparametric regression.

Point 1: practical implications of the theoretical work

Reply: One guidance for designing the neural network is to utilize a smaller-sized neural network for approximating score functions corresponding to larger values of $t$. This is motivated by the smoothing effect of Gaussian noise injection during the forward diffusion. So, we may consider a neural network class whose size diminishes over time, such as the piece-wise constant complexity neural network class described in equation (1):

$S(y,x,t)=\sum_{i=1}^{\mathcal{I}} S_i(y,x, t) \cdot \mathbf{1}\left(t_{i-1} \leq t<t_{i}\right) ,$

where for each $i \in [\mathcal{I}]$, $S_i$ is a ReLU neural network whose size decreases as $i$ increases. A simulation study included in Appendix A examines the performance of this piecewise NN structure in toy examples. Please see our global comments for details. The size of each ReLU neural network should be carefully chosen to achieve a balance between approximation error and estimation error: larger sizes increase the capacity of the NN, thus reducing the approximation error, but they also introduce larger random fluctuation in the training objective and increase the estimation error. The optimal size of each ReLU neural network should achieve a balance between these two errors, and will grow with $n$. Explicit expressions for the optimal sizes are provided in Theorem 2.

Point 2: Implications of assumption D.

Reply: Sorry for the confusion. Here the assumption means that there exists a positive constant $r_0$, so that for any $x_0\in \mathcal{M}$, an encoder-decoder structure can be locally defined over $x \in \mathbb{B}_{\mathcal{M}_X}\left(x_0, r_0\right)$. Here, the incorporation of the terminologies "encoder" and "decoder" is primarily intended to enhance clarity. The formal mathematical definition of a manifold is defined through local charts/parametrizations, which can be interpreted as an encoder-decoder pair and enable us to define the manifold smoothness with respect to the response and covariate ($\beta_Y$ and $\beta_X$).