On Statistical Rates and Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)

Reviewer's Question 1: I have reviewed Appendix C. How does the proof for DiTs in this paper differ from that presented in [Chen et al., 2023a]?

Response: Thanks for the question. Here are some clarifications.

• To prove score approximation (Theorem 3.1), our approach utilizes the universal approximation of the Transformer network, while the work [Chen23] relies on the approximation capabilities of the ReLU neural network.
• To prove score estimation (Corollary 3.1.1) and distribution estimation (Corollary 3.1.2), our work uses the covering number of the Transformer network, while the work [Chen23] uses the covering number of the ReLU neural network.
• Our work provides the analysis of the computational limits for all possible efficient DiT algorithms/methods for both forward inference and backward training under the strong exponential time hypothesis, while the work [Chen23] does not consider the computational limit analysis.

Reviewer's Question 2: The assumption of a low-dimensional linear latent space in the context of DiTs is somewhat unclear to me. Is this assumption widely accepted within the DiTs community, or has it been adopted primarily to simplify technical analysis and proofs?

Response: Thank you for your question. Let us provide some additional details.

The assumption is widely accepted within the DiTs community. In practice, existing works [Peebles23, Ma24] use the DiTs with an autoencoder to compress input data into a low-dimensional latent space, where the autoencoder can be a simple linear layer. This aligns with the low-dimensional linear latent space assumption.

Reviewer's Question 3: Additionally, is there potential for the proof to be generalized beyond this assumption? My concerns are further reinforced by the nontrivial gap suggested in Corollary 3.1.1.

Response: Grateful for your question. Here's some further explanation.

Yes. We can generalize our proof beyond the low-dimensional linear latent space assumption by setting the matrix $B$ in Assumption 2.1 as an identity matrix. However, the linear subspace assumption leads to a more robust conclusion, suggesting that the latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.

---

We hope these points address the reviewer's questions.

We're open to any further questions or clarifications you might have about our work. Thank you!

===

• [Chen23] Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning (ICML), 2023.
• [Peebles23] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2023.
• [Ma24] Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.03048, 2024.

Reviewer's Comment 1: It is meaningful to further elucidate why existing DiTs models are difficult to train.

Response: Thanks for your comment. We'd like to clarify a few points.

This relates to the high-dimensional latent data representation, which increases both approximation and estimation errors. To demonstrate this, we can generalize our proof to the setting without the low-dimensional linear latent space assumption by using an identity matrix for matrix $B$ in Assumption 2.1, implying a large $d_0$. According to Theorem 3.1, and the proof details of Corollaries 3.1.1 and 3.1.2 (page 48 and 52), we see that the errors in score approximation, score estimation, and distribution estimation all depend on $d_0$. A larger $d_0$ leads to greater errors.

Reviewer's Question 1: How will these theoretical proofs contribute to the exploration of fast training and sampling algorithms?

Response: Thanks for the question.

In essence, our hardness results provides necessary conditions for designing efficient methods:

• The latent dim should be small enough $d=O(\log L)$ (Thm 4.1 & Prop 4.1, 4.2)
• Normalization of $K,Q,V$ in DiT attention heads is beneficial for performance and efficiency. For example:
  • For efficiency inference: $\max{\|| W_KA_1\||,\||W_QA_2\||,\||W_{OV}A_3\|| }\le B$ with $B=o(\sqrt{\log L})$ (Prop 4.2)
  • For efficiency training: $\max{\||W_KA_1\||,\||W_QA_2\||,\||W_{OV}A_3\|| }\le \Gamma$ with $\Gamma=o(\sqrt{\log L})$ (Thm 4.1)

We want to emphasize that these conditions are necessary but not sufficient. Sufficient conditions should depend on the detailed designs of specific methods.

We hope these points address the reviewer's concerns. We have revised the latest version of our paper accordingly.

We're open to any further questions or clarifications you might have about our work. Thank you!

Thank you for your response.

In summary, the paper is overly theoretical, and I would appreciate it if it could include some toy examples. For instance, it could demonstrate improvements in DiTs by adhering to the design principles for d and QKV you mentioned in R2, and show a certain performance boost as a result. I would be glad to see such experiments incorporated into the paper.

Considering the limited evaluation as a criterion for scoring, I will deduct one point. I look forward to further discussions that can alleviate my concerns.

Reviewer's Comment 1: The paper is too technical (e.g. there are many long formal definitions), and the results are hard to understand for a non-expert in the area. The formulations of the theorems contain a lot of parameters, it makes them very hard to read. I suggest that you may write down informal versions that would be clear and would illustrate the result (with a reference to the formal version in the appendix).

Reviewer's Question 1: Could you write here some informal versions of the theorems (or some simple corollaries) and the clear comparison with prior works? I'll be happy to increase the score if I see that the results are better than the state of the art.

Response: Thanks for the suggestion. We agree adding an informal version of our results will be beneficial for the general audience. We have added them in introduction accordingly, which are quoted as follows:

• Theorem 1.1 (Informal Version of Theorem 3.1). For any approximation error $\epsilon>0$ and any initial data distribution $P_0$ under Assumptions 2.1 to 2.3, there exists a DiT score network such that for any time $t\in [T_0,T]$ ($T$ is the stopping time in forward process, and $T_0$ is the early stopping time in backward process), the upper bound of the approximation error is $\epsilon \cdot \sqrt{d_0}/\sigma(t)$, where $\sigma(t)=1-e^{-t}$.

• Corollary 1.1.1 (Informal Version of Corollary 3.1.1). Under Assumptions 2.1 to 2.3 and using $\epsilon \in (0,1)$, we choose the characterized structural configuration of DiT score network in the proof details of Theorem 1.1. Let n denote the sample size, then with probability $1-1/\mathrm{poly}(n)$, the sample complexity for the score estimation is $\mathcal{O}(\xi(n, \epsilon))$, where $\xi(n, \epsilon)=C_1 \cdot 2^{\epsilon^{-C_2}} \cdot n^{-0.5}+C_3 \cdot \epsilon^2+n^{-1}$, and $C_1, C_2, C_3$ are positive constants.

• Corollary 1.1.2 (Informal Version of Corollary 3.1.2). With the estimated DiT score network in Corollary 1.1.1, we denote $\hat{P}_{T_0}$ as the generated distribution at time $T_0$. We have the following with probability $1-1/\mathrm{poly}(n)$.

1. The distribution estimation error of $\hat{P}_{T_0}$ in the latent subspace (see Assumption 2.1) is $\mathcal{O}(\sqrt{\xi(n, \epsilon)})$.
2. The estimated distribution $\hat{P}_{T_0}$ in the orthogonal subspace degenerates to a point mass at origin as $T_0 \rightarrow 0$.

• Theorem 1.2 (Informal Version of Theorem 4.1). Assuming SETH and all numerical values are in $O(\log L)$ encoding, there exists an algorithm for approximating gradient computation of optimizing DiT loss up to $1/\mathrm{poly}(L)$ accuracy that runs in $L^{1+o(1)}$ time, if certain norm bound conditions are satisfied.

• Proposition 1.1 (Informal Version of Proposition 4.1). Assuming SETH, the existence of sub-quadratic time algorithm for approximating DiT inference depends on the norm bounds of $K,Q,V$ of the attention heads in DiT.

• Proposition 1.2 (Informal Version of Proposition 4.2). Assuming SETH, $L^{1+o(1)}$ time DiT inference is possible.

---

Reviewer's Comment 2: In addition, while the comparison with prior works is present, it is unclear to me and seems to be imprecise. It would be nice to see something like (for example) "we have this error bound, while all prior works [..., ...] have asymptotically worse errors".

Response: Thanks for the comment. We acknowledge the current draft is not precise enough and have made modifications accordingly. We quote them as follows:

• line 179 After Theorem 3.1:

Remark 3.1 (Comparing with Existing Works.) We are the first to prove the score approximation capability of DiT, while the prior theoretical works about DiT [Benton24, Wibisono24] merely assume that the score function is well approximated.

• line 200 After Corollary 3.1.1:

Remark 3.2 (Comparing with Existing Works.) We are the first to provide a sample complexity for DiT, while the prior theoretical works on the sample complexity of diffusion models [Zhu23, Chen23] only focus on ReLU-based diffusion models and does not include the attention mechanism.

• line 223 After Corollary 3.1.2:

Remark 3.5 (Comparing with Existing Works.) We are the first to provide a distribution estimation for DiT, incorporating the tail behavior assumption of the latent variable distribution (Assumption 2.2). However, the prior work [Oko23] does not address DiT and relies on assumptions about the initial data distribution that are far from empirical realities.

These are all "1st known" DiT results (i.e., SOTA?).

For computational limits, since we are the first analysis on DiT, there is no prior work to compare with.

---

We hope the revisions and clarifications provided in this response address the reviewer's concerns.

We look forward to further feedback and discussion. Thank you for your time and effort!

===

• [Benton24] Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear convergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
• [Wibisono24] Andre Wibisono, Yihong Wu, and Kaylee Yingxi Yang. Optimal score estimation via empirical bayes smoothing. arXiv preprint arXiv:2402.07747, 2024.
• [Zhu23] Zhenyu Zhu, Francesco Locatello, and Volkan Cevher. Sample complexity bounds for score-matching: Causal discovery and generative modeling. Advances in Neural Information Processing Systems (NeurIPS), 2023.
• [Chen23] Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning (ICML), 2023.
• [Oko23] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning (ICML), 2023.

Dear Reviewer jLPK,

As the discussion period coming to its end, we want to check if our rebuttal has addressed your concerns.

Please let us know if you have any further questions or need clarification. Thank you!

Best regards,

Authors

Reviewer's Question 1: Can the linear subspace assumption on input data be relaxed to a more general one, such as the manifold data? Since it's more natural to consider the intrinsic geometric structures of data.

Response: Thanks for the question. Yes, but not trivial. Here are some clarifications:

• If we consider that the initial data is manifold without the subspace assumption, we can relax our proof to the manifold data by using an identity matrix for matrix $B$ in Assumption 2.1.
• If we consider that the low-dimensional latent representation is manifold with the subspace assumption, the price to pay is that we no longer have the score function decomposition (Lemma 2.1). We need new techniques to characterize the rates.

In addition, the linear subspace assumption yields stronger results. The [Chen23] with the subspace assumption obtained a sharper rate than that of [Oko23] on both score estimation and distribution estimation.

Reviewer's Question 2: One of the backgrounds of this work comes from the quadratic complexity of transformer blocks with respect to sequence length. While in literature, many alternative attention mechanisms that have a linear complexity in sequence length have been proposed (though not in the context of diffusion). I'm wondering whether it is possible to evaluate these methods' efficiency in your framework.

Response: Thanks for the question. We would like to clarify a few points.

Yes.

For the statistical limits (approximation and estimation theory), we can generalize our proof by deriving the universal approximation capability and the covering number of the attention mechanism with linear complexity.

For the computational limits, these alternative attention variants (efficient or not) have already fallen into our framework. This is because we consider “all possible” algorithms in Problem 1. That is, any attempt to compute the DiT gradient faster is considered in our analysis. This includes using different attention mechanisms as score estimators.

Reviewer's Question 3: In the right-hand side of equation (3.1), it seems that the second term does not involve and thus will not goes to zero even when goes to infinity. Does it mean that the estimation can never be consistent?

Response: Appreciate your question. Here are the clarifications.

Yes. This stems from the double exponential factor $2^{\epsilon^{-2L}}$ mentioned in Remark 3.4, where $\epsilon$ denotes any given approximation error. We plan to explore the possibilities noted in Remark 3.4 to avoid the double exponential factor. By doing so, we can choose $\epsilon$ as a function of the variable $n$ to balance the first and second terms on the right-hand side of the equation (3.1). This approach is expected to lead to consistent estimation.

---

Reviewer's Comment 1: The practical insights of the provably efficient criteria (Section 4) are relatively scarce.

Response: Thanks for your comment. We acknolodge the importance of enhancing the applicability of our findings. Here are some clarifications.

Our provably efficient criteria offer some insights for the design of more efficient methods. Specifically, we demonstrate that:

• The latent dim should be small enough $d=O(\log L)$ (Thm 4.1 & Prop 4.1, 4.2)
• Normalization of $K,Q,V$ in DiT attention heads is beneficial for performance and efficiency. For example:
  • For efficiency inference: $\max{\|| W_KA_1\||,\||W_QA_2\||,\||W_{OV}A_3\|| }\le B$ with $B=o(\sqrt{\log L})$ (Prop 4.2)
  • For efficiency training: $\max{\||W_KA_1\||,\||W_QA_2\||,\||W_{OV}A_3\|| }\le \Gamma$ with $\Gamma=o(\sqrt{\log L})$ (Thm 4.1)

We want to emphasize that these conditions are necessary but not sufficient. Sufficient conditions should depend on the detailed designs of specific methods.

We hope these points address your concern and enhance the utility of our criteria in practical settings.

Reviewer's Comment 2: A few noticeable typos were encountered in the paper.

Response: Thanks for the comment. We apologize for the typos. In response, we have conducted 3 more rounds of proofreading and fixed all typos identified in our latest version. This includes:

• line 109, we corrected “denosing” to “denoising”.
• line 120, we corrected “knonw” to “known”.
• line 155, we corrected “a” to “an”.
• line 192, we corrected “an” to “a”.
• line 228, we corrected “depend” to “depends”.
• line 242, we deleted the extra “full”.
• line 253, we corrected “analyze” to “analyzing”.
• line 303, we corrected “exists” to “exist”.
• line 320, we corrected “motivate” to “motivates”.
• line 322, we corrected “do” to “does”.
• line 954, we corrected “subset” to “subsets”.
• line 977, we corrected “part” to “parts”.
• line 989, we corrected “a” to “an”.
• line 1168, we added “at” after “arrive”.

Your feedback is greatly appreciated and has been essential in enhancing the clarity and readability of our paper. Thank you for your help.

---

We hope the revisions and clarifications provided in this response address the reviewer's concerns.

We welcome additional feedback and look forward to further discussions. Thank you for your time and valuable inputs!

---

• [Chen23] Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning (ICML), 2023.
• [Oko23] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning (ICML), 2023.