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

Authors: Rong Tang, Lizhen Lin, Yun Yang
Primary area: learning theory
Keywords: conditional distribution estimation, diffusion models, distribution regression, generative models, manifold, minimax rate

Abstract:
We consider a class of conditional forward-backward diffusion models for conditional generative modeling, that is, generating new data given a covariate (or control variable). To formally study the theoretical properties of these conditional generative models, we adopt a statistical framework of distribution regression to characterize the large sample properties of the conditional distribution estimators induced by these conditional forward-backward diffusion models. Here, the conditional distribution of data is assumed to smoothly change over the covariate. In particular, our derived convergence rate is minimax-optimal under the total variation metric within the regimes covered by the existing literature. Additionally, we extend our theory by allowing both the data and the covariate variable to potentially admit a low-dimensional manifold structure. In this scenario, we demonstrate that the conditional forward-backward diffusion model can adapt to both manifold structures, meaning that the derived estimation error bound (under the Wasserstein metric) depends only on the intrinsic dimensionalities of the data and the covariate.

Decision: accept (poster)
Reviewer scores: 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730430146832

thread_type:
official_review

<thread_object>
creator: reviewer_yrrd

<content>
summary:
This paper uses distribution regression to estimate the conditional distribution induced by conditional forward-backward diffusion models. Authors utilize neural network to estimate the key points and also allow for high-dimensional X and Y with low-dimensional structures. Solid theoretical results are provided.

strengths:
1. Authors introduce SDE in a simple way and made it easy to connect SDE to distribution regression.

2. The paper extends the unconditional distribution estimator to the conditional one and provides theoretical results. Theories are given in an elegant way that they clearly explain the connections to the theories of unconditional distribution estimator, mean regression problems, and nonparametric regression under l2 error. Theoretical results are interesting and inspiring.

weaknesses:
The paper mentions that it provides practical guidance to design the neural network, including the network size, smoothness level and so on. Though it is a theoretical work, some empirical results would significantly demonstrate the usefulness of the paper. I would recommend authors run some simulations to show how the theoretical results guide the design of neural networks. For example, under what scenarios, the errors are optimally controlled?

The theoretical results are interesting.  Can the authors highlight the difficulties and summarize the new ways of proof if there are? This could be a huge contribution to the statistical learning community. I didn't see the challenges from my end and all seem to be standard.

questions:
What is p_data on Line 122?

Authors need to check all the parenthetical and narrative citations throughout the whole paper. Many were mis-used. For example, Line 316 should be parenthetical citation.

Definition of neural network class is already given in Tang and Yang, 2024 (AISTATS). Authors need to reference it or directly call it from the citation.

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

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




Rong Tang and Yun Yang. Adaptivity of diffusion models to manifold structures. 27th International Conference on Artificial Intelligence and Statistics, 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661791676

thread_type:
official_review

<thread_object>
creator: reviewer_2ase

<content>
summary:
- The paper presents a derivation of the estimation error bound for the diffusion models with sampling guided by the covariate regressor. The authors closely build upon the methodology from Oko et al 2023 https://proceedings.mlr.press/v202/oko23a/oko23a.pdf who explored the statistical learning theory for diffusion models via score matching estimators and proved nearly minimax optimal estimation rates in total variation and Wasserstein metric of order 1 for fully connected neural-networks (including suggestions on generalisation to U-nets and ConvNets). 

- The presented paper aims to introduce the statistical learning theory for conditional diffusion models using the distribution regression of the covariate structure guiding the sampling. Authors investigate the properties of the estimators using the finite sample analysis of the statistical error bound under various metrics, the relations of latent lower dimensional data structures and the smoothness properties of the data. 

- The main contribution is: 
   1) Given the assumptions on the size of the architecture,  the estimation error bound of the conditional diffusion estimator is minimax optimal under total variation and Wasserstein distance of order 1, authors include the discussion on the generalization to the unconditional case as a special case, 
   2) if there is a low-dimensional latent space structure within the high-dimensional data manifold (say within the feature space or feature and label space), the resulting conditional distribution estimators adapts to the lower-dimensional latent space.

strengths:
- The paper is mathematically well written. The notation is consistent and precise. Definitions and proofs are provided in the extensive Appendix. These can be useful for deriving and investigating other novel properties and algorithms.
- The problem is relevant and well chosen and the previous work is well mapped. 
- There is an interesting potential applications for measuring the quality of the conditional diffusion estimators and guided  sampling. The paper also  provides a background for the reasoning about the selection of the model architecture.
- The interesting contribution is in analyzing the latent dimension representations of the data and the adaptability of the higher dimensional estimator.
- The results provide insight into the impact of the smoothness of the underlying data manifold on the performance of the estimator.

weaknesses:
1) Paper is missing motivational examples and practical experiments to demonstrate how the presented theory can be used on real-life problems. The results are theoretically interesting and there is no issue with the soundness of the arguments. Nevertheless, without a practical demonstration on how to use the derived results, the main contribution of the paper may be harder to access to the wider machine learning community.

2) While it is great to see the methodology related to the proofs as a part of the main paper, it would be great to see particular direct applications in the main body of the paper.

3) The usual approach to estimate the gradients of log p_t(y_t|y) (see line 200) is to use the gradient of unconditional model and add gradient of log-likelihood of the ‘reverse predictor' which detects ‘the covariate values’ x based on the features constructed from whom the unconditional values ‘y_t’. Authors decided to use marginal forward diffusion which results in empirical estimation of the joint density between the covariate X (the guiding covariates) and Y (the unconditional diffusion). The discussion on incorporating computational efficiency of this approach in comparison to the 'reverse predictor' approach is missing.
4) The proven Theorems assume the initial constraints of the sizes of the architecture such as depth, width, norm of the parameters, etc. These constraints are set as functionals of the order of magnitude of the data sample $n$. It is not well discussed what happens if these constraints are not satisfied.

questions:
1) Are the covariates X also located in the Besov space support?
2) The marginals M_X and M_Y are assumed to be  [-1,1]  D-X, resp. D_Y dimensional cube. How does the scaling of the data to the [-1, 1] cube impact some other intrinsic features, e.g. latent translations? What type of scaling transformations would you recommend?
3) Can you provide guidelines on how you estimate $\alpha_X$ and $\alpha_Y$ or do you assume these values?
4) How does presented results in Theorem 2 relate to latent diffusion models?

5) p.7 Remark 1: Does this remark imply that the kernel estimator can well approximate the performance of the conditional diffusion based estimator?

6) It would be great if authors can  include experiments for different embedded lower dimensional structures, and demonstrate how the presented theory supports and navigates the experimentally obtained results. Using even simple synthetic geometric objects would be great.

7) Can you provide more discussion on the implications of the different regimes discussed in Remark 3? In particular how would you select the smoothness parameters $\alpha_x$ to investigate the transition phases?

8 )   How does Lemma 2 helps to assess the generalization error?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667056937

thread_type:
official_review

<thread_object>
creator: reviewer_x4xs

<content>
summary:
Previous work has studied the min-max optimality of score-based diffusion models in estimating the density of an unknown distribution, and in estimating the intrinsic dimension of data lying on a submanifold of a sample space. 

This work extends these results to the conditional modelling setting, where a conditioning variate changes the distribution being modeled. The distribution is assumed to vary smoothly with the conditioning variate.

These results show that very similar rates of convergence of the total variation and Wasserstein distances between the true and  estimated data distributions hold to the unconditional modeling setting, with additional terms depending on the dimensionality of the conditioning variate. In the limit of the conditioning variate having dimension zero, the same rates as previous works on the unconditional setting are obtained.

strengths:
- The contributions are, while possibly not surprising, are meaningful and interesting results extending the theoretical analysis of diffusion models.
- The paper is generally very well written and easy to follow, given the technical nature of the content. In particular the remarks on the theorems make understanding the implications of them more accessible.

weaknesses:
- The practical implications of the theoretical work presented are not completely clear to me. The authors mention that the results suggest how to design networks, but from my understanding on theorem 2, this suggests simply that larger networks produce better bounds?
- The implications of assumption D are unclear to me. The variable $r_0$ is undefined. Is this condition meant to hold for any $r_0$? If so then this implies a need for the decoding to be smooth globally, with I believe would limit topological structure of the low-dimension data manifold. If this is local condition, the is this not a consequence of the smooth manifold structure of $\mathcal{M}$?

questions:
Please see the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731991996723

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global comments

comment:
**Examples to illustrate the contribution.**  We have included a simulation in our paper, which can be found in Appendix A of the supplementary material. In this toy example, we consider a   response $Y$ generated by the process $$Y = ((1 + 1.5 \cos(x)) \cos(\phi), (1 + 1.5 \sin(x)) \sin(\phi), 1.5 \sin(x)),$$  and  aim at estimating  $Y|x$ and $Y|\phi$.   We employ two neural network classes to approximate the conditional score functions: a standard ReLU NN and a piecewise ReLU NN. For the piecewise ReLU NN, the total time horizon is partitioned into $5$ intervals, with a ReLU NN used to model the conditional score within each interval (c.f. Section 2.3), which aligns with the structure suggested and employed in our theoretical analysis. The width of the hidden layers is selected to ensure comparability of training parameters between the two types of conditional score families.

Through simulations, we observe that for different $x$ (or $\phi$) values, the generated response $Y$ from the fitted conditional diffusion model will concentrate around different ellipse (or tilted ellipse). This indicates that the conditional diffusion model effectively captures both the covariate information and the underlying manifold structure. Additionally, we assess the Maximum Mean Discrepancy (MMD) distance between the generated data and the true conditional distribution across various covariate values, as well as the average across all covariate values. Consistent with our theoretical findings, incorporating the piecewise structure into the neural network results in a noticeably reduced MMD distance. Below, we provide the MMD results for your convenience.


 | Condition          | Piecewise NN | NN    | Condition            | Piecewise NN | NN    |
|--------------------|--------------|-------|----------------------|--------------|-------|
| **$x = 0$**        | 0.0023       | 0.0032| **$\phi = 0$**       | 0.0002       | 0.0034|
| **$x = 0.5$**      | 0.0013       | 0.0014| **$\phi = 0.5$**     | 0.0009       | 0.0013|
| **$x = 1$**        | 0.0007       | 0.0017| **$\phi = 1$**       | 0.0014       | 0.0028|
| **$x \sim Unif[0,1]$** | 0.0021   | 0.0031| **$\phi \sim Unif[0, 2\pi]$** | 0.0008 | 0.0015|

*Table: The table presents the MMD between the conditional diffusion model estimator and the target conditional distribution. The first three rows show the MMD for the conditional distributions $Y \mid x = 0$, $Y \mid x = 0.5$, and $Y \mid x = 1$ as well as $Y \mid \phi = 0$, $Y \mid \phi = 0.5$, and $Y \mid \phi = 1$. The last row displays the expected conditional MMD, where the expectation is taken with respect to $x$ sampled from a uniform distribution on $[0,1]$ and $\phi$ sampled from uniform over $[0,2\pi]$.*

 

We will also add a discussion in our main text to emphasize this simulation: 

 

*We have also conducted a simulation study (see Appendix A) to demonstrate the effectiveness of this theoretically guided neural network architecture compared to a standard single ReLU neural network (across both space and time). In this experiments,  we consider cases where, given the covariate $X$, the response $Y$ is supported on different (tilted) ellipses depending on the values of the covariate.   Consistent with our theoretical findings,  the simulations show that incorporating  the piecewise structure into the neural network results in a more accurate estimation of the conditional distribution.*
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies conditional diffusion models for generative modeling, focusing on generating data given a covariate. Using a statistical framework of distribution regression, the authors analyze the large-sample properties of conditional distribution estimators under the assumption that the conditional distribution changes smoothly with the covariate. They derive a minimax-optimal convergence rate under the total variation metric, consistent with existing literature. The analysis is further extended to scenarios where both data and covariates lie on low-dimensional manifolds, showing that the model adapts to these structures. The resulting error bounds under the Wasserstein metric depend only on the intrinsic dimensionalities of the data and covariates. This paper has received unanimous support from the reviewers. Therefore, I recommend acceptance.