On Statistical Rates and  Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)
Authors: Jerry Yao-Chieh Hu
Weimin Wu
Zhuoru Li
Sophia Pi
Zhao Song
Han Liu
Venue: NeurIPS 2024 poster (2024)
Keywords: DiT
Score Matching
Diffusion Transformer
Efficiency
Universal Approximation
Convergence Analysis
Diffusion Generative Model
DDPM
Primary area: diffusion_based_models
TL;DR: We study the statistical and computational limits of Diffusion Transformers (DiTs).

Abstract: We investigate the statistical and computational limits of latent **Di**ffusion **T**ransformers (**DiTs**) under the low-dimensional linear latent space assumption. Statistically, we study the universal approximation and sample complexity of the DiTs score function, as well as the distribution recovery property of the initial data. Specifically, under mild data assumptions, we derive an approximation error bound for the score network of latent DiTs, which is sub-linear in the latent space dimension. Additionally, we derive the corresponding sample complexity bound and show that the data distribution generated from the estimated score function converges toward a proximate area of the original one.
Computationally, we characterize the hardness of both forward inference and backward computation of latent DiTs, assuming the Strong Exponential Time Hypothesis (SETH). For forward inference, we identify efficient criteria for all possible latent DiTs inference algorithms and showcase our theory by pushing the efficiency toward almost-linear time inference. For backward computation, we leverage the low-rank structure within the gradient computation of DiTs training for possible algorithmic speedup. Specifically, we show that such speedup achieves almost-linear time latent DiTs training by casting the DiTs gradient as a series of chained low-rank approximations with bounded error.
Under the low-dimensional assumption, we show that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.

On Statistical Rates and Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)
Jerry Yao-Chieh Hu
Weimin Wu Zhuoru Li
Sophia Pi
Zhao Song
Han Liu
'Center for Foundation Models and Generative AI, 'Department of Computer Science, *Department of Statisties and Data Science, Northwestern University, Evanston, II. 60208, USA
Simons Institute for the Theory of Computing, UC Berkeley, Berkeley, CA 94720, USA
{ jiu, wwm J @u. nottwweteerr.. edd;
magic. 1ww.xkdeeuukeaaaaaaa.aeoo com;
edu
Abstract
We investigate the statistical and computational limits of latent Diffusion Transformers (DiTs) under the low-dimensional linear latent space assumption. Statistically, study the universal approximation and sample complexity of the DīTs score function, as well as the distribution recovery property of the initial data. Specifically, under mild data assumptions, we derive approximation error bound for the score network of latent DiTs, which is sub-linear in the latent space dimen- sion. Additionally, we derive the corresponding sample complexity bound and show that the data distribution generated from the estimated score function converges toward proximate the hardness of both forward inference and backward computation of latent DiT', assuming the Strong Exponential Time Hypothesis (SETH). For forward inference, we identify efficient criteria for all possible latent DiTs inference algorithms and showcase our theory by pushing the efficiency toward almost-lineat time inference. For backward computation, we leverage the low-rank structure within the gradient computation of DiTs training for possible algorithmic speedup. Specifically, we show that such speedup achieves almost- -inear time latent DiTs training by casting the DiTs gradient as series of chained low-rank approximations with bounded error. Under the oowdiimesioona assumption, we show that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DìT's have the potential to bypass the challenges associated with the high dimensionality of initial data.
of the original one. Computationally, we characterize
Introduction
We investigate the statistical and computational limits of latent diffusion transformers (DiTs), assum- ing the data is supported practical but also timely. On one hand, DiTs have demonstrated revolutionary success in generative AI and digital creation by using Transformers as score networks [Esser et al., 2024, Ma et al., 2024, Chen et al., 2024a, Mo et al., 2023, Peebles and Xie, 2023]. On the other hand, they require significant computational resources [Liu al., 2024], making them challenging industrial labs. Therefore, it is natural to ask whether it is possible to make them lighter and faster without sacrificing performance. Answering these questions requires a fundamental understanding of the DiT architecture. This work provides timely theoretical analysis of the fundamental limits of DiT' architecture, aided by the analytical feasibility provided by the low-dimensional data assumption.
unknown -iw-dimessional linear subspace. This analysis is not only
train outside of specialized
Equal contribution. Version: January 3, 2025. Future updates arXiv.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Empirically, Latent Diffusion is go-to design for effectiveness and computational efficiency [Rom- bach al., 2022, Liu al., 2021, Pope et al., 2021, Su and Wu, 2018]. Theoretically, it is capable of hosting the assumption of low-dimensional data structure (see Assumption 2.1 for formal definition) for detailed analytical characterization [Chen et al., 2023, Bortoli, 2022]. In essence, diffusion models with low-dimensional data structures manifest natural lower-dimensional diffusion process through an encoder/decoder within a robust and informative latent representation feature space [Rombach et al., 2022, Pope et al., 2021]. Such lower-dimensional diffusion improves computational efficiency by reducing data complexity without sacrificing essential information [Liu et al., 2021]. With this assumption, Chen et al. [2023] decompose the score function of U-Net based diffusion models into on-support and orthogonal components. This decomposition allows for the characterization of the distinct behaviors of the two components: the on- on-support component facilitates latent distribution learning, while the orthogonal component facilitates subspace recovery.
In our work, we utilize low-dimensional data structure assumption to explore statistical and com- putational limits of latent DiT's Our analysis includes the characterizations of statistical rates and provably efficient criteria. Statistically, we pose two questions and provide a theory to characterize the statistical rates of latent DiT under the assumption of a low-dimensional data:
What is the approximation limit of using transformers to approximate the DiT score function, particularly in the low-dimensional data subspace?
Question 1.
How accurate is the estimation limit for such score estimator in practical training scenarios? With the score estimator, how well can diffusion transformers recover the data distribution?
Question 2.
Computationally, the primary challenge of DiT lies in the transformer blocks' quadratic complexity. This computational burden applies to both inference and training, even with latent diffusion. Thus, it is essential to design algorithms and methods to circumvent this (?(L)) where L is the latent DiT sequence length. However, there are no formal results to support and characterize such algorithms. To address this gap, we pose the following questions and provide a fundamental theory to fully characterize the complexity of latent DiT under the -ww-dimeniional linear subspace data assumption:
Question 3. Is it possible to improve the ((L²) time complexity with a bounded approximation error for both forward and backward passes? What is the computational limit for such an improvement?
Contributions. We study the fundamental limits of latent DIT. Our contributions are threefold:
Score Approximation. We address Question 1 by characterizing the approximation limit of matching the DiT score function with transformer-basee score estimator. Specifically, under mild data assumptions, we derive an approximation error bound for the score network, sub-linear in the latent space dimension (Theorem 3.1). These results not only explain the expressiveness of latent DiT (under mild assumptions) but also provide guidance for the structural configuration of the score network for practical implementations (Theorem 3.1).
Score and Distribution Estimation. We address Question by exploring the limitations of score and distribution estimations of latent DiT's in practical training scenarios. Specifically, we provide a sample complexity bound for score estimation (Theorem 3.2), using norm-based covering number bound of transformer architecture. Additionally, we show that the learned score estimator is able to recover the initial data distribution (Corollary 1.2.
Provably Efficient Criteria and Existence of Almost Linear Time Algorithms. We address Question 3 by providing provably efficient criteria for latent DiTs in both forward inference and backward computation/training
For forward inference, we characterize all possible efficient DiT algorithms using a norm-hased efficiency threshold for both conditional and unconditional generation (Proposition 4.1). Efficient algorithms, including almost-linear time algorithms (Propo- sition 4.2), are possible only below this threshold. For backward computation, we prove the existence of almost-linear time DiT training algorithms (Theorem 4. by utilizing the inherent low-rank structure in DiT gradients through a chained low-rank approximation.
Interestingly, both our statistical and computational results are dominated by the subspace dimen- sion under the dimennsiiii assumption, suggesting that latent DiT can potentially bypass the challenges associated with the high dimensionality of initial data.
Organization. Section 2 includes background on score decomposition and Transformer-basee score networks. Section includes DiTs' statistical rates. Section 4 includes DiTs' provably efficient criteria. Section 5 includes concluding remarks. We defer discussions of related works to Appendix C.
2

---

Notations. We use lower case letters to denote vectors, e.g., 2 E R". | and || denote its Euclidean norm and Infinite norm respectively. We use upper case letters to de- note matrix, e.g. 7 E RdxL || 1 and || denote the 2-norm, operator norm and Frobenius norm respectively.
|| p.y denotes the P, q-norm where the p-norm is over columns and g-norm is over rows. Given function f. let م = ( (l and Ilpp ((l / x yll2). With a distribution P. we denote [l Up [l the (P)) norm. Let f&P be a pushforward measure, i.e., for any measurable 0 (f) P(f P(f-1(())) (S2)). We use for (conditional) Gaussian density functions.
Background
This section reviews the ideas we built on, including an overview of diffusion models (Section 2.1). the score decomposition under the linear latent space assumption (Section 2.2), and the transformer backbone in DiT (Section 2.3).
2.1 Score-Matching Denoising Diffusion Models
We briefly review forward process, backward process and score matching in diffusion models.
Forward and Backward Process. In the forward process, Diffusion models gradually add noise to the original data zo E RP, and To ~ Po- Let Hy denote the noisy data at the timestamp t, with marginal distribution and destiny P. and Pe- The conditional distribution ((z follows N (B(t)=nn, Ip). where S(t) exp(- wss)ds/2), (tt) 1 (t),,, and w(t) > is nondecreasing weighting function. In practice, the forward process terminates at a large enough T such that Pr is close (o, Tp). In the hackward process, we obtain ye by reversing the forward process. The generation of 1/ depends on the score function V log P.).). However, this is unknown in practice, we use a score estimator sw t) to replace log Pe(-), where sw sw(, t) is usually a neural network with parameters W. See Appendix D.I for the details.
Score Matching. To estimate the score function, we use the following loss
/' (() [1s, v log (Tx))) de, To
min W
where (t) is the weight function, and To is a small value to stabilize training and prevent score function from blowing up [Vahdat et al., 2021]. However, it is hard to compute V log p+(-) with available data samples. Therefore, we minimize the equivalent denoising score matching objective
min [" (t)) [Ex [|ss (2t, Vz. log ( zo 20) dt. W To
(2.1)
where (2))))))) is the transition kernel, then Va, log thr = (8())zo 1//x///// /o()).
To train the parameters W in the score estimator sw (-, t), we use the empirical version of (2.1). We select n i..dd data samples { ~ Po. and sample time t; (1 < I n) uniformly from interval [To, T). Given Zo,i, sample Ft, from N((((t)) (t.)) The empirical loss is
((W) 1 Σ ||w ( (,c,,,4) 4) =o.ill²
(2.2)
For convenience of notation, we denote population loss L(W) = Ep, [[W)]
2.2 Score Decomposition Linear Latent Space
In this part, we review the score decomposition in [Chen et al., 2023]. We consider that the D- dimensional input data supported on do di-dimensional subspace, where do < D.
Assumption 2.1 (Low--Dieessional Linear Latent Space). Let x denote the initial data at = 0. has latent representation via Bh, where B E IR. Dxda is an unknown matrix with orthonormal columns. The latent variable & E do follows the distribution Ps with density function ps.
Remark 2.1. By "linear latent space," we mean that each entry of a given latent vector is linear combination of the corresponding input, i.e., x = Bh. This is also known as the "low-dimensional data" assumption in literature [Chen et al., 2023].
3

---

Let T and denote the perturbed data and its associated latent variable at > 0. respectively. Based on the low-dimensional data structure assumption, we have the following score decomposition theory: on-support score 84 B x. 1) and orthogonal score 8 (x,().
Lemma 2.1 (Score Decomposition, Lemma I of [Chen et al.. 2023]). Let data x = Bh follow Assumption 2.1. The decomposition of score function V log pr (3)
v PT) BV logrl pl (h) (ID- BB x/o(1),
h BB,
(2.3)
(h.4)
(7.6)
where (h) = &h (hih))p (h)dh, vix (-1h) is the Gaussian density function of ((B(t)), 3(t) e/2 and (t) 1 e- We restate the proof in Appendix D.2 for completeness.
Additionally, our theoretical analysis is based on two following assumptions as in [Chen et al., 2023].
Assumption 2.2 (Tail Behavior of Pt). The density function Ph > 0 is twice continuously differen- tiable. Moreover, there exist positive constants Au, A1, A2 such that when |Ah|| Au, the density function (h)h) (h) < (2x)-dd/22 A exp( A2 A211/1//22.
Assumption 2.3 (L,, Lipschitz of s+ (h, The on-support score function s (h, is Le+- Lipschitz E RRdo for any t E [0, T).
2.3 Score Network and Transformers
In this part, we introduce the score network architecture and Transformers. Transformers are the backbone of the score network in DIT. By Assumption 2.1, h B x E Rdo with do D.
(Latent) Score Network. Following [Chen al., 2023], rearrange (2.3) into
V (x) ((o(L)V B x) BT x))((()) (t).
(2.4)
c a.t): Rd (Ta.7)
We use WB Dxdo to approximate B Dxdo, and neural network f f(W B H, t) approximate
q( B F, t). We adopt the following score network class for diffusion in latent space (i.e., in Te R
S sw (i, Wof ()/ (/o W {WB.f)}
(2.5)
where the columns in Wp are orthogonal, f RRd To, T Rd is neural network. In this work, we focus on the diffusion transformers (DiTs), i.e. using Transformer for / [Peebles and Xie, 2023].
Transformers. A Transformer block consists of a self-attention layer and a feed-forward layer, with both layers having skip connection. We use n"m, Rd - Rdd to denote Transformer block. Here r and m are the number of heads and head size in self-attention layer, and 7 is the hidden dimension in feed-forwant layer. Let X E d be the model input, then have the model output
Attn Att = X + Σ ==1 W Softmax ((W& wxx).
(2.6)
FF o Attn(X) Attn(X) + W2 ReLU(W, Attn(X) b, b21 where W W W. E Rm W Rdxm Wi Riidd W2 E Rdxl bi E Re,, Rdd
(2.7)
In our work, we use Transformer networks with positional encoding E RdxL We define the Transformer networks as the composition of Transformer blocks
T..m [TT
pdxL _ | is composition of blocks r..m.ss..
For example, the following is a Transformer network consisting K blocks and positional encoding (x( = (K) Attn (K) o .F..F))))
(*) + E).
(2.8)
3 Statistical Rates of Latent DiTs with Subspace Data Assumption
In this section, we analyze the statistical rates of latent DiTs. Section 3.1 introduces the class of latent DiT score networks. In Section 3.2, we prove the approximation limit of matching the DiT score function with the score network class, and characterize the structural configuration of the score network when a specified approximation error required. Following this, Section 3.3, utilizing the characterized structural configuration, prove the score and distribution estimation for latent DiTs.

---

-1/o7
Latent Encoder
Reshape Layer
Transformer
Revensed Reshape Layer
Lateut Decoder
1/o(()
Wz
R(-)
fT E Tramd
We
R"
Rdh
R
R
R²
(a..))
J= f R
Figure I: Overview of DiT Score Network Architecture sw sw(.,), (-,t). W denotes the linear layer from the input data space the linear latent space. f(-) -H- fT R(() denotes the transformer network fr(+) with reshaping layer R(-), where fr(-) E Tp..m. WB denotes the linear layer from the linear latent space to the input data space. (f) denote the variance of the conditional distribution ((r,a zo)-
3.1 DiT Score Network Class
Here, we provide the details about DiT score network class used in our analysis. (2.5), f is network with Transformer as the backbone, and (7,, R&d [To denotes the input data. Following [Peebles and Xie, 2023], DiT uses time point / to calculate the scale and shift value in the Transformer backbone, and it transforms input picture into sequential version. To achieve the transformation, we introduce reshape layer.
Definition 3.1 (DiT Reshape Layer R(-)). Let R(.) Ro Rd be reshape layer that transforms the dn-dimensional imput into a L matrix. Specifically, for any do = x image input, R(-) converts it into sequence representation with feature dimension d p² (where p 2 2) and sequence length == (i/p)² Besides, define the corresponding reverse reshape (flatten) layer R (-) R
Ro the inverse of R(-). By du dL, R. R are associative w.r.t their input.
To simplify the self-attention block (2.6), let Wov = W and Wik = (Wk) W
Definition 3.2 (Transformer Network Class Tp T..m..), We define the Transformer network class as
Tp T. (K, CT, Cov Cov, CQq Ckq; C.. 22,00 ,Cr, Cr, ITT), satisfying the constraints
Model architecture with K blocks: fr(X) FF(K) Attn o... ....F))) Attn (X):
Model output bound: supx [X
CT:
Parameter bound in Attn ( (W 112,00 I cor,, ) S Cov, |W S
O ||Wk <. CкQ. E 11 I CC,Vi [K];
Parameter bound FF || 2 S S
||Wll S Cr, Vj E (2) [K] Lipschitz of fr: (X,) (X2)ll I |1 - X2lle vX,, X2 E Rd
Definition 3.3 (DiT Score Network Class Syr.m. (Figure 1)). We denote S Symma as the DiT score network class in (2.5), replacing f with At o fr R, and fr from the Transformer class Tr.m.l
3.2 Score Approximation of DiT
Here, we explore the approximation limit of latent DiT score network class Symmm. under linear latent space assumption. Recall that P. is the distribution of zt, (t) is the variance of P(x: mo). do the dimension of latent space, L the sequence length of transformer input, the stopping time in forward process, To is the early stopping time in backward process, and La, is the Lipschitz coefficient of on-support score function. Then we have the following Theorem 3.1.
Theorem 3.1 (Score Approximation of DiT). For any approximation error > and any data distribution Po under Assumptions 2.1 to 2.3, there exists a DiT score network sw from Syz., (defined in Definition 3.2), where W (WB. Fr.. such that for any [To, 71, we have:
(:,() V log (2(P)) Se. (d)/
where (t)) = e et and the upper bound of hyperparameters in S Sy2.1.4 are
K O(E -21). Cr (doLe, dd /T/ log(1/e))
5

---

cov- (1/6)0(1, Cov (1/6)011, oke (1/e)0(1, CKQ (1/c) CE. - ( - = (1/6)00), Cr = (1/6)0m, I.T =0 (dol.s)).
Proof Sketch. Our proof is built the key observation that there is tail behavior of the low-dimensional latent variable distribution Pr. (Assumption 2.2). Recall that V log (ō) = Bq(h, tt/a(t) - -//a(t), where h = B x (defined in (2.4)). By taking W = B, our aim reduces to construct transformer network to approximate qh, t). To achieve this, we firstly approximate q(h,t) with compact-supportee continuous function, based on the tail behavior of Ph. Then we construct transformer to approximate the compact-supporter continuous function using the universal approximation capacity of transformer [Yun et al., 2020J. See Appendix F.1 for detailed proof. D
Intuitively, Theorem 3.1 indicates the capability of the transformer-based score network to approx- imate the score function with precise guarantees. Furthermore, Theorem 3.1 provides empirical guidance for the design choices of the score network when specified approximation error is required.
Remark 3.1 (Comparing with Existing Works). Theoretical analysis of DiTs is limited. Previous works that do not specify the model architecture assume that the score estimator is well-apppoximmted [Benton et al., 2024, Wibisono et al., 2024). To the best of our knowledge, this work is the first to present an approximation theory for DiTs, offering the estimation theory Theorem 3.2 and Corol- lary 3.2.1 based on the estimated score network, rather than a perfectly trained one.
Remark 3.2 (Latent Dimension Dependency). Theorem 3.1 suggests that the approximation capacity and Transformer network size primarily depend on the latent variable dimension do d L. This indicates that DiTs can potentially bypass the challenges associated with the high dimensionality of initial data by transforming imput data into -iwddmmnniioaal latent variable.
3.3 Score Estimation and Distribution Estimation
Besides score approximation capability, Theorem 3.1 also characterizes the structural configuration of the score network for any specific precision, e.g., K, Cg. CF, etc. This characterization enables further analysis of the performance of score network in practical scenarios. In Theorem 3.2, provide a sample complexity bound for score estimation. In Corollary 3.2.1, 1, show that the learned score estimator is able to recover the initial data distribution.
Score Estimation. To derive a sample complexity for score estimation using Sy..4, we rewrite the score matching objective in (2.2) W argmin wes, S. [(sw), W (WB.
Theorem 3.2 shows that as sample size n + oo, sw (w())) (s,() convergences to V llggp()).
Theorem 3.2 (Score Estimation of DiT). Under Assumptions 2.1 to 2.3, we choose S72.1.1 as in
Theorem 3.1 using (0, 1) and L 1, With probability 1 1/poly(n), have
To |w (- log I (P(P))
T
dt - = 1
11/3TTTT 2(1/0)²*
1
nI33TOT
ToT
(3.1)
where hides the factors related to D, do, d, Lay, and log n.
Proof. See Appendix F.2 for a detailed proof.
Intuitively, Theorem 3.2 shows a sample complexity bound for score estimation in practice.
Remark 3.3 (Comparing with Existing Works). [Zhu al., 2023] provides a sample complexity for simple ReLU-based diffusion models under the assumption of an accurate score estimator. To best of our knowledge, we are the first to provide sample complexity for DiTs, based on the learned score network in Theorem 3.1 and the quantization (piece-wise approximation) approach for transformer universality [Yun al., 2020]. Furthermore, our first term shows convergence rate of 1/T, outperforming [Chen et al. 2023], in which the first term is independent of

---

Remark 3.4 (Double Exponential Factor and Inconsistent Convergence). Theorem 3.2 reports an explicit result on sample complexity bounds for score estimation of latent DiTs: a double exponential factor 2(/e) in the first term. We remark that this arises from the required depth K is O(c-2)) and the norm of required weight parameters is (()) as shown in Theorem 3.1, assuming the universality of transformers requires dense layers [Yun et al., 2020]. This double exponential factor causes inconsistent convergence with respect to sample size n, as its large value prevents setting as function of balance the first and second terms in (3.1). This motivates to rethink transformer universality and explore new proof techniques for DiTs, which we leave for future work.
Definition 3.4. For later convenience, define E(n, e)),,,L) n7 2(1/e)
n +
Distribution Estimation. In practice, DiTs generate data using the discretized version with step size μr, see Appendix D.I for details. Let Pr. be the distribution generated by sw using the discretized backward process in Theorem 3.2. Let PA and pt, be the distribution and density function of on-support latent variable h at To- We have the following results for distribution estimation.
Corollary 3.2.1 (Distribution Estimation of DīT, Modified From Theorem 3 of [Chen et al., 2023].
(1xxs)] }). where co is the minimum eigenvalue of Ep, [Ah ]. Let T = O(log n), To O(min (co, 1/Ls,
With the estimated DiT score network $W in Theorem 3.2, we have the following with probability
1 - 1/poly(n).
(i) The accuracy to recover the subspace B is [WWW3 BB = (E(n, L)/o)-
(ii) With the conditions KL (L(PAll (0, Id.)) < co, there exists an orthogonal matrix U paxd such that we have the following upper bound for the total variation distance
TV(P), (W (WBU) Pro) Ö(v E(n, L) log n)),
(3.2)
where δ hides the factor about D, do, d, La, log n. and T To- and (WnU)) Pr. denotes the pushforward distribution.
(iii) For the generated data distribution Pr. the orthogonal pushforward (I Wn )3 Pro is ((0, Σ), where Σ < aTo for a constant a >
B
Proof. See Appendix F.3 for a detailed proof.
Intuitively, Corollary 3.2.1 shows the estimation results in 3 parts: (i) the accuracy of recovering the subspace B; (ii) the estimation eπror between Pr, and Ph and (iii) the vanishing behavior of Pr, in the orthogonal space. These indicate that the learned score estimator is capable of recovering the initial data distribution. Notably, Corollary 3.2.1 is agnostic to the specifics of {(n, c, L).
Remark 3.5 (Comparing with Existing Works). Oko et al. [2023] analyze the distribution estimation under the assumption that the initial density is supported on [-1, 1] and smooth in the boundary. Our Assumption 2.2 demonstrates greater practical relevance. This suggests that our method of distribution estimation aligns more closely with empirical realities.
Remark 3.6 (Subspace Recovery Accuracy). (i) of Corollary 3.2.1 confirms that the subspace learned by DiTs. The error is proportional to the sample complexity for score estimation and depends on the minimum eigenvalue of the covariance of Ph.
4 Provably Efficient Criteria
Here, we analyze the computational limits of latent DiT' under low-dimensional linear subspace data assumption (i.e., Assumption 2.1). The hardness of DiT models ties both forward and backward passes of the score network in Definition 3.3. We characterize them separately.
4.1 Computational Limits of Backward Computation
Following Section 2, suppose we have n i..dd data samples (ro, 1 ~ Pa. and time 1.p (1 S
i < n) uniformly sampled from [To, T). For each data Zo,; E R., we sample Izs, E RD. from N [[.((t.))oo.. (t.,).p) Let MiaAR--300 be the inverse transformation of W WAR-1()), and denote
7

---

Yo,= = (WAR ( E BAXI, We rewrite the empirical denoising score satremng loss (2.2) as
dx I.
72
11 dxL
(Fr(
T0.i
WA R- (RT veco
Yo.
YL
dox1
Dxda
dox1
(4.1)
For efficiency, it suffices to focus on just transformer attention heads of the DiT score network due to their dominating quadratic time complexity in both passes. Thus, we consider only single layer attention for /T. to simplify our analysis. Further, we consider the following simplifications:
(SO) To prove the hardness of (4.1) for both full gradient descent and stochastic mini-batch gradient descent methods, suffices to consider training on a single data point.
(S1) For the convenience of our analysis, we consider the following expression for attention mech- anism. Let X,Y RdxL_ Let WK, Wo, Wy E Rexd be attention weights such that Q WoX E RdxL K WK X RexL and V Wy X RexL, We write attention mechanism of hidden size and sequence length I. as
Att(X)= = (WoWvX) exp(X WXXTTWWoX E RdxL,
(4.2)
V multiplication
K-Q multiplication
1z). Here, exp(-) is entry-wise exponential function, i.e., exp(A),, i. exp(Ai.) for any matrix diag(-) converts a vector into diagonal matrix with the vector's entries on the diagonal, and I. the length- L all ones vector.
with D = ding(exp(XWog))) )
(S2) Since V multiplication is linear in weight while K K-Q multiplication is exponential in weights, we only need to focus on the gradient update of K-Q multiplication. Therefore, for efficiency analysis of gradient, it is equivalent to analyzing a reduced problem with fixed Wo WoWy. const..
(S3) To focus on the DIT, consider the low- low-dimensional linear encoder WA A to be pretrained and to not participate in gradient computation. This aligns with common practice [Rombach et al., 2022] and is justified by the trivial computation cost due to the linearity of 2 (S4) To further simplify, introduce A1, A2, Ag R and W Raxd via
Wn R
(R(W (210))
Yo.
(By (SO), (S1) and (S2)
=XERXXA
YERdxz
WWR WoWy
D-1
exp
X
WWWo
-
=WovERexd :
=AT
-WRdaa =A
(4.3)
Notably, A2, As, Y,Y are constants w.r.t. training above loss with gradient updates.
Therefore, simplify the objective of training DiT into
Definition 4.1 (Training Generic DiT Loss). Given A,, A2, Aa, Y E Rdx and Wov. W e Rdxd following (S4), Training DiT with & loss on single data point X, E Idd is formulated as
"Wov exp A W AA) - r)
min Lo min W
(4.4)
2 Here D = diag(exp AT W A _
Remark 4.1 (Conditional and Unconditional Generation). Lo is generic. If A1 + A2 E Rdxt.
Definition 4.1 reduces to cross-attention in DiT score net (for conditional generation). If A1 Au Rdxt. Definition 4.1 reduces to self-attention DiT score net (for unconditional vanilla generation).
We introduce the next problem to characterize all possible gradient computations of optimizing (4.4).
"The gradient computation is linear in WA. and hence the computation w.r.t WA is cheap and upper- bounded by L (dly(d) time in straightforward way.
8

---

Problem (Approximate DiT Gradient Computation AADrGGCIL,dTT.. Given AI,Az, Aa, E RdxL, Let E > O. Assume all numerical values are in ((log(L) ()) encoding. Let loss function Lo follow Definition 4.1. The problem of approximating gradient computation of optimizing empirical DiT loss (4.4) is to find an approximated gradient matrix (OS) E Rdxd such that || matrix A.
ac aW
I 1/poly(L) L). Here, ||A|] = max,,) Assy for any
In this work, we aim to investigate the computational limits of all possible efficient algorithms of ADITGC with c = 1/poly(L). Yet, the explicit gradient of DiT denoising score matching loss (4.4) is too complicated to characterize ADITGC. To combat this, we make the following observations.
(01) Let 91(+) WARR Rd + Rd 92(-) Att(-) R 1. - Rdxt, and ga(+) =
R(W RDD - RaxL such that ga(x)) X for x E RD (with D> do dT.).
(02) Vectorization of fy. For the ease of presentation. we use notation flexibly that fr to denote both a matrix in R and vector in R. in the following analysis. This practice does not affect correctness. The context in which JT is used should clarify whether it refers to a matrix vector. Explicit vectorization follows Definition D.1.
(03) Linearity of g1. By linearity of W R (.), fr(-) RdL.
we treat 91 as a matrix in Rd xdl. acting on vector
dgs From Therefore, we have 16=1191 192(9s) rIl such that its gradient involves dCg = 91 aw
dw above, only need to focus proving the computation time and error control of term deg aw for gradient w.r. W. Luckily, with tools from fine-grained complexity theory [Alman and Song, 2023, 2024,,b,c] and tensor trick (see Appendix D.3), we prove the existence of almost-linear time algorithms for Problem in the next theorem. Let vec(W) = W for any matrix W following Definition D.I.
Theorem 4.1 (Existence of Almost-Linear Time Algorithms for ADITGC), Suppose all numerical values are in (log b-tt encoding. Let max(|) Wov mallmax: Ax max- [ Allmmx) I P. There exists a L ((1)(1) time algorithm to solve ADITGC(L,, I, d = O(log L), ol log L) (i.e., Problem 1) with loss Ly from Definition 4.1 up to 1/poly(L) accuracy. particular, this algorithm
outputs gradient matrices ((W) E Rddd such that (W)
ac BW max S 1/poly(L).
Proof Sketch. Our proof is built the key observation that there exist low-rank structures within the DiT training gradients. Using the tensor trick [Diao et al., 2019, 2018] and computational hardness results of attention [Hu al., 2024b, Alman and Song, 2023], we approximate DîT training gradients with a series of low-rank approximations and carefully match the multiplication dimensions so that the computation of dgg forms a chained low-rank approximation. We complete the proof by aw
demonstrating that this approximation is bounded by 1/poly(L) error and requires only almost- linear time. See Appendix G.2 for a detailed proof.
Remark 4.2. We remark that Theorem 4.1 is dominated by the relation between and d, hence by the subspace dimension do dL. A smaller do makes Theorem 4.1 more likely to hold.
4.2 Computational Limits of Forward Inference
Since the inference of matching diffusion models is a forward pass of the trained score estimator sw, the computational hardness of DiT ties the munsformer-baseee score network,
sw (A1, A2, Ag) W (Wov Ag D D- D-eexp exp
(ATWK WqA2)),
(4.5)
dxL
LXL
Lxs
LxL
following notation in Definition 4.1. For inference, we study the following approximation problem. Notably, by Remark 4.1, (4.5) subsumes both conditional and unconditional DiT inferences.
Problem 2 (Approximate DiT Inference ADTTI I,, T, Ōr)). Let dr > 0 and B > 0. Given A1, Ag, Ag E _ and Wov, WK, We E ddd with guarantees that (WovAAslls I
B, || Aillo < B and [[WoAells S B, we aim to study approximation problem
See Assumption 2.

---

ADITI d,, L, B,ā), that approximates sw (A1, A2, A3) with a vector ž E Rdo (with do d- L) such that || WAR (WovAsD- exp(A, W W2A22) mux I óp. Here, || = max.,j Aij for any matrix A.
By (02) and (03), make observation that Problem 2 is just special case of [Alman and Song, 2023]. Hence, we characterize the all possible efficient algorithms for ADITI with next proposition.
Proposition 4.1 (Norm-Basee Efficiency Phase Transition), Let || Σ B, Aıllx Σ B and WW As I B with B = O vlgg L). Assuming SETH (Hypothesis 1), for every q > 0. there are constants C, Ca C. > such that: there is no (n²2-))-tim (sub-quadratic) algorithm for the problem ADITI(L, C log L, B L, L L.
Remark 4.3. Proposition 4.1 suggests efficiency threshold for the upper bound of || A1 oc" [|W4 oot WWovAAll..c Only below this threshold are efficient algorithms for Problem 2 possible.
Moreover, there exist almost-linear DiT inference algorithms following [Alman and Song, 2023].
Proposition 4.2 (Almost- Linear Time DiT Inference). Assuming SETH, the DiT inference problem ADITI(L,4 d = (log I., = olo L), 5f = 1 (T)) can be solved in J, [11+ time.
Remark 4.4. Proposition 4.2 a special case of Proposition 4.1 under the efficiency threshold.
Remark 4.5. Propositions 4.1 and 4.2 are dominated by the relation between L and d, hence by the subspace dimension do = dL. A smaller do makes Propositions 4.1 and 4.2 more likely to hold.
5 Discussion and Concluding Remarks
We explore the fundamental limits of latent DiTs with 3 key contributions. First, we prove that transformers are universal approximators for the score functions in DiTs (Theorem 3.1), with approximation capacity and model size dependent only the latent dimension, suggesting DiTs can handle high-dimensional data challenges. Second, we show that Transformer-basee score estimators converge the true score function (Theorem 3.2), ensuring the generated data distribution closely approximates the original (Corollary 3.2.1). Third, we provide provably efficient criteria (Proposition 4. and prove the existence of almost-linear time algorithms for forward inference (Proposition 4.2) and backward computation (Theorem 4.1). Our computational results hold for both unconditional and conditional generation of DiTs (Remark 4. 4). These results highlight the potential of latent DiTs to achieve both computational efficiency and robust performance in practical scenarios.
Practical Guidance from Computational Results. Section 4 analyzes the computational feasibility and identifies all possible efficient DiT algorithms/methods for both forward inference and backward training. These results provide practical guidance for designing efficient methods:
The latent dimension should be sufficiently small: d = Olog L) (Theorem 4.1, Propositions 4.1 and 4.2).
Normalization of K, Q, and V in DiT attention heads enhances performance and efficiency: For efficient inference: max [[W A1ll, A2ll, [Woo S B with (Proposition 4.2) and A1, A2, Ag being the input data associated with K, Q. V. For efficient training: max A. ||Wa [Wov I with I olViog (The-
ol o(yoo L)
orem 4.1).
We remark that these conditions necessary but not sufficient; sufficient conditions depend on the specific design of the methods used. This is due to the best- or worst-case nature of hardness results.
Limitations and Future Direction. As discussed in Remark 3.4, the double exponential factor in our explicit sample complexity bound (Theorem 3.2) suggests a possible gap in our understanding of transformer universality and its interplay with DiT architecture. This motivates us to rethink transformer universality and explore new proof techniques for DiTs, which we leave for future work. Besides, due to its formal nature, this work does not provide immediate practical implementations. However, we expect that our findings provide valuable insights for future diffusion generative models.
Post-Acceptance Note [October, 29, 2024]. A follow-up work by Hu et al. [2024f] alleviates the double exponential factor and achieves minimax optimal statistical rates for DiTs under Hölder smoothness data assumptions.
10

---

Broader Impact
This theoretical work aims to shed light on the foundations of diffusion generative models and is not anticipated have negative social impacts.
Acknowledgments
JH would like to thank to Minshuo Chen, Sophia Pi, Yi-Chen Lee, Yu-Chao Huang, Yibo Wen, Damien Jian, Jialong Li, Zijia Li, Tim Tiz-Kit Lau, Chenwei Xu, Dino Feng and Andrew Chen for enlightening discussions on related topics; Ting-Chur Liu for pointing out typos; and the Red Maple Family for support. The authors would like to thank the anonymous reviewers and program chairs for constructive comments.
JH is partially supported by the Walter P. Murphy Fellowship. HL is partially supported by NIH R01LMM372201, AbbVie and Dolby. The content is solely the responsibility of the authors and does necessarily represent the official views of the funding agencies.
References
Beatrice Achilli, Enrico Ventura, Gianluigi Silvestri, Bao Pham, Gabriel Raya, Dmitry Krotov, Carlo Lucibello, and Luca Ambrogioni. Losing dimensions: Geometric memorization in generative diffusion. arXiv preprint arXiv:2410.08727, 2024.
Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances in Neural Information Processing Systems (NeurIPS), 36, 2023.
Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix soft- max attention to kronecker computation. In The Twelfth International Conference on Learning Representations (ICLR), 2024a.
Josh Alman and Zhao Song. The fine fine-grained complexity of gradient computation for training large language models. In NeurIPS. arXiv preprint arXiv:2402.04497, 2024b.
Josh Alman and Zhao Song. Fast rope attention: Combining the polynomial method and fast fourier transform. manuscript, 2024c.
Luca Ambrogioni. search of dispersed memories: Generative diffusion models are associative memory networks. arXiv preprint arX9v:2.009. 17290, 2023.
Fan Bao, Chongxuan Li, Yue Cao, and Jun Zhu. All are worth words: vit backbone for score-based diffusion models. In NeurIPS 2022 Workshop on Score-Based Methods, 2022.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear con- vergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. Transactions Machine Learning Research, 2022. ISSN 2835-8856.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:187711 2020.
Junsong Chen, Jincheng YU, Chongjian GE, Lewei Yao, Enze Xie, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. SSalphaa: Fast training of diffusion transformer for photorealistic text-to-image synthesis. In The Twelfth International Conference on Learning Representations (ICLR), 2024a.
Minshuo Chen, Xingguo Li. and Tuo Zhao. On generalization bounds of a family of recurrent neural networks. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics (AISTATS), volume 108, pages 1233-1243, 2020a.
Minshuo Chen, Wenjing Liao, Hongyuan Zha, and Tuo Zhao. Distribution approximation and statisti- cal estimation guarantees of generative adversarial networks. arXiv preprint arXiv:2002.0:3: 03938 2020b.
11

---

Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on ow-diimensional data. International Conference on Machine Learning (ICML), pages 4672-4712 PMLR, 2023.
Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Ap- plications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771, 2024b.
Marek Cygan, Holger Dell, Daniel Lokshtanov, Dániel Marx, Jesper Nederlof, Yoshio Okamoto, Ramamohan Paturi, Saket Saurabh, and Magnus Wahlström. On problems as hard as enf-sat. ACM Transactions Algorithms (TALG), 12(3):1-2 2016.
Huaian Diao, Zhao Song, Wen Sun, and David Woodruff. Sketching for kronecker product regression and p-splines. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 1299-1308. PMLR, 2018.
Huaian Diao, Rajesh Jayaram, Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for kronecker product regression and low rank approximation. Advances in neural information processing systems (NeurIPS), 32, 2019.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self aelf-attention mechanisms. In International Conference on Machine Learning (ICML), pages 5793-5831. PMLR, 2022.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024.
Luciano Floridi and Massimo Chiriatti. Gpt-3: Its nature, scope, limits, and consequences. Minds and Machines, 30:681-699 2020.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pages 23164-23173, 2023a.
Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. fast optimization view: Reformulating single layer attention in I1m based on tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint arXiv:2309.07416 2023b.
Yeqi Gao, Zhao Song, and Shenghao Xie. In-context learning for attention scheme: from single soft- max regression to multiple softmax regression via tensor trick. arXiv preprint arXiv:2307.02419, 2023c.
Jiuxiang Gu, Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Tensor attention training: Provably efficient learning of higher-order transformers. arXiv preprint arXiv:2405. 16411, 2024.
Jiaqi Guan, Xiangxin Zhou, Yuwei Yang, Yu Bao, Jian Peng, Jianzhu Ma, Qiang Liu, Liang Wang, and Quanquan Gu. Decompdiff: diffusion models with decomposed priors for structure- based drug design. arXiv preprint ar arXiv:2403.07902, 2024.
Alexander Havrilla and Wenjing Liao. Understanding scaling laws with statistical and approximation theory for transformer neural networks on intrinsically low-dimensional data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Benjamin Hoover, Hendrik Strobelt, Dmitry Krotov, Judy Hoffman, Zsolt Kira, and Duen Horng Chau. Memory in plain sight: survey of the uncanny resemblances between diffusion models and associative memories. arXiv preprint arXiv: 16755 2023.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint 006.02200600685,. 2021.
12

---

Jerry Yao- Vao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu, Bo- Yu Chen, and Han Liu. On sparse modemn hopfield model. In Thirty-sevent Conference on Neural Information Processing Systems (NeurIPS), 2023.
Jerry Yao-Chieh Hu, Pei-Hsuan Chang, Haozheng Luo, Hong-Yu Yu Chen, Weijian Li, Wei-Po Wang, and Han Liu. Outlier-efficient hopfield layers for large transformer-based models. In Forty-first International Conference on Machine Learning (ICML), 2024a.
Jerry Yao-Chich Hu. Thomas Lin, Zhao Song, and Han Liu. On computational limits of modern hopfield models: A fine aine-grinned complexity analysis. In Forty-first International Conference Machine Learning (ICML), 2024b.
Jerry Yao-Chieh Hu, Maojiang Su, En-Jui Kuo, Zhao Song. and Han Liu. Computational limits of low-rank adaptation (Iora) for transformer-based models. arXiv preprint arXiv:2406.03) 36, 2024c.
Jerry Yao-Chieh Hu, Wei-Po Wang, Ammar Gilani, Chenyang Li, Zhao Song, and Han Liu. Funda- mental limits of prompt tuning transformers: Universality, capacity and efficiency. arXiv preprint arXiv:2411.16525 2024d.
Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hopfield models: Transformer-compaiinn dense associative memories as spherical codes. In Tiihty-eighth Conference on Neural Information Processing Systems (NeurIPS), volume 38, 2024e.
Jerry Yao-Chieh Hu, Weimin Wu, Yi-Chen Lee, Yu-Chao Huang, Minshuo Chen, and Han Liu, On statistical rates of conditional diffusion transformers: Approximation, estimation and minimax optimality. arXiv preprint 11215:2411.17522, 2024f.
Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat, Journal of Computer and System Sciences, 62(2):367-375, 2001.
Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. Dnabert: pre-trained bidirectional encoder representations from transformers model for dna-language in genome. Bieinformatics, 37 (15):2112-2120) 2021.
Haotian Jiang and Qianxiao Li. Approximation theory of transformer networks for sequence modeling. arXiv preprint arXiv:2305. 18475, 2023.
Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? arXiv preprint arXiv:2307. 14023, 2023.
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations (ICLR), 2022.
Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Zhuoyan Xu, and Junze Yin. Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers. arXiv 2024a.
preprint arXiv:2405.0522
Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Yufa Zhou Beyond linear approx- imations: A novel pruning approach for attention matrix. arXiv preprint arXiv:2410.1 12661, 2024b.
Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Multi-layer transformers gradient can be approximated in almost linear time. arXiv preprint arXiv:2408. /3233, 2024c.
Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Tensor attention training: Provably efficient learning of higher-order transformers. arXiv preprint urXiv:2405.16411, 2024d.
Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, Jianfeng Gao, Lifang He, and Lichao Sun. Sora: A review on background, technology, limitations, and opportunities of large vision models, 2024.
Zhonghua Liu, Yue Lu, Zhihui Lai, Weihua Ou, and Kaibing Zhang. Robust sparse low-rank embedding for image dimension reduction. Applied Soft Computing, 113:107907, 2021.
13

---

Zhengxiong Luo, Dayou Chen, Yingya Zhang, Yan Huang, Liang Wang, Yujun Shen, Deli Zhao, Jingren Zhou, and Tieniu Tan. Videofusion: Decomposed diffusion models for high- high-quality video generation. In 2023 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10209-10218 IEEE, 2023.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden- and Saining Xie. Sit: Exploring flow and diffusion-bant generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.087. 2024.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv: 02010. 2023.
Shentong Mo, Enze Xie, Ruihang Chu, Lanqing Hong. Matthias Niessner, and Zhenguo Li. Dit-3d: Exploring plain diffusion transformers for 3d shape generation. Advances in Neural Information Processing Systems (NeurIPS), 36, 2023.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:21 2021.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning (ICML), pages 26517-26582. PMLR, 2023.
OpenAL. Sora: A video generative model based on transformer diffusion. OpenAl Research, 2024. Accessed: 08/16/2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 4205, 2023.
Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning, arXiv preprint arXiv:2104.08894 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1(2):3, 2022.
Hubert Ramsauer, Bernhard Schafl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitmer, Milena Pavlovic, Geir Kjetil Sandve, al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02211 2020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition (CVPR), pages 10684-10695, 2022.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems (NeurIPS), 32, 2019.
Bing Su and Ying Wu. Learning low-dimensional temporal representations. In International Conference on Machine Learning (ICML), pages 4761-4770. PMLR, 2018
Wenpin Tang and Hanyang Zhao. Score-based diffusion models via stochastic differential equations a technical tutorial. arXiv preprint wrXiv:2402.07487 2024.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 11287-11302, 2021.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Diffusion language models are versatile protein learners. arXiv preprint 18857,2402.18567, 2024a.
14

---

Yan Wang, Lihao Wang, Yuning Shen, Yiqun Wang, Huizhuo Yuan, Yue Wu, and Quanquan Gu. Protein conformation generation via force-guided se (3) diffusion models. arXiv preprint 140888, 2024b.
Yihan Wang. Jatin Chauhan, Wei Wang, and Cho-Jui Hsieh. Universality and limitations of prompt tuning. Advances in Neural Information Processing Systems (NeurIPS), 36, 2023.
Andre Wibisono, Yihong Wu, and Kaylee Yingxi Yang. Optimal score estimation via empirical bayes smoothing. arXiv preprint arXiv:2402.07744 2024.
Virginia Vassilevska Williams. On some fine- aine-griiedd questions in algorithms and complexity. In Proceedings of the international congress of mathematicians: Rio de janeiro 2018, pages 3447- 3487. World Scientific, 2018.
Dennis Wu, Jerry Yao-Chieh Hu, Teng- Yun Hsiao, and Han Liu. Uniform memory retrieval with larger capacity for modern hopfield models. In Forty-first International Conference on Machine Learning (ICML), 2024a.
Dennis Wu, Jerry Yao-Chich Hu, Weijian Li, Bo-Yu Chen, and Han Liu. STanhop: Sparse tan- dem hopfield model for memory-enhance time series prediction. In The Twelfth International Conference on Learning Representations (ICLR), 2024b.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are trans- formers universal approximators of sequencc-to-sqquemee functions? In International Conference on Learning Representations (ICLR), 2020.
Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar, Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306. 09305, 2023.
Xiangxin Zhou, Xiwei Cheng, Yuwei Yang. Yu Bao, Liang Wang, and Quanquan Gu, Decompopt: Controllable and decomposed diffusion models for structure-based molecular optimization. arXiv preprint arXiv:2403. 13829, 2024a.
Xiangxin Zhou, Dongyu Xue, Ruizhe Chen, Zaixiang Zheng, Liang Wang, and Quanquan Gu. Antigen-specific antibody design via direct energy- -aasey-based preference optimization. arXiv preprint arXiv:2403.16576, 2024b.
Zhihan Zhou, Yanrong Ji, Weijian Li, Pratik Dutta, Ramana Davuluri, and Han Liu. Dnabert-2: Effi- cient foundation model and benchmark for multi-species genome. arXiv preprint arXiv:2306.15000 15006, 2023.
Zhihan Zhou, Weimin Wu, Harrison Ho, Jiayi Wang, Lizhen Shi, Ramana V Davuluri, Zhong Wang, and Han Liu. Dnabert-s: Learning species-aware dna embedding with genome foundation models. ArXiv, 2024c.
Zhenyu Zhu, Francesco Locatello, and Volkan Cevher. Sample complexity bounds for score-matee matchnn Causal discovery and generative modeling. Advances in Neural Information Processing Systems (NeurIPS), 36, 2023.
15

---

Appendix
More Discussions on Low-Dimensional Linear Latent Space
17
Notation Table
18
C Related Works
19
D Supplementary Theoretical Background
21
D.I Diffusion Models
21
D.2 Proof of Lemma 2.1
21
D.3 Preliminaries: Strong Exponential Time Hypothesis (SETH) and Tensor Trick
23
E More Background and Auxiliary Lemmas: Universal Approximation of Transformers via Piecewise Approximation
E.1
Piccewise-Constant Function Approximates Compact-Supported Continuous Function
25
E.2
Modified Transformer Approximates Piecewise-constant Function
26
E.2.1
Quantization by Modified Feed-forward Layers
27
E.2.2
Contextual Mapping by Modified Self-attention Layers
28
E.2.3
Map to the Desired Output by Modified Feed-forward Layers
29
E.3
Standard Transformers Approximate Modified Transformers
E,4
All Together: Standard Transformers Approximate Compact-supported Continuous Functions
E.5
Supplementary Proofs
30
E.5.1
Preliminaries
30
E.5.2
Proof of Lemma E.2
32
E.5.3
Proof of Lemma E.4
32
E.5.4
Proof of Lemma E.5
37
E.5.5
Proof of Lemma E.7
39
F Proofs of Section 3
40
E.1
Proof of Theorem 3.1
40
E.1.1
Auxiliary Lemmas for Theorem 3.1.
40
F.1.2
Main Proof of Theorem 3.1
F2
Proof of Theorem 3.2
46
F2.1
Auxiliary Lemmas for Theorem 3.2
46
F.2.2
Proof of Theorem 3.2
................................
47
E3
Proof of Corollary 3.2.1
................................
52
F.3.1
Auxiliary Lemmas
52
F.3.2
Main Proof of Corollary 3.2.1
54
G Proofs of Section 4
55
G.I Auxiliary Theoretical Results for Theorem 4.1
55
G.1.1 Low-Rank Decomposition of DiT Gradients
55
G.1.2 Low-Rank Approximations of Building Blocks Part f()), (()) and c(.)
58
G.1.3 Low-Rank Approximations of Building Blocks Part II: p(.)
59
G.2 Proof of Theorem 4.1
16

---

A More Discussions on Low-Dimensional Linear Latent Space
Our analysis is based on the low-dimensional linear latent space assumption (Assumption 2.1). Here we further discuss this in light of our theoretical results
Our results are more general and extend beyond Assumption 2.1. In addition to the case where do D, our theoretical results apply do = D, our results still hold by setting B as the identity matrix Ip. Namely, our results hold after removing the linear subspace assumption.
other settings: do D and do > D. Especially, for
Statistically, for score approximation, score estimation, and distribution estimation, the upper bounds depend on the dimension of the latent variable do, other than A smaller do allows for reduced model size to achieve a specified approximation error compared larger one (Theorem 3.1). Additionally, with smaller do. both score and distribution estimation errors are reduced relative scenarios with larger ones (Theorem 3.2 and Corollary 3.2.1).
Computationally, smaller do benefits the provably efficient criteria (Proposition 4.1, almost-linear time algorithms for forward inference (Proposition 4.2) and backward computation (Thcorem 4.1).
17

---

B Notation Table
We summarize our notations in the following table for easy reference.
Table 1: Mathematical Notations and Symbols
Symbol
Description
Euclidean norm, where is vector Infinite norm, where is vector 2-norm, where Z is a matrix Operator norm, where Z is matrix Frobenius norm, where Z is a matrix p. q-norm, where Z is a matrix 12--norm where f function
=
1 2
op
12
F
[2 p.g.
|f(x))
[1(x)) LP(())
nP) where f is function and P distribution Lipschitz-norm, where is funetion
[f( 1
f&P
Pushforward measure, where is function and P is distribution
n
Sample size
Data point in original data space, R.D Latent variable in low-dimensional subspace, h Rdo
h
Ph
The destiny function of h
B
The matrix with orthonormal columns to transform x, where B E Dxdo Perturbed data variable at t >
h
h= B
T
Stopping time in the forward process of Diffusion model Stopping time in the backward process of Diffusion model Discretized step size in backward process
To
P
The density function of x for at time t The density function of h at time t (Conditional) Gaussian density function
d
Input dimension each token in the transformer network of DiT Token length in the transformer network of DiT
X
Sequence input of transformer network in DiT, where X R Position encoding, where E RdxL
E
R(-)
Reshape layer in DiT, R(-) 12 + RdxL
WB
The orthonormal matrix to approximate B, where WB E R Dxdo
18

---

C Related Works
Organization. In the following, we first discuss recent developments in DiTs. Then, we discuss the main technique of our statistical results: the universality (universal approximation) of transformer. Next, we discuss recent theoretical developments in diffusion generative models. Lastly, we discuss other aspects of transformer in foundation models beyond diffusion models.
Diffusion Transformers. Diffusion [Ho et al., 2020] and score-based generative models [Song and Ermon. 2019] have been particularly successful as generative models of images, video and biomedical data [Nichol et al., 2021, Ramesh et al., 2022, Liu al., 2024, Zhou et al., 2024a.b, Wang et al., 2024a,b]. Recently, transformer-basee attention in research. The U-ViT model [Bao et al., 2022] incorporates transformer blocks into U-net architecture, treating all inputs as tokens. In contrast, DiT [Peebles and Xie, 2023] utilizes a straightforward, non-hierarchical transformer structure. Empirically, diffusion transformers (DiT's [Peebles and Xie, 2023] have emerged as a significant advancement (e.g.. SoRA [OpenAI, 2024, Liu et al., 2024] from OpenAl), effectively combining the strengths of transformer architectures and diffusion-based approaches. Models like MDT [Gao et al., 2023a] and MaskDiT [Zheng et al., 2023] improve the training efficiency of DiT by applying a masking strategy.
diffusion models have garnered significant
Universality and Memory Capacity of Transformers. The universality of transformers refers to their ability serve universal approximators. This means that transformers theoretically models any sequence-to-see to-sequence function to a desired degree of accuracy. Yun et al. [2020] establish that transformers can universally approximate sequence-to-sequee sequeeeeee functions by stacking numerous layers of feed-forward functions and self-attention functions. In a different approach, Jiang and Li [2023] affirm the universality of transformers by utilizing the Kolmogorov- representation Theorem. Most recently, Kajitsuka and Sato [2023] show that transformers with one self-attention layer is universal approximator.
The memory capacity of a transformer is a practical measure to test the theoretical results of the transformer's 's universality, by ensuring the model can handle necessary context and dependencies. By memory capacity, we refer to the minimal set of parameters such that the model (i.e., transformer) approximates all input-output pairs in the training dataset with bounded error. Several works address the memory capacity of transformers. Kim et al. [2022] show that transformers with Ö L+ L) parameters are sufficient to memorize N length-L and dimension-d sequence-to-seqqeee
data points by constructing a contextual mapping with (L) attention layers. Mahdavi et al. [2023] show that a multi-head- -atenntitoo with h heads is able to memorize (hL) examples under linear independence data assumption. Kajitsuka and Sato [2023] show that single layer transformer with (N Ld + d²) parameters is able to memorize N length-L and dimension-d sequence-to wowennceeto-sqqueee sequenne data points by utilizing the connection between the softmax function and Boltzmann operator. Hu et al. [2024d], Wang et al. [2023] extend the results of (Kajitsuka and Sato, 2023, Yun et al., 2020] to prompt tuning and discuss the memorization of the data sequences. Another line of research establishes a different kind of memory capacity for transformers by connecting transformer attention with dense associative memory models (modern Hopfield models) [Hu et al., 2024a, 2023, Wu et al., 2024a,b, Ramsauer et al., 2020]. Notably, they define memory capacity as the smallest number of (length- L and dimension-d) data points the model (transformer attention) is able to store and derive exponential-in bigh-probability capacity lower bounds. In particular, Hu et al. [2024e] report tight exponential scaling of capacity with feature dimension from the perspective of spherical codes.
Our work is motivated by and builds on [Yun et al., 2020] to bridge the transformer's function approximation ability with data distribution estimation. While we do not address the memorization of DiT's (or diffusion models in general), recent studies on dense associative models suggest viewing pre- trained diffusion generative models as associative memory models [Achilli et al., 2024, Ambrogioni, 2023, Hoover et al., 2023]. We plan to explore this aspect in future work.
Theories of Diffusion Models. In addition to empirical success, there has been several theoretical analysis about diffusion models [Chen et al., 2024b, Tang and Zhao, 2024]. Chen et al. [2023] studies score approximation, estimation, and distribution recovery of U-Net based diffusion models. Benton et al. [2024] provide convergence bounds linear in data dimensions, assuming accurate score function approximation. Zhu et al. [2023], Wibisono et al. [2024] provide statistical sample complexity bounds for matching under the similar assumptions. Oko et al. [2023] analyze the distribution
19

---

estimation under the assumption that the initial density is supported on [-1, and smooth in the boundary.
Among these works, our work is built on and closest to [Chen et al., 2023], as both assume the data has low-dimensional structure". However, work differs in three key aspects. First, beyond the simple ReLU networks considered in [Chen et al., 2023], we provide the first score approximation analysis for DiTs with transformer-based score estimator. Second, our work is the first provide the statistical rates of DiTs (score and distribution estimation) based on transformer universality [Yun al., 2020] and norm-based converging number bound [Edelman et al., 2022], supporting the practical success of DiT's [Esser et al., 2024, Ma et al.. 2024]. Lastly, our work provides the first comprehensive analysis the computational limits and all possible efficient DiT algorithms/methods for both forward inference and backward training. This offers timely insights into the empirical computational inefficiency of DiTs [Liu et al., 2024] and guidance for future DiT architectures.
Transformers in Foundation Models: Transformer-Bine
Pretrained Models. Transformer- based pretrained models utilize attention mechanisms to process sequential data, enabling the learning of contextual relationships for tasks like natural language understanding and generation. These models encompass three types: ancoder-based, decoder-based, and diffusion transformers. Encoder- based transformers, such as DNABERT [Zhou al., 2024c, 2023, et al., 2021], employ bidirectional attention to extract feature representations DNABERT shows great potential to capture complex patterns genome sequences and improve tasks such as gene prediction. Decoder-based transformers generate output sequences from encoded information using unidirectional attention, such ChatGPT [Radford al., 2019, Floridi and Chiriatti, 2020, Brown et al., 2020] for natural language. The diffusion transformers generate sequence toward target distribution, such SoRA [Liu et al., 2024] and Videofusion [Luo et al., 2023] for video generation and DecompDiff [Guan et al., 2024] for drug design. In our paper, we present early exploration of the statistical and computational limits of diffusion transformer models.
'Recent work by Havrilla and Liao [2024] examines the generalization and approximation of transformers under Hölder smoothness and oov-dieeniiona subspace assumptions.
20

---

D Supplementary Theoretical Background
In this section, provide some further background. We show the details about the forward and backward process in Diffusion Models in Appendix D.I. Besides, we give the details of the proof about the score decomposition in Appendix D.2.
D.1 Diffusion Models
Forward Process. Diffusion models gradually add noise the original data in the forward process. We describe the forward process as the following SDE
dæ, w()xx,d vw(t)dW, 21 E IRD.
(D.1)
where z0
Po. (We) a standard Brownian motion, and w(t) > is nondecreasing weighting function. Let Pt and Pr denote the marginal distribution and destiny of xz. The conditional distribution P(r) follows (3(())=o, DD) where B(t) = exp Jo (ss)ds/22) and (t) (t).
In practice, (D.1) terminates at large enough such that Pr is close (0,1D))
Backward Process. process satisfies
We obtain the backward process 2T by reversing (D.1). The backward
dy: 1 (T t)ye + w(T -1)
llggprrr -t) dt ()dW,
(D.2)
where the score function Vlogp(()) is the gradient of log probability density function of zt, and W. is reversed Brownian motion. However, V (-og and Pr are both unknown in (D.2). To resolve this, we use score estimator sw(,t) f) to replace V log pr(),, where (w () is usually neural network with parameters W. Secondly, we replace Pr by the standard Gaussian distribution. Consequently, we obtain the following SDE
di WT t)ū, + wT - t)sw (Ur, T -rt) t) dt Vw(T - t)dW,, e. To ~ N(O, ID). (D.3)
In practice, we use discrete schemes of (D.3) to generate data, following [Song and Ermon, 2019]. We denote the discretization step size. For E [Riu, (k 1)pi,, we have
1
diy
(T (T -t)Te, W( t)ss t T kj)) dt Vw((T tjdW..
(D.4)
D.2 Proof of Lemma 2.1
Here restate the proof of [Chen et al., 2023, Lemma 1] for completeness.
Proof. Recall x = Bh by Assumption 2.1 with x E ERD, B R Dxdo and E Rd.
By the forward process (D.1), we have
((( =
v1(3 /1(( 3h)pm(h)dh
(D.5)
where
[1)(t) - 2 (1)
( Bh) [2mh((t)-PP)) D/2 exp
(D.6)
the Gaussian transition kernel.
21

---

Then we write the score function as
(P((T)
V (o))((F)
(D.7)
Pr (T
V J (3 Bh)ps (h)dh
(By pluging in pe pc()))
J (T Bh)pk (h)dh
( Bh) (h)dh
(By interchanging f with V)
J [e(T (h)pk
where the last equality holds since (/ Bh) is continuously differentiable æ. Plugging (D.6) into (D.7), we have
v log (T)
[2x((1)]-D22 V/ (h)ch
/ o(t) (B(1) - x) exp
(13(4) x112
ph (h)dh.
20(t)
We then decompose above score function by projecting of into Span(B) i.e., replacing -x with -BB (LD -BB J=:
v log
(TTh(1]]DD) D/2
Jull (h)pp(()jdl
/ (t))
(B - BB 8) (ID BB ) E exp [1,(t) 01/2
pa (h dh.
20(()
Absorbing the factor of [2th(t)] D/2 into the Gaussian kernel (
Bh), we have
v log Pe (x)
J (x Bh)pa (h)dh /
[2)((t)] D/2
[1(1)Bh- m1
(8(1)B BBTT x) exp
ph (h)dh
(tt)
20(()
((t) (ID (TD- --BBB) (/((1)) |
(h)pk(h)dh
(T Bh)pa
1
1
( (h)ph(h)dh o(t) (B(t) BB x) ve(E (h)pp(h)dh)
o(t) (ID BB F.
==
To further simplify 84, decompose (> Bh)
(x Bh)
-2rhh(t)]-D/2 D/2 exp
1
BB(t) m12)
(t)(t)
1 (t [1B(t)Bh - BBT x - (ID BB )x1((]
-22 exp
-12mh(t)]-D/2 D/2 exp
20(t) (11((t)Bh- BB 2112 + |1 - BBB F
=
- BB
22

---

[2 D/2 exp
(t( (13(t)Bh- BB =B2 1I
1
BB 2
(B((()) - 13 in (p)n(B) while (To BB* ) orthogonal to Span (B)
||8(1) -B x
[1 -BB x
[2xh((t]]do// exp
[2th((t]-(D-do)// exp
20(
(1)
(0 BB)))) (since B has orthonormal columns)
where both 1/2 (B fi) and the (LLD - BBTBBBBB )z) are Gaussian.
Plugging the (x Bh) = the (BTx x hb) vh ((ID BBT into 84, we obtain
(z = C o(t) (B(t) BB Z) ( ( BB (z)T)))p(( (h)ch
C(D BB )x)
a(t) (B(1)B). BBT x) (B (h)dh
J 11)p / o(t) (8(t)Bh - BB ( T Ø [h)p,h)dd
where C = (LD BB (. h))p())dh]
Notably, 4. depends only the projected data B x. Therefore,
able replace 8 (z, t) with 8+ (BT E, t). The benefit is that the dimension do of the first input in 8+ (B" a, ) is much smaller.
Lastly, by denoting h = B I such that V we (h h)) (8(t)h h) (t), we arrive at
V ()) dh J The 1))p
8+ ( x,,) B
BV logpl (B
(T) = f (()
This completes the proof.
D.3 Preliminaries: Strong Exponential Time Hypothesis (SETH) and Tensor Trick
Here we present the ideas we built upon for Section 4.
Strong Exponential Time Hypothesis (SETH). Impagliazzo and Paturi [2001] introduce the Strong Exponential Time Hypothesis (SETH) as a stronger form the P * conjecture. It suggests that our current best SAT algorithms are optimal and is popular conjecture for proving fine-grained lower bounds for a wide variety algorithmic problems [Cygan al., 2016, Williams, 2018].
Hypothesis (SETH), For every c > 0, there is a positive integer k 2 3 such that k-SAT on formulas with n variables cannot be solved ((2(1=e)n) time, even by a randomized algorithm.
Tensor Trick for Computing Gradients. The tensor trick [Diao et al., 2019, 2018] is instrument to compute complicated gradients in a clean and tractable fashion. We start with some definitions.
Definition D.1 (Vectorization). For any matrix X L.xd we define X vec (X) e Edd such that X1, = (i-1))a+ for all i e [L] and j E [d].
Definition D.2 (Matrixization). For any vector X E Idd, we define mat(X) = X such that X1, mat (X) X ((-1)4+ for all [L] and [d]. namely mat(.) vec (.))
Definition D.3 (Kronecker Product), Let A E xd and B Rloxd, We define the Kronecker product of A and B as A @ B E Rluloxddd such that (A & Ain Bio with i, E [Lo ja E [da], ds E [Lo].jo E [db]-
-1d,+ is equal to
23

---

Definition D.4 (Sub-Block of a Tensor). For any E Rlaxd and B E RLxxd, let A =- A B E Rhads Xd,,d, For any E [L]], we define A E IR FaLaxda Xdada be the j-th La dads sub-block A.
Lemma D.1 (Tensor Trick [Diao et al., 2019, 2018]), For any A R. R Rlaxd and X E &dd,, it holds vec (A X = (AT @ B )X Rhat.
To showcase the tensor trick, let's consider a (single data point) attention following [Gao al., 2023b,c]. Setting D = diag (exp(""W?wxx))
and W = wx W E Rdddd we have
Co = Wv X
D
expl S
1
(D.8)
dxd ERexL ER
EREXL
ERax4
Proposition D.1 (Definition 4.7 of [Gao et al., 2023b]). By Definition D.3 and Definition D.4,
we identify Djj = (exp A 11) E R for all j [L]. with A X @ X R1 and W R Therefore, for each i E [L] and i [d]. it holds Lo ΣE 111 ( exp W. .XWvl-))) y3s)
The elegance of Proposition D.I emerges when we vectorize the weights into vectors W, Wy. making the gradient computations (e.g., dCo/WW and dCa/W,) more tractable by avoiding complex matrix or tensor derivatives. This approach systematically simplifies the handling of chain-rule terms in the gradient computation of losses like Co.
Fine-Grained Complexity for Transformer. Many recent works also utilize similar techniques from fine-grained complexity analyze transformer architectures. Alman and Song [2023, 2024b], Liang et al. [2024d], Alman and Song [2024a] explore the computational feasibility of inference and training for standard softmax and tensor attention. Liang layer training results from [Alman and Song, 2024b] to deep transformer models. [Liang et al., 2024a] extend [Alman and Song, 2024b] to provide a fast attention gradient approximation based on Fourier transform. [Liang et al., 2024b] extend [Alman and Song, 2024b] to sparse attention matrix. Hu al. [2024d] study the computational limits of inference and training prompt-tuning for pretrained transformers. Hu et al. [2024c] study the computational limits of LoRA [Hu et al., 2021] in transformers, identifying norm- norm-boure bound conditions for efficient LoRA training and proving the existence of nearly linear-time LoRA algorithms.
[2024c] extend the single-
Our work closest to [Alman and Song, 2024b, 2023]. Our forward inference computational results build on [Alman and Song, 2023]. Our backward training computational results are related to [Alman and Song, 2024b] but include additional analysis on reshaping and latent embedding.
24

---

E
More Background and Auxiliary Lemmas: Universal Approximation of Transformers via Piecewise Approximation
Here, we review the universal approximation transformers following [Yun et al., 2020).
Our goal is to reproduce the results of (Yun et al., 2020] and use or modify them as auxiliary lemmas for proofs Section (i.e., Appendix F.)
We start with their central result and prove it in the rest of the section.
Lemma E.1 (Universal Approximation of Transformers, Theorem [Yun et al., 2020]). Let > For any given compact-supported continuous function f R R 1, there exists transformer network fr E T,2,,44 such that
1/2
[10((X)
Sc.
Proof Overview. We use the following proof strategy:
Step 1. We show that the piccewise-constant function is able to approximate compact-supported continuous function in Appendix E.I.
Step 2. We define modified self-attention and feed-forward layers to construct the modified transformer. We show that the modified transformer is able to approximate piecewise-constant function Appendix E.2.
Step 3. We show that the standard transformer in Appendix E.3 is able to approximate the modified transformer.
We provide details of Step 1. in Appendix E.I, Step 2. in Appendix E.2, and Step 3. in Appendix E.3. Then we summarize our results in Appendix E.4.
E.1 Piecewise-Constant Function Approximates Compact-Supported Continuous Funetion
In this subsection, we show that the piecewise-constant function is able to approximate compact- supported continuous function.
We start with the definition of the compact-supported continuous functions of interest.
Assumption E.1. Without loss of generality, we assume that the target function in discussion is supported on [0, 1 1]dx1, We denote the [0, 1]dx -supporred continuous functions as F.
We introduce the notion of grid and cube for the compact support [0, 1
Definition E.1 (Grid and Cube with Width 6). Given a grid width 6, let Go == {0, ......I - 5)xx1. denote the set of grids within [0, 1,xL, For grid point G = (Gze(d),eE()) G8. we denote its associated cube as
So S 8k [Gj.A: Gj,E 6) C [0, 1
Each cube So represents hyper rectangular in the multi-dimennional space |0, 11aX4, constructed discretize the space into smaller subspaces.
We introduce the notion of piecewise-constant fuention class w.r.t. the [0, _jdx L supported continuous function class F.
Definition E.2 (Piecewise-Consse Constant Function Class). Let fa denote the piesewise constant function of grid width 6, and Il {.} denote the indicator function. For each G' E Go. and any matrix Ag E RdxL define the piesewise-constant function class
F(6) = fo X Eceg. AG 1{X Sc),Ao E RdxL)
(E.1)
25

---

We recall that for a given sequence-to-sequeee function f.
|ff = - 1/2
We approximate the compact-supported function with a piecewise-constant function in the next lemma.
Lemma E.2. (Lemma 8 of [Yun et al., 2020]) For any given f F and e/3 O, can find
5* 0, such that there exists fa- E F((*) satisfying |I - fō- 11 S c/33.
Proof. See Appendix E.5.2 for a detailed proof.
E,2 Modified Transformer Approximates Piecewise-constant Function
this subsection, we define modified self-attention and feed-forward layers to construct the modified transformers, We use the modified transformers to approximate the piecewise-constant function.
Definition E.3 (Modified Transformer Networks). The modified transformer network Tr.ma includes two modifications to the standard transformer network Tp.m..
Modified attention layer: Replace Softmax operator with Hardmax operator 0 (-).
Modified feed-forward layer: Replace ReLU(-) with an activation function E V. Here, V denotes the set of all piecewise linear functions with at most three pieces and at least one constant.
We approximate F(8) with this modified transformer networks Tr.m.l
Lemma E.3 (Modified from Proposition 4 of [Yun et al., 2020J). For each fs E F(6), there exists JT.a E 22.1.1 such that [| - ST.cll ((d/2).
Proof Sketch. Given 8, and for any grid G E Gs. we have grid set Go and the cube Sç. Our proof follows steps:
Quantization. For all X pdxt, we quantize it to a finite set:
If E Sg C [0, 1]dxL, we quantize to the element E G5. If & [0, 1]dx we quantize it to element out of Gs.
Mapping. For any G E Go, map it to the desired output Ag.
For Quantization, we achieve this by series of modified feed-forward layers. We show this in Appendix E. For Mapping, follow two steps:
For any G + G' E Go. we use a "contextual mapping" 9e(-) (defined as Definition E.4). The series of
mapping maps all the elements in (G) and qe to different values. Then, modified self-ateention layers to achieve "contextual mapping". We show this in Appendix E.2.2.
Definition E.4 (Contextual Mapping). Consider finite set Gs E RdxL A map defines contextual mapping if the map satisfies the following:
Go - R x
For any G E Gs. the entries in ge (G) are all distinct.
For any G A G' E Go. all entries of ge (G) and (G')") are distinct.
For any G E Gs. this Appendix E.2.3.
series of modified feed-forward layers to map (G) to Ac. We show
26

---

Remark E.1. Our proof differs from (Yun et al., 2020] in one aspect: Although [Yun et al., 2020, Proposition 4] outlines proof for transformer networks without positional encoding and sketches the proof for networks with it, we provide a detailed proof for the latter to support our proof.
E.2.1 Quantization by Modified Feed-forward Layers
We use series of modified beew-fowward layers in Tr.m. element G of the following grid:
quantize an input X E gdxL to an
p
1- 8)
where J > > 0 is a large number to be determined later. We achieve this via two steps.
Step 1: Map the element out of [0, 1) to J.
We e, represent the standard unit vector where the i-th element is 1. For the i-th row of X, we define the following feed-forve forward layer to achieve our aim.
Definition E.5 forward Layer 1), The vector e; acts as the weight parameters, and (i(+) acts as the activation function in the feed-forward layer
(i()) { (-t -t--J, J, fort 0r 1, 0, otherwise.
X X + (e,
(E.2)
We take = 1 as example to give the specific calculation. Let X (wij)dxL then we have
FF(X)
(G()))) ((x12))
(x
Si (2),1) (1(212)
( 0
0
X
:
In the first row of X, the above layer transforms the element that is out of [0, 1) to -J. We stack the above layers together for i = 1, 1....... d. If the element of is out [o, 1), the series of layers maps it to J.
Step 2: Map the element in 1) to (0, 5, 26, 1- 8).
respectively. We define the following layer. Definition E.6 (Feed-forward Layer 2). The vector e; acts as the weight parameters and (2(-) acts as the activation function in the feed-forward layer
For the i-th of X, take 0, 1, 1..., 1/8 1/,.................
0, <<0 or l 6, 0 6.
X + x + tecole - ko1,)), (o(t)
(E.3)
-t,
We take = and A: example. We give the following specific calculation
FF(X) = x +
(62(01.1 - 5) (2 6)
Ç2(01,1 8))
(2(11,1 6)
62(21,2-1 -)
(e( 5) 0
0
= X +
:
:
0
0
27

---

In the first row of X, the above layer transforms the element in [6, 26 to 6. We stack the above layers together for i 1, ......... dd and 0, o..,..... 1/8 1. If the element of X [kō, (k: + 1)6), the series layers maps it to kö.
Combining the above two parts, achieve our goal with d/ô + d feed-forward layers. We denote the d/5 + d series layers as fr,
E.2.2 Contextual Mapping by Modified Self-attention Layers
In our attention layers, we use the following positional encoding E E RdxL
2
L
1
2
L
E
(E.4)
:
2
L
According to Appendix E.2.1, the output of fr is in the grid J, 0, ..... .,. 87ax1, For any X this grid, the first column X E is in
(-J,0,8, 6-100,
and the second column is in
The results are similar in the other columns.
For i 0, ...., L 1, we use the following notation:
[ 1 j], = fi J,i,i 2 1 -6..
Then, define the grid 9 the following.
Definition E.7 (Grid 9 ). We add E to all the grid points in Gs to generate the modified grid G defined as follows:
9 00 8: 1 813 [1 8: 515 x LL 8 8]3.
Next, show that the modified attention layer computes contextual mapping (Definition E.4) for gg. For = 1,2, L 1, we use the following notation:
8 6) {i, + 8, 25.
Lemma E.4 (Modified from Lemma 6 of [Yun et al., 2020]). We consider the following subset of G3:
G = 0::::: 5)d 11:0: 2 8)dt x x L - 1 8/2
Assume that and 8-1 > 2. Then, there exist a function fT, RdxL - d composed of
5-d modified attention layers (Definition E.3), a vector Rdd and two constants 4, fr E IR:
(0 t,), such that (G) u fT, (G), G E G satisfies the following properties:
1. For any G E G. all the entries of (G( are distinct.
2. For any different G, G' E all the entries of (G). gG()") are distinct.
3. For any G E Go. the entries (G) are in [4i, 4...
4. For any G E - \ Gs. all the entries of (G) are outside [t, tr...
28

---

Proof. See Appendix E.5.3 for a detailed proof.
Remark E.2. Our proof differs from [Yun et al., 2020] in one aspect: The original [Yun et al., 2020, Lemma 6] does not include positional encoding (E.4). Although Yun et al. [2020] sketches the proof for networks with (E.4) in the attention layer input, we detail the proof.
E.2.3 Map to the Desired Output by Modified Feed-forward Layers
Next, we show that series feed-forware layers map the output of modified attention layers fT, to the desired output of function for.
Lemma E.5 (Lemmaa 7 of [Yun al., 2020J). There exists function fT,03 RdxL d composed of (LL(1/8)ál LL!) modified feed-forward layers, such that
AG
if 9a, ifG E 9
(T
1sxL
\ Gg..
Proof. See Appendix E.5.4 for detailed proof.
In conclusion, we have the following lemma for the required number of layers in the modified transformer.
Lemma E.6 (Total Number of Layers). From the proof of Lemma E.3, if want achieve approximation error (84/2) by the modified transformer, need O(8-1) modified ferw-forward
layers in fT,,1, O68 d) modified self-attention layers in fr. fc, and (88-4() modified feed-forward layers in FT.c3.
Proof. By the proof of Lemma E.3, we complete the proof.
E.3 Standard Transformers Approximate Modified Transformers
In this subsection, we show that standard neural network layers are able to approximate the modified self-attention layers and the modified feed-forward layers (Definition E.3). We have the following Lemma E.7.
Lemma E.7 (Lemma of [Yun et al., 2020]), For each JT.c E T and any > 0, there exists fr E 2,1,1 such that [[ fT.ell,3 I c/3.
Proof. See Appendix E.5.5 for a detailed proof.
E.4 All Together: Standard Transformers Approximate Compact-supported Continuous Functions
We summarize the results of Lemmas E.2, E.3 and E.7. Then prove Lemma E.I. Furthermore, to achieve the approximation error in Lemma E.1, we take å = O(e//d) Lemma E.3.
29

---

E,5
Supplementary Proofs
We first present two preliminary concepts: selective shift operation and bijective column ID mapping in Appendix E.
Then we show
Proof of Lemma E.2 in Appendix E.5.2
Proof of Lemma E.4 in Appendix E.5.3
Proof of Lemma E,5 in Appendix E.5.4
Proof of Lemma E.7 in Appendix E.5.5
E.5.1
Preliminaries
Here, we give the definition of two preliminary concepts: selective shift operation and bijective column ID mapping.
Selective Shift Operation. This operation refers to shifting certain entries of the input selectively. To achieve this, consider the following function E(:-) Rdx + R
((X;bo) equ Xo [(u T X) (u TuTx X bq [,)]
(E.5)
where X E RdxL = (1,0,0, _) E Rd, and bo E R. E Rdd is vector to be determined.
To see the output, we consider the j-th column of T XoH [(u x)T (uT TX bbo1,]]]
If X.. > bq. it calculates argmax of u X;
If T., <
bq. it calculates argmin of u X.
All rows of E(X; hQ) except the first row are zero. We consider the j-th entry of the first row in b(X;bo), which is denoted as E(X; bo)1j. Then for all j [L], we have
E(X;bo) = = Xon [(uJ X)T (u X.J - bo)l = mmax u X., (ming u X. if u X < bq-
if f X.j > bo
From this observation, we define a function parametrized by bo and 1/2 (with bo < 6/2)
<(X; (/2) E(X;bo) 5(X:
(E.6)
Then have
max T X.,x - ming u X. if bo < < =x < B
c(X;bo, 19)1.
(0.
others.
We define an attention layer of the form X + + + EXX; R2). For any column X if be Xx. < b/g its first coordinate X1,; is shifted up by max; X ming u "X while all the other coordinates stay untouched. We call this the selective shift operation because can choose ba and i/ shift certain entries of the input selectively.
Bijective Column ID Mapping.
We consider the input G E G (Definition E.7), We use
J= L + dl. and u (1,0-1,8 8 dd11)
(E.7)
For any j [4], we have the following two conclusions:
30

---

If Gr,y 2 0 for all i for
[d], i... G. [ - 1::::
then we have
u G [ō :0:00;0 + +d++11_ d+1 6]. where dj = - 8 ( d+1 )
(E.8)
The mapping G. + T G., maps the elements in [j- j to [8, 6: 83 + 0 -d+1 - 8].
This is a bijection.
If there exists E [d] such that Gi.j JJ + j. then
8 18-dt1 - 8- < <0.
u G S -3L0 (j ()-
1-8
(E.9)
We say that u G.g gives the "column ID" for each possible value of G E Ei - 1: 0: 8:5 87 Remark E.3 (lllustration of Bijection Properity). For the bijection property, we give the following illustration. Let G.j- (gi ". 3 and G.j (2jj.0aj." -aaa) If u G.j u G. and G.j. * G.g. we deduce
(g) - 91j) + (925 (2j) + + 8 d+1 (gdj 9dj) 0.
(E.10)
Because Gy # G.j. then there exists a k (k < d), such that gkj # Jas and g&و = Digli > k). We have
0 10-++1 (gks -mx)) 8-k+2
However,
(9)s D1s) + ... + 8- (gh-1s -x-1,))
I 101s - 0ıjl + +... + 10k+t k+2 (gx-1,5 ~x-1,)
+ 8 8) (1 5) ....
8-+
This contradicts with (E.10). Thus prove the property of bijection.
31

---

E.5.2 Proof of Lemma E.2
Proof of Lemma E.2. We restate the proof from [Yun al., 2020] for completeness.
By the nature of the compact-supportee continuous function, is uniformly continuous.
Because I1-1loo is equivalent to ||-11p when the number of entries are finite, we have the following by the definition of uniform continuity.
For any e/3 > 0, there exists a 0* > O, such that for any x, Y E RdxL, and || - oY < <d", we have (I](X) f(YIl, < e/33.
Then we perform the following steps following Definitions E.1 and E.2:
We create grid Gs- by choosing grid width 8*, We also create cube So with respect to G Gs..
For any grid point G G5*, define CG E Sg as the center point of the cube Sg.
We define piecewise-constant function (X) ELeg, ((CG) Sc).
For any E Sc. we have ||X - Callo < 8 According to the uniform continuity, we drive
lX f8- (X)Il = llf(X) f(Co)ll < e/3.
This implies that |lf - fs- Ia E/3 and completes the proof.
E.5.3
Proof of Lemma E.4
We give the proof Lemma E.4 by constructing the network to satisfy the requirements.
Proof of Lemma E.4. Recall the selective shift operation in Appendix F.5.1. The overall idea of the construction includes two steps:
Step 1: For each j E [L], we stack 5-d attention layers. For E 1ō 5 0, + +8-4+1 - (E.8) in the increasing order, use the attention layer as
(E.11)
The total number of layers is L8-d, These layers cast G Go L different entries required by Property of Lemma E.4.
Step 2: We add an extra single-head attention layer with the following attention part
L8 -6-((++)d--(((:0))
(E.12)
This layer achieves global shifting and casts different G e Go to unique elements required by the Property of Lemma E.4.
The two operations map To and 9 \ I3 to different sets, as required by Property and Property 4 of Lemma E.4. The bounds l and 4, are calculated then.
Then, we give detailed proof by showing the impact of the two steps and verifying the four properties of Lemma E.4. We achieve this by making a category division of 9
Category G G.. all entries in the point are between and L
Category G e 5f G. the point G has least one entry that equals J.
Let (1,0-1 108--18-22 8-2 8 d+.+). Recall that 6, (j -1) -----11 - 1) for any j E [r.] in (E.8).
32

---

Category 1. We denote 9j = u G.... then we have 91 < 92 < <... 9L. The first 8-d layers
sweep the set (6, 8 : 6, + 8- - 8 [L] and apply selective shift operation on each element the set. This means that selective shift operation will be applied to 9i first, then 92, followed by 9s. and so on.
The First Shift Operation. In the first selective shift operation with g going through [61 ö ô1 + 8-d+1 - 6). the (1, 1)-th entry of G (i.e., G1,1) is shifted by the operation, while the other entries are left untouched. The updated value G1,1 is
-I,1=1, GI,1 + max (u (.,G..)) min (u G.A))
G1,1 + 801 91).
Therefore, the output of the layer after the operation is
(G.1 G.22 G.
Let 91 == o² G. We have
91 G1,1 + Σe 8- 1-it G1,1
1=2
= G1,1 + (0 - 91) an + Σ
G-1+1
i=2
91 + ( -9)) 91).
Then we deduce gr. < 31 because
91 = 91 5-
5- - 91)
(gt.
8 8 d+1
2 10+8 8-d
(L 1
-8- 8--+11 + 0)
(By (E.8)
8
+ + (L 1) 5-d+1
- rd
(L 11-1) 1)
8 8-
8
8
2 pd.
1)
+ 8
1
8-d+1
(L 1)
+ 8 --111 d+1
8
gL.
(By 8 < 1 and (E.8)
Thus, after updating, we have max T (G.. G.
G.L) = max(gi192. 92..........91
]66 = I1.
and the new minimum is 92-
The Second Shift Operation. In the second selective shift operation with g going through [8g : 8: 8g + 8 5-4+1 5), the (1, -)th entry of G (i.e., G1,2) is shifted by the operation, while the other entries left untouched. The updated value G1,2 is
G G1,2 + 8- 92)
G1,2 d 92) + 5 2dl 91).
Therefore, the output of the layer after the operation is
(ō.,1 G.22
G.,))
33

---

We have
92 = u T G2
92 + 8 (91 92) + 000 -11).
Then we deduce Φ1 < 92. because
91 + 8 (gL - 91) < 92 + ( 92) 8-
[g -91)
d ()(22 91) < 8 18 () - -gi).
(By 5- > 1 and g2 > 92)
Thus, after updating, we have
max u (G.a T.2² G.,)) max (0r, 92. ................................
91 .9. 92.
and the minimum is 9s-
Repeating the Process. By repeating this process, we show that the j-th shift operation shifts G1,5 by 8 "0 - 9j). Then we have
9j u G..
= Jj +
- - (2--+1) + (gr
(gj=
+ - 91).
We deduce Øj-1 < Øj holds for 2 S L, because
カ1 9j
9--1 + Σ
bd+d - - 1+1 + -8- (-(i-1)d (oL - 91)
km
9j +
kdd (g - -gg-++1) + 8 (61 -91)
1-1
kd 1)(0)-441-----))) - (8--))00-- 1) - g1),
where the last inequality holds because
i-1 Σ
kd+d ffi-b+1 - -gi-k)
k=1
< < 1)d
[gj- gÿ-N
8- 11-11 (gL
g1)-
Therefore, after the j-th selective shift operation, Øj is the new maximum among (Ø....., Ij. (jtti...,,99. and 9j+1 is the new minimum.
After L Shift Operations. After the whole L shift operations, the input G is mapped to new point G. where u. G = (01 Øz bound of gr, we have the following lemma.
ΦL) and ÿ < 52 < Tr. For the lower and upper
Lemma E.8 (Lemma 10 of [Yun et al., 2020]). ÿL = u T G. satisfies the following bounds:
8 1l(-11)
1) s タL L8 (L+1jd
34

---

Also, the mapping from (91 92
gt) to IL is one-to-one mapping.
Global Shifting by the Last Layer We note that after the above L shift operations, there is another attention layer with attention part Lô L8- (L+1)dd-1(((;)) Since 9 91, it adds the following to each entry in the first row of G:
18-(L+1jd-1 max u G. Lb L8 9L.
k
The output of this layer is defined to be the function GG). In summary, for any Go. [d], and j [L], we have
GL + o% = Gaj
(T)
if sisss d,
where 83 = Ek= For any Lg and j [L],
j-1 8 ( -
-j- -
8 (9L 9)) + L8 -8 9L.
T fT. (T) = - + 1-5 (L+1)d 9L.
Next, we check the Property 1, Property and Property of Lemma E.4.
Checking Property 1 of Lemma E.4. Given any G G.. already prove that
Φ1 < 92 < ... < gL.
All of them are distinct.
Checking Property 2 of Lemma E.4. Note that the upper bound on ğL from Lemma E.8 also holds for other 9j (j E [L 1)). For all j e [L], we have
L8 L8- "r S u (G) < L8- gr. + L8 (L-+1)))d
By Lemma E.8, two different G, G' E Go are mapped different ğL and Hh. and they differ at least by This means that the following two intervals are guaranteed to be disjoint:
[L- (1+1)d-1 -2. L8 (L+ 1)d
9L L8 L8- (E++))))))
[LS [L88((41] Vy, 1----81--111 (L+1)d-1 "y L85 1-------00 (1+1)d)
Thus, the entries of u fG. and u fT (G')) are all distinct. Now, we finish showing that the mapping /T (-) uses (1/5)² + attention layers to implement contextual mapping on G8.
Checking Property 3 of Lemma E.4. Given Lemma E.8 and u" (G) [Lō (Lo-(L+1)d-17 'r. L8-(4+17)) (1+1))-11 "r. + L8 (8-(4)))d), for any Gs, we have
u fT.) 2 L8 18-2
a (T,) I28 1 2(L+1)d-1 L5 L5- (4+1)d
This proves that TT T are between 4 and t,, where
t L8 2(L+1)d 1),,
35

---

f = 2 + Ld (8+14)
Category Now we check the Property 4 of Lemma E.4. For the input points G E 9 \ 96. note that the point G has at least one entry that equals to -J k, [L - 11. Let gj = G.j Recall that whenever column G.j has an entry that equals to - k. k [T. 1}, we have 9j 0. Without loss of generality, assume that 91
Because the selective shift operation is applied to each element of [0 : 8 òn + 0 8- 6] and is not applied to negative values, thus we have min; u G., = < 0. 91 never gets shifted upwards and remains the minimum for the whole time.
All gj 's are Negative. When all 9j 's are negative, selective shift operation never shifts the input G. Thus G G. Recall that u urG, G. < O for all j [L]. The last layer with attention part L8 L8-(L++)) E(: 0) adds L8-(I+1)d-1 ming u T G. < 0 to each entry in the first row of G. This makes G remain negative. Therefore, fT,22(G) satisfies u (T)))(G)., 0 < t for all j [4].
Not All Øj 's are Negative. Now consider the case where at least one is positive. Suppose that there are k positive elements and they satisfy Iú < 9tg <... Dix Thus selective shift operation does not affect gs. where E
[L] \ fi....., i It shifts 9s, by
8 (maxu u G. min u G.,a)
2 5 (2L88 di - (L - 8- -
8- - 8 - -8- + lik - 1)
(By (E.9))
1 -
1
8-d+1 - 8
-3LL8 8 (L ix)
1 1-5
2 24..225
(By 2)
2L0 228-((11)4
The next shift operations shift Dig....> Bi, by an even larger amount. Therefore, at the end of the first L(1/0)d layers, we have L8- (L+iid s gh ... S The and .j 0 for j [L] \ (in...., 8.
Then, we shift by the last layer. The last layer with attention part Lg - acts differently for negative and positive 9J 's. (i). For negative 9j 's, it adds the following to gj. [L] fi.,..,,
L8 (L+1)d-1 min u G., = LG- 91 < O.
This term pushes them further the negative side. (ii). For positive gi 's, it adds
Lö 18-(1+1)d- max u T G = 1.5-(T+1)4d Ti. 2 2-22--2(22-))) 2(1++))
Thus they are all greater than or equal to (12882(L+1)d++ Note that
2 (L+1)dd11 tr, where t, L28 ( + 18-(4+1)))
Then the final output fr.e2(G) (G) satisfies u fr.e(G):.) $ [tr,ty], for all j [L]. This completes the verification of Property 4 of Lemma F.4.
In conclusion, we need O(L8- layers of modified self-attention layer to obtain our approximation. This completes the proof.
36

---

E.5.4 Proof of Lemma E.5
Proof of Lemma E.5. We restate the proof from (Yun et al., 2020] for completeness.
Note that [ (1/6 + 1)dk < co, so the output of (g has finite number of distinct real values. Let M be the upper bound of all these possible values. By the construction of fT,c2, M
Construct the Layers: fi, JT,ea(fT,c2(
OaxL if 9 \ Gs. According to Lemma E.4, for all j E [L], we have u. (G)c(()) E [t,, t,] i. G. and T fT)) & [t,t, if E 93 G... Due to this property, add the following feed-forward layer.
Definition E.8 freed-forward Layer 3). The vectors u and 1 act as the weight parameters, and (a(-) acts the activation function in the feew-forward layer.
0
ift E [4, 4r]
X x (M + 1)L (x), x), Ss(t)
(E.13)
1
if E 41,4).
Case for E 9 G8. We have (5(u") JT.ch (G)) T. Thus, all the entries of the input are shifted by -M and become strictly negative.
Case for G E G... We have Ca(u" (G)) = 0 so the output stays the same as the fT,2 (G).
With the input fT (G), Gs. then Ca(u (G)) OF. Thus, the output stays the same as the input. If G E 9 \ Gs, then (ı(uT (G)) = 1[. Thus, all the entries of the input are shifted by- -M 1 and become strictly negative.
Next, we map those negative entries to zero. For 1,2,
add the following layer:
Definition E.9 (Feed-forward Layer 4). The vectors u and e₁ act as the weight parameters and (a(-) acts the activation function in the feed-ford beed-fowward layer.
X x ecaa((es) X), (a(1) {-t ift
(F.14)
if 2
After these d layers, the output for G 9 G a zero matrix, while the output for G Go remains (G).
Construct the Layers: JT, (T,a((TT=2(G)) Ag if E G.. Each different G is mapped to J. unique numbers T (T), which are at least apart from each other. We map each unique number to the
corresponding output column follows. We choose one G E Gs.. For each u fr fT, j E [L]. we add the following feed-forward layer.
Definition E.10 (Fread-orwarr Layer 5). The vectors and e, act as the weight parameters, and (-)- acts the activation function the feed-forward layer.
X +X + ((Ag).s (T. Colu Tx x aT T (G.
(E.15)
(1( = {1 -8/2 8/2,
(E.
others.
Case for G 9 Go. Recall that the input X of this layer is fr.2(G). If X is zero matrix, which is the case for G E 95 No. we have u T X OT. Then u X-u (T.ee (G) I -4.1. Since t > 6/2, 2, the output remains the same X.
37

---

Case for G E Gg. Let the input X be fT), where G 9s is not equal to G. According to the Property 2 of Lemma E.4 and given a j E [L], u fr.a(G): (k [4]) differs from u (T) by at least 8. Then we have
fT,a'T. u (T, = -
Thus the imput is left untouched. If G G. then
(s(u fT. u ( = (c))
Thus shift the j-th column of fT, (T)
(T) + ((Ag)ss- - - ((Ao).. - -
In other word, this layer maps the column (G) to (Ag)aj. without affecting any other columns.
For each E G.. we defer that we need one layer for each unique value of u (T)))()).g) Note that there are ( such numbers, so we use (68-d2) layers to finish our construction. This completes the proof.
38

---

E.5.5 Proof of Lemma E.7
Proof of Lemma E.7. We restate the proof from [Yun al., 2020] for completeness.
The proof follows two steps: (i) Approximate the modified self-attention layers. (ii) Approximate the modified feed-forward layers.
Step 1: Approximate the Modified Self-Attention Layers.
We achieve this by approximating the Softmax operator with the Hardmax operator Given a matrix X R we have
(AX) OX),
X0.
The operator is the only difference between the normal and the modified self-attention layers. We approximate the modified self-attention layer in Tr.m,l by the normal self-attention layer with the same number of heads r and head size m.
Step2: Approximate the Modified Feed-Forward Layers.
We achieve this by approximating the activation function in with four ReLU functions. From Definition E.3, recall that v denotes three-piecewise functions with at least a constant piece. We consider the following Ç E
b1
x < C1: if c1 S < C2 if 0g S x;
((x)
02 b2
0g by
where a2, 03, b1, b3, c1, 02 R, and Ci c2.
We approximate ((x) by ((x) composed of four ReLU functions:
((x) =b 0241 + b2 bj ReLU(x ci + 82
8201 + by - b1
ReLU(x c1)
c
0302 + by a2(c2 E) -
02
6
0zl2 + by uz(02 -6) b2
ReI.U(x 02)
03
61
if < Ci - E,
(agc1 + by - b1)(x - cı)/e aze1 + b2
if - S <
02 + b2
if < 02 6,
(ag02 az(c2 - t) b2)(xx o2)/e a302 + by
if
<02,
0g + by
if S x.
Ase 0, we approximate ((x) by ((x). The activation function is the only difference between the normal and modified feed-forward layers. We approximate the modified feed-forward layer in Tr.m.l by the normal one.
Thus, for any fr.e T2,1,1, there exists a function fr E T,2,1,4 to approximate fr.--
This completes the proof.
39

---

Proofs of Section
Our proofs are motivated by the approximation and estimation theory of U-Net-based diffusion models in [Chen et al., 2023]. We use transformer networks' universal approximation theory in Appendix E and the covering number proceed with our proof. Specifically, derive the approximation error bound in Appendix F.1 and the corresponding sample complexity bound in Appendix F.2. Then we show that the data distribution generated from the estimated score function converges toward proximate area of the original one in Appendix F.3.
F1 Proof of Theorem 3.1
Here we present some auxiliary theoretical results Appendix F.I.1 to prepare for our main proof of Theorem 3.1. Then derive the approximation error bound of DiTs (i.e., the proof of Theorem 3.1) in Appendix F.1.2.
F1111 Auxiliary Lemmas for Theorem 3.1.
We restate some auxiliary lemmas and their proof' from [Chen al., 2023] for later convenience.
Lemma F.1 (Lemma 16 of [Chen et al., 2023]). Consider probability density function ps (h) xxp(-CI)112/2) for h E Rd and constant C 0. Let rh be fixed radius. Then holds
2d/
JINI,TM Pa (h)dd -
(d 1) 1o (op(-Crr//2)
2d do/2
J1N1227 |M1l2p (h)dh <
do exp (
CT +
Lemma F.2 (Lemma of [Chen et al., 2023]). Suppose Assumption 2.2 holds and g defined as:
q(h, 1) /
heb, px) dh, = B B
(h hp (h)
Given e 0, with rh = c log(d//T) + log(1/e)) for an absolute constant c, it holds
19(h, Th) ra s, for E [To. TT.
Lemma F3 (Theorem of [Chen et al., 2023]). We denote
T(rx)= = sup agto.11eejo.rr.
sup
With q(h,t) = hhae (h))))) dh (H)))) and pa satisfies Assumption 2.2, we have coarse upper bound for T(ri):
(1 3(t)
(1+
1+ (t)
Vdorn) -e do)
T(r)) 0
Let
o(1)
Lemma F.4 (Lemma 10 of [Chen et al., 2020b]). For any given > 0, and L-Lipschitz function defined on [o, 1,d,, there exists continuous function constructed by trapezoid function, such that
|
40

---

Moreover, the Lipschitz continuity of is bounded:
[f(x) [w)) s 10dgL|la vll2 for any y, [0. 1j
F.1.2 Main Proof of Theorem 3.1
Proof of Theorem 3.1. With V logpt (h) = =Baa 84 5h,h, t), we have the following in (2.4)
(h. o(t log (T) +Ba ottBB ((tt)BB(+,A0 (h, +
(F.1)
We proceed as follows:
Step 1. Approximate qh,, t) with a compact-supported continuous function f(h, t).
Step 2. Approximate f(h,t) t) with transformer network.
Step 1. Approximate qh, t) with a Compact-supported Continuous Function f(h, 8). We partition de into compact subset H₃ |bll|b/12 S rh and its complement H2, where rh is to be determined later. We approximate q(h, t) on the two subsets respectively and then prove f's continuity. Such a step achieves an estimation error of dag between q(h, 1) and f(h, (). We show the main proof here.
Approximation on H2 x [To 2). For any > O, we take = do log(do/Tb)) log e). From Lemma F.2, we have
2 Tn) 2A)))) Se for te [T..
So we set F(h,t) = on H2 [To.T].
Approximation on H1 [To. TJ. On H1 (oo, T), we approximate q(h, ()) by approximating each coordinate qu (h,&) f) respectively, where q(h, t) [qi (h, t), g2(,, t), qdo (h, t)]. We rescale the input by y/ (h+rr, /)/rr, and L' L/"T. Then the transformed input space is [0, 1] (To/ / 1]. We implement such transformation by single feed-forward layer.
By Assumption 2.3, on-support score s (h, is I,, -Liscchit in T.. This implies qh, t) is (1 + La, )-Lipschitz in h. When taking the transformed inputs, g(v .1)) q(2rny) Th 1, TT') becomes (1 + T.x, )-Lipschitz in y Similarly, each coordinate (v/ L) is also 2r Tag
Lipschitz in y'. Here we take Ln- = 1 1 Ley.
Besides, g(v',t') (') is
with respect to ł, where
T(rk) = sup tE(To. T Te
sup
at ql. t)
We have a coarse upper bound for 7 (Fw) in Lemma F.3. We restate it here for convenience
((+))= :0 (0 (2)( 3(t) T.e
a(t) 11 Vidoa) o ( Log Th T do)
In conclusion, each ()), is Lipsichitz continuous. So can apply Lemma F.4 to determine Fu(v', ()) for approximating each coordinate. We concatenate fi's together and construct f
[f Fa,]" According to the construction in Lemma F.4 and for any given achieve
[V/ 2()',()))l Se.
sup V x (Tbb/.T
41

---

Considering the input rescaling (i.e., h 3/ and t - t'), we obtain:
The constructed function is Lipschitz continuous in h. For any h1, h2 Hi and holds
[To, T,
[ - f(22, 1) I 10d h1 h2ll2-
(F.2)
The function is also Lipschitz in t. For any t1, 22 E [To. 7 and | S rhe holds
F1, (1) F(h, 107(rx)lltz tell2-
Due to the fact that the construction of f(h, 1) is based on trapezoid function, we have J(h, t) = for [h | Th and any E (To, Thus, the two parts of f(h, t) can be joined together. To be more specific, the above Lipschitz continuity in extends to the whole RR:
Approximation Error Analysis under L Norm. The L approximation error of f can be decomposed into two terms:
(1a(h. 1) F(h,t) (P(p))
=)(()(htt) - F(h, { <TAllee0rpy + [lq( >wll
The second term in the RHS above has already been bounded with the selection of rh:
(l)(h, [ rR)
ra3ll rell2e Se.
The first term is bounded by:
[1(((T, (T, <rallle((p)
$ do
[f((',t) ,(vl,
sup =.
dde.
Then we obtain
[a(h, - F(h, (lraaep) S (Vdo + 1)c.
If substitute with e/2, we obtain that the approximation error of F(h, t) is vdge.
Step 2. Approximate f(h, t) by a Transformer. This step is based on the universal approximation of transformers for the compact-supported continuous function in Lemma E. DiT uses time point t to calculate the scale and shift value in the transformer backbone [Peebles and Xie, 2023]. It also transforms an input picture into a sequential version. We ignore time point in the notation of the transformer network in DIT. Recall the reshape layer R(.) in Definition 3.1, we consider using f(.) == R- fr o R(-) to approximate Fe(-) == f(.,t), t), where fr E Tp 22,1,4
Overall Approximation Error. R-1 fr o R(-). We denote
With Lemma E.I, we approximate S.()) with 7()
H = R(h).
We have
1/2
= (1 [.(()) - f(h) (1
[T((h) F(h) La(P))
42

---

- (I
1/2
Ro Ro R R (H) dh
1/2
Fr [ (1)
Ro
p
S
(FF3)
Along with Step 1, we obtain
o((h, t) f( ( (e(py) & (la(h,,t F(h, LP(PP) + [(h, - F(h) (2((P)) S (1+ doje..
The constructed approximator to Vlogp( is & (BJ(BTan x. =))a((), and the approxi- mation error is
1 voo for any t E T. -
v 11ggp(()s$()))))) (P(P))
o1)
Settling-down of Hyperparameters. We settle down the hyperparameters to configure network here. We refer to Appendix E.2 for some of the following calculations.
1. Model Architecture Depth K.
From Lemma E.6, we have K ((1/8) To achieve e-error approximation, we set 8 O (22/d) according to Lemma E.3. Thus we obtain
K = K=
(F.4)
2. Lipchitz Upperbound for Transformer: Ly.
We denote f.) R Fr R. (-). We get the Lipshitz upper bound for fr T T in the following way
|| [ (H1) Fr (Hz)
$ Fr (H1) Fr, (Ha)
+ =IFF.. (H1) - Fr. (H2)
+ +I (H2) - Fr (H2)
$ 2e + | (H1) F.R (H2)llF
(By (F3))
$ 26 + 10d L,e, 1 - H2ll
(By (F:2))
Then we get
T.T - O
(F.5)
3. Model Output Bound for Sy.1..
For the output of the constructed transformer (-). according to Lemma E.5, we have Fr(0)
O, where O OZXL- Thus, with the Lipschitz upperbound ((d) we have [l (H))l
Oddobs rn), where || S rh. With Tı. = el /d (og(d//// TT) + log(1/e)), we obtain
CT = or= ddo log(do/ + log(1/e))
(F.6)
4. Model Parameters Bound: C2 Cov, C2ae .Q Cxq, Cg.
By definition, we have:
[[(Wbv) T s co
[( s Cov, | SOko |[ s CkQ
43

---

where = 1, For simplicity, we omit i hereafter, which does not affect our discussion.
Recall that | denotes the 2, co-norm, where the 2-norm is over columns and co-norm over rows. By the construction of modified attention layers (E.11) and (E.12) in Appendix E.5.3, we consider Wov to have the largest norm, i.e.,
1
8 3-441
0
Wov L3 18-(1+1)d-1
1
We give the following upper bounds
Wovll2.0 Ldd (L+2( = =0 (8-La)
(F.7)
d-1
Worll22 = sup |=
Wov Lo 1.5-(1+1)d-)
Σ - (6
(F.8)
By (E.11) and (E.12) Appendix E.5.3, and the self-attention layers in Appendix E.5.5, we consider WKQ to have the largest norm, i.e.,
8-1
8-1 8-2
8-d
WkQ== =
(1.0-1.... ... .0..,...ttt))
0-dt1)=
8-d+1
d+1
8 d
5 2d+2
Then we have
W kQll,,0
8-22 2d42 ((8-2d).
(F.9)
[Wkell2- = sup |
[|xo0ll2 = 2d+2 = 0(8-2))
(F.10)
We substitute 8 with (e2/d) (according to Appendix E.4) and get:
car (1)))
Cov (1/c)
C2.00
KK (1/6)011)
CкQ (1/))00)
From the construction of positional encoder (E.4) in Appendix E.2, have
0
L - L-1
0
E
;
: :
0
L-1
We deduce
ET
-V(
2,00
44

---

Thus we have
C3 = O
(F.11)
5. Parameters Bound in Feed Forward Layers: C CF.
Recall the construction of modified forw-forward layers in the proof of Lemma E.4, which includes Definitions E.5, E.6 and E.8 to E.10. With the approximation by normal feed-forward layers in Appendix E.5.5, we consider the weight parameters with the largest norm in the feed-forward layers, i.e.,
1 i
1 8-1
8-d+1
1 8 8-1
g-d+1
Wi =
1 (1,0
, 5-d+1) =
E BRadd
1 8-1
6-d+1
1 8-1
d-d+1
Then we have
d-1
C² =
Σ 1=0
8-22 =0
(F.12)
(1/c)O())
(By setting O(E1/4) according to Appendix E.4)
and
Cp
sup ==
|W = O (6-4)
(F.13)
(1)))
(By setting = -((2/d) according to Appendix E.4)
This completes the proof.
45

---

F.2 Proof of Theorem 3.2
Here present the auxiliary theoretical results about the covering number of transformer networks in Appendix F.2.1. 1. The results are based on [Edelman et al., 2022, Theorem A.17J. Then we derive the sample complexity bound of DiTs (i.e., the proof of Theorem 3.2) in Appendix F.2.
F.2.1 Auxiliary Lemmas for Theorem 3.2
Lemma F.5 (Lemma 15 of [Chen et al., 2023]). Let & be a bounded function class. Then there exists
constant such that the output of any g E 9 Rd >> [0, bj is bounded by b. Let 21, 22,
#2," .* ; Rdd
be i.i.d. random variables. For any E (0, 1), S 1, and > 0, we have
(1 + 3/a)B
(e 1-11oo)))) + (2 (2+aae) a)c) s
P sup geg
g(2)) (1 [)))
log
3n
1 1+ ((;)>>
(1+ 6/a)B
N11.l,,1-11oo (2+ 5d.
P sup E [g(2)] gEG
log
3n
8
Now, we give the definition of the covering number as follows.
Definition F11 (Covering Number). Given function class F and data distribution P. Sample data points (x,) from P. For any > 0, the covering number (c,,F, {X, IH-1) is the smallest size of collection (a cover) C E F, such that for any f E F, there exists a satisfying
max -(x.)) Se.
Furthermore, we define the covering number with respect to the data distribution as
N(e, F1 11-1)= -
sup (X.))_>PPP
((
Then we give the covering number of the transformer networks.
Lemma F.6 (Modified from Theorem A.17 of [Edelman et al., 2022]). Let T (K, Cr, Cov Cov, 22.0 CKQ. 02,00 ,CF,CE, LT) represent the class of functions of K-layer transformer blocks satisfying the norm bound for matrix and Lipsichitz property for feed-forward layers. Then for all data point 12 S Cx. have
log.N(co: Tr.m CT, c2rr Cov, Cov, C20, CKQ, C2 , CE, (T)
log(nL)
Σ
3
( (0200) d3 (2( rm3 (oo))) $
2
where a == [Ifc(Cr))Cov( 4CKo) + -CE.
Remark F.I. We modify [Edelman et al., 2022, Theorem A.17] in seven aspects:
1. We not consider the last linear layer the model, which converts each column vector of the transformer output to scalar. Therefore, ignore the item related to the last linear layer in [Edelman et al., 2022, Theorem A.17].
2. We do not consider the normalization layer in our model. Because the normalization layer IInmm (.) in the original proof only ensures that III II (X1) ( S
|| I X2112,0. ignoring this layer does not change the result.
3. Our activation function is ReLU. Thus, we replace the Lipschitz upperbound of activate function by 1.
46

---

4. We consider the positional encoding (E.4). Then we need to replace the upperbound Cx for the inputs with the upperbound Cx CF. Besides, for multi-layer transformer, the original conclusion in [Edelman et al., 2022, Theorem A.17] uses as the upperbound for the 2, co-norm of inputs. We incorporate the upperbound for the inputs into the result stated Lemma F.6.
5. We use (2.7) as the feed-forward layer, including two linear layers and residual layer. Thus, we replace the original upperbound for the norm of weight matrix with the upperbound for the norm of La + W1 in Lemma F.6. In the following, we use to estimate the log-covering number, thus we ignore the item for La here for converience. This is the same for the self-attention layer.
6. We multi-head attention, and incorporate the number of heads into our result, which is similar to [Edelman et al., 2022, Theorem A. 1.].
7. In our work, we use the transformer T2.1,4 i.e., T = 2, m =
F.2.2 Proof of Theorem 3.2
Proof of Theorem 3.2. Our proof is built on [Chen et al., 2023, Appendix B.2]. For one data sample, we define the empirical score matching loss objective (2.1) as follows
1
l"
V., log (xa//o))))) sT (oxs, f)ll
((x;8)) T To
Then we define [($w) = =E ((1;; ww
Following [Chen et al., 2023, Appendix B.2], for any (0, 1), we have
((s
Ctrino (1 + 2
aurnn ($
[(s@) - Ctrune ( (s +(1 +
inf [(sw), ESNN
(T)
(17)
(III)
where
Atrunc (8 = ExxPo
(runc((;;$ ) = Eg ((==)))(1()))) srat] > B.
We denote
" 4CT(CT + D) 2
To)),
2x = o (V (dd log log do log Cr log log(n/8)
Then we have
S (T To)
(F.14)
For any 8 > O, according to Lemma F.5, the following holds for term (I) with probability 1-8 8,
N( (T--Toll-
(T' (T-To)((-)) ((C+Te)) log(t /oy S (l-112)
(1) 0 (1 (1+6//)(C7+12 (T -To)
log
+ (2 + a)j
where 0 is constant, and 0 will be determined later.
47

---

We set
2 TT To)
1
where < is constant to determined later. Remark F2 (Selection Criteria of r). We have two criteria:
Recall that the covering number used in our setting is .V (T (1-1b)(-1) STZ..,, S132.,,11-111) Thus, we must ensure 2 7/. According to (F.14), we consider satisfying the condition 2 (nToo(T To)) Therefore, we consider 0 < 1
(CT++x)log[[])
For the exponent of (T To), although selecting a value smaller than I is possible, we find that the convergence rate with respect to 7' is dominated by the 1/T term appearing later in the second term of (8). Therefore, we continue to consider the exponent to be
Then we have
(1 (1+6/a) (C7 + log nTo(T To)
ra) N ro " 8,2,1,, 11-112)
+ 2a n TTo To)
(1) =
5
with probability 5. 3.
Following the proof structure of term (11) in [Chen et al., 2023, Appendix B.2], we have
(m)) - ( To -rre
For any > 0, let sw be the transformer network approximator the score function in Theorem 3.1. For the term (IIT), we have
(III) < ((TT) (1+ + a) ctruno ( +(1 + a) Ctrune (sw)
(177)
(1TT))
For any 8 >>0 O, according to Lemma F.5 and given that sw is a fixed function, the following holds for term (III)1 with probability 1 J.
[1)==01 1
( - To)
(1 3/a) (C7
1 7 log
nTo(T
Following the proof structure of term (I/T)2 in [Chen et al., 2023, Appendix B.2], we have
de
(IIT); =
+ Cs.
TT To)
where C3 a constant.
Putting (I), (II), and (1II) together and setting = e², then we have
(,0)(1)) - Plogp:
Tm
(C4 + +
N (n' (T) 1, nー doe²
TnToTT - -To) log
(F.15)
5
TT To)
with probability 1- 38.
48

---

Covering Number of Sya.a. The next step is to calculate the covering number of ST3.1.. STa.a consists of two components: (i) Matrix WB with orthonormal columns; (ii) Network function fr.
Suppose we have WBI, WB2 and f1, such that [[B - WBal I 81 and sup || 33r +v Dlog 7T (lfi(z,t) fe(2,t)ll S 82. where fi R fri R. f2
R o fr2 o R. Then we have
sup S3r (T..)
IsW (7,
1
WWILi(WB10,) - W22 1
at))
sup S0r, Dlog (T
1
|| [ ( f)
sup S3r Doo DIE(To,))
(tt)
WWnnf(W B2 W2nf2(Wing B2 1) [[Wnrf(WW22,.)) W WWa2
$ att) (LTO (3r + Dlog D) + 82 + 6
(F.16)
where L. upper bounds the Lipschitz constant of fT.
For the set {WB E Dxda ||Walla I 1}, its d covveriing number is (1 + 2//5//1) Ddo [Chen al., 2020a, Lemma 8]. The ó2-covering number of f needs further discussion as there is reshaping process in network. The input is reshaped from ì E Rd to H E Rdx and
|11| STZ
| Ta.
Thus we have
()( 52(h, 0)11 < 62
sup + (6.21
|31p Dlo [| ST2(H))l < 82.
sup
1,4 in Lemma F6,. We Then we follow the covering number sequence-to-sequenc transformer 12,1, get the following 82 covering number
log(nL)
K
dà (cZ,0) $ +d3 (2(C))Cov + (10r) $) 8))
Σ
52
where
a; = [o CE).
jai
According to the (F.4, (F5), (F.7,, (F88), (F.9), (F.10,, (F.12), (F.13), (F.11) and (F.6 in Ap- pendix F.1.2, we derive the following with ((2/)) (Appendix E.4) and d 4 (Theorem 3.1):
K = 0c LT ). chr O(de th1, Cov =0(e" -0c-44)), CR O(c4), 4), CkQ = 0 4) o2 = O(c Cr= =r=0(c === (T,31))
(F.17)
Cr = O (dober Vdo log(dg/ To) + log(1/e)) Ts = - (
(V log log + log Cr + log(n/8))
Each element of the input data is within [0, 1], as shown in Appendix E.
49

---

For any 83 > 0, we get the oog-coveingg number of 1,2.1.4
log (ds, T2 -11-11)) =0 (
L d2 log(nL)
83
(28k log(L//)) de log(nL)
=0(1)
83
According to (F.15), we adopt the following value for 63 in our setting
83
(C + To (T/T)
According to [Chen et al., 2023, Appendix B.2], the log-covering number of Spa..
log.N 63, S S7, 11-12)
6CT ITT do
Volo 1))
28K log(L/e) d² log(nI)
= =0 2Ddo log 1+
Tod3
(By (F.16)
0
112 Da
Da LgLL log((L))
(By (F.17))
(122211004 de log(nL))
o
(By (1/e)** > 8 log L/t))
2
De 22211110) Ded dg
By ignoring the log factors)
Substituting the log-covering number into (F.15), we have
1
l'
Vlogpt() (EAp de
T
To To
0
C2 +
(logN((s) Sy2.24, l1.1 + log(1/8)) +
do To - To)
2
(TTo(T To)
TT To)
C7 +2
1
do
(log.N S Saan l1.) + (og(1/0)) +
(F.18)
enToT
ToT
ToT
1st term
2ud term
Recall the following parameters:
C OdBL, do logg(do/To) + (og((/))),
22 = (dg log do + log CT log.l/////////o 0)),
8: probability error,
approximation error,
n: sample size,
To T/2,
D,d, do > 1: feature dimension,
L > 1: sequence length,
50

---

du = L. L.d.
L.y: Lipschitz coefficient.
Ignoring the log factors and poly d, du, Ls ), the first term in (F.18) becomes
1
1
2(1/e)
n
ToT
The second term simplified
ToT
Thus, the final bound is
20/
+
ToT è
n1-26 26
ToT
To balance the first and second terms with respect to n, we select b 1/3. Therefore, we give the final bound as
1
20
+
1
2
n1/3
TOT
n1/3TOT
ToT
This completes the proof.
51

---

F.3 Proof of Corollary 3.2.1
Our proof is built [Chen et al., 2023, Appendix CJ. The main difference between work and [Chen et al., 2023] is our score estimation error in Theorem 3.2. Consequently, only the subspace emror and the total variation distance differ from [Chen al., 2023, Theorem 3].
First, we introduce the ground truth backward SDE and the learned backward SDE of the latent variable. Recall from (D.2), Yr. denotes the backward process. We denote the backward latent variable by h = B yt- Since we write the time index explicitly, we drop the y, h notation for > 0.
Following [Chen et al., 2023, Appendix C.2], have the following ground truth backward process
[ 2 hr + V log Vlog -/ at d W.),
dhj
where Wr denotes the reversed Wiener process (standard Brownian motion) at time t (see Section for more details).
We define PA as the ground truth marginal distribution of h7.-
For the learned process Vr. we consider hi = W.T Vr. For any orthogonal matrix U E do, we define the U transformed version of h5 as TU U Uh. Then the backward SDE for hi- is
a = 11 + b (he (T, TT - 1] di + +ddTTWWI
where
SU. h Uht
1 1- +UT f(Ult t)). o(t)
We define P as the estimated marginal distribution h7 from above continuous SDE.
The discretized backward SDE of hTp"" -+.U is
dh. =, kky
[T
- (hpp T kw) dt + CaUTTTTTT i [ky,, (I + 1)1).
We define dh. To
as the estimated marginal distribution of h.4 from above discrete SDE.
Next, we present the auxiliary theoretical results in Appendix F.3.1 to prepare our main proof of Corollary 3.2.1. Then we give detailed proof of Corollary 3.2.1 in Appendix F.3.2.
F.3.1 Auxiliary Lemmas
Here include few auxiliary lemmas from [Chen et al., 2023] without proofs. Recall the definition of Lipschitz norm: for given function f. Ilf(() l = (lf - f(y)ll//1a v1l2))
Lemma E7 (Lemma 3 of [Chen et al., 2023]), Assume that the following holds
En-P, log [ < Csh.
AminE [hh [hhh1 2
ExePy [1b 1 S Ch.
where Amin denotes the smallest eigenvalue. We denote
Elo(h,t) t)] ''
1 Ex P a ddt). JTo o2(1)
(F.19)
Let To < minn(2lo(dd//sa)) (co) co) and T max log (Cz/C///////C (o)ddo), 1}. Suppose we have
E Wof(T t) Bq x, k <
(F.20)
52

---

Then we have
WBWW BB BBT O(TTb/oo)
and there exists orthogonal matrix U &R such that:
EllUT t(Uh,t) - g(h,t)))))))))) t)
To [T log To max I + Can] max; (l((,.t) Ch Co
= |
cn
Lemma E88 (Lemma of [Chen et al., 2023]). Assume that Pa. is sub- Gaussaan and that ff f(h, t) and v log (T) are Lipschitz continuous with respect to and t. For any orthogonal matrix U E RR ddd define
1 - o(t) a(t) ( (UTh+))
t (1,1)
Assume that we have the latent score matching error-bound
[" BiJny || (h, t) V log (T) dt S Clatent (T -T);;
where Elatent > 0. Then we have the following latent distribution estimation error for the continuous backward SDE:
TV PP) Pb:)
Gatten (T- (T) (Pal 14.)) exp(-T).
where PA is the marginal distribution of the generated hy, using the continuous backward SDE. To
Furthermore, let PTo denote the marginal distribution the generated hT, using the discretized backward SDE. Then we have the following latent distribution estimation error for the discretized backward SDE
TV
P Pho
Elatent (T - TT) + V KLL (0, exp(-T) + v Edis (T To),
where
max; [[
m,x [1, 1)
2 7?²
Edis
(26)
T
max, I15(;;))llipp
n/² max
E |holl? .d + nda.
(To)
and is the step size in the backward process.
Lemma F.9 (Lemma 6 of [Chen et al., 2023]). Consider the following discretized SDE with step size y satisfying T To KrØ for some Kr E N+.
1 dy: =
2 oTT kp)) yk + dB,, for tE [kp, (k 1)H),
where Yo N(O, T). Then, for T > and To x² Se(To u).
have yr-To N (0,,21) with
53

---

Lemma F.10 (Lemma 10 [Chen et al., 2023]), Assume that V log ph (h) is L, -ipschhii Then we have EA V/ log Ph ( - doLn.
F.3.2 Main Proof of Corollary 3.2.1
Proof. Recall the estimation enror in Theorem 3.2 is c, where
5(n, c, L) = n//33 2(1/c)
+ n/3 + c².
Proof of (i). By the definition (F.19) and the estimation enror in Theorem 3.2, the error bound in (F.20) equals to {(n, e, L) (TT)) in Lemma F.7. By Lemma F.10, we set Csh. do T...
Then, we have
WEW BB" 1 o [(n, L
CO
By substituting the value of E(n, e, L) and T = Ologn) into the bound above, we deduce
WBW BBB 1 o cont/3 20100)2
c/n 60
Proof of (ii). Recall that max, (10((,)).. I LT. Furthermore, according to Lemma F.7 and
Lemma F.10, we have
BlluT S(Uh, (h, (Etnten -To)),
where
5(n, c, L).
To [T
I.3 Ch co
log To)do L7 doLn)
Gistent
TTo
Co
Following the proof structure in [Chen al., 2023, Appendix C.4], get
["
U f(U h
Blu f(Uh,t) q(h, t)l1
ET
V logg (h)
dt
(t))
To
Elatent
TT To).
Following the proof structure in [Chen et al., 2023, Appendix C.4] and setting T = Ologgn), we obtain
TV(PA) To: To
=Ō -((v
(VVenuen Clatent (T-To)
1
1 + e²
2(1/e)a2
log
1//3
1/3
where hides the factor about D, da, d, La log n, and T - To
di.dis (UWB). Pro, where Pr, is the distribution generated by s using the By definition, Pro discretized backward process. This completes the proof of the total variation distance.
Proof' of (iii). We apply Lemma E9 due to our score decomposition. With the marginal distribution at time T To and observing μ < To. obtain the last property.
This completes the proof.
54

---

G Proofs of Section 4
Our proofs are motivated by the observation of low-rank gradient decomposition transformer like models [Alman and Song, 2024b, Gu et al., 2024]. With our simplifications and observations made in Section 4, we utilize the fine-grained complexity results of transformer and attention [Hu et al., 2024b, Alman and Song, 2024a,b] and tensor trick (Lemma D. and [Diao et al., 2019, 2018)) to proceed our proofs. Specifically, we approximate DīT training gradients with a series low-rank approximations in Appendices G.1.1 to G.1.3, and carefully match the multiplication dimensions so that the computation of des forms a chained low-rank approximation in Appendix G.2. JW
G.1 Auxiliary Theoretical Results for Theorem 4.1
Here we present some auxiliary theoretical results to prepare our main proof of the Existence of almost-linear Time Algorithms for ADITGC Theorem 4.1.
G.1.1 Low-Rank Decomposition of DiT Gradients
We start by some definitions. Recall that W Rddd and W. E Rd denotes the vectorization of W paxå following Definition D.I.
Definition G.1. Let A1, A2 E RdxL be matrices. Suppose A A & A RL Define Ajo E gxxd as an L x de sub-block of There are L such sub-blocks total. For each jo E [L]. define the function u v(W), Rdd - R by u(W), exp(Aj, W) E R.4.
Definition G.2. Let A,, A2 E RdxL be two matrices. Suppose A = Af @ AT E RL' Define Ajo E L2x as an L x d² sub-block of A. There are L such sub-blocks in total. For every index jo [L]. consider the function a a(W) jo mdd IR defined by a(W),, (exp(A,, W 1
Lx1
LX1
Definition G.3. Suppose that a (W) E IR and j RR_ are defined as in Definitions G.1 and G.2, respectively. For a fixed jo E [T.]. consider the function f(W) oR + RR defined by
f f(W)
(WW) u(
scalar
LX1
Define f(W) R as the matrix where the ju-th row is (f(W),.)T.
Definition G.4. For every in E [d], define the function h(W )io Rdd + R2 by
h(Wov)i
Ag
(Wov)
Lxd
dx1
Here, Wov dxd denotes the matrix representation of Wov Rdd and (Wov) represents the ig-th column of T Define h(Wov) E Lxd as the matrix where the in column is h(W i
Definition G.5. For each jo E [L]. denote f(W); E g& as the normalized vector defined by Definition G.3. For each io E [d]. (Wo )i is defined as per Definition G.4. For every pair (jo. do) E [L] x [d] define the function e(W) Edd x d + IR by
c(W) (f(W),, jo h h(W ov)i.)) -YT 10.
55

---

where (Wov)jono is the element at the (jo, 1o) position of the matrix Wov E BR c(-) has matrix form
W
((W)
Lxd
Lxd
With the tensor trick (Appendix D.3), compute the gradient dag
of the DiT loss follows:
dg dw
Jo.i0 W
(G.1)
dW
ja=1 to=1
(G.1) presents a neat decomposition of dga Each term is easy enough to handle. Thus, we arrive at aw
the following lemma. Let Z(i, =] and Z|:j] be the i-th row and j-th column of matrix Z.
Lemma G.1 (Low-Rank Decomposition DIT Gradient). Let matrix A1, A2, A3. W, Wov, and loss function C follow Definition 4.1, and A A A It holds
( diag (S(W),) (W)ss h(W ov)ig.
(
dg2 dW
Σ
(W)oo 1o.10 Ajo (W)
(G.2)
(I)
Proof. Let Z[i, and [[ the i-th row and j-th column of matrix Z.
With DiT loss Definition 4.1, we have
L d
dg aw
Ja.i0 W
2
jo=1
i=1
d
de(W)) jo.i0 dW to
30 (W
10.80
dW
===1 1=1
d
(((W), jor h(W )io)
J0 W J0.50
By Definition G.5)
dW
dW
æ=1
d
df(W) jn
aw Jo c(W)
n( )io
dW,
== ==1
da -1 (W) u(W)
dw 20 c (WW)
h h(Wovl) )ia
(By Definition G.3)
dW,
æ= 1=1
d
to c (WW
a(W) du(W))
in dW, in dW.
da())
u (W) J0 Massh(W5 (Wov)
dW
1==1
du(W)
(WW) da(WW)
Ja W
Je.oo Jo.50
(o(W)
u ( jo: h h(oov)s
dw
dW,
jo
dW,
===1
By chain rule)
For each jo E [L]. we have
d (A,, W)
dW
A dw,
= (A,o) [..
dW 40
56

---

Therefore, for each jo E [L], we have
du(W),
dexp (Ajo W)
(By Definition G.1)
dW
dW
d A W
cxp (A W)
By entry-wise product rule)
dW
2 & (WW)
By Definition G.I again)
Similarly,
da(W), j0
d (u(W)) Nor 1t)
(By Definition G.2)
dW io
dw,
(A @ w(W 1z)
By entry-wise product rulc)
(Auul,)), 1), uW)).
(y Definition G.I again)
Putting all together, we have
dg2 (W) jo
dW
= ((h(Woo)) ((W)s) - (h(Wov)). ((W) (Asol, ((W)s))))))))) do.
where
(h (h(W )in Agols, f(W),.) (h(Wov)sn J(W)so)
Ajo (diag (f()),)) f(W),o (WW) b(Wov
This completes the proof.
Observe (G.2) carefully. We see that (I) is diagonal and (II) is low-rank. This provides hint for algorithmic speedup through low-rank approximation: If we approximate the other parts with low-rank approximation and carefully match the multiplication dimensions, we might formulate the computation of dgg as chained low-rank approximation. aW
Surprisingly, such an approach makes computing (G.2) as fast as in almost-linear time. To proceed, we further decompose (G.2) according to the chain-rule in the next lemma, and then conduct approximation cemm-by-trrm.
To facilitate our proof, it's convenient to introduce the following notations.
Definition G.6 (q(-)). Define c(W) E IR Lxd as specified in Definition G.5 and h hWovv) E R gLxd as described in Definition G.4. Define (W) E R by
(W) = c(W) h(Wov)
Lxd
dxL
addition, g.W), denotes the jo-th row of q(W), transposed, making an 1 vector.
Definition G.7 (p(-).p1(-), P2())). For each index jo E [L]. we define p(W),, E R as follows:
pWW), jo (diag(((W)s, ff-ff((w)) f(W),
9(W) jo-
57

---

We define p(W) RLxL such that J. forms the jo-th row of p(W). In addition, for every index Jo E [L], we define pi (W) P2 (W), R
pi io = diag (s (W) g (WW) Jor
P2 == (WW) / (W),, ((W) jo:
such that pWW pi (W) P2(W)
p(-) allows us express
dga in neat form: aw
Lemma G.2. Define the functions f(W) E RLXL, c(W) E RdxL, h(Wov) E Rd RIX,, and p(W) E R as specified in Definitions G.3 to G.7, respectively. Let A1, A2 E R be two given matrices, and define A = Al @ AT. Define d2 according to (01), and g2 W jo. be as described in (G.1). It holds
(W) E
dg2 dW
vec (AIP(W)A))
(G.3)
Proof. By definitions, (G.1) gives
d(g2)io
(G.4)
dWin
= Cjov5o
((f()), @ A h(W h(Wov)iol
(f(W) jo h(W ov )io))o) (f(W) Asorto))
Jo. dlag(f(W)) MWWov)o
=A s(Wso SWNT hooo
(By (a 2 b) = a diag(b)c for a, b, R*)
Therefore, (G.4) becomes
d d( d(g2)jonie jonie Cja.io (A)..
AT..t ((W),, (WW), N(W )io)
dW, io
Cjoio AJ. (diag(f(()) - ((W), I(W) 1. (WWoV
(G.5)
Then, by definitions of q(-).((())))))))) p(-), we complete the proof.
G.1.2 Low-Rank Approximations of Building Blocks Part 1: ((),,(()), and c(-)
The definitions of p. Pt. P2, and Lemma G.2 show that the DiT training gradient des entry-wise products f. 9. and c. Therefore, we approximate these with dimension-matchhe dga itself becomes a low-rank approximation. In the following
aw involves
low-rank approximations, computing aW sections, we present low-rank approximations for f. and c.
Lemma G.3 (Approximate ((-), Modified from [Alman and Song, 2023]). Let r ol o(vlog L)
and k [,0(1), Let A1, A2, E RaxL W Baxd and f(W) D-1 exp A1 X A2)
with
D diag (exp (A] W Ag) 1z) follows Definitions G.1 G.3 and G.5. If max (1) W
max S
1' < 1,, then there exist two matrices U1, V1 E 4xki such that ||U,K f(W)) S
e/poly In addition, takes I.(+((1) time construct U, and Vi.
Proof. By [Alman and Song, 2023, Theorem 3), we complete the proof.
Lemma G.4 (Approximate e(-)). Assume all numerical values in Olog L) bits. Let O L) and c(W) E IR follows Definition G.5. There exist matrices U1, V E L2 such that V. h(Wov) Y -c(W) max I </poly(L).
58

---

Proof of Lemma C.4.
V h(Wov) - -Y -c(W -V--(W
(Wor)
h(Wor) -Y - -((W))(Wov) - rT) max (By Definition G.5)
max
= =1[[un - -f(W)] h(Wov)
<//poly(L).
(By [Alman and Song. 2023, Theorem 31)
Lemma G.5 (Approximate q(-)). Let l22 Ld(), c(-) R follow Definition G.5 and let
(W) = (WW E RRLL (follow Definition G.6). There exist two matrices U2, V/2 Rlxk2 such that 122/ ((W) max I </poly(LL. In addition, it takes I, [++)( time to construct U2, 1/2
Proof of Lemma G.5, Our proof is built on [Alman and Song, 2023, Lemma D.3].
Let q(-) denote an approximation to q()).
By Lemma G.4, V h(Wov) Y approximates cl (W) up to accuracy 1/poly(L..
Thus, by setting (WW = h(Wov)
YY", y)) we find a low-rank form for q(-):
q (W)
h(ov)) (h(Wor)) V.U h(Wov)
such that
(W q( W max ((Wor) (UIV, ((Wov) N(Wov)
max
d dlln(Wov))ll.. |12 h(Wor) -Y - -c(W)
c/poly(L).
By k1, d = (1),, compute (h(Wov)) Vi U. takes only L [++((1) time. This completes the
dxL.
Lxki XyXI
proof.
G.1.3 Low-Rank Approximations of Building Blocks Part II: p(-)
Now, we use the low-rank approximations of f, to construct low-rank approximations for pr pr()). p.(),.()))
Lemma G.6 (Approximate p1 (-)). Let k2 I,((1), Suppose U1, Vi R approximates
f(W) I L2L such that | f(W)
max I c/poly(L), and U2, V/2 E L approximates max S </poly(L.). Then there exist two matrices Ug. V3 E I lx such that ||UsVT Pı PP(WW) </poly(L). In addition, it takes L [+((1) time to construct U3, Vz.
the (WW E R such that |U2V - q(W)
Proof of Lemma G.6. By tensor trick, we construct U3. V/ tensor products of U1, Vi and U2, V2, respectively, while preserving their low-rank structures. Then, we show the low-rank approximation of (-) with bounded error by Lemma G.3 and Lemma G.5.
Let @ be column-wise Kronecker product such that A B = [A[,1] & B|.1] ... ||] ky] & Bi. ki]] E LL for TR. Lx B E Lxkk
Let 7(W)= = V.V and ?(W) V denote matrix-multillicttion approximations to f(W) and (WW,, respectively.
59

---

Licks
Lxck
Lxkk
Lxkk
For the case of presentation, let U3
U2
and V/3
Vi
2
V22
It holds
|UgV3 - Pi (W)l max
|UsV/s f(W) O (WW) max
By pr (W) f(W) o (WW)
| (2) (V2 V22) f(W) o ((W)
max
= [[(UGNT]) @ (U2VVT - -(WW) O (WW)
max
= (1F(W) o ((W) - ((W) ④ ((W) max
< (1(W)) @ (()) j(W (WW) Imax
[ o (WW) - S(W) o ((W)
max
Se/poly(L)
Se/poly(L)
<</poly(L).
(By Lemma G.3 and Lemma G.S)
Computationally, by k, I,0(1), computing U3 and V33 takes I, [+(( time. This completes the proof.
Lemma G.7 (Approximate (-)). Let k1, k2, K4 = L11), Let p2W E g follow Definition G.7 such that its ju-th column is (2))) = j( T
jo 9 (W) jo for each jo [L]. Suppose U1, E Lx approximates the f(X) such that V| ((W) max S </poly(L.). and U2. V2 E Riikk approximates the g() E IR 1.x such that |12V ((W) max I c/poly(LL) Then there exist matrices U4, V4 E TR. Lxks such that [UAV P20)| max I c/poly(L). In addition, it takes [+((1) time to construct U4, V4.
Proof of Lemma G.7. From Definition G.7,
(IT) (2(W = (W),, jo f (W)J ((W),.
For (I), show its low-rank approximation by observing the low-rank-preserringg property of the multiplication between f(-) and g(-) (from Lemma G.3 and Lemma G.5). For (II), we show its low-rank approximation by the low-rank structure of f()) and (I).
Part (I). We define function r(W) Rdd + RC such that the jo-th component r(W) jo (f(W);.) T (WW), for all ju E [L]. Let F(W) denote the approximation of r(W) via decomposing into f(-) and y(:):
F(W),. = (F(W) ior 7(W) (UIV)) lio. [(U]()) Lo, :I]
-iijo VG 1/2 ()2
(G.6)
Lxka
for all Jo E [L]. This allows us to write P2 (W) = f(W) diag(r(W)) with diag(F(W)) denoting a diagonal matrix with diagonal entries being components of F(W).
Part (II). With r()), we approximate p2(-) with (2( = J(W) diag(F(W)) as follows.
Since F(W) has low rank representation, and diag(F((W) is a diagonal matrix, f2(.) has low-rank representation by definition. Thus, we set (2( V with U, Un and V diag(F(W)) Then, we bound the approximation error
U
P2) max
60

---

(122(W) P2) max
[F(1W) Fo( j( (o
- max JoEIL]
+ [[WW)..F( F(W' ((W), (WW) jo
max joE[L]
max
By triangle incquality)
<e/poly(D).
Computationally, computing V2 1/2 takes L time by ki, k2 1o()), Once we have V/ 1/2 precomputed, (G.6) only takes k2) time for each jo E [4]. Thus, the total time is O (Lk)k2) = ()+o(1)) Since U1 and V1 takes [+((( time to construct and 142 diag(F(WW)
V1 also takes
LXL
Lxkg
[,1+o(1) time, U4 and V4 takes 1,1+((1) time to construct. This completes the proof.
G.2 Proof of Theorem 4.1
Proof of Theorem 4.1. By the definitions of matrices p(-), P1(-) and pz(-) (Definition G.7), we have p(W) P1 (W) - PP)
By Lemma G.2, we have
dg2 vec (AIP(W)A3)
(G.7)
dw
To show the existence of (o))) algorithms for DiT backward computation Problem 1, prove fast low-rank approximations for Aipı (W)AJ and AIP2(W))4 as follows.
Let Pi (W', P2 (W) denote the approximations to pi (W), p2 (W), respectively. By Lemma G.6, it takes L 1++((1) time to construct U3, V/3 E R Lxka such that
A15 A = V3U AT.
Then, computing A
AT takes I, +(1+ due to the fact that d, k3 I,((1).
dx Lx
log X Lxd
Therefore, total running time for A1p (W) AT is L. [0(I) = [11+(11).
For the same reason (by Lemma G.7), total running time for A1p2(W)AL is L- ( = L1+0(1),
Lastly, we have
Og2 G(W))
OW
max
||vec (A) vec (A))(W)AT) max
(By Lemma G.2)
= (A)) (AIP(W)A)) max
(By definition, ||
max = maxi, |Agl for any matrix A)
S (A,,, (W) (W)) AT) Il (A (W)
P2 P2( (W)) AT) mm
(By Definition G.7 and triangle incquality)
< 11 (W) P1 (W)) + (I(2
(W))
(By the sub-multiplicative property of || 11ol.c))
<e/poly(L.).
(By Lemma G.6 and Lemma G.7)
Set = IL). We complete the proof.
61

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our contributions and scope Section and Section are reflected by the claims in abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specificatio. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., the approach was only tested a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Yes. include our proofs in the appendix and have made every effort ensure the correctness of our theoretical results.
62

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might specific model and empirical evaluation, it may be necessary
suffice, if the contribution to either make it possible for others to replicate the model with the same dataset, or provide access the model. general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., the case of large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but should be possible for other researchers to have some path to reproducing verifying the results.
the case of
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification:
This is formal theory work without experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
63

---

Please see the NeurIPS code and data submission guidelines (https://nips.cr public/quide SrodeSubmit se sionPolicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this central to the contribution (e.g.. for a new spen-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:/ /ttp///nppp...// cc/ public/guiden des CodeeeooeSSiiii Sulmi ss Looooooooioooooo iyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits. hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: This is a formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI. if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
64

---

8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, every respect, with the NeurIPS Code of Ethics https:/ /////euurpppp wttpp/biccctthiccssllelele.e..ed
Answer: [Yes]
Justification: Yes. We follow the code of ethics in this work.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g., there special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This theoretical work aims to shed light on the foundations of diffusion generative models and anticipated have negative social impacts.
Guidelines:
The answer NA means that there is no societal impact the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g.. deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will foundational research and not tied to par- ticular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
65

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g,, gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
This is a formal theory work without experiments.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi rswi thcode. com/oooooooooooooo /daaaaatsssssse has curated licenses for some datasets. Their licensing guide can help determine the license of dataset. For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged reach out the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/codemmodel as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
66

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This is formal theory work without experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
67


Reviews:
[{"review_id": "m2yt2bMXkd", "reviewer": "Reviewer_LxhB", "summary": "This paper establishes the statistical rates and provably efficient criteria of Latent Diffusion Transformers (DiTs). Specifically, there are three main theoretical results:\n* the approximation error bound for the transformer-based score estimator,\n* the sample complexity bound for score estimation,\n* and the provably efficient criteria for latent DiTs in both forward inference and backward training.\n\nThese results closely rely on the low-dimensional linear subspace assumption on input data.", "questions": "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. \n2. 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. \n3. In the right-hand side of equation (3.1), it seems that the second term does not involve $n$ and thus will not goes to zero even when $n$ goes to infinity. Does it mean that the estimation can never be consistent?", "limitations": "The limitations of this work have been stated in Section 5 and are left for future work as claimed by the authors.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Given the popularity of generative AI these days, the theoretical understanding of DiTs are of great interest to the community.\n2. Most of the concepts and the results are clearly presented.", "weaknesses": "1. The practical insights of the provably efficient criteria (Section 4) are relatively scarce.\n2. A few noticeable typos were encountered in the paper. For example, \n* At line 120 on page 3, 'knonw' should be corrected to 'known' in the sentence 'This is also knonw as ...';\n* At line 242 on page 7, there is an extra 'full'.\n3. Including simulation results on error bounds against related parameters could enhance the credibility and reliability of theoretical results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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.\n\n**Response:**\nThanks for the question. Yes, but not trivial. Here are some clarifications:\n* 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. \n* 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.\n\nIn 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.\n\n> **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. \n\n**Response:**\nThanks for the question. We would like to clarify a few points.\n\nYes. \n\nFor 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. \n\nFor 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.  \n\n> **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?\n\n**Response:**\nAppreciate your question. Here are the clarifications.\n\nYes.  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.\n\n---\n\n> **Reviewer's Comment 1:** The practical insights of the provably efficient criteria (Section 4) are relatively scarce.\n\n**Response:**\nThanks for your comment. We acknolodge the importance of enhancing the applicability of our findings. Here are some clarifications.\n\nOur provably efficient criteria offer some insights for the design of more efficient methods. Specifically, we demonstrate that:\n\n* The latent dim should be small enough $d=O(\\log L)$ (Thm 4.1 & Prop 4.1, 4.2)\n* Normalization of $K,Q,V$ in DiT attention heads is beneficial for performance and efficiency. For example: \n  * 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)\n  * 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)\n\nWe want to emphasize that these conditions are necessary but not sufficient. Sufficient conditions should depend on the detailed designs of specific methods.\n\nWe hope these points address your concern and enhance the utility of our criteria in practical settings.\n\n> **Reviewer's Comment 2:** A few noticeable typos were encountered in the paper.\n\n**Response:**\nThanks 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:\n* `line 109`, we corrected “denosing” to “denoising”.\n* `line 120`, we corrected “knonw” to “known”.\n* `line 155`, we corrected “a” to “an”.\n* `line 192`, we corrected “an” to “a”.\n* `line 228`, we corrected “depend” to “depends”.\n* `line 242`, we deleted the extra “full”.\n* `line 253`, we corrected “analyze” to “analyzing”.\n* `line 303`, we corrected “exists” to “exist”.\n* `line 320`, we corrected “motivate” to “motivates”.\n* `line 322`, we corrected “do” to “does”.\n* `line 954`, we corrected “subset” to “subsets”.\n* `line 977`, we corrected “part” to “parts”.\n* `line 989`, we corrected “a” to “an”.\n* `line 1168`, we added “at” after “arrive”.\n\nYour feedback is greatly appreciated and has been essential in enhancing the clarity and readability of our paper. Thank you for your help.\n\n---\n\n\nWe hope the revisions and clarifications provided in this response address the reviewer's concerns.\n\nWe welcome additional feedback and look forward to further discussions. Thank you for your time and valuable inputs!\n\n---\n\n* [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.\n* [Oko23] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning (ICML), 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mdyr3QaGOb", "reviewer": "Reviewer_jLPK", "summary": "The paper studies the statistical and computational limits of latent diffusion transformers.", "questions": "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.", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The results seem to be new and non-trivial (though I'm not an expert in the field, so I might have a wrong impression).", "weaknesses": "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).\n\nIn 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\".\n\nUPDATE: Increased the score and the confidence after the rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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).  \n\n>  **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.\n\n**Response:**\nThanks 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:\n\n* > **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 \n$\\epsilon \\cdot \\sqrt{d_0}/\\sigma(t)$, \nwhere $\\sigma(t)=1-e^{-t}$.\n\n* > **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 \n$\\mathcal{O}(\\xi(n, \\epsilon))$, \nwhere $\\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.\n\n* > **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)$. \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)})$.\n> 2. The estimated distribution $\\hat{P}_{T_0}$ in the orthogonal subspace degenerates to a point mass at origin as $T_0 \\rightarrow 0$.\n\n* > **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.\n\n* > **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.\n\n* > **Proposition 1.2 (Informal Version of Proposition 4.2).** Assuming SETH, $L^{1+o(1)}$ time DiT inference is possible.\n\n---\n\n> **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\".\n\n\n**Response:**\nThanks for the comment. We acknowledge the current draft is not precise enough and have made modifications accordingly. We quote them as follows:\n\n* `line 179` After Theorem 3.1: \n> **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. \n* `line 200` After Corollary 3.1.1:\n> **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.\n* `line 223` After Corollary 3.1.2:\n> **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.\n\n**These are all \"1st known\" DiT results (i.e., SOTA?).**\n\nFor computational limits, since we are the first analysis on DiT, there is no prior work to compare with.\n\n\n\n---\nWe hope the revisions and clarifications provided in this response address the reviewer's concerns.\n\nWe look forward to further feedback and discussion. Thank you for your time and effort!\n\n===\n\n* [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.\n* [Wibisono24] Andre Wibisono, Yihong Wu, and Kaylee Yingxi Yang. Optimal score estimation via empirical bayes smoothing. arXiv preprint arXiv:2402.07747, 2024.\n* [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.\n* [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.\n* [Oko23] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning (ICML), 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7nGCc2MRgV", "reviewer": "Reviewer_cP7q", "summary": "This paper explores the statistical and computational limits of latent DiTs:\n\n1. It proves that Transformers are sufficient as universal approximators for the score function in DiTs, with their approximation capacity depending on the latent dimension.\n2. Transformer-based score estimators converge to the true score function, indicating that the Transformer architecture is adequate for estimating the original data distribution.\n3. It provides provably efficient criteria to demonstrate the existence of almost-linear time algorithms for forward inference and backward computation, offering a theoretical basis for the efficient training and inference of DiTs.", "questions": "see weaknesses.", "limitations": "The paper has outlined its limitations and broader impact.", "rating": 6, "confidence": 1, "soundness": 4, "presentation": 1, "contribution": 3, "strengths": "The advantages of the paper are as follows:\n1. This paper derives an approximation error bound for the score network that is sub-linear in the latent space dimension. This finding not only explains the expressiveness of latent DiTs (under mild assumptions) but also offers guidance for the structural configuration of the score network in practical implementations.\n2. This paper prove that learned score estimator is able to recover the initial data distribution, which provides a theoretical basis for the feasibility of using neural networks to estimate the score.\n3. This paper proves the existence of almost-linear time DiT training algorithms for forward inference and backward computation, providing a theoretical foundation for efficient training and inference.\n4. All statistical and computational results are analyzed in a low-dimensional subspace, demonstrating the feasibility of using VAE for dimensionality reduction in latent DiTs.", "weaknesses": "The weaknesses of the paper are as follows:\n1. It is meaningful to further elucidate why existing DiTs models are difficult to train.\n2. How will these theoretical proofs contribute to the exploration of fast training and sampling algorithms?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Reviewer's Comment 1:** It is meaningful to further elucidate why existing DiTs models are difficult to train. \n\n**Response:**\nThanks for your comment. We'd like to clarify a few points.\n\nThis 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.\n\n> **Reviewer's Question 1:** How will these theoretical proofs contribute to the exploration of fast training and sampling algorithms?\n\n**Response:**\nThanks for the question. \n\nIn essence, our hardness results provides necessary conditions for designing efficient methods:\n* The latent dim should be small enough $d=O(\\log L)$ (Thm 4.1 & Prop 4.1, 4.2)\n* Normalization of $K,Q,V$ in DiT attention heads is beneficial for performance and efficiency. For example: \n  * 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)\n  * 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)\n\nWe want to emphasize that these conditions are necessary but not sufficient. Sufficient conditions should depend on the detailed designs of specific methods.\n\nWe hope these points address the reviewer's concerns. We have revised the latest version of our paper accordingly.\n\nWe're open to any further questions or clarifications you might have about our work. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ifksjr8php", "reviewer": "Reviewer_Zkkk", "summary": "In this paper, the authors explore the statistical and computational limits of latent Diffusion Transformers (DiTs) under the assumption of a low-dimensional linear latent space. Their contributions include an approximation error bound for the DiTs score function, which is sub-linear in the latent space dimension, as well as a sample complexity bound demonstrating convergence of the data distribution generated from the estimated score function. Additionally, they identify efficient criteria for forward inference and backward computation, achieving almost-linear time complexity for both processes.", "questions": "See above.", "limitations": "See above.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The authors derived approximation error bounds for transformer-based score estimators in latent DiTs, providing practical structural guidance. \n- The paper also provided sample complexity bounds for score estimation and demonstrated recovery of initial data distribution.\n- The authors provided efficiency criteria and characterized efficient algorithms for latent DiTs, including almost-linear time algorithms for both forward inference and training.", "weaknesses": "- I have reviewed Appendix C. How does the proof for DiTs in this paper differ from that presented in [Chen et al., 2023a]? 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? 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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]?\n\n**Response:**\nThanks for the question. Here are some clarifications.\n\n* 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. \n* 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. \n* 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.\n\n> **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? \n\n**Response:**\nThank you for your question. Let us provide some additional details.\n\nThe 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.\n\n> **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.\n\n**Response:**\nGrateful for your question. Here's some further explanation.\n\nYes. 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.\n\n---\n\nWe hope these points address the reviewer's questions.\n\nWe're open to any further questions or clarifications you might have about our work. Thank you!\n\n===\n\n* [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.\n* [Peebles23] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2023.\n* [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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The submission focuses on statistical and computational bounds for latent Diffusion Transformers (DiTs), alongside provably efficient criteria for both forward inference and backward training. The results are obtained under the assumption that input data derive from linear latent low-dimensional space. The paper is quite dense and technical. The discussion phase helped to partially improve its form and the results might have limited practical impact: concerning these points, authors are encouraged to add numerical evidence of the obtained bounds, if possible. On the other hand, the novelty of the results and the relevance of the considered architecture are of great theoretical interest, as acknowledged during the reviewing process, and therefore I support the acceptance of the manuscript.
Decision: Accept (poster)