Title: On Statistical Rates and  Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)
Venue: NeurIPS 2024
Accepted: yes
Keywords: DiT

Score Matching

Diffusion Transformer

Efficiency

Universal Approximation

Convergence Analysis

Diffusion Generative Model

DDPM
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)

Abstract

We investigate the statistical and computational limits of latent Diffusion Transformers (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.

1 Introduction

We investigate the statistical and computational limits of latent diffusion transformers (DiTs), assuming the data is supported on an unknown low-dimensional linear subspace. This analysis is not only 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 et al., 2024], making them challenging to train outside of specialized 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 a timely theoretical analysis of the fundamental limits of DiT architecture, aided by the analytical feasibility provided by the low-dimensional data assumption.

Empirically, Latent Diffusion is a go-to design for effectiveness and computational efficiency [Rombach et al., 2022, Liu et 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 a 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-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 computational limits of latent DiTs. 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:

Question 1. What is the approximation limit of using transformers to approximate the DiT score function, particularly in the low-dimensional data subspace?

Question 2. How accurate is the estimation limit for such a score estimator in practical training scenarios? With the score estimator, how well can diffusion transformers recover the data distribution?

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 2 ) 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 low-dimensional linear subspace data assumption:

Question 3. Is it possible to improve the Ω( L 2 ) 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 a transformer-based 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 2 by exploring the limitations of score and distribution estimations of latent DiTs 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 3.2.1).

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-based efficiency threshold for both conditional and unconditional generation (Proposition 4.1). Efficient algorithms, including almost-linear time algorithms (Proposition 4.2), are possible only below this threshold. For backward computation, we prove the existence of almost-linear time DiT training algorithms (Theorem 4.1) 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 dimension under the low-dimensional 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-based score networks. Section 3 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.

/negationslash

Notations. We use lower case letters to denote vectors, e.g., z ∈ R D . ‖ z ‖ 2 and ‖ z ‖ ∞ denote its Euclidean norm and Infinite norm respectively. We use upper case letters to denote matrix, e.g., Z ∈ R d × L . ‖ Z ‖ 2 , ‖ Z ‖ op , and ‖ Z ‖ F denote the 2 -norm, operator norm and Frobenius norm respectively. ‖ Z ‖ p,q denotes the p, q -norm where the p -norm is over columns and q -norm is over rows. Given a function f , let ‖ f ( x ) ‖ L 2 := ( ∫ ‖ f ( x ) ‖ 2 2 d x ) 1 / 2 , and ‖ f ( · ) ‖ Lip = sup x = y ( ‖ f ( x ) -f ( y ) ‖ 2 / ‖ x -y ‖ 2 ) . With a distribution P , we denote ‖ f ‖ L 2 ( P ) = ( ∫ P ‖ f ( x ) ‖ 2 2 d x ) 1 / 2 as the L 2 ( P ) norm. Let f /sharp P be a pushforward measure, i.e., for any measurable Ω , ( f /sharp P )(Ω) = P ( f -1 (Ω)) . We use ψ for (conditional) Gaussian density functions.

2 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 x 0 ∈ R D , and x 0 ∼ P 0 . Let x t denote the noisy data at the timestamp t with marginal distribution and destiny as P t and p t . The conditional distribution P ( x t | x 0 ) follows N ( β ( t ) x 0 , σ ( t ) I D ) , where β ( t ) = exp( -∫ t 0 w ( s )d s/ 2) , σ ( t ) = 1 -β 2 ( t ) , and w ( t ) > 0 is a nondecreasing weighting function. In practice, the forward process terminates at a large enough T such that P T is close to N (0 , I D ) . In the backward process, we obtain y t by reversing the forward process. The generation of y t depends on the score function ∇ log p t ( · ) . However, this is unknown in practice, we use a score estimator s W ( · , t ) to replace ∇ log p t ( · ) , where s W ( · , t ) is usually a neural network with parameters W . See Appendix D.1 for the details.

Score Matching. To estimate the score function, we use the following loss

where γ ( t ) is the weight function, and T 0 is a small value to stabilize training and prevent score function from blowing up [Vahdat et al., 2021]. However, it is hard to compute ∇ log p t ( · ) with available data samples. Therefore, we minimize the equivalent denoising score matching objective

where ψ t ( x t | x 0 ) is the transition kernel, then ∇ x t log ψ t ( x t | x 0 ) = ( β ( t ) x 0 -x t ) /σ ( t ) .

To train the parameters W in the score estimator s W ( · , t ) , we use the empirical version of (2.1). We select n i.i.d. data samples { x 0 ,i } n i =1 ∼ P 0 , and sample time t i (1 ≤ i ≤ n ) uniformly from interval [ T 0 , T ] . Given x 0 ,i , we sample x t i from N ( β ( t i ) x 0 ,i , σ ( t i ) I D ) . The empirical loss is

For convenience of notation, we denote population loss L ( W ) = E P 0 [ ̂ L ( W )] .

2.2 Score Decomposition in Linear Latent Space

In this part, we review the score decomposition in [Chen et al., 2023]. We consider that the D -dimensional input data x supported on a d 0 -dimensional subspace, where d 0 ≤ D .

Assumption 2.1 (Low-Dimensional Linear Latent Space) . Let x denote the initial data at t = 0 . x has a latent representation via x = Bh , where B ∈ R D × d 0 is an unknown matrix with orthonormal columns. The latent variable h ∈ R d 0 follows the distribution P h with a density function p h .

Remark 2.1. By 'linear latent space,' we mean that each entry of a given latent vector is a 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].

Let x and h denote the perturbed data and its associated latent variable at t > 0 , respectively. Based on the low-dimensional data structure assumption, we have the following score decomposition theory: on-support score s + ( B /latticetop x, t ) and orthogonal score s -( x, t ) .

Lemma 2.1 (Score Decomposition, Lemma 1 of [Chen et al., 2023]) . Let data x = Bh follow Assumption 2.1. The decomposition of score function ∇ log p t ( x ) is

where p h t ( h ) := ∫ ψ t ( h | h ) p h ( h )d h , ψ t ( ·| h ) is the Gaussian density function of N ( β ( t ) h, σ ( t ) I d 0 ) , β ( t ) = e -t/ 2 and σ ( t ) = 1 -e -t . 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 P h ) . The density function p h > 0 is twice continuously differentiable. Moreover, there exist positive constants A 0 , A 1 , A 2 such that when ‖ h ‖ 2 ≥ A 0 , the density function p h ( h ) ≤ (2 π ) -d 0 / 2 A 1 exp( -A 2 ‖ h ‖ 2 2 / 2) .

Assumption 2.3 ( L s + -Lipschitz of s + ( h, t ) ) . The on-support score function s + ( h, t ) is L s + -Lipschitz in h ∈ R d 0 for any t ∈ [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 /latticetop x ∈ R d 0 with d 0 < D .

(Latent) Score Network. Following [Chen et al., 2023], we rearrange (2.3) into

We use W B ∈ R D × d 0 to approximate B ∈ R D × d 0 , and a neural network f ( W /latticetop B x, t ) to approximate q ( B /latticetop x, t ) . We adopt the following score network class for diffusion in latent space (i.e., in h ∈ R d 0 )

where the columns in W B are orthogonal, f : R d 0 × [ T 0 , T ] → R d 0 is a neural network. In this work, we focus on the diffusion transformers (DiTs), i.e., using Transformer for f [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 τ r,m,l : R d × L → R d × L to denote a Transformer block. Here r and m are the number of heads and head size in self-attention layer, and l is the hidden dimension in feed-forward layer. Let X ∈ R d × L be the model input, then we have the model output

where W i K , W i Q , W i V ∈ R m × d , W i O ∈ R d × m , W 1 ∈ R l × d , W 2 ∈ R d × l , b 1 ∈ R l , b 2 ∈ R d .

In our work, we use Transformer networks with positional encoding E ∈ R d × L . We define the Transformer networks as the composition of Transformer blocks

For example, the following is a Transformer network consisting K blocks and positional encoding

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 is required. Following this, in Section 3.3, utilizing the characterized structural configuration, we prove the score and distribution estimation for latent DiTs.

W

/latticetop

B

Latent Encoder

W

B

Latent Decoder

R

(

·

)

f

T

∈ T

r,m,l

Transformer Network

R

-

1

(

·

)

Reversed

⊕

x

∈

R

D

R

d

0

R

d

×

L

R

d

0

R

D

1

/σ

(

t

)

R

d

×

L

-

1

/σ

2

t

f

=

R

-

1

◦

f

T

◦

R

:

R

d

0

→

R

d

×

L

s

W

(

x, t

)

3.1 DiT Score Network Class

Here, we provide the details about DiT score network class used in our analysis. In (2.5), f is a network with Transformer as the backbone, and ( h, t ) ∈ R d 0 × [ T 0 , T ] denotes the input data. Following [Peebles and Xie, 2023], DiT uses time point t to calculate the scale and shift value in the Transformer backbone, and it transforms an input picture into a sequential version. To achieve the transformation, we introduce a reshape layer.

Definition 3.1 (DiT Reshape Layer R ( · ) ) . Let R ( · ) : R d 0 → R d × L be a reshape layer that transforms the d 0 -dimensional input into a d × L matrix. Specifically, for any d 0 = i × i image input, R ( · ) converts it into a sequence representation with feature dimension d := p 2 (where p ≥ 2 ) and sequence length L := ( i/p ) 2 . Besides, we define the corresponding reverse reshape (flatten) layer R -1 ( · ) : R d × L → R d 0 as the inverse of R ( · ) . By d 0 = dL , R,R -1 are associative w.r.t. their input.

To simplify the self-attention block in (2.6), let W i OV = W i O W i V and W i KQ = ( W i K ) T W i Q .

Definition 3.2 (Transformer Network Class T r,m,l p ) . We define the Transformer network class as

Model architecture with K blocks: f T ( X ) = FF ( K ) ◦ Attn ( K ) ◦ · · · FF (1) ◦ Attn (1) ( X ) ;

Model output bound: sup X ‖ f T ( X ) ‖ 2 ≤ C T ;

Parameter bound in Attn (i) : ∥ ∥ ( W i OV ) /latticetop ∥ ∥ 2 , ∞ ≤ C 2 , ∞ OV , ∥ ∥ ( W i OV ) /latticetop ∥ ∥ 2 ≤ C OV , ∥ ∥ W i KQ ∥ ∥ 2 , ∞ ≤ C 2 , ∞ KQ , ∥ ∥ W i KQ ∥ ∥ 2 ≤ C KQ , ∥ ∥ E /latticetop ∥ ∥ 2 , ∞ ≤ C E , ∀ i ∈ [ K ] ;

Parameter bound in FF (i) : ∥ ∥ W i j ∥ ∥ 2 , ∞ ≤ C 2 , ∞ F , ∥ ∥ W i j ∥ ∥ 2 ≤ C F , ∀ j ∈ [2] , i ∈ [ K ] ;

Lipschitz of f T : ‖ f T ( X 1 ) -f T ( X 2 ) ‖ F ≤ L T ‖ X 1 -X 2 ‖ F , ∀ X 1 , X 2 ∈ R d × L .

Definition 3.3 (DiT Score Network Class S T r,m,l p (Figure 1)) . We denote S T r,m,l p as the DiT score network class in (2.5), replacing f with R -1 ◦ f T ◦ R , and f T is from the Transformer class T r,m,l p .

3.2 Score Approximation of DiT

Here, we explore the approximation limit of latent DiT score network class S T r,m,l p under linear latent space assumption. Recall that P t is the distribution of x t , σ ( t ) is the variance of P ( x t | x 0 ) , d 0 is the dimension of latent space, L is the sequence length of transformer input, T is the stopping time in forward process, T 0 is the early stopping time in backward process, and L s + 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 /epsilon1 > 0 and any data distribution P 0 under Assumptions 2.1 to 2.3, there exists a DiT score network s W from S T 2 , 1 , 4 p

̂ (defined in Definition 3.2), where ̂ W = { ̂ W B , ̂ f T } , such that for any t ∈ [ T 0 , T ] , we have:

where σ ( t ) = 1 -e -t , and the upper bound of hyperparameters in S T 2 , 1 , 4 p are

Proof Sketch. Our proof is built on the key observation that there is a tail behavior of the low-dimensional latent variable distribution P h (Assumption 2.2). Recall that ∇ log p t ( x ) = Bq ( h, t ) /σ ( t ) -x/σ ( t ) , where h = B /latticetop x (defined in (2.4)). By taking ̂ W B = B , our aim reduces to construct a transformer network to approximate q ( h, t ) . To achieve this, we firstly approximate q ( h, t ) with a compact-supported continuous function, based on the tail behavior of P h . Then we construct a transformer to approximate the compact-supported continuous function using the universal approximation capacity of transformer [Yun et al., 2020]. See Appendix F.1 for a detailed proof.

Intuitively, Theorem 3.1 indicates the capability of the transformer-based score network to approximate the score function with precise guarantees. Furthermore, Theorem 3.1 provides empirical guidance for the design choices of the score network when a 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-approximated [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 in Theorem 3.2 and Corollary 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 d 0 = d × L . This indicates that DiTs can potentially bypass the challenges associated with the high dimensionality of initial data by transforming input data into a low-dimensional 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,C E , C F , etc. This characterization enables further analysis of the performance of score network in practical scenarios. In Theorem 3.2, we provide a sample complexity bound for score estimation. In Corollary 3.2.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 S T 2 , 1 , 4 p , we rewrite the score matching objective in (2.2) as ̂ W ∈ argmin s W ∈S T 2 , 1 , 4 p ̂ L ( s W ) , ̂ W = { ̂ W B , ̂ f T } .

Theorem 3.2 shows that as sample size n →∞ , s W ( · , t ) convergences to ∇ log p t ( · ) .

Theorem 3.2 (Score Estimation of DiT) . Under Assumptions 2.1 to 2.3, we choose S T 2 , 1 , 4 p as in Theorem 3.1 using /epsilon1 ∈ (0 , 1) and L > 1 , With probability 1 -1 / poly( n ) , we have

where ˜ O hides the factors related to D,d 0 , d, L s + , 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 et al., 2023] provides a sample complexity for simple ReLU-based diffusion models under the assumption of an accurate score estimator. To the best of our knowledge, we are the first to provide a 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 et al., 2020]. Furthermore, our first term shows a convergence rate of 1 /T , outperforming [Chen et al., 2023], in which the first term is independent of T .

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 (1 //epsilon1 ) 2 L in the first term. We remark that this arises from the required depth K is O ( /epsilon1 -2 L ) , and the norm of required weight parameters is (1 //epsilon1 ) O (1) 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 /epsilon1 as a function of n to balance the first and second terms in (3.1). This motivates us to rethink transformer universality and explore new proof techniques for DiTs, which we leave for future work.

Distribution Estimation. In practice, DiTs generate data using the discretized version with step size µ , see Appendix D.1 for details. Let ̂ P T 0 be the distribution generated by s ̂ W using the discretized backward process in Theorem 3.2. Let P h T 0 and p h T 0 be the distribution and density function of on-support latent variable h at T 0 . We have the following results for distribution estimation.

Corollary 3.2.1 (Distribution Estimation of DiT, Modified From Theorem 3 of [Chen et al., 2023]) . Let T = O (log n ) , T 0 = O (min { c 0 , 1 /L s + } ) , where c 0 is the minimum eigenvalue of E P h [ hh /latticetop ] . With the estimated DiT score network s ̂ W in Theorem 3.2, we have the following with probability 1 -1 / poly( n ) .

(i) The accuracy to recover the subspace B is ∥ ∥ W B W /latticetop B -BB /latticetop ∥ ∥ 2 F = ˜ O ( ξ ( n, /epsilon1,L ) /c 0 ) .

(ii) With the conditions KL ( P h || N (0 , I d 0 )) < ∞ , there exists an orthogonal matrix U ∈ R d × d such that we have the following upper bound for the total variation distance

where ˜ O hides the factor about D,d 0 , d, L s + , log n , and T -T 0 . and ( W B U ) /latticetop /sharp ̂ P T 0 denotes the pushforward distribution.

(iii) For the generated data distribution ̂ P T 0 , the orthogonal pushforward ( I -W B W /latticetop B ) /sharp ̂ P T 0 is N (0 , Σ) , where Σ /precedesequal aT 0 I for a constant a > 0 .

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 error between ̂ P T 0 and P h T 0 ; and (iii) the vanishing behavior of ̂ P T 0 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, /epsilon1,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] D 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 is learned by DiTs. The error is proportional to the sample complexity for score estimation and depends on the minimum eigenvalue of the covariance of P h .

4 Provably Efficient Criteria

Here, we analyze the computational limits of latent DiTs under low-dimensional linear subspace data assumption (i.e., Assumption 2.1). The hardness of DiT models ties to 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.i.d. data samples { x 0 ,i } n i =1 ∼ P d , and time t i 0 (1 ≤ i ≤ n ) uniformly sampled from [ T 0 , T ] . For each data x 0 ,i ∈ R D , we sample x t i 0 ∈ R D from N ( β ( t i 0 ) x 0 ,i , σ ( t i 0 ) I D ) . Let ( W A R -1 ( · )) † be the inverse transformation of W A R -1 ( · ) , and denote

Y 0 ,i := ( W A R -1 ) † ( x 0 ,i ) ∈ R d × L . We rewrite the empirical denoising score-matching loss (2.2) as

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 a single layer attention for f T , to simplify our analysis. Further, we consider the following simplifications:

(S0) To prove the hardness of (4.1) for both full gradient descent and stochastic mini-batch gradient descent methods, it suffices to consider training on a single data point.

(S1) For the convenience of our analysis, we consider the following expression for attention mechanism. Let X,Y ∈ R d × L . Let W K , W Q , W V ∈ R s × d be attention weights such that Q = W Q X ∈ R d × L , K = W K X ∈ R s × L and V = W V X ∈ R s × L . We write attention mechanism of hidden size s and sequence length L as

with D := diag(exp ( XW Q W T K X T ) 1 L ) . Here, exp( · ) is entry-wise exponential function, i.e., exp( A ) i,j = exp( A i,j ) for any matrix A , diag( · ) converts a vector into a diagonal matrix with the vector's entries on the diagonal, and 1 L is the lengthL all ones vector.

(S2) Since V multiplication is linear in weight while 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 W O W V X = const..

(S3) To focus on the DiT, we consider the low-dimensional linear encoder W 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 W A 2 .

(S4) To further simplify, we introduce A 1 , A 2 , A 3 ∈ R s × L and W ∈ R d × d via

Notably, A 1 , A 2 , A 3 , X, Y are constants w.r.t. training above loss with gradient updates.

Therefore, we simplify the objective of training DiT into

Definition 4.1 (Training Generic DiT Loss) . Given A 1 , A 2 , A 3 , Y ∈ R d × L and W OV , W ∈ R d × d following (S4), Training a DiT with /lscript 2 loss on a single data point X,Y ∈ R d × L is formulated as

Here D := diag(exp ( A /latticetop 1 WA 2 ) 1 n ) ∈ R L × L .

Remark 4.1 (Conditional and Unconditional Generation) . L 0 is generic. If A 1 = A 2 ∈ R d × L , Definition 4.1 reduces to cross-attention in DiT score net (for conditional generation). If A 1 = A 2 ∈ R d × L , Definition 4.1 reduces to self-attention in DiT score net (for unconditional vanilla generation).

/negationslash

We introduce the next problem to characterize all possible gradient computations of optimizing (4.4).

Problem 1 (Approximate DiT Gradient Computation (ADITGC ( L, d, Γ , /epsilon1 ) )) . Given A 1 , A 2 , A 3 , Y ∈ R d × L . Let /epsilon1 > 0 . Assume all numerical values are in O (log( L )) -bits encoding. Let loss function L 0 follow Definition 4.1. The problem of approximating gradient computation of optimizing empirical DiT loss (4.4) is to find an approximated gradient matrix ˜ G ( W ) ∈ R d × d such that ∥ ∥ ˜ G ( W ) -∂ L ∂W ∥ ∥ max ≤ 1 / poly( L ) . Here, ‖ A ‖ max := max i,j | A ij | for any matrix A .

In this work, we aim to investigate the computational limits of all possible efficient algorithms of ADITGC with /epsilon1 = 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.

(O1) Let g 1 ( · ) := W A R -1 ( · ) : R d × L → R d 0 , g 2 ( · ) := Att( · ) : R d × L → R d × L , and g 3 ( · ) := R ( W /latticetop A · ) : R D → R d × L such that g 3 ( x ) = X for x ∈ R D (with D > d 0 = dL ).

(O2) Vectorization of f T . For the ease of presentation, we use notation flexibly that f T to denote both a matrix in R d × L and a vector in R dL in the following analysis. This practice does not affect correctness. The context in which f T is used should clarify whether it refers to a matrix or a vector. Explicit vectorization follows Definition D.1.

(O3) Linearity of g 1 . By linearity of W A R -1 ( · ) , we treat g 1 as a matrix in R d 0 × dL acting on vector f T ( · ) ∈ R dL .

Therefore, we have L 0 = ‖ g 1 · [ g 2 ( g 3 ) -Y ] ‖ 2 2 , such that its gradient involves d L 0 d W = g 1 d g 2 d W . From above, we only need to focus on proving the computation time and error control of term d g 2 d W for gradient w.r.t W . Luckily, with tools from fine-grained complexity theory [Alman and Song, 2023, 2024a,b,c] and tensor trick (see Appendix D.3), we prove the existence of almost-linear time algorithms for Problem 1 in the next theorem. Let vec( W ) := W for any matrix W following Definition D.1.

Theorem 4.1 (Existence of Almost-Linear Time Algorithms for ADITGC) . Suppose all numerical values are in O (log L ) -bits encoding. Let max( ‖ W OV A 3 ‖ max , ‖ W K A 1 ‖ max , ‖ W Q A 2 ‖ max ) ≤ Γ . There exists a L 1+ o (1) time algorithm to solve ADITGC ( L p , L, d = O (log L ) , Γ = o ( √ log L )) (i.e., Problem 1) with loss L 0 from Definition 4.1 up to 1 / poly( L ) accuracy. In particular, this algorithm outputs gradient matrices ˜ G ( W ) ∈ R d × d such that ∥ ∥ ˜ G ( W ) -∂ L ∂W ∥ ∥ max ≤ 1 / poly( L ) .

Proof Sketch. Our proof is built on 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 et al., 2024b, Alman and Song, 2023], we approximate DiT training gradients with a series of low-rank approximations and carefully match the multiplication dimensions so that the computation of d g 2 d W forms a chained low-rank approximation. We complete the proof by demonstrating that this approximation is bounded by a 1 / poly( L ) error and requires only almostlinear time. See Appendix G.2 for a detailed proof.

Remark 4.2. We remark that Theorem 4.1 is dominated by the relation between L and d , hence by the subspace dimension 3 d 0 = dL . A smaller d 0 makes Theorem 4.1 more likely to hold.

4.2 Computational Limits of Forward Inference

Since the inference of score-matching diffusion models is a forward pass of the trained score estimator s W , the computational hardness of DiT ties to the transformer-based score network,

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 ADITI ( d, L, Γ , δ F ) ) . Let δ F > 0 and B > 0 . Given A 1 , A 2 , A 3 ∈ R d × L , and W OV , W K , W Q ∈ R d × d with guarantees that ‖ W OV A 3 ‖ ∞ ≤ B , ‖ W K A 1 ‖ ∞ ≤ B and ‖ W Q A 2 ‖ ∞ ≤ B , we aim to study an approximation problem

ADITI ( d, L, B, δ F ) , that approximates s W ( A 1 , A 2 , A 3 ) with a vector ˜ z ∈ R d 0 (with d 0 = d · L ) such that ∥ ∥ ˜ z -W A R -1 ( W OV A 3 D -1 exp ( A /latticetop 1 W /latticetop K W Q A 2 ))∥ ∥ max ≤ δ F . Here, ‖ A ‖ max := max i,j | A ij | for any matrix A .

By (O2) and (O3), we make an observation that Problem 2 is just a special case of [Alman and Song, 2023]. Hence, we characterize the all possible efficient algorithms for ADITI with next proposition.

Proposition 4.1 (Norm-Based Efficiency Phase Transition) . Let ‖ W Q A 2 ‖ ∞ ≤ B , ‖ W K A 1 ‖ ∞ ≤ B and ‖ W OV A 3 ‖ ∞ ≤ B with B = O ( √ log L ) . Assuming SETH (Hypothesis 1), for every q > 0 , there are constants C, C a , C b > 0 such that: there is no O ( n 2 -q ) -time (sub-quadratic) algorithm for the problem ADITI ( L, d = C log L, B = C b √ log L, δ F = L -C a ) .

Remark 4.3. Proposition 4.1 suggests an efficiency threshold for the upper bound of ‖ W K A 1 ‖ ∞ , ‖ W Q A 2 ‖ ∞ , ‖ W OV A 3 ‖ ∞ . 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, d = O (log L ) , B = o ( √ log L ) , δ F = 1 / poly( L )) can be solved in L 1+ o (1) time.

Remark 4.4. Proposition 4.2 is 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 d 0 = dL . A smaller d 0 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 on the latent dimension, suggesting DiTs can handle high-dimensional data challenges. Second, we show that Transformer-based score estimators converge to 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.1) 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.1). 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 = O (log 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 training: max {‖ W K A 1 ‖ , ‖ W Q A 2 ‖ , ‖ W OV A 3 ‖} ≤ Γ with Γ = o ( log L ) (Theorem 4.1).

-For efficient inference: max {‖ W K A 1 ‖ , ‖ W Q A 2 ‖ , ‖ W OV A 3 ‖} ≤ B with B = o ( log L ) (Proposition 4.2) and A 1 , A 2 , A 3 being the input data associated with K,Q,V . √

We remark that these conditions are 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.

Broader Impact

This theoretical work aims to shed light on the foundations of diffusion generative models and is not anticipated to 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 Tsz-Kit Lau, Chenwei Xu, Dino Feng and Andrew Chen for enlightening discussions on related topics; Ting-Chun 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 R01LM1372201, AbbVie and Dolby. The content is solely the responsibility of the authors and does not 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 softmax attention to kronecker computation. In The Twelfth International Conference on Learning Representations (ICLR) , 2024a.

Josh Alman and Zhao Song. The 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. In manuscript , 2024c.

Luca Ambrogioni. In search of dispersed memories: Generative diffusion models are associative memory networks. arXiv preprint arXiv:2309.17290 , 2023.

Fan Bao, Chongxuan Li, Yue Cao, and Jun Zhu. All are worth words: a 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 convergence 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 on 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:1877-1901, 2020.

Junsong Chen, Jincheng YU, Chongjian GE, Lewei Yao, Enze Xie, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-$\alpha$: 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 statistical estimation guarantees of generative adversarial networks. arXiv preprint arXiv:2002.03938 , 2020b.

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) , pages 4672-4712. PMLR, 2023.

Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Applications, 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 cnf-sat. ACM Transactions on Algorithms (TALG) , 12(3):1-24, 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-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-694, 2020.

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 23164-23173, 2023a.

Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulating single layer attention in llm based on tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint arXiv:2309.07418 , 2023b.

Yeqi Gao, Zhao Song, and Shenghao Xie. In-context learning for attention scheme: from single softmax regression to multiple softmax regression via a 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 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: A survey of the uncanny resemblances between diffusion models and associative memories. arXiv preprint arXiv:2309.16750 , 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 arXiv:2106.09685 , 2021.

Jerry Yao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu, Bo-Yu Chen, and Han Liu. On sparse modern hopfield model. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS) , 2023.

Jerry Yao-Chieh Hu, Pei-Hsuan Chang, Haozheng Luo, Hong-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-Chieh Hu, Thomas Lin, Zhao Song, and Han Liu. On computational limits of modern hopfield models: A fine-grained complexity analysis. In Forty-first International Conference on Machine Learning (ICML) , 2024b.

Jerry Yao-Chieh Hu, Maojiang Su, En-Jui Kuo, Zhao Song, and Han Liu. Computational limits of low-rank adaptation (lora) for transformer-based models. arXiv preprint arXiv:2406.03136 , 2024c.

Jerry Yao-Chieh Hu, Wei-Po Wang, Ammar Gilani, Chenyang Li, Zhao Song, and Han Liu. Fundamental 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-compatible dense associative memories as spherical codes. In Thirty-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 arXiv: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. Bioinformatics , 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 preprint arXiv:2405.05219 , 2024a.

Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Yufa Zhou. Beyond linear approximations: A novel pruning approach for attention matrix. arXiv preprint arXiv:2410.11261 , 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.13233 , 2024c.

Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Tensor attention training: Provably efficient learning of higher-order transformers. arXiv preprint arXiv: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.

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-quality video generation. In 2023 IEEE/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-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.08740 , 2024.

Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv:2306.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:2112.10741 , 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.

OpenAI. Sora: A video generative model based on transformer diffusion. OpenAI 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 4195-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, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical textconditional 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 Holzleitner, Milena Pavlovic, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217 , 2020.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference 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 arXiv: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 arXiv:2402.18567 , 2024a.

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 arXiv:2403.14088 , 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.07747 , 2024.

Virginia Vassilevska Williams. On some fine-grained questions in algorithms and complexity. In Proceedings of the international congress of mathematicians: Rio de janeiro 2018 , pages 34473487. 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-Chieh Hu, Weijian Li, Bo-Yu Chen, and Han Liu. STanhop: Sparse tandem hopfield model for memory-enhanced 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 transformers universal approximators of sequence-to-sequence 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-based preference optimization. arXiv preprint arXiv:2403.16576 , 2024b.

Zhihan Zhou, Yanrong Ji, Weijian Li, Pratik Dutta, Ramana Davuluri, and Han Liu. Dnabert-2: Efficient foundation model and benchmark for multi-species genome. arXiv preprint arXiv:2306.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-matching: Causal discovery and generative modeling. Advances in Neural Information Processing Systems (NeurIPS) , 36, 2023.

Appendix

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 d 0 < D , our theoretical results apply to two other settings: d 0 = D and d 0 > D . Especially, for d 0 = D , our results still hold by setting B as the identity matrix I D . Namely, our results hold after removing the linear subspace assumption.

Statistically, for score approximation, score estimation, and distribution estimation, the upper bounds depend on the dimension of the latent variable d 0 , other than d . A smaller d 0 allows for a reduced model size to achieve a specified approximation error compared to a larger one (Theorem 3.1). Additionally, with a smaller d 0 , both score and distribution estimation errors are reduced relative to scenarios with larger ones (Theorem 3.2 and Corollary 3.2.1).

Computationally, smaller d 0 benefits the provably efficient criteria (Proposition 4.1, almost-linear time algorithms for forward inference (Proposition 4.2) and backward computation (Theorem 4.1).

B Notation Table

We summarize our notations in the following table for easy reference.

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 et al., 2024, Zhou et al., 2024a,b, Wang et al., 2024a,b]. Recently, transformer-based diffusion models have garnered significant attention in research. The U-ViT model [Bao et al., 2022] incorporates transformer blocks into a 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 (DiTs) [Peebles and Xie, 2023] have emerged as a significant advancement (e.g., SoRA [OpenAI, 2024, Liu et al., 2024] from OpenAI), 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.

Universality and Memory Capacity of Transformers. The universality of transformers refers to their ability to serve as universal approximators. This means that transformers theoretically models any sequence-to-sequence function to a desired degree of accuracy. Yun et al. [2020] establish that transformers can universally approximate sequence-to-sequence 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-Albert representation Theorem. Most recently, Kajitsuka and Sato [2023] show that transformers with one self-attention layer is a universal approximator.

The memory capacity of a transformer is a practical measure to test the theoretical results of the transformer'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 a bounded error. Several works address the memory capacity of transformers. Kim et al. [2022] show that transformers with ˜ O ( d + L + √ NL ) parameters are sufficient to memorize N lengthL and dimensiond sequence-to-sequence data points by constructing a contextual mapping with O ( L ) attention layers. Mahdavi et al. [2023] show that a multi-head-attention with h heads is able to memorize O ( hL ) examples under a linear independence data assumption. Kajitsuka and Sato [2023] show that a single layer transformer with O ( NLd + d 2 ) parameters is able to memorize N lengthL and dimensiond sequence-to-sequence 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,b,e, 2023, Wu et al., 2024a,b, Ramsauer et al., 2020]. Notably, they define memory capacity as the smallest number of (lengthL and dimensiond ) data points the model (transformer attention) is able to store and derive exponential-ind high-probability capacity lower bounds. In particular, Hu et al. [2024e] report a 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 DiTs (or diffusion models in general), recent studies on dense associative models suggest viewing pretrained 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 score-matching under the similar assumptions. Oko et al. [2023] analyze the distribution

estimation under the assumption that the initial density is supported on [ -1 , 1] D 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 a low-dimensional structure 4 . However, our 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 a transformer-based score estimator. Second, our work is the first to provide the statistical rates of DiTs (score and distribution estimation) based on transformer universality [Yun et al., 2020] and norm-based converging number bound [Edelman et al., 2022], supporting the practical success of DiTs [Esser et al., 2024, Ma et al., 2024]. Lastly, our work provides the first comprehensive analysis of 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-Based Pretrained Models. Transformerbased 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: encoder-based, decoder-based, and diffusion transformers. Encoder-based transformers, such as DNABERT [Zhou et al., 2024c, 2023, Ji et al., 2021], employ bidirectional attention to extract feature representations DNABERT shows great potential to capture complex patterns of genome sequences and improve tasks such as gene prediction. Decoder-based transformers generate output sequences from encoded information using unidirectional attention, such as ChatGPT [Radford et al., 2019, Floridi and Chiriatti, 2020, Brown et al., 2020] for natural language. The diffusion transformers generate a sequence toward a target distribution, such as 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 an early exploration of the statistical and computational limits of diffusion transformer models.

D Supplementary Theoretical Background

In this section, we provide some further background. We show the details about the forward and backward process in Diffusion Models in Appendix D.1. 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 to the original data in the forward process. We describe the forward process as the following SDE

where x 0 ∼ P 0 , ( W t ) t ≥ 0 is a standard Brownian motion, and w ( t ) > 0 is a nondecreasing weighting function. Let P t and p t denote the marginal distribution and destiny of x t . The conditional distribution P ( x t | x 0 ) follows N ( β ( t ) x 0 , σ ( t ) I D ) , where β ( t ) = exp ( -∫ t 0 w ( s )d s/ 2 ) and σ ( t ) = 1 -β 2 ( t ) . In practice, (D.1) terminates at a large enough T such that P T is close to N (0 , I D ) .

Backward Process. We obtain the backward process y t := x T -t by reversing (D.1). The backward process satisfies

where the score function ∇ log p t ( · ) is the gradient of log probability density function of x t , and W t is a reversed Brownian motion. However, ∇ log p t ( · ) and P T are both unknown in (D.2). To resolve this, we use a score estimator s W ( · , t ) to replace ∇ log p t ( · ) , where s W ( · , t ) is usually a neural network with parameters W . Secondly, we replace P T by the standard Gaussian distribution. Consequently, we obtain the following SDE

In practice, we use discrete schemes of (D.3) to generate data, following [Song and Ermon, 2019]. We use µ > 0 to denote the discretization step size. For t ∈ [ kµ, ( k +1) µ ] , we have

D.2 Proof of Lemma 2.1

Here we restate the proof of [Chen et al., 2023, Lemma 1] for completeness.

Proof. Recall x = Bh by Assumption 2.1 with x ∈ R D , B ∈ R D × d 0 and h ∈ R d 0

By the forward process (D.1), we have

where

is the Gaussian transition kernel.

.

Then we write the score function as

where the last equality holds since ψ t ( x | Bh ) is continuously differentiable in x . Plugging (D.6) into (D.7), we have

We then decompose above score function by projecting of x into Span( B ) , i.e., replacing -x with -BB /latticetop x -( I D -BB /latticetop ) x :

Absorbing the factor of [2 πh ( t )] -D/ 2 into the Gaussian kernel ψ t ( x | Bh ) , we have

To further simplify s + , we decompose ψ t ( x | Bh ) as

where both ψ t ( B /latticetop x | h ) and ψ t ( ( I D -BB /latticetop ) x ) are Gaussian.

Plugging

ψ

t

(

x

|

Bh

) =

ψ

t

(

B

/latticetop

x

|

h

)

ψ

t

(

(

I

D

-

BB

/latticetop

)

x

)

into

s

+

, we obtain

where C := [ ψ t (( I D -BB /latticetop ) x ) ∫ ψ t ( B /latticetop x | h ) p h ( h )d h ] -1 .

Notably, s + depends only on the projected data B /latticetop x . Therefore, we are able to replace s + ( x, t ) with s + ( B /latticetop x, t ) . The benefit is that the dimension d 0 of the first input in s + ( B /latticetop x, t ) is much smaller.

Lastly, by denoting h = B /latticetop x such that ∇ h ψ t ( h | h ) = ( β ( t ) h -h ) ψ t ( h | h ) /σ ( t ) , we arrive at

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 of the P = NP conjecture. It suggests that our current best SAT algorithms are optimal and is a popular conjecture for proving fine-grained lower bounds for a wide variety of algorithmic problems [Cygan et al., 2016, Williams, 2018].

/negationslash

Hypothesis 1 (SETH) . For every /epsilon1 > 0 , there is a positive integer k ≥ 3 such that k -SAT on formulas with n variables cannot be solved in O (2 (1 -/epsilon1 ) n ) time, even by a randomized algorithm.

Tensor Trick for Computing Gradients. The tensor trick [Diao et al., 2019, 2018] is an instrument to compute complicated gradients in a clean and tractable fashion. We start with some definitions.

Definition D.1 (Vectorization) . For any matrix X ∈ R L × d , we define X := vec ( X ) ∈ R Ld such that X i,j = X ( i -1) d + j for all i ∈ [ L ] and j ∈ [ d ] .

Definition D.2 (Matrixization) . For any vector X ∈ R Ld , we define mat( X ) = X such that X i,j = mat( X ) := X ( i -1) d + j for all i ∈ [ L ] and j ∈ [ d ] , namely mat( · ) = vec -1 ( · ) .

Definition D.3 (Kronecker Product) . Let A ∈ R L a × d a and B ∈ R L b × d b . We define the Kronecker product of A and B as A ⊗ B ∈ R L a L b × d a d b such that ( A ⊗ B ) ( i a -1) L b + i b , ( j a -1) d b + j b , is equal to A i a ,j a B i b ,j b with i a ∈ [ L a ] , j a ∈ [ d a ] , i b ∈ [ L b ] , j b ∈ [ d b ] .

Definition D.4 (Sub-Block of a Tensor) . For any A ∈ R L a × d a and B ∈ R L b × d b , let A := A ⊗ B ∈ R L a L b × d a d b . For any j ∈ [ L a ] , we define A j ∈ R L b × d a d b be the j -th L b × d a d b sub-block of A .

Lemma D.1 (Tensor Trick [Diao et al., 2019, 2018]) . For any A ∈ R L a × d a , B ∈ R L b × d b and X ∈ R d a × d b , it holds vec ( A /latticetop XB ) = ( A /latticetop ⊗ B /latticetop ) X ∈ R L a L b .

To showcase the tensor trick, let's consider a (single data point) attention following [Gao et al., 2023b,c]. Setting D := diag ( exp ( X T W T K W Q X ) 1 L ) and W := W K W T Q ∈ R d × d , we have

Proposition D.1 (Definition 4.7 of [Gao et al., 2023b]) . By Definition D.3 and Definition D.4, we identify D j,j := 〈 exp ( A j W ) , 1 L 〉 ∈ R for all j ∈ [ L ] , with A := X ⊗ X ∈ R L 2 × d 2 and W ∈ R d 2 . Therefore, for each j ∈ [ L ] and i ∈ [ d ] , it holds L 0 = ∑ L j =1 ∑ d i =1 1 2 (〈 D -1 j,j exp ( A j W ) , XW V [ · , i ] 〉 -Y j,i ) 2 .

The elegance of Proposition D.1 emerges when we vectorize the weights into vectors W,W V , making the gradient computations (e.g., d L 0 / W and d L 0 / W V ) 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 L 0 .

Fine-Grained Complexity for Transformer. Many recent works also utilize similar techniques from fine-grained complexity to 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 et al. [2024c] extend the singlelayer 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 et al. [2024d] study the computational limits of inference and training in prompt-tuning for pretrained transformers. Hu et al. [2024c] study the computational limits of LoRA [Hu et al., 2021] in transformers, identifying norm-bound conditions for efficient LoRA training and proving the existence of nearly linear-time LoRA algorithms.

Our work is 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.

E More Background and Auxiliary Lemmas: Universal Approximation of Transformers via Piecewise Approximation

Here, we review the universal approximation of 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 of Section 3 (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 3 of [Yun et al., 2020]) . Let /epsilon1 > 0 . For any given compact-supported continuous function f : R d × L → R d × L , there exists a transformer network f T ∈ T 2 , 1 , 4 p , such that

Proof Overview. We use the following proof strategy:

Step 1. We show that the piecewise-constant function is able to approximate compact-supported continuous function in Appendix E.1.

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 in 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.1, 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 Function

In this subsection, we show that the piecewise-constant function is able to approximate compactsupported 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] d × L . We denote the set of [0 , 1] d × L -supported continuous functions as F .

We introduce the notion of grid and cube for the compact support [0 , 1] d × L .

Definition E.1 (Grid and Cube with Width δ ) . Given a grid width δ , let G δ := { 0 , δ, . . . , 1 -δ } d × L denote the set of grids within [0 , 1] d × L . For a grid point G = ( G j ∈ [ d ] ,k ∈ [ L ] ) ∈ G δ , we denote its associated cube as

Each cube S G represents a hyper rectangular in the multi-dimensional space [0 , 1] d × L , constructed to discretize the space into smaller subspaces.

We introduce the notion of piecewise-constant fucntion class w.r.t. the [0 , 1] d × L -supported continuous function class F .

Definition E.2 (Piecewise-Constant Function Class) . Let f δ denote the piesewise constant function of grid width δ , and 1 {·} denote the indicator function. For each G ∈ G δ , and any matrix A G ∈ R d × L , we define the piecewise-constant function class as

We recall that for a given sequence-to-sequence function f ,

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 /epsilon1/ 3 > 0 , we can find a δ /star > 0 , such that there exists a f δ /star ∈ F ( δ /star ) satisfying ‖ f -f δ /star ‖ L 2 ≤ /epsilon1/ 3 .

Proof. See Appendix E.5.2 for a detailed proof.

E.2 Modified Transformer Approximates Piecewise-constant Function

In 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 T r,m,l p includes two modifications to the standard transformer network T r,m,l p :

Modified attention layer: Replace Softmax operator with Hardmax operator σ H ( · ) .

Modified feed-forward layer: Replace ReLU( · ) with an activation function ζ ∈ Ψ . Here, Ψ denotes the set of all piecewise linear functions with at most three pieces and at least one constant.

We approximate F ( δ ) with this modified transformer networks T r,m,l p .

Lemma E.3 (Modified from Proposition 4 of [Yun et al., 2020]) . For each f δ ∈ F ( δ ) , there exists a f T ,c ∈ T 2 , 1 , 1 such that ‖ f δ -f T ,c ‖ 2 = O ( δ d/ 2 ) .

Proof Sketch. Given δ , and for any grid G ∈ G δ , we have a grid set G δ and the cube S G .

Our proof follows two steps:

Quantization. For all X ∈ R d × L , we quantize it to a finite set:

-If X ∈ S G ⊂ [0 , 1] d × L , we quantize it to the element G ∈ G δ .

-If X / ∈ [0 , 1] d × L , we quantize it to an element out of G δ .

Mapping. For any G ∈ G δ , we map it to the desired output A G .

For Quantization , we achieve this by a series of modified feed-forward layers. We show this in Appendix E.2.1.

For Mapping , we follow two steps:

/negationslash

For any G = G ′ ∈ G δ , we use a 'contextual mapping' q c ( · ) (defined as Definition E.4). The mapping maps all the elements in q c ( G ) and q c ( G ′ ) to different values. Then, we use a series of modified self-attention layers to achieve 'contextual mapping'. We show this in Appendix E.2.2.

Definition E.4 (Contextual Mapping) . Consider a finite set G δ ∈ R d × L . A map q c : G δ → R 1 × L defines a contextual mapping if the map satisfies the following:

-For any G ∈ G δ , the entries in q c ( G ) are all distinct.

/negationslash

-For any G = G ′ ∈ G δ , all entries of q c ( G ) and q c ( G ′ ) are distinct.

For any G ∈ G δ , we use a series of modified feed-forward layers to map q c ( G ) to A G . We show this in Appendix E.2.3.

Remark E.1. Our proof differs from [Yun et al., 2020] in one aspect: Although [Yun et al., 2020, Proposition 4] outlines a 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 a series of modified feed-forward layers in T r,m,l p to quantize an input X ∈ R d × L to an element G of the following grid:

where J > L > 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 use e i to represent the standard unit vector where the i -th element is 1 . For the i -th row of X , we define the following feed-forward layer to achieve our aim.

Definition E.5 (Feed-forward Layer 1) . The vector e i acts as the weight parameters, and ζ 1 ( · ) acts as the activation function in the feed-forward layer

We take i = 1 as an example to give the specific calculation. Let X = ( x i,j ) d × L , then we have

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 , 2 , . . . , d . If the element of X is out of [0 , 1) , the series of layers maps it to J .

Step 2: Map the element in [0 , 1) to { 0 , δ, 2 δ, . . . , 1 -δ }

.

For the i -th row of X , we take k = 0 , 1 , . . . , 1 /δ -1 respectively. We define the following layer.

Definition E.6 (Feed-forward Layer 2) . The vector e i acts as the weight parameters and ζ 2 ( · ) acts as the activation function in the feed-forward layer

We take i = 1 and k = 1 as an example. We give the following specific calculation

In the first row of X , the above layer transforms the element in [ δ, 2 δ ] to δ . We stack the above layers together for i = 1 , 2 , . . . , d and k = 0 , 1 , . . . , 1 /δ -1 . If the element of X is in [ kδ, ( k +1) δ ] , the series layers maps it to kδ .

Combining the above two parts, we achieve our goal with d/δ + d feed-forward layers. We denote the d/δ + d series layers as f T ,c 1 .

E.2.2 Contextual Mapping by Modified Self-attention Layers

In our attention layers, we use the following positional encoding E ∈ R d × L

According to Appendix E.2.1, the output of f T ,c 1 is in the grid {-J, 0 , δ, . . . , 1 -δ } d × L . For any X in this grid, the first column of X + E is in

The results are similar in the other columns.

For i = 0 , 1 , . . . , L -1 , we use the following notation:

Then, we define the grid G + δ as the following.

Definition E.7 (Grid G + δ ) . We add E to all the grid points in G δ to generate the modified grid G + δ , defined as follows:

Next, we show that the modified attention layer computes contextual mapping (Definition E.4) for G + δ . For i = 1 , 2 , . . . , L -1 , we use the following notation:

Lemma E.4 (Modified from Lemma 6 of [Yun et al., 2020]) . We consider the following subset of G + δ :

Assume that L ≥ 2 and δ -1 ≥ 2 . Then, there exist a function f T ,c 2 : R d × L → R d × L composed of δ -d +1 modified attention layers (Definition E.3), a vector u ∈ R d , and two constants t l , t r ∈ R ( 0 < t l < t r ), such that q c ( G ) := u /latticetop f T ,c 2 ( G ) , G ∈ G + δ satisfies the following properties:

For any G ∈ ˜ G δ , all the entries of q c ( G ) are distinct.

For any different G,G ′ ∈ ˜ G δ , all the entries of q c ( G ) , q c ( G ′ ) are distinct.

For any G ∈ ˜ G δ , all the entries of q c ( G ) are in [ t l , t r ] .

For any G ∈ G + δ \ ˜ G δ , all the entries of q c ( G ) are outside [ t l , t r ] .

and the second column is in

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 a series of feed-forward layers map the output of modified attention layers f T ,c 2 to the desired output of function f δ /star .

Lemma E.5 (Lemma 7 of [Yun et al., 2020]) . There exists a function f T ,c 3 : R d × L → R d × L composed of O ( L (1 /δ ) dL /L !) modified feed-forward layers, such that

Proof. See Appendix E.5.4 for a 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 we want to achieve a approximation error O ( δ d/ 2 ) by the modified transformer, we need O ( δ -1 ) modified feed-forward layers in f T ,c 1 , O ( δ -d ) modified self-attention layers in f T ,c 2 , and O ( δ -dL ) modified feed-forward layers in f T ,c 3 .

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 9 of [Yun et al., 2020]) . For each f T ,c ∈ T 2 , 1 , 1 p and any /epsilon1 > 0 , there exists f T ∈ T 2 , 1 , 4 p such that ‖ f T -f T ,c ‖ L 2 ≤ /epsilon1/ 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 we prove Lemma E.1.

Furthermore, to achieve the /epsilon1 approximation error in Lemma E.1, we take δ = O ( /epsilon1 2 /d ) in Lemma E.3.

E.5 Supplementary Proofs

We first present two preliminary concepts: selective shift operation and bijective column ID mapping in Appendix E.5.1.

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, we consider the following function ξ ( · ; · ) : R d × L → R d × L

where X ∈ R d × L , e 1 = (1 , 0 , 0 , · · · , 0) /latticetop ∈ R d , and b Q ∈ R . u ∈ R d is a vector to be determined.

To see the output, we consider the j -th column of u /latticetop Xσ H [ ( u /latticetop X ) /latticetop ( u /latticetop X -b Q 1 /latticetop n ) ] :

If u /latticetop X : ,j > b Q , it calculates argmax of u /latticetop X ;

If u /latticetop X : ,j < b Q , it calculates argmin of u /latticetop X .

All rows of ξ ( X ; b Q ) except the first row are zero. We consider the j -th entry of the first row in ξ ( X ; b Q ) , which is denoted as ξ ( X ; b Q ) 1 ,j . Then for all j ∈ [ L ] , we have

From this observation, we define a function parametrized by b Q and b ′ Q (with b Q < b ′ Q )

Then we have

We define an attention layer of the form X → X + ξ ( X ; b Q , b ′ Q ) . For any column X : ,j , if b Q < u /latticetop X : ,j < b ′ Q , its first coordinate X 1 ,j is shifted up by max k u /latticetop X : ,k -min k u /latticetop X : ,k , while all the other coordinates stay untouched. We call this the selective shift operation because we can choose b Q and b ′ Q to shift certain entries of the input selectively.

Bijective Column ID Mapping. We consider the input G ∈ G + δ (Definition E.7). We use

For any j ∈ [ L ] , we have the following two conclusions:

If G i,j ≥ 0 for all i ∈ [ d ] , i.e., G : ,j ∈ [ j -1 : δ : j -δ ] d , then we have

The mapping G : ,j → u /latticetop G : ,j maps the elements in [ j -1 : δ : j -δ ] d to [ δ j : δ : δ j + δ -d +1 -δ ] . This is a bijection.

If there exists i ∈ [ d ] such that G i,j = -J + j , then

/negationslash

We say that u /latticetop G : ,j gives the 'column ID' for each possible value of G : ,j ∈ [ j -1 : δ : j -δ ] d . Remark E.3 (Illustration of Bijection Properity) . For the bijection property, we give the following illustration. Let G : j = ( g 1 j , g 2 j , · · · , g dj ) /latticetop and G : j = ( g 1 j , g 2 j , · · · , g dj ) /latticetop . If u /latticetop G : j = u /latticetop G : j and G : j = G : j , we deduce

/negationslash

/negationslash

Because G : j = G : j , then there exists a k ( k < d ) , such that g kj = g kj and g ij = g ij ( i > k ) . We have

This contradicts with (E.10). Thus we prove the property of bijection.

However,

E.5.2 Proof of Lemma E.2

Proof of Lemma E.2. We restate the proof from [Yun et al., 2020] for completeness.

By the nature of the compact-supported continuous function, f is uniformly continuous.

Because ‖·‖ ∞ is equivalent to ‖·‖ F when the number of entries are finite, we have the following by the definition of uniform continuity.

For any /epsilon1/ 3 > 0 , there exists a δ /star > 0 , such that for any X,Y ∈ R d × L , and ‖ X -Y ‖ ∞ < δ /star , we have ‖ f ( X ) -f ( Y ) ‖ F < /epsilon1/ 3 .

Then we perform the following steps following Definitions E.1 and E.2:

We create a grid G δ /star by choosing grid width δ /star . We also create cube S G with respect to G ∈ G δ /star .

For any grid point G ∈ G δ /star , we define C G ∈ S G as the center point of the cube S G .

We define a piecewise-constant function f δ /star ( X ) = ∑ L ∈G δ /star f ( C G ) 1 { X ∈ S G } .

For any X ∈ S G , we have ‖ X -C G ‖ ∞ < δ /star . According to the uniform continuity, we drive

This implies that ‖ f -f δ /star ‖ L 2 < /epsilon1/ 3 and completes the proof.

E.5.3 Proof of Lemma E.4

We give the proof of Lemma E.4 by constructing the network to satisfy the requirements.

Proof of Lemma E.4. Recall the selective shift operation in Appendix E.5.1. The overall idea of the construction includes two steps:

Step 1: For each j ∈ [ L ] , we stack δ -d attention layers. For g ∈ [ δ j : δ : δ j + δ -d +1 -δ ] (E.8) in the increasing order, we use the attention layer as

The total number of layers is Lδ -d . These layers cast G ∈ ˜ G δ to L different entries required by Property 1 of Lemma E.4.

Step 2: We add an extra single-head attention layer with the following attention part

This layer achieves a global shifting and casts different G ∈ ˜ G δ to unique elements required by the Property 2 of Lemma E.4.

The two operations map ˜ G δ and G + δ \ ˜ G δ to different sets, as required by Property 3 and Property 4 of Lemma E.4. The bounds t l and t r are calculated then.

Then, we give a 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 G + δ :

Category 1: G ∈ ˜ G δ , all entries in the point G are between 0 and L -δ .

Category 2: G ∈ G + δ \ ˜ G δ , the point G has at least one entry that equals to -J .

Let u = (1 , δ -1 , δ -2 , . . . , δ -d +1 ) . Recall that δ j = ( j -1)( δ -δ -d +1 ) / ( δ -1) for any j ∈ [ L ] in (E.8).

Category 1. We denote g j := u /latticetop G : ,j , then we have g 1 < g 2 < · · · < g L . The first δ -d layers sweep the set [ δ j : δ : δ j + δ -d +1 -δ ] , j ∈ [ L ] and apply selective shift operation on each element in the set. This means that selective shift operation will be applied to g 1 first, then g 2 , followed by g 3 , and so on.

The First Shift Operation. In the first selective shift operation with g going through [ δ 1 : δ : δ 1 + δ -d +1 -δ ] , the (1 , 1) -th entry of G (i.e., G 1 , 1 ) is shifted by the operation, while the other entries are left untouched. The updated value ˜ G 1 , 1 is

Therefore, the output of the layer after the operation is

Let ˜ g 1 := u T ˜ G : , 1 . We have

Then we deduce g L < ˜ g 1 , because

Thus, after updating, we have

and the new minimum is g 2 .

The Second Shift Operation. In the second selective shift operation with g going through [ δ 2 : δ : δ 2 + δ -d +1 -δ ] , the (1 , 2) -th entry of G (i.e., G 1 , 2 ) is shifted by the operation, while the other entries are left untouched. The updated value ˜ G 1 , 2 is

Therefore, the output of the layer after the operation is

We have

Then we deduce ˜ g 1 < ˜ g 2 , because

Thus, after updating, we have

and the new minimum is g 3 .

Repeating the Process. By repeating this process, we show that the j -th shift operation shifts G 1 ,j by δ -d ( ˜ g j -1 -g j ) . Then we have

We deduce ˜ g j -1 < ˜ g j holds for all 2 ≤ j ≤ L , because

where the last inequality holds because

Therefore, after the j -th selective shift operation, ˜ g j is the new maximum among { ˜ g 1 , . . . , ˜ g j , g j +1 , . . . , g L } and g j +1 is the new minimum.

After L Shift Operations. After the whole L shift operations, the input G is mapped to a new point ˜ G , where u /latticetop ˜ G = ( ˜ g 1 ˜ g 2 . . . ˜ g L ) and ˜ g 1 < ˜ g 2 < · · · < ˜ g L . For the lower and upper bound of ˜ g L , we have the following lemma.

Lemma E.8 (Lemma 10 of [Yun et al., 2020]) . ˜ g L = u /latticetop ˜ G : ,L satisfies the following bounds:

Also, the mapping from ( g 1 g 2 · · · g L ) to ˜ g L 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δ -( L +1) d -1 ξ ( · ; 0) . Since 0 < ˜ g 1 < · · · < ˜ g L , it adds the following to each entry in the first row of ˜ G :

The output of this layer is defined to be the function f T ,c 2 ( G ) .

In summary, for any G ∈ ˜ G δ , i ∈ [ d ] , and j ∈ [ L ] , we have

where δ + j = ∑ j -1 k =1 δ -kd ( g j -k -g j -k +1 ) + δ -jd ( g L -g 1 ) + Lδ -( L +1) d -1 ˜ g L .

For any G ∈ ˜ G δ and j ∈ [ L ] ,

All of them are distinct.

Checking Property 2 of Lemma E.4. Note that the upper bound on ˜ g L from Lemma E.8 also holds for other ˜ g j ( j ∈ [ L -1] ). For all j ∈ [ L ] , we have

By Lemma E.8, two different G,G ′ ∈ ˜ G δ are mapped to different ˜ g L and ˜ g ′ L , and they differ at least by δ . This means that the following two intervals are guaranteed to be disjoint:

Thus, the entries of u /latticetop f T ,c 2 ( G ) and u /latticetop f T ,c 2 ( G ′ ) are all distinct.

Now, we finish showing that the mapping f T ,c 2 ( · ) uses (1 /δ ) d +1 attention layers to implement a contextual mapping on ˜ G δ .

Checking Property 3 of Lemma E.4. Given Lemma E.8 and u /latticetop f T ,c 2 ( G ) : ,j ∈ [ Lδ -( L +1) d -1 ˜ g L , Lδ -( L +1) d -1 ˜ g L + Lδ -( L +1) d ) , for any G ∈ ˜ G δ , we have

This proves that all u /latticetop f T ,c 2 ( L ) : ,j are between t l and t r , where

Next, we check the Property 1, Property 2 and Property 3 of Lemma E.4.

Checking Property 1 of Lemma E.4. Given any G ∈ ˜ G δ , we already prove that

Category 2. Now we check the Property 4 of Lemma E.4. For the input points G ∈ G + δ \ ˜ G δ , note that the point G has at least one entry that equals to -J + k, k ∈ [ L -1] . Let g j := u /latticetop G : ,j . Recall that whenever a column G : ,j has an entry that equals to -J + k, k ∈ [ L -1] , we have g j < 0 . Without loss of generality, assume that g 1 < 0 .

Because the selective shift operation is applied to each element of [0 : δ : δ L + δ -d +1 -δ ] and is not applied to negative values, thus we have min k u /latticetop G : ,k = g 1 < 0 . g 1 never gets shifted upwards and remains the minimum for the whole time.

All g j 's are Negative. When all g j 's are negative, selective shift operation never shifts the input G . Thus ˜ G = G . Recall that u /latticetop ˜ G : ,j < 0 for all j ∈ [ L ] . The last layer with attention part Lδ -( L +1) d -1 ξ ( · ; 0) adds Lδ -( L +1) d -1 min k u /latticetop ˜ G : ,k < 0 to each entry in the first row of ˜ G . This makes ˜ G remain negative. Therefore, f T ,c 2 ( G ) satisfies u /latticetop f T ,c 2 ( G ) : ,j < 0 < t l for all j ∈ [ L ] .

Not All g j 's are Negative. Now consider the case where at least one g j is positive. Suppose that there are k positive elements and they satisfy g i 1 < g i 2 < · · · < g i k . Thus selective shift operation does not affect g i , where i ∈ [ L ] \ { i 1 , . . . , i k } . It shifts g i 1 by

The next shift operations shift g i 2 , . . . , g i k by an even larger amount. Therefore, at the end of the first L (1 /δ ) d layers, we have Lδ -( L +1) d ≤ ˜ g i 1 ≤ · · · ≤ ˜ g i k , and ˜ g j < 0 for all j ∈ [ L ] \ { i 1 , . . . , i k } .

Then, we shift G by the last layer. The last layer with attention part Lδ -( L +1) d -1 ξ ( · ; 0) acts differently for negative and positive ˜ g j 's. (i). For negative ˜ g j 's, it adds the following to ˜ g j , j ∈ [ L ] \ { i 1 , . . . , i k } :

This term pushes them further to the negative side. (ii). For positive ˜ g i 's, it adds

Thus they are all greater than or equal to 2 L 2 δ -2( L +1) d +1 . Note that

Then the final output f T ,c 2 ( G ) satisfies u /latticetop f T ,c 2 ( G ) : ,j / ∈ [ t l , t r ] , for all j ∈ [ L ] . This completes the verification of Property 4 of Lemma E.4.

In conclusion, we need O ( Lδ -d ) layers of modified self-attention layer to obtain our approximation. This completes the proof.

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 |G + δ | = (1 /δ +1) dL < ∞ , so the output of f T ,c 2 ( G + δ ) has finite number of distinct real values. Let M be the upper bound of all these possible values. By the construction of f T ,c 2 , M > 0 .

Construct the Layers: f T ,c 3 ( f T ,c 2 ( G )) = 0 d × L if G ∈ G + δ \ ˜ G δ . According to Lemma E.4, for all j ∈ [ L ] , we have u /latticetop f T ,c 2 ( G ) : ,j ∈ [ t l , t r ] if G ∈ ˜ G δ , and u /latticetop f T ,c 2 ( G ) : ,j / ∈ [ t l , t r ] if G ∈ G + δ \ ˜ G δ . Due to this property, we add the following feed-forward layer.

Definition E.8 (Feed-forward Layer 3) . The vectors u and 1 L act as the weight parameters, and ζ 3 ( · ) acts as the activation function in the feed-forward layer.

Case for G ∈ G + δ \ ˜ G δ . We have ζ 3 ( u /latticetop f T ,c 2 ( G )) = 1 /latticetop L . Thus, all the entries of the input are shifted by -M -1 and become strictly negative.

Case for G ∈ ˜ G δ . We have ζ 3 ( u /latticetop f T ,c 2 ( G )) = 0 /latticetop L , so the output stays the same as the f T ,c 2 ( G ) .

With the input f T ,c 2 ( G ) , if G ∈ ˜ G δ , then ζ 3 ( u /latticetop f T ,c 2 ( G )) = 0 /latticetop L . Thus, the output stays the same as the input. If G ∈ G + δ \ ˜ G δ , then ζ 3 ( u /latticetop f T ,c 2 ( G )) = 1 /latticetop L . 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 i = 1 , 2 , · · · , d , we add the following layer:

Definition E.9 (Feed-forward Layer 4) . The vectors u and e i act as the weight parameters and ζ 4 ( · ) acts as the activation function in the feed-forward layer.

After these d layers, the output for G ∈ G + δ \ ˜ G δ is a zero matrix, while the output for G ∈ ˜ G δ remains f T ,c 2 ( G ) .

Construct the Layers: f T ,c 3 ( f T ,c 2 ( G )) = A G if G ∈ ˜ G δ . Each different G is mapped to L unique numbers u /latticetop f T ,c 2 ( G ) , which are at least δ apart from each other. We map each unique number to the corresponding output column as follows. We choose one G ∈ ˜ G δ . For each u /latticetop f T ,c 2 ( G ) : ,j , j ∈ [ L ] , we add the following feed-forward layer.

Definition E.10 (Feed-forward Layer 5) . The vectors u and e i act as the weight parameters, and ζ 4 ( · ) acts as the activation function in the feed-forward layer.

Case for G ∈ G + δ \ ˜ G δ . Recall that the input X of this layer is f T ,c 2 ( G ) . If X is a zero matrix, which is the case for G ∈ G + δ \ ˜ G δ , we have u /latticetop X = 0 /latticetop L . Then u /latticetop X -u /latticetop f T ,c 2 ( G ) : ,j 1 /latticetop L < -t l 1 L . Since t l > δ/ 2 , the output remains the same as X .

Case for G ∈ ˜ G δ . Let the input X be f T ,c 2 ( G ) , where G ∈ ˜ G δ is not equal to G . According to the Property 2 of Lemma E.4 and given a j ∈ [ L ] , u /latticetop f T ,c 2 ( G ) : ,k , ( k ∈ [ L ]) differs from u /latticetop f T ,c 2 ( G ) : ,j by at least δ . Then we have

Thus the input is left untouched.

If G = G , then

Thus we shift the j -th column of f T ,c 2 ( G ) to

In other word, this layer maps the column f T ,c 2 ( G ) : ,j to ( A G ) : ,j , without affecting any other columns.

For each G ∈ ˜ G δ , we defer that we need one layer for each unique value of u /latticetop f T ,c 2 ( G ) : ,j . Note that there are O ( δ -dL ) such numbers, so we use O ( δ -dL ) layers to finish our construction.

This completes the proof.

E.5.5 Proof of Lemma E.7

Proof of Lemma E.7. We restate the proof from [Yun et 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 σ S with the Hardmax operator σ H . Given a matrix X ∈ R d × L , we have

The operator is the only difference between the normal and the modified self-attention layers. We approximate the modified self-attention layer in T r,m,l p 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, we recall that Ψ denotes three-piecewise functions with at least a constant piece. We consider the following ζ ∈ Ψ :

where a 2 , a 3 , b 1 , b 2 , b 3 , c 1 , c 2 ∈ R , and c 1 < c 2 .

We approximate ζ ( x ) by ˜ ζ ( x ) composed of four ReLU functions:

As /epsilon1 → 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 T r,m,l p by the normal one.

Thus, for any f T ,c ∈ T 2 , 1 , 1 p , there exists a function f T ∈ T 2 , 1 , 4 p to approximate f T ,c .

This completes the proof.

F Proofs of Section 3

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 to proceed with our proof. Specifically, we 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 a proximate area of the original one in Appendix F.3.

F.1 Proof of Theorem 3.1

Here we present some auxiliary theoretical results in Appendix F.1.1 to prepare for our main proof of Theorem 3.1. Then we derive the approximation error bound of DiTs (i.e., the proof of Theorem 3.1) in Appendix F.1.2.

F.1.1 Auxiliary Lemmas for Theorem 3.1.

We restate some auxiliary lemmas and their proofs from [Chen et al., 2023] for later convenience.

Lemma F.1 (Lemma 16 of [Chen et al., 2023]) . Consider a probability density function p h ( h ) = exp ( -C ‖ h ‖ 2 2 / 2 ) for h ∈ R d 0 and constant C > 0 . Let r h > 0 be a fixed radius. Then it holds

Lemma F.2 (Lemma 2 of [Chen et al., 2023]) . Suppose Assumption 2.2 holds and g is defined as:

Given /epsilon1 > 0 , with r h = c ( √ d 0 log( d 0 /T 0 ) + log(1 //epsilon1 ) ) for an absolute constant c , it holds

Lemma F.3 (Theorem 1 of [Chen et al., 2023]) . We denote

With q ( h, t ) = ∫ hψ t ( h | h ) p h ( h ) / ( ∫ ψ t ( h | h ) p h ( h )d h )d h and p h satisfies Assumption 2.2, we have a coarse upper bound for τ ( r h ) :

Lemma F.4 (Lemma 10 of [Chen et al., 2020b]) . For any given /epsilon1 > 0 , and L -Lipschitz function g defined on [0 , 1] d 0 , there exists a continuous function f constructed by trapezoid function, such that

Moreover, the Lipschitz continuity of f is bounded:

F.1.2 Main Proof of Theorem 3.1

Proof of Theorem 3.1. With ∇ log p h t ( h ) = B /latticetop s + ( h, t ) , we have the following in (2.4)

We proceed as follows:

Step 1. Approximate q ( h, t ) with a compact-supported continuous function f ( h, t ) .

Step 2. Approximate f ( h, t ) with a transformer network.

Step 1. Approximate q ( h, t ) with a Compact-supported Continuous Function f ( h, t ) . We partition R d 0 into a compact subset H 1 := { h | ∥ ∥ h ∥ ∥ 2 ≤ r h } and its complement H 2 , where r h 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 √ d 0 /epsilon1 between q ( h, t ) and f ( h, t ) . We show the main proof here.

Approximation on H 2 × [ T 0 , T ] . For any /epsilon1 > 0 , we take r h = c ( √ d 0 log( d 0 /T 0 ) -log /epsilon1 ) . From Lemma F.2, we have

So we set f ( h, t ) = 0 on H 2 × [ T 0 , T ] .

Approximation on H 1 × [ T 0 , T ] . On H 1 × [ T 0 , T ] , we approximate q ( h, t ) by approximating each coordinate q k ( h, t ) respectively, where q ( h, t ) = [ q 1 ( h, t ) , q 2 ( h, t ) , · · · , q d 0 ( h, t )] . We rescale the input by y ′ = ( h + r h 1 ) / 2 r h and t ′ = t/T . Then the transformed input space is [0 , 1] d 0 × [ T 0 /T, 1] . We implement such a transformation by a single feed-forward layer.

By Assumption 2.3, on-support score s + ( h, t ) is L s + -Lipschitz in h . This implies q ( h, t ) is (1 + L s + ) -Lipschitz in h . When taking the transformed inputs, g ( y ′ , t ′ ) = q (2 r h y ′ -r h 1 , T t ′ ) becomes 2 r h (1 + L s + ) -Lipschitz in y ′ . Similarly, each coordinate g k ( y ′ , t ) is also 2 r h (1 + L s + ) -Lipschitz in y ′ . Here we take L h = 1 + L s + .

Besides, g ( y ′ , t ′ ) is Tτ ( r h ) -Lipsichitz with respect to t , where

We have a coarse upper bound for τ ( r h ) in Lemma F.3. We restate it here for convenience

In conclusion, each g k ( y ′ , t ) is Lipsichitz continuous. So we can apply Lemma F.4 to determine f k ( y ′ , t ) for approximating each coordinate. We concatenate f i 's together and construct f = [ f 1 , . . . , f d 0 ] /latticetop . According to the construction in Lemma F.4 and for any given /epsilon1 , we achieve

Considering the input rescaling (i.e., h → y ′ and t → t ′ ), we obtain:

-The constructed function is Lipschitz continuous in h . For any h 1 , h 2 ∈ H 1 and t ∈ [ T 0 , T ] , it holds

-The function is also Lipschitz in t . For any t 1 , t 2 ∈ [ T 0 , T ] and ∥ ∥ h ∥ ∥ 2 ≤ r h , it holds

Due to the fact that the construction of f ( h, t ) is based on trapezoid function, we have f ( h, t ) = 0 for ∥ ∥ h ∥ ∥ 2 = r h and any t ∈ [ T 0 , T ] . Thus, the two parts of f ( h, t ) can be joined together. To be more specific, the above Lipschitz continuity in h extends to the whole R d 0 .

Approximation Error Analysis under L 2 Norm. The L 2 approximation error of f can be decomposed into two terms:

The second term in the RHS above has already been bounded with the selection of r h :

The first term is bounded by:

Then we obtain

If we substitute /epsilon1 with /epsilon1/ 2 , we obtain that the approximation error of f ( h, t ) is √ d 0 /epsilon1 .

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.1. 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 t in the notation of the transformer network in DiT. Recall the reshape layer R ( · ) in Definition 3.1, we consider using f ( · ) := R -1 ◦ f T ◦ R ( · ) to approximate f t ( · ) := f ( · , t ) , where f T ∈ T 2 , 1 , 4 p .

Overall Approximation Error. With Lemma E.1, we approximate f t ( · ) with ̂ f ( · ) := R -1 ◦ ̂ f T ◦ R ( · ) . We denote

We have

Along with Step 1, we obtain

The constructed approximator to ∇ log p t ( x ) is s ̂ W = ( B ̂ f ( B /latticetop x, t ) -x ) /σ ( t ) , and the approximation error is

Settling-down of Hyperparameters. We settle down the hyperparameters to configure our network here. We refer to Appendix E.2 for some of the following calculations.

Model Architecture Depth K .

From Lemma E.6, we have K = O ((1 /δ ) dL ) . To achieve /epsilon1 -error approximation, we set δ = O ( /epsilon1 2 /d ) according to Lemma E.3. Thus we obtain

Lipchitz Upperbound for Transformer: L T .

We denote f t,R ( · ) = R ◦ f t ◦ R -1 ( · ) . We get the Lipshitz upper bound for ̂ f T ∈ T 2 . 1 . 4 p in the following way

Then we get

Model Output Bound for S T 2 , 1 , 4 p .

For the output of the constructed transformer ̂ f T ( · ) , according to Lemma E.5, we have ̂ f T ( O ) = O , where O = 0 d × L . Thus, with the Lipschitz upperbound O ( d 0 L s + ) , we have ‖ ̂ f T ( H ) ‖ F = O ( d 0 L s + r h ) , where ‖ H ‖ F ≤ r h . With r h = c ( √ d 0 log( d 0 /T 0 ) + log(1 //epsilon1 )) , we obtain

Model Parameters Bound: C 2 , ∞ OV , C OV , C 2 , ∞ KQ , C KQ , C E

. By definition, we have:

where i = 1 , 2 . For simplicity, we omit i hereafter, which does not affect our discussion.

Recall that ‖ Z ‖ 2 , ∞ denotes the 2 , ∞ -norm, where the 2 -norm is over columns and ∞ -norm is over rows. By the construction of modified attention layers (E.11) and (E.12) in Appendix E.5.3, we consider W OV to have the largest norm, i.e.,

We give the following upper bounds

By (E.11) and (E.12) in Appendix E.5.3, and the self-attention layers in Appendix E.5.5, we consider W KQ to have the largest norm, i.e.,

Then we have

We substitute δ with O ( /epsilon1 2 /d ) (according to Appendix E.4) and get:

From the construction of positional encoder (E.4) in Appendix E.2, we have

We deduce

Thus we have

5. Parameters Bound in Feed Forward Layers: C 2 , ∞ F , C F .

Recall the construction of modified feed-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.,

Then we have

and

This completes the proof.

F.2 Proof of Theorem 3.2

Here we present the auxiliary theoretical results about the covering number of transformer networks in Appendix F.2.1. The results are based on [Edelman et al., 2022, Theorem A.17]. 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

LemmaF.5 (Lemma 15 of [Chen et al., 2023]) . Let G be a bounded function class. Then there exists a constant b such that the output of any g ∈ G : R d 0 ↦→ [0 , b ] is bounded by b . Let z 1 , z 2 , · · · , z n ∈ R d 0 be i.i.d. random variables. For any δ ∈ (0 , 1) , a ≤ 1 , and c > 0 , we have

Now, we give the definition of the covering number as follows.

Definition F.1 (Covering Number) . Given a function class F and a data distribution P . Sample n data points { X i } n i =1 from P . For any /epsilon1 > 0 , the covering number N ( /epsilon1, F , { X i } n i =1 , ‖·‖ ) is the smallest size of a collection (a cover) C ∈ F , such that for any f ∈ F , there exists a ̂ f ∈ C satisfying

Furthermore, we define the covering number with respect to the data distribution as

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 r,m,l p ( K,C T , C 2 , ∞ OV , C OV , C 2 , ∞ KQ , C KQ , C 2 , ∞ F , C F , C E , L T ) 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 ‖ X ‖ 2 , ∞ ≤ C X , we have

Remark F.1. We modify [Edelman et al., 2022, Theorem A.17] in seven aspects:

We do not consider the last linear layer in the model, which converts each column vector of the transformer output to a scalar. Therefore, we ignore the item related to the last linear layer in [Edelman et al., 2022, Theorem A.17].

We do not consider the normalization layer in our model. Because the normalization layer ∏ norm ( · ) in the original proof only ensures that ‖ ∏ norm ( X 1 ) -∏ norm ( X 2 ) ‖ 2 , ∞ ≤ ‖ X 1 -X 2 ‖ 2 , ∞ , ignoring this layer does not change the result.

Our activation function is ReLU . Thus, we replace the Lipschitz upperbound of activate function by 1.

We consider the positional encoding (E.4). Then we need to replace the upperbound C X for the inputs with the upperbound C X + C E . Besides, for multi-layer transformer, the original conclusion in [Edelman et al., 2022, Theorem A.17] uses 1 as the upperbound for the 2 , ∞ -norm of inputs. We incorporate the upperbound for the inputs into the result stated in Lemma F.6.

We use (2.7) as the feed-forward layer, including two linear layers and a residual layer. Thus, we replace the original upperbound for the norm of weight matrix with the upperbound for the norm of I d + W 2 W 1 in Lemma F.6. In the following, we use O to estimate the log-covering number, thus we ignore the item for I d here for converience. This is the same for the self-attention layer.

We use multi-head attention, and incorporate the number of heads τ into our result, which is similar to [Edelman et al., 2022, Theorem A.12].

In our work, we use the transformer T 2 , 1 , 4 p , i.e., τ = 2 , m = 1 .

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

Then we define L ( s ̂ W ) = E x ∼ P 0 [ /lscript ( x ; s ̂ W ) ] .

Following [Chen et al., 2023, Appendix B.2], for any a ∈ (0 , 1) , we have

where

We denote

Then we have

For any δ > 0 , according to Lemma F.5, the following holds for term ( I ) with probability 1 -δ ,

where c ≤ 0 is a constant, and ι > 0 will be determined later.

We set where C 3 is a constant.

Putting ( I ) , ( II ) , and ( III ) together and setting a = /epsilon1 2 , then we have

with probability 1 -3 δ .

where 0 < b ≤ 1 is a constant to be determined later.

Remark F.2 (Selection Criteria of τ ) . We have two criteria:

Recall that the covering number used in our setting is N ( ( T -T 0 )( ι -η ) ( C T + r x ) log( T/T 0 ) , S T 2 , 1 , 4 p , ‖·‖ 2 ) . Thus, we must ensure ι ≥ η . According to (F.14), we consider ι satisfying the condition ι ≥ ( nT 0 ( T -T 0 )) -1 . Therefore, we consider 0 < b ≤ 1 .

For the exponent of ( T -T 0 ) , although selecting a value smaller than 1 is possible, we find that the convergence rate with respect to T is dominated by the 1 /T term appearing later in the second term of (F.18). Therefore, we continue to consider the exponent to be 1.

Then we have

with probability 1 -δ .

Following the proof structure of term ( II ) in [Chen et al., 2023, Appendix B.2], we have

For any /epsilon1 > 0 , let s W be the transformer network approximator to the score function in Theorem 3.1. For the term ( III ) , we have

For any δ > 0 , according to Lemma F.5 and given that s W is a fixed function, the following holds for term ( III ) 1 with probability 1 -δ ,

Following the proof structure of term ( III ) 2 in [Chen et al., 2023, Appendix B.2], we have

Covering Number of S T 2 , 1 , 4 p . The next step is to calculate the covering number of S T 2 , 1 , 4 p . S T 2 , 1 , 4 p consists of two components: (i) Matrix W B with orthonormal columns; (ii) Network function f T .

Suppose we have W B 1 , W B 2 and f 1 , f 2 , such that ‖ W B 1 -W B 2 ‖ F ≤ δ 1 and sup ‖ x ‖ 2 ≤ 3 r x + √ D log D,t ∈ [ T 0 ,T ] ‖ f 1 ( x, t ) -f 2 ( x, t ) ‖ 2 ≤ δ 2 , where f 1 = R -1 ◦ f T 1 ◦ R,f 2 = R -1 ◦ f T 2 ◦ R . Then we have

where L T upper bounds the Lipschitz constant of f T .

For the set { W B ∈ R D × d 0 : ‖ W B ‖ 2 ≤ 1 } , its δ 1 -covering number is ( 1 + 2 √ d 0 /δ 1 ) Dd 0 [Chen et al., 2020a, Lemma 8]. The δ 2 -covering number of f needs further discussion as there is a reshaping process in our network. The input is reshaped from h ∈ R d 0 to H ∈ R d × L , and

Thus we have

Then we follow the covering number of sequence-to-sequence transformer T 2 , 1 , 4 p in Lemma F.6. We get the following δ 2 -covering number

where

According to the (F.4), (F.5), (F.7), (F.8), (F.9), (F.10), (F.12), (F.13), (F.11) and (F.6) in Appendix F.1.2, we derive the following with δ = O ( /epsilon1 2 /d ) (Appendix E.4) and d = 4 (Theorem 3.1):

Each element of the input data is within [0 , 1] , as shown in Appendix E.

For any δ 3 > 0 , we get the log-covering number of T 2 , 1 , 4 p ,

According to (F.15), we adopt the following value for δ 3 in our setting

According to [Chen et al., 2023, Appendix B.2], the log-covering number of S T 2 , 1 , 4 p is

Substituting the log-covering number into (F.15), we have

Recall the following parameters:

C 2 T = O ( d 2 0 L 2 s + d 0 log( d 0 /T 0 ) + log(1 //epsilon1 )) ,

r 2 x = O ( d 0 log d 0 +log C T +log ( n/δ ) ) ,

δ : probability error,

/epsilon1 : approximation error,

n : sample size,

T 0 < T/ 2 ,

D,d,d 0 > 1 : feature dimension,

L > 1 : sequence length,

d 0 = L · d ,

L s + : Lipschitz coefficient.

Ignoring the log factors and poly( D,d,d 0 , L S + ) , the first term in (F.18) becomes

The second term is simplified to

Thus, the final bound is

To balance the first and second terms with respect to n , we select b = 1 / 3 . Therefore, we give the final bound as

This completes the proof.

F.3 Proof of Corollary 3.2.1

Our proof is built on [Chen et al., 2023, Appendix C]. The main difference between our work and [Chen et al., 2023] is our score estimation error in Theorem 3.2. Consequently, only the subspace error and the total variation distance differ from [Chen et 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), y t denotes the backward process. We denote the backward latent variable by h ← t = B /latticetop y t . Since we write the time index explicitly, we drop the y, h notation for t > 0 .

Following [Chen et al., 2023, Appendix C.2], we have the following ground truth backward process

where W t denotes the reversed Wiener process (standard Brownian motion) at time t (see Section 2 for more details).

We define P h T 0 as the ground truth marginal distribution of h ← T 0 .

For the learned process ˜ y t , we consider ˜ h ← t = W /latticetop B ˜ y t . For any orthogonal matrix U ∈ R d 0 × d 0 , we define the U transformed version of ˜ h ← t as ˜ h ← ,U t = U /latticetop ˜ h ← t . Then the backward SDE for ˜ h ← ,U t is

where

We define ̂ P h T 0 as the estimated marginal distribution of ˜ h ← ,U T 0 from above continuous SDE.

The discretized backward SDE of ˜ h ← ,U T 0 is

We define ̂ P h, dis T 0 as the estimated marginal distribution of ˜ h ⇐ ,U T 0 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 a detailed proof of Corollary 3.2.1 in Appendix F.3.2.

F.3.1 Auxiliary Lemmas

Here we include a few auxiliary lemmas from [Chen et al., 2023] without proofs. Recall the definition of Lipschitz norm: for a given function f , ‖ f ( · ) ‖ Lip = sup x = y ( ‖ f ( x ) -f ( y ) ‖ 2 / ‖ x -y ‖ 2 ) .

/negationslash

Lemma F.7 (Lemma 3 of [Chen et al., 2023]) . Assume that the following holds

where λ min denotes the smallest eigenvalue. We denote

Let T 0 ≤ min { 2 log( d 0 /C sh ) , 1 , 2 log( c 0 ) , c 0 } and T ≥ max { 2 log( C h /d 0 ) , 1 } . Suppose we have

Then we have

and there exists an orthogonal matrix U ∈ R d 0 × d 0 , such that:

Lemma F.8 (Lemma 4 of [Chen et al., 2023]) . Assume that P h is sub-Gaussian and that f ( h, t ) and ∇ log p h t ( h ) are Lipschitz continuous with respect to h and t . For any orthogonal matrix U ∈ R d 0 × d 0 , we define

Assume that we have the latent score matching error-bound

where /epsilon1 latent > 0 . Then we have the following latent distribution estimation error for the continuous backward SDE:

where ̂ P h T 0 is the marginal distribution of the generated h T 0 using the continuous backward SDE. Furthermore, let ̂ P h, dis T 0 denote the marginal distribution of the generated h T 0 using the discretized backward SDE. Then we have the following latent distribution estimation error for the discretized backward SDE

where

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 µ satisfying T -T 0 = K T µ for some K T ∈ N + ,

where y 0 ∼ N(0 , I ) . Then, for T > 1 and T 0 + µ ≤ 1 , we have y T -T 0 ∼ N ( 0 , σ 2 I ) with σ 2 ≤ e ( T 0 + µ ) .

Lemma F.10 (Lemma 10 in [Chen et al., 2023]) . Assume that ∇ log p h ( h ) is L h -Lipschitz. Then we have E h ∼ P h ‖∇ log p h ( h ) ‖ 2 2 ≤ d 0 L h .

F.3.2 Main Proof of Corollary 3.2.1

Proof. Recall the estimation error in Theorem 3.2 is ξ ( n, /epsilon1,L ) / ( TT 0 ) , where

Proof of (i). By the definition of (F.19) and the estimation error in Theorem 3.2, the error bound in (F.20) equals to ξ ( n, /epsilon1,L )( T -T 0 ) / ( TT 0 ) in Lemma F.7. By Lemma F.10, we set C sh = d 0 L h . Then, we have

By substituting the value of ξ ( n, /epsilon1,L ) and T = O (log n ) into the bound above, we deduce

Proof of (ii). Recall that max t ‖ f ( · , t ) ‖ Lip ≤ L T . Furthermore, according to Lemma F.7 and Lemma F.10, we have

where

Following the proof structure in [Chen et al., 2023, Appendix C.4], we get

Following the proof structure in [Chen et al., 2023, Appendix C.4] and setting T = O (log n ) , we obtain

where ˜ O hides the factor about D,d 0 , d, L s + , log n , and T -T 0

By definition, ̂ P h, dis T 0 = ( UW B ) /latticetop /sharp ̂ P T 0 , where ̂ P T 0 is the distribution generated by s ̂ W using the discretized backward process. This completes the proof of the total variation distance.

Proof of (iii). We apply Lemma F.9 due to our score decomposition. With the marginal distribution at time T -T 0 and observing µ /lessmuch T 0 , we obtain the last property.

This completes the proof.

G Proofs of Section 4

Our proofs are motivated by the observation of low-rank gradient decomposition in 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.1 and [Diao et al., 2019, 2018]) to proceed our proofs. Specifically, we approximate DiT training gradients with a series of low-rank approximations in Appendices G.1.1 to G.1.3, and carefully match the multiplication dimensions so that the computation of d g 2 d W forms a chained low-rank approximation in Appendix G.2.

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 ∈ R d × d and W ∈ R d 2 denotes the vectorization of W ∈ R d × d following Definition D.1.

Definition G.1. Let A 1 , A 2 ∈ R d × L be two matrices. Suppose A = A /latticetop 1 ⊗ A /latticetop 2 ∈ R L 2 × d 2 . Define A j 0 ∈ R L × d 2 as an L × d 2 sub-block of A . There are L such sub-blocks in total. For each j 0 ∈ [ L ] , define the function u ( W ) j 0 : R d 2 → R L by u ( W ) j 0 := exp( A j 0 W ) ∈ R L .

Definition G.2. Let A 1 , A 2 ∈ R d × L be two matrices. Suppose A = A /latticetop 1 ⊗ A /latticetop 2 ∈ R L 2 × d 2 . Define A j 0 ∈ R L × d 2 as an L × d 2 sub-block of A . There are L such sub-blocks in total. For every index j 0 ∈ [ L ] , consider the function α ( W ) j 0 : R d 2 → R defined by α ( W ) j 0 := 〈 exp( A j 0 W ) ︸ ︷︷ ︸ L × 1 , 1 L ︸︷︷︸ L × 1 〉 .

Definition G.3. Suppose that α ( W ) j 0 ∈ R and u ( W ) j 0 ∈ R L are defined as in Definitions G.1 and G.2, respectively. For a fixed j 0 ∈ [ L ] , consider the function f ( W ) j 0 : R d 2 → R L defined by

Define f ( W ) ∈ R L × L as the matrix where the j 0 -th row is ( f ( W ) j 0 ) /latticetop .

Definition G.4. For every i 0 ∈ [ d ] , define the function h ( W OV ) i 0 : R d 2 → R L by

Here, W OV ∈ R d × d denotes the matrix representation of W OV ∈ R d 2 , and ( W OV ) /latticetop ∗ ,i 0 represents the i 0 -th column of W /latticetop OV . Define h ( W OV ) ∈ R L × d as the matrix where the i 0 -th column is h ( W OV ) i 0 .

Definition G.5. For each j 0 ∈ [ L ] , we denote f ( W ) j 0 ∈ R L as the normalized vector defined by Definition G.3. For each i 0 ∈ [ d ] , h ( W OV ) i 0 is defined as per Definition G.4. For every pair ( j 0 , i 0 ) ∈ [ L ] × [ d ] , define the function c ( W ) j 0 ,i 0 : R d 2 × R d 2 → R by

where ( W OV ) j 0 ,i 0 is the element at the ( j 0 , i 0 ) position of the matrix W OV ∈ R L × d . c ( · ) has matrix form

With the tensor trick (Appendix D.3), we compute the gradient d g 2 d W of the DiT loss as follows:

(G.1) presents a neat decomposition of d g 2 d W . Each term is easy enough to handle. Thus, we arrive at 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 of DiT Gradient) . Let matrix A 1 , A 2 , A 3 , W, W OV , Y and loss function L follow Definition 4.1, and A := A /latticetop 1 ⊗ A /latticetop 2 . It holds

Proof. Let Z [ i, · ] and Z [ · , j ] be the i -th row and j -th column of matrix Z .

For each j 0 ∈ [ L ] , we have

Therefore, for each j 0 ∈ [ L ] , we have

Similarly,

Putting all together, we have

where

This completes the proof.

Observe (G.2) carefully. We see that (I) is diagonal and (II) is low-rank. This provides a 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 d g 2 d W as a chained low-rank approximation.

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 the approximation term-by-term.

To facilitate our proof, it's convenient to introduce the following notations.

Definition G.6 ( q ( · ) ) . Define c ( W ) ∈ R L × d as specified in Definition G.5 and h ( W OV ) ∈ R L × d as described in Definition G.4. Define q ( W ) ∈ R L × L by

In addition, q ( W ) /latticetop j 0 denotes the j 0 -th row of q ( W ) , transposed, making it an L × 1 vector.

Definition G.7 ( p ( · ) , p 1 ( · ) , p 2 ( · ) ) . For each index j 0 ∈ [ L ] , we define p ( W ) j 0 ∈ R n as follows:

We define p ( W ) ∈ R L × L such that p ( W ) /latticetop j 0 forms the j 0 -th row of p ( W ) . In addition, for every index j 0 ∈ [ L ] , we define p 1 ( W ) j 0 , p 2 ( W ) j 0 ∈ R L as

such that p ( W ) = p 1 ( W ) -p 2 ( W ) .

p ( · ) allows us to express d g 2 d W in a neat form:

Lemma G.2. Define the functions f ( W ) ∈ R L × L , c ( W ) ∈ R d × L , h ( W OV ) ∈ R d × L , q ( W ) ∈ R L × L , and p ( W ) ∈ R L × L as specified in Definitions G.3 to G.7, respectively. Let A 1 , A 2 ∈ R d × L be two given matrices, and define A = A /latticetop 1 ⊗ A /latticetop 2 . Define g 2 according to (O1), and let g 2 ( W ) j 0 ,i 0 be as described in (G.1). It holds

Proof. By definitions, (G.1) gives

Therefore, (G.4) becomes

Then, by definitions of q ( · ) , p ( · ) , we complete the proof.

G.1.2 Low-Rank Approximations of Building Blocks Part I: f ( · ) , q ( · ) , and c ( · )

The definitions of p , p 1 , p 2 , and Lemma G.2 show that the DiT training gradient d g 2 d W involves entry-wise products of f , q , and c . Therefore, if we approximate these with inner-dimension-matched low-rank approximations, computing d g 2 d W itself becomes a low-rank approximation. In the following sections, we present low-rank approximations for f , q , and c .

Lemma G.3 (Approximate f ( · ) , Modified from [Alman and Song, 2023]) . Let Γ = o ( √ log L ) and k 1 = L o (1) . Let A 1 , A 2 , ∈ R d × L , W ∈ R d × d and f ( W ) = D -1 exp ( A /latticetop 1 X A 2 ) with D = diag ( exp ( A /latticetop 1 WA 2 ) 1 L ) follows Definitions G.1 to G.3 and G.5. If max (∥ ∥ A /latticetop 1 W ∥ ∥ max ≤ Γ , ‖ A 2 ‖ max ) ≤ Γ , then there exist two matrices U 1 , V 1 ∈ R L × k 1 such that ∥ ∥ U 1 V /latticetop 1 -f ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) . In addition, it takes L 1+ o (1) time to construct U 1 and V 1 .

Proof. By [Alman and Song, 2023, Theorem 3], we complete the proof.

LemmaG.4 (Approximate c ( · ) ) . Assume all numerical values are in O (log L ) bits. Let d = O (log L ) and c ( W ) ∈ R L × d follows Definition G.5. There exist two matrices U 1 , V 1 ∈ R L × k 1 such that ∥ ∥ U 1 V /latticetop 1 h ( W OV ) -Y /latticetop -c ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) .

Lemma G.5 (Approximate q ( · ) ) . Let k 2 = L o (1) , c ( · ) ∈ R L × d follow Definition G.5 and let q ( W ) := c ( W ) h ( W OV ) T ∈ R L × L (follow Definition G.6). There exist two matrices U 2 , V 2 ∈ R L × k 2 such that ∥ ∥ U 2 V /latticetop 2 -q ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) . In addition, it takes L 1+ o (1) time to construct U 2 , V 2 .

Our proof is built on [Alman and Song, 2023, Lemma D.3].

Let ˜ q ( · ) denote an approximation to q ( · ) .

By Lemma G.4, U 1 V /latticetop 1 h ( W OV ) -Y approximates c ( W ) up to accuracy /epsilon1 = 1 / poly( L ) .

Thus, by setting ˜ q ( W ) = h ( W OV ) ( U 1 V /latticetop 1 h ( W OV ) -Y ) /latticetop , we find a low-rank form for ˜ q ( · ) :

such that

By k 1 , d = L o (1) , compute ( h ( W OV )) /latticetop ︸ ︷︷ ︸ d × L V 1 ︸︷︷︸ L × k 1 U /latticetop 1 ︸︷︷︸ k 1 × L takes only L 1+ o (1) time. This completes the proof.

G.1.3 Low-Rank Approximations of Building Blocks Part II: p ( · )

Now, we use the low-rank approximations of f, q, c to construct low-rank approximations for p 1 ( · ) , p 2 ( · ) , p ( · ) .

Lemma G.6 (Approximate p 1 ( · ) ) . Let k 1 , k 2 = L o (1) . Suppose U 1 , V 1 ∈ R L × k 1 approximates f ( W ) ∈ R L × L such that ∥ ∥ U 1 V /latticetop 1 -f ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) , and U 2 , V 2 ∈ R L × k 2 approximates the q ( W ) ∈ R L × L such that ∥ ∥ U 2 V /latticetop 2 -q ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) . Then there exist two matrices U 3 , V 3 ∈ R L × k 3 such that ∥ ∥ U 3 V /latticetop 3 -p 1 ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) . In addition, it takes L 1+ o (1) time to construct U 3 , V 3 .

Proof of Lemma G.6. By tensor trick, we construct U 3 , V 3 as tensor products of U 1 , V 1 and U 2 , V 2 , respectively, while preserving their low-rank structures. Then, we show the low-rank approximation of p 1 ( · ) with bounded error by Lemma G.3 and Lemma G.5.

Let /circledivide be column-wise Kronecker product such that A /circledivide B := [ A [ · , 1] ⊗ B [ · , 1] | . . . | A [ · , k 1 ] ⊗ B [ · , k 1 ]] ∈ R L × k 1 k 2 for A ∈ R L × k 1 , B ∈ R L × k 2 .

Let ˜ f ( W ) := U 1 V T 1 and ˜ q ( W ) := U 2 V T 2 denote matrix-multiplication approximations to f ( W ) and q ( W ) , respectively.

For the case of presentation, let U 3 = L × k 1 ︷︸︸︷ U 1 /circledivide L × k 2 ︷︸︸︷ U 2 and V 3 = L × k 1 ︷︸︸︷ V 1 /circledivide L × k 2 ︷︸︸︷ V 2 . It holds

Computationally, by k 1 , k 2 = L o (1) , computing U 3 and V 3 takes L 1+ o (1) time. This completes the proof.

LemmaG.7 (Approximate p 2 ( · ) ) . Let k 1 , k 2 , k 4 = L o (1) . Let p 2 ( W ) ∈ R L × L follow Definition G.7 such that its j 0 -th column is p 2 ( W ) j 0 = f ( W ) j 0 f ( W ) /latticetop j 0 q ( W ) j 0 for each j 0 ∈ [ L ] . Suppose U 1 , V 1 ∈ R L × k 1 approximates the f( X ) such that ∥ ∥ U 1 V /latticetop 1 -f ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) , and U 2 , V 2 ∈ R L × k 2 approximates the q ( W ) ∈ R L × L such that ∥ ∥ U 2 V /latticetop 2 -q ( W ) ∥ ∥ max ≤ /epsilon1/ poly( L ) . Then there exist matrices U 4 , V 4 ∈ R L × k 4 such that ∥ ∥ U 4 V /latticetop 4 -p 2 () ∥ ∥ max ≤ /epsilon1/ poly( L ) . In addition, it takes L 1+ o (1) time to construct U 4 , V 4 .

Proof of Lemma G.7. From Definition G.7,

For (I), we show its low-rank approximation by observing the low-rank-preserving property of the multiplication between f ( · ) and q ( · ) (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 a function r ( W ) : R d 2 → R L such that the j 0 -th component r ( W ) j 0 := ( f ( W ) j 0 ) /latticetop q ( W ) j 0 for all j 0 ∈ [ L ] . Let ˜ r ( W ) denote the approximation of r ( W ) via decomposing into f ( · ) and q ( · ) :

for all j 0 ∈ [ L ] . This allows us to write p 2 ( W ) = f ( W ) diag( r ( W )) with diag( ˜ r ( W )) denoting a diagonal matrix with diagonal entries being components of ˜ r ( W ) .

Part (II). With r ( · ) , we approximate p 2 ( · ) with ˜ p 2 ( W ) = ˜ f ( W ) diag( ˜ r ( W )) as follows.

Since ˜ f ( W ) has low rank representation, and diag( ˜ r ( W )) is a diagonal matrix, ˜ p 2 ( · ) has low-rank representation by definition. Thus, we set ˜ p 2 ( W ) = U 4 V T 4 with U 4 = U 1 and V 4 = diag( ˜ r ( W )) V 1 . Then, we bound the approximation error

Computationally, computing V /latticetop 1 V 2 takes L 1+ o (1) time by k 1 , k 2 = L o (1) . Once we have V /latticetop 1 V 2 precomputed, (G.6) only takes O ( k 1 k 2 ) time for each j 0 ∈ [ L ] . Thus, the total time is O ( Lk 1 k 2 ) = L 1+ o (1) . Since U 1 and V 1 takes L 1+ o (1) time to construct and V 4 = diag( ˜ r ( W )) ︸ ︷︷ ︸ L × L V 1 ︸︷︷︸ L × k 1 also takes

L 1+ o (1) time, U 4 and V 4 takes L 1+ o (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 ( · ) , p 1 ( · ) and p 2 ( · ) (Definition G.7), we have

To show the existence of L 1+ o (1) algorithms for DiT backward computation Problem 1, we prove fast low-rank approximations for A 1 p 1 ( W ) A /latticetop 2 and A 1 p 2 ( W ) A /latticetop 2 as follows.

Let ˜ p 1 ( W ) , ˜ p 2 ( W ) denote the approximations to p 1 ( W ) , p 2 ( W ) , respectively.

By Lemma G.6, it takes L 1+ o (1) time to construct U 3 , V 3 ∈ R L × k 3 such that

Then, computing A 1 ︸︷︷︸ d × L U 3 ︸︷︷︸ L × k 3 V /latticetop 3 ︸︷︷︸ k 3 × L A /latticetop 2 ︸︷︷︸ L × d takes L 1+ o (1) due to the fact that d, k 1 k 3 = L o (1) .

Therefore, total running time for A 1 p 1 ( W ) A /latticetop 2 is L · L o (1) = L 1+ o (1) .

For the same reason (by Lemma G.7), total running time for A 1 p 2 ( W ) A /latticetop 2 is L · L o (1) = L 1+ o (1) . Lastly, we have

Set /epsilon1 = 1 / poly( L ) . We complete the proof.

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 in Section 3 and Section 4 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 contributions made in the paper and important assumptions and limitations. A 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 is 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 a 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 (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 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. We include our proofs in the appendix and have made every effort to ensure the correctness of our theoretical results.

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 or the supplemental material, but if they appear in 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 proofs 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 to reproduce the main experimental 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 a formal theory work without experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

If 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 or 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 suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In 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., in the case of a large language model), releasing of a 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 a new model (e.g., a 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 to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: This is a formal theory work without experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-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://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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 to 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 a 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 a 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 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is 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).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

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 to reproduce 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 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 experimental runs as 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, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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.

If 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., if there is a 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 is not anticipated to have negative social impacts.

Guidelines:

The answer NA means that there is no societal impact of 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 of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as 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 or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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, or scraped datasets)?

Answer: [NA]

Justification: This is a formal theory work without experiments.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to usage guidelines or restrictions to access the model or 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 or 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 a 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 of the asset is used and, if 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 a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to 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 a formal theory work without experiments.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions 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 is used.

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, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This is a formal theory work without experiments.

Guidelines:

The answer NA means that the paper does not 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 to 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 a 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 in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

STRENGTHS:
- The authors derived approximation error bounds for transformer-based score estimators in latent DiTs, providing practical structural guidance. 
- The paper also provided sample complexity bounds for score estimation and demonstrated recovery of initial data distribution.
- 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.

QUESTIONS:
See above.

LIMITATIONS:
See above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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?\n\nQUESTIONS:\nsee weaknesses.\n\nLIMITATIONS:\nThe paper has outlined its limitations and broader impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-15"}

comments: SUMMARY:
The paper studies the statistical and computational limits of latent diffusion transformers.

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).

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".

UPDATE: Increased the score and the confidence after the rebuttal.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe limitations of this work have been stated in Section 5 and are left for future work as claimed by the authors.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-27"}