Diffusion Transformer Captures Spatial-Temporal Dependencies: A Theory for Gaussian Process Data
Authors: Hengyu Fu
Zehao Dou
Jiawei Guo
Mengdi Wang
Minshuo Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion model
transformers
sequential data
spatial-temporal dependency
sample complexity
Primary area: generative models
TL;DR: We analyze the theory of diffusion transformers in capturing spatial-temporal dependencies.

Abstract: Diffusion Transformer, the backbone of Sora for video generation, successfully scales the capacity of diffusion models, pioneering new avenues for high-fidelity sequential data generation. Unlike static data such as images, sequential data consists of consecutive data frames indexed by time, exhibiting rich spatial and temporal dependencies. These dependencies represent the underlying dynamic model and are critical to validate the generated data. In this paper, we make the first theoretical step towards bridging diffusion transformers for capturing spatial-temporal dependencies. Specifically, we establish score approximation and distribution estimation guarantees of diffusion transformers for learning Gaussian process data with covariance functions of various decay patterns. We highlight how the spatial-temporal dependencies are captured and affect learning efficiency. Our study proposes a novel transformer approximation theory, where the transformer acts to unroll an algorithm. We support our theoretical results by numerical experiments, providing strong evidence that spatial-temporal dependencies are captured within attention layers, aligning with our approximation theory.

Published as a conference paper at ICLR 2025
DIFFUSION TRANSFORMER CAPTURES SPATIAL- TEMPORAL DEPENDENCIES: A THEORY FOR GAUS- SIAN PROCESS DATA
Hengyu Fu'
Zehao Dou!
Jiawei Guo
Mengdi Wang
Minshuo Chen 1
ABSTRACT
Diffusion Transformer, the backbone of Sora for video generation, successfully scales the capacity of diffusion models, pioneering new avenues for high-fidelity sequential data generation. Unlike static data such as images, sequential data consists of consecutive data frames indexed by time, exhibiting rich spatial and temporal dependencies. These dependencies represent the underlying dynamic model and critical to validate the generated data. In this paper, make the first theoretical step towards bridging diffusion transformers for capturing spatial-temporal dependencies. Specifically, establish score approximation and distribution estimation guarantees of diffusion transformers for learning Gaussian process data with covariance functions of various decay patterns. We highlight how the spatial-temporal dependencies are captured and affect learning efficiency. Our study proposes a novel transformer approximation theory, where the transformer acts to unroll an algorithm. We support our theoretical results by numerical experiments, providing strong evidence that spatial-temporal dependencies are captured within attention layers, aligning with our approximation theory.
INTRODUCTION
Diffusion models have emerged powerful new technology for generative AI, which widely adopted in computer vision and audio generation (Song and Ermon, 2019; Song et al., 2020; Ho et al., 2020; Zhang et al., 2023), sequential data modeling (Alcaraz and Strodthoff, 2022; Tashiro al., 2021; Tian et al., 2023), reinforcement learning and control (Pearce et al., 2023; Hansen-Estrude et al., 2023; Zhu al., 2023; Ding and Jin, 2023), as well as computational biology (Xu et al., 2022; Guo et al., 2023). The basic functionality of diffusion models is to generate new samples replicating essential characteristics in the training data.
Diffusion models generate new samples by sequentially transforming Gaussian white noise. Each step of the transformation is driven by a so-called "score function", which is parameterized by neural network. In order train the score neural network, diffusion models utilize forward process to produce noise corrupted data and the score neural network attempts to remove the added noise. In early implementations of diffusion models, the score neural network is typically chosen as the U-Net (Ronneberger et al., 2015). Afterwards, a few works also demonstrate the capability of using transformers score neural network (Peebles and Xie, 2023; Wu et al., 2024; Bao al., 2023). Throughout the paper, we adopt the terminology in (Peebles and Xie, 2023) to name diffusion models with transformers as diffusion transformers.
Recently, the astounding success of applying diffusion models in dynamic (sequential) data, including video generation (Gupta et al., 2023; Liu al., 2024b) and financial data augmentation (Gao et al., 2024), strengthens the seemingly unlimited potentiality of diffusion models. These models advocate transformers over the traditional U-Net for parameterizing the score function. high-level intuition is that video data comprises rich spatial and temporal dependencies, induced by the underlying dynamics of objects and background. For example, the movement of an object should be continuous along
WWok done at Stanford University. Email: wwwchh018ccaccc
I Yale University. Email: zehao. www.daulleee..ooo 2 edu
Work done Northwestern University. Email: guokaku6668 (aae:666@maa....mmmm
S Princetoon University. Email:
edu
Northwestem University. Email: mi nshuo. chen@northwesterr.

---

Published as a conference paper at ICLR 2025
learned temporal dependency shown estimated kerrel of sumples
spatial dependency step
truth tomporal dependency
shown in kernel Gauasiou Process
ebown la covariacio matrii truth at 20th
learned at 20ch
learned 70ch
70th
300
300
126
20 200
time step
Figure I: Diffusiin transformer learns spatial-tempraa dependencies. The diffusion transformer is trained with data sampled from stationary Gaussian process consisting of 128 time steps. At each time step, the data dimension is 8. We obtain 1000 generated samples at each time step. The left large heat map demonstrates the estimated temporal correlation (see Appendix F for the estimation method) in the process between different time steps, which aligns well with the ground truth on the right. The smaller heat maps are the estimated covariance matrix of data at a single time step, which demonstrate the spatial dependencies in data. They also align well with the ground truth.
the time. These dependencies resonate well with the aeff-attention mechanism in transformers for capturing token-wise correlation, suggesting benefits for learning with sequential data. We illustrate diffusion transformer accurately learning teriiall dependencies in Figure
Despite the empirical success, there lacks rigorous understanding of diffusion transformers for sequential data modeling. Different from static data, sequential data consists of series of dependent data frames. The data sequence can be extensively long. For instance, one-minute video would contain over 1440 image frames, and intraday data financial applications may be even longer.
Therefore, naïvely treating the entire sequential data solely as high-dimensional data without con- sidering inner spatial-temport correlations will lead a large dimensionality dependence and inefficient learning. As a result, existing results of diffusion models for static data can provide few insights on the following fundamental questions:
Can diffusion transformers efficiently capture spatial-temporal dependencies in sequential data? dependencies affect the learning efficiency?
If yes, how do
We answer these questions for the first time by studying using diffusion transformers for learning Gaussian process data. Besides its simplicity, Gaussian process exhibits intriguing and salient properties. Firstly, Gaussian process data can be high-dimenssonaa, highlighting the influence of data dimensionality in diffusion transformer. Secondly, the spatial-temporal dependencies are the defining quantities of a Gaussian process. This necessitates an effective learning of these dependencies. In fact, Gaussian process encode wide variety of complicated correlations in real-world applications (Seeger, 2004; Williams and Rasmussen, 2006). For instance, Brownian motion falls into the category of Gaussian process for describing particle movements in fluid. Gaussian process is also powerful statistical tool for regression, classification and forecasting problems (Banerjee et al., 2013; Deringer al., 2021; Chen et al., 2021; Borovitskiy et al., 2021).
Contributions Our results show that by construction, transformers can adapt to the spatial-temporal dependencies as to promote the learning efficiency. Furthermore, we show sample complexity bounds of diffusion transformers, demonstrating the influence of the decay of correlation in the sequential data. We summarize our contributions follows.
We propose novel score function approximation scheme for Gaussian process data, which represents the score function by a gradient descent algorithm (Lemma 1). Then we construct transformer architecture to unroll the gradient descent algorithm in Theorem I. We particularly highlight that the attention layer effectively captures the spatial-temporal dependencies. Meanwhile, the decay pattern of those dependencies influences the approximation efficiency.
Built upon our score function approximation theory, we establish the first sample complexity bound for diffusion transformers in learning sequential data (Theorem 2). We show that the generalization error scales with 1/ /n, where n is the sample size. Our generalization error also demonstrates the influence of dependency decay speed and the length of sequences.
2

---

Published as a conference paper at ICLR 2025
We provide numerical results to support our theory by showing the learning performance under ious settings. More interestingly, we demonstrate that a well-trained diffusion transformer reproduces the ground-truth ppatial-temporda dependencies accurately within attention layer, emphasizing the applicability of our theoretical insights.
Our theories are the first explain how diffusion transformers model sequential data, while most existing theories of diffusion models focus on static data. Due to space limit, we provide a discussion about related works and our technical novelty to Appendix B.
Notation We use bold letters to denote vectors and matrices. For a vector v, we denote its Euclidean norm as l|vll2- For matrix A, we denote its operator, Frobenius norm ||A||2 and || p. respectively. Moreover, we denote || = m AAajl We denote the condition number of a positive definite matrix A by ((A) Amaax (A)/ Amin (A), where Amax and Amin denote the
maximum and minimum eigenvalues. We denote f S if f I Cg holds for a constant C' > O.
GAUSSIAN PROCESS AND DIFFUSION TRANSFORMER
In this section, we formalize data modeling and sampling problem with Gaussian process data. Meanwhile, we briefly introduce diffusion processes and transformer architectures.
Gaussian Process We denote [X,))Ejo,HH as a continuous-time Gaussian process in the time interval [0, H.. The process lives in the d-dimensional Euclidean space, i.e., X, Rd for any E [0, H]. A defining property of Gaussian process is that for any finite collection of time indices h1, hi,..., hv with N E N the joint distribution of (Xh, ..... Xny } is still Gaussian. As particular example, for a fixed time h, the marginal distribution of X, is Gaussian.
To fully describe a Gaussian process, need the concept of mean and covariance functions. Roughly speaking, a mean function (h) = E x characterizes the expected evolution trend of the process. A covariance function [(h1 h2)) [X,,,, ((h1) )) (h2))) I captures the correlation between two time indices in the process. Note that, when h1 = h2 h, the covariance function computes the covariance matrix of Xx. Remarkably, the covariance function determines many basic properties of the continuous-time process, such as its stationarity, periodicity, and smoothness (Williams and Rasmussen, 2006). In our later study, we reveal an intimate connection between the behavior of the covariance function to the learning efficiency of diffusion transformer.
Throughout the paper, we focus on Gaussian processes whose covariance functions only depend on the gap between time indices. Accordingly. we reparameterize the covariance function as [(h,, h22) (h1, E2 where ((,.) is a scalar-output function and Σ = Cov [x]] (identical for any h).
Sequential Data Sampled from Gaussian Process In real-world scenarios, an underlying continuous-time process is often perceived by sequence of data sampled at discrete times. For example, a video typically consists of 24 to 30 image frames per second. When played back, these frames appear seamless the human eye, which cannot distinguish between individual frames as if the video is continuous. Following the same spirit, we denote hy..... ......,,y hy for a sufficiently large N E N as a uniform grid on the interval [0, H] and form a discrete sequence {X....... X observed at those time indices from an underlying Gaussian process. By the definition of Gaussian process, if we stack X X Xby, consecutively as a vector RdN it follows a Gaussian distribu- tion. The mean is fa " MN]] with μa; = ((h;)) and the covariance is block-wise matrix represented as T 3 E, where
....:
h1)
k(),NN) :
THE
TINE
:
and T+E=
h1),, hy) hx,,
INIE INNE
Here, @ the matrix Kronecker product. Notably, I captures the temporal dependency between time indices and E captures the spatial dependency of entries in X,. Since hı,..., hy form uniform grid, I is a symmetric Toeplitz matrix with entries taking at most N different values.
Gaussian process data provides explicit description of the spatial-temporal dependencies, but still remains highly relevant to real-world diffusion models. These models utilize pre-trained Variational AutoEncoder (VAE) to map data into low-dimensional latent representation (Wang et al., 2023; Blattmann et al., 2023). The typical prior distribution of the low- www-immeniioall representation in VAEs is assumed to be Gaussian. Empirical results have demonstrated the effectiveness of Gaussian latent prior for sequential data modeling in some variants of VAEs (Casale et al., 2018; 8; Fortuin et al., 2020). Our study aligns well with these empirical observations.
3

---

Published as a conference paper at ICLR 2025
In a learning setting, collect i..d. realizations of the discrete sequence (Xiy..... X denoted as D generating new samples mimicking the distribution of the discrete sequence. The subtlety here is that naively learning the joint distribution of (X
We aim to use a diffusion transformer for learning and
} is subject to a large dimension factor N, heavily exaggerating the problem dimension and jeopardizing the learning efficiency. Fortunately, will show that the behavior of the covariance function may induce benign temporal dependencies and largely promote the sample complexity.
Diffusion Processes A diffusion model generates new data by progressively removing noise using the so-called "score function". We adopt continuous-time perspective for brief review of diffusion models. Interested readers may refer to recent surveys for a comprehensive exposure (Chen et al., 2024; Tang and Zhao, 2024; Chan, 2024). A diffusion model utilizes forward and a backward process for training and sample generation, respectively:
1 dx; - X,tt + dWW, for t [0, T] and Xo Po.
(Forward)
dx; X Vlogpr- dt dw,, for t E [0, 7] and X N(O, I). (Backward)
Here, W, and W, are independent Wiener processes and T is finite terminal diffusion timestep. The initial distribution of Xo is Pa. which is also the clean data distribution, and we denote P; as the marginal distribution of X, Since we are corrupting Po by Gaussian noise, Pt has density function Pt. Thus, V log Pt in the backward process is recognized as the score function, which is unknown and requires learning. Typically, the score function will be parameterized by neural network and trained by optimizing a loss function, which we will introduce in Section 5. When generating new samples, simulate discretized version of the backward process using the learned score function.
For Gaussian process data, we interpret X, as a vector in I by stacking N observations. We term each observation as a patch. Equivalently, the forward process is to add independent Gaussian noise to each patch simultaneously. However, the backward process cannot be decomposed according to patches, as the score function encodes their correlation; see Section for a detailed discussion on the structure of V log Pt.
Transformer Architecture A transformer comprises a series of blocks and each block encompasses
a multi-head attention layer and a feedforward layer. Let Y = [y yNv E Ti DxN be the (column) stacking matrix of N patches. In transformer block, multi-head attention computes
Attn =
M
V a Y(Q T)
where V Q", and Km are weight matrices of corresponding sizes in the m-th attention head, and is an activation function. The attention layer is followed by feedforward layer, which computes
FFN Y + W, ReLU(W, Y bb1 + 1
Here, WT, W2 are weight matrices, bj and b2 are offset vectors, 1 denotes a vector of ones, and the ReLU activation function applied entrywise. For our study, the raw input to transformer N patches of d-dimensional vectors and diffusion timestep t in the backward process. We refer to T(D, L. M, B, R) as a transformer architecture defined by
T(D, L, M, B, (f f fout (FFN, Attny) o... (FFN, o Attnn) fin Attn; uses entrywise ReLU activation for i = I....,,, L, number of heads in each Attn is bounded by M, the Frobenius norm of each weight matrix is bounded by B, the output range [| bounded by R).
See Figure 2 for illustration of fi and fout- For attention layers, we consider ReLU activation for technical convenience and postpone a discussion with softmax activation to Appendix D.5.
REPRESENT SCORE FUNCTION AS THE LAST ITERATE OF GRADIENT DESCENT
Diffusion transformer learns the sequential data distribution through estimating the score function. In this section, study how transformers effectively represent the score function of our Gaussian process data. We are particularly interested understanding how can transformers capture the spatial-temporal dependencies via multi-head attention.
4

---

Published as a conference paper at ICLR 2025
R Dx
Transfomer Block
Mead Attertion
Fend
Y1
y N
Bd N
(Ik))))))))))) VVy
X1
X
Enbeddiig
W2 Belu
oon R
Time (
Head N [0(("))) (x))
Reakdual Connectioon
Residral Connertion
Figure 2: Transformer architecture. Here fin is a linear layer to lift input patch to RD, which appends the input raw data with time index embedding and other useful information. After pussing through L transformer blocks, Jour projects each patch into the data original dimension TRd and clip the oulput range by R. We allow the output range be diffusion timestep dependent (denoted Re).
3.1 SCORE FUNCTION FOR GAUSSIAN PROCESS DATA
Recall from Section 2 that the joint distribution of our N-patch Gaussian process data is still Gaussian. We can show that at diffusion timestep (0, 7]. given V, [xI,...... N, in the backward process, the score function is
- (a @ E - aap) for Ot = et and 0; vI- (2) where μ the mean vector, and a and σr are determined by the diffusion forward process. We defer the derivation to Appendix C.I. Observe that the score function for the i-th patch depends on the evolution of patches, reflecting temporal dependencies in sequential data. Moreover, the influence of each patch is determined by the covariance function Fış = 7(h, h2)) In the extreme case of T = 1 = j}. i.c., there is no correlation between patches, the score function reduces (a I @ Σ + I Gaussian distribution N(J,, Σ) for each patch.
Vlog VlogP:(v.)
arrf),, which isolates patches and reproduces the score function of the
Although the score function assumes the closed-fooe form expression in (2), effectively representing may suffer from difficulties. In fact, there are dependencies among N patches and the correlation encoded by the inverse covariance matrix (o²T @ Σ + - Directly representing using a transformer is subject to high complexity and large dimension depending on N. To overcome the challenge, we resort to an algorithm unrolling perspective and leveraging the attention mechanism in transformers. The key insight is to approximate the score function as the last iterate of a gradient descent algorithm, where the algorithm can be efficiently implemented by a transformer.
3.2 SCORE FUNCTION IS THE OPTIMIZER OF A CONVEX FUNCTION AND GRADIENT DESCENT FINDS IT
We consider a fixed diffusion timestep E (0, T]. It is straightforward to check by the first-order optimality that V log.:(v.) is the minimizer of the following quadratic objective function,
V log (v.. argmin (s(s) 2 T (a
Σ) + of + (v. aaft)
(3)
SERAN
Examining (3) reveals that the objective function strongly convex as long as 0. Moreover importantly, the formulation (3) is free of matrix inverse and the optimal solution can be found by gradient descent algorithm. Specifically, in the k-tth iteration, gradient descent for L, computes
(k++) ( t7 ((a)(T @ Σ) + +o (*)) + +ve - (ve-a.)))
(4)
where 77, is a proper step size. Comparing to (2), the gradient descent iteration avoids the matrix inverse, and we can further decompose the update (4) according to each patch. With well-conditioned covariance matrix r @ Σ, the gradient descent algorithm converges exponentially fast for approx- imating the ground-truth score function. Moreover, could substitute r in (4) by any of its approximation I' We quantify the representation error after sufficient gradient descent iterations. Lemma (Gradient Descent Iterate Approximates the Score Function). For an arbitrarily fixed t E [0, T] and vi, given an error tolerance e > and any integer J < N, if I with F, = T 1 < J} is positive semidefinite, then running gradient descent in (4) with suitable step size for K = OKs log (1/c)) iterations gives rise to
8 (v)) V Vlogp((v)) <
+
Σ
V1 of
aspelle
Σ r |!j
G,G representation eroor
Eg correlation truncation enror
5

---

Published as a conference paper at ICLR 2025
where K2 k< Σ) ofl) is the condition number of a @ Σ) + o
The proof is deferred to Appendix C.2. To interpret the lemma, we first consider J = N, which implies &2 = 0 and & recovers the typical convergence guarantee of gradient descent algorithm: for strongly convex and smooth function, gradient descent algorithm converges exponentially fast.
Controlling the Length of Dependencies The special treatment in Lemma 1 concentrates on the truncation length J. On the high level, J defines the maximum length of temporal dependencies aim to model in the score function. In particular, instead working with (4), we consider a surrogate gradient descent iteration driven by replacing I by T, which neglects correlations beyond length J. We discuss sufficient conditions for ensuring I' being positíve semidefinite after Assumption 1. Introducing I' incurs the truncation error ε2, whose magnitude depends on the pattern of temporal dependencies. Apparently, with long-horizer dependencies, the truncation eπror &g tends to be large. However, advantage appears in the presence of decaying dependencies, since we can neglect faint correlation to promote the representation and learning efficiency.
3.3 BOUNDING CORRELATION TRUNCATION ERROR &g
In many dynamical systems, temporal dependencies decay rather fast function of the time gap. The decay pattern may lead to a well controlled truncation error &g As the decay is determined by the covariance function Gaussian process, we impose the following assumption.
Assumption Let el,.... eN E Rd be di-dimensional time embedding of hı, hy with |le: l = r for each such that there exists a positive and increasing function f with |le, e,ll2 f(i jI) 2 cli - -jl for an absolute constant > O. Further, the covariance function y satisfies
(h, h2) exp (- lle: - e,, for V [1.2] and 8> 0.
Firstly, Assumption I says that the time embedding (a.k.a. position embedding) preserves the gap be- tween real times. For example, common time embedding used sequential data modeling (Vaswani et al., 2017) is sinusoidal transformations, where e; [r siint /C), cos(2in /C)] R2 for positive radius r and a large constant C > 0. We can check that ||e: - e5l1, 2r sin jlt/C)
4rli j is positive and approximately linearly increasing for sufficiently large C.
Secondly, the covariance function decays exponentially fast and the speed is controlled by the exponent and the bandwidth f. Large E or small indicates that the correlation between different time indices decays relatively slowly. Thus, the sequential data has some long-horizen morghhorzzon dependencies. The range of includes the well-known quadratic-exponentii (Gaussian) covariance function (v 2). Moreover, when v = 1, the covariance function coincides with the correlation in Brownian motion. Varying v E [1, 2] capture abundant temporal dependency patterns.
Besides, under Assumption 1, we can provide a sufficient condition for ensuring î being positive semidefinite for any J. Specifically, when f S c", diagonal entry has a larger magnitude than the sum of the magnitudes of off diagonal entries. This implies that for any truncation length J, always positive semidefinite. See Remark in Appendix C.3 for formal justification. Apparently, requiring & < c² is not necessary and as long as the covariance function decays sufficiently fast, I can be positive semidefinite. We now establish the following corollary to demonstrate reasonable choice of J.
symmetric diagonally dominant, i.e., the
Corollary (Correlation Truncation with Decay). Suppose Assumption I holds with f S c", For any Amin (T), under the setup of Lemma 1, by setting J = = (((LlogNN/(co))))) holds that
(v)) - -V Vlog < 20, |v: axp4l12
The proof is deferred to Appendix C.3, where we keep the & < c condition for technical convenience. This should not considered restrictive, as our theory holds as long as the truncation of I at length J is positive semidefinite. Corollary shows that the truncation error Eg can be controlled at the same order of E,. The truncation length J is logarithmically dependent on the full length of the sequence, indicating that we focus relatively bandwidth f. This observation the key to promote the representation and learning efficiency of diffusion transformer.
dependencies proportional to the

---

Published as a conference paper at ICLR 2025
4 APPROXIMATION THEORY OF SCORE FUNCTION USING TRANSFORMERS
This section devotes to establishing a transformer approximation theory of the score function. Differ- ent from the existing universal approximation theories (Cybenko, 1989; Yarotsky, 2018), we construct a transformer to unroll the gradient descent algorithm for representing the score function. We show this perspective leads to an efficient approximation in the following theorem.
Theorem 1 (Score Approximation by Transformers). Suppose Assumption 1 holds with & < c². Given any to T) and a small c < Amin (T), there exists transformer architecture T(D, L, M, B, Re) such that with proper weight parameters, it yields an approximation $ to the score function V logge with
/ ((v)))
v log (log S So for all E [to,
The transformer architecture satisfies
D OOd de), L - ( log (Nd/((rr)))) M
B = (o(oo(NNd/(eom))))) 2 Nd((2 |> bo)), R₂ = =oog(NV/((oa)))NNddo)
The proof is deferred to Appendix D.I. Here we observe that the approximation error depends on Tee which matches existing results for studying score approximation using neural networks (Oko et al., 2023; Chen et al., 2023a; Tang and Yang, 2024). Yet we remark that our algorithm unrolling approach very different from these existing works. Our results also hold for softmax activated transformer architectures, which is discussed in Appendix D.5. We remark on other interpretations.
Score
log.)(v.)). arg min aERON
C((s) = 2
(af 60 Σ) + +1 + (v. -a114
Fuuction
Tre (s (0))
Gradient
00
r
(K)
Duscent
Lemma
Truncation Error
-n (s
Correlation Truncation
r
(K)
Approxinate One GD Iteration
Theorem
Input X
Tip Linear Offset
Embedding e;
5-1
Transformer
Algorithm
Iterate (0)
Multiplication Module
Trausformer Block
Transformer
(K)
Block
Unrolling
Mean P
Atta
FFN
Embedding
Input
Mean
Repent Times
Figure 3: Construction of score function approximation using transformer. By rewriting the score function as the optimizer of quadratic objective function, we use gradient descent algorithm to approximate the optimizer. We allow correlation truncation manipulate the maximum length of temporal dependencies to model in Lemma 1. Then we construct a transformer architecture to unroll the gradient descent algorithm for score approximation in Theorem 1. Each gradient descent iteration realized by a multiplication module followed by two transformer blocks. In the first transformer block, its attention layer calculates correlation iú utilizing time embedding. The second transformer block calculates the linear offset -maoos 1T(2, axj4a in (4).
Glimpse of Transformer Architecture The constructed transformer architecture demonstrated in Figure 3. To achieve approximation, the transformer unrolls the gradient descent algorithm. For realizing a single step, a multiplication module calculates time dependent rescaling parameters, e.g., a2 and 0, Then two transformer blocks implement the iteration in (4). As be seen, the raw input is lifted into a higher dimensional vector, containing time embedding and other useful information. It is worth mentioning that the total number of transformer blocks is proportional to
7

---

Published as a conference paper at ICLR 2025
the condition number Kt. For sharply decaying covariance functions, Keg is constant. However, for slowly decaying covariance functions, Kt, can be large, indicating the fundamental difficulty of capturing long-horizon dependencies. While explicitly bounding Kt, for finite N goes beyond the current technical limit, we discuss its asymptotic behavior in Appendix G.4.
What Is Represented in Self-Attention We zoom into our constructed transformer architecture in Figure 3 to understand the role of multi-head attention layer in Transformer Block 1. Here the attention layer is constructed with proper Q and K matrices, so that it finds the correlation between a pair of time embedding e; and e,. Specifically. it calculates the inner product e, e, for approximating the correlation coefficient ((,, h2). Interestingly, our numerical results in Section 6 support this construction, providing evidence that a well-trained diffusion transformer puts large weight corresponding to time embedding. We suspect that our construction provides practical insights on how correlation is learned in attention layers.
SAMPLE COMPLEXITY OF DIFFUSION TRANSFORMER
Given a properly transformer architecture, this section studies the sample complexity of diffusion transformer for learning Gaussian process data. As mentioned in Section 2, the training of diffusion transformer is estimate the score function. Conceptually, we can use quadratic loss,
argmin So Ev.-P, t|s(v,, t) V log pa (v.) dt.
(5)
However, this loss function is not directly implementable due to V Plogpt being unknown and numerical instability when approaches zero. Therefore, we consider the following loss,
argmingeT 1r T. E Be s(vv..) aivo Vi
dt.
Here, to an early-stopping time and Civa-Yi substitutes the unknown score function llggP Pt. The equivalence of (6) to (5) established in the seminal works (Vincent, 2011; Hyvärinen and Dayan, 2005). When given the collected data D {xi.j....,, XN.j we replace the population expectation Ev, in (6) by a sample empirical distribution. We denote vj = ,x. x x[,] as the stacking vector data sequence. Then the estimated score function can written as an empirical risk minimizer,
S E argmin + f- Ev.-Noovv
s(v., |
Ve
dt.
2
To generate new samples, diffusion transformer uses Sn the backward process. Correspondingly; we denote the distribution learned by such a diffusion transformer as P. We bound the divergence of P to our ground data distribution Po in the following theorem.
Theorem (Sample Complexity of Diffusion Transformer). Suppose Assumption holds with È < c" Assume there exists a constant C such that C-1 Amiu (Σ) < Amax (Σ) < C and || S C. We choose the transformer architecture T(D, M, B, R2) as in Theorem with 1/n. By setting the terminal diffusion timestep T log(n), considering P generated by the empirical risk minimizer S, we have
2 Nd
ED TV P)
p)){3 S
log
Se+1 2v
ndntt
where Po is a perturbed data distribution satisfying (PO,, Po) SdttoNd
The proof is deferred to Appendix E. We remark that the emergence of Po owes to the early-stopping in score estimation, which is obtained by evolving Po along the forward process for timestep (o- The optimal choice on lo depends on the condition number Keg- For fast decay covariance functions,
Kig (1) and can choose = 1/n and the generalization error is in the order of Ö(VeNdd/n) In this case, we have a relatively weak dependence on N, demonstrating the efficiency diffusion transfommers in sequential data modeling. On the other hand, when Ks, is large, i.c., with the presence of long-horizee horizon dependencies, the learning efficiency suffers from heavier dependence on N. We experiment with various decaying patterns in Section and show the corresponding performance.
8

---

Published as a conference paper at ICLR 2025
6 NUMERICAL RESULTS
6.1 EXPERIMENTS ON GAUSSIAN PROCESS DATA
In this section, we conduct experiments diffusion transformer for learning synthetic Gaussian process data. We test various aspects of our theory, including learning efficiency and the capture of spatoal-temporal dependencies.
Experiment Setup We consider synthetic Gaussian process data with to the setting where the covariance matrix is generated by Σ A A E R for a Gaussian random matrix A E RSxs, We set the mean vector la = O for simplicity. The covariance function is chosen as in Assumption 1 with ?(ha, h2) exp [h, hyl" /2)- We set & and as hyperparameters and report their influences on the learning. We generate E { {1000, 3200, 10000, 32000, 100000} sequences from the Gaussian process as our training data in different settings. The diffusion trans- former is implemented based on the DiT (Peebles and Xie, 2023) code base. We set the number of transformer blocks to be 12 throughout all the experiments. We modify the patchify module in DiT to cope with our Gaussian process data. Additional implementation details can be found in Appendix F.
and N = 128. We stick
Influence of Covariance Function Decay on Capturing Spatial-Temporal Dependencies We study the influence of covariance functions on the performance of diffusion transformers. We vary the exponent / and bandwidth as follows: 1) we keep & = 64, while v varies in (1, 21/4 21/44,22/2222/ 2//2,23/4 2): 2) we keep y = 2, while l varies in 14,,,,4,256, 1024). After training, we generate 1000 samples at each time step for performance evaluation. The metric is a relative error e of the estimated sample covariance matrix to its ground-truth (see definition in Appendix F), which is reported in Figure 4(a).
roladive for 64 and
relative error for and various
Q'K
10
y=1
10a
1-24
y= 222
2-256
10a
10a
16
64
Q₂ K
Q2
2
t-
10a
10
20
10
10a
Q.
Q.
10a
10 10P
10
10th
16
(a) Relative error under different n, M,
(b) Query-Key matrix
Figure 4: In panel (a), we observe that the relative error decreases as the sample size increases. Meanwhile, larger or smaller i leads to better performance, supporting our generalization bound in Theorem 2. In panel (b), we split Q to Q = Que Qel where Qu E g16x32 and Qe E IR 16x33 corresponds to time embedding e,, We also split the key matrix K as JK, Ke]. The sub-block Qº Ke has dominant weights compared to other sub-blocks.
Transformer's Query-Key Matrices Coincides with Our Approximation Theory We dive into transformer blocks to understand how attention layer captures dependencies. Inside a transformer block, the input is a concatenation of data vector and corresponding time embedding written as |2.,,, ] E 323 where e; E IR is the time embedding and Zi, E RR is the output of the patchify module in DiT, obtained by a linear transformation on raw data at diffusion timestep t. We plot the heat map of query and key matrices in the 5th transformer block in Figure 4(b). Plots for other transformer blocks are provided in Appendix F, As can be seen, the interaction between time embedding (bottom-right block) is dominant, which aligns with our approximation theory for constructing the transformer architecture.
To further under- stand how DIT captures the temporal dependencies, i.e., matrix I for Gaussian process data, we plot the evolution of score matrices (QY) KY in the attention layers at different steps in the backward diffusion process. Besides, we demonstrate the gradual change of score matrices in different attention layers. Specifically, Figure 5 presents that with progressive denoising in the backward process, the attention score matrix becomes more and more similar to the ground truth T. In addition, in the first few attention layers, e.g., the first and the second layers, the dependencies are not well-structured and as predicted by our theory, these layers are still realizing some transformations on their inputs. Starting from the 3rd layer, the attention score matrices gradually exhibit the pattern of the ground
Backward Diffuusion Process Unveils the Kernel Matrix in Attention Scores

---

Published as a conference paper at ICLR 2025
truth temporal dependencies. In further subsequent layers, observe that the learned pattern of temporal dependencies is kept. We refer readers to Figure 12 in Appendix for complete plot of score matrices in each attention layer.
1- 50
1- 100
t- 200
r- 400
1-
500
1- 1000
[-
1- 50
1- 100
t- 200
t- 400
1- 500
1- (000
Layer
Layer
Lager
Layer
Figure 5: We demonstrate score matrices in different attention layers and at different backward denoising steps. The learned temporal dependencies gain more and more clarity as the denoising in the backward process proceeds. Meanwhile, observe that the temporal dependencies are well captured starting from the 3rd layer.
6.2 EXPERIMENTS ON SEMI-SYNTHTTIC VIDEO DATA
To further demonstrate the capability of diffusion transformers capturing spatial-temporal de- pendencies, we consider learning 2D motions of ball. The motion is described by sequence of gray-scale image frames of resolution 64 x 64, which characterizes the ball that starts moving toward random direction in a cube and bounces back when hitting wall. Because of the bouncing-back mechanism, the dynamic of the ball goes beyond the class of Gaussian process and exemplifies more complex dependencies with abrupt changes. We train latent diffusion model (Rombach et al., 2022), where we first generate 20000 image frames for training a 2D Variational Autoencoder (VAE). The 2D VAE sets the latent representation in TR... Once the VAE is trained, we fix it and generate independently 10000 motion videos, each consisting of 240 image frames. The diffusion transformer trained on the latent representations the generated videos. As shown Figure 6, we observe that spatial-temporal dependencies well captured. We provide generated video frames in Appendix F, which features appealing con--consittency and satisfies the bouncing-beed -bak mechanism.
Coveranto malrs geomett
Tre-covaricc EEE2
0
512
23
455
108
0,
0.
20
20 30
80
(a) Learned spatial-temporal dependencies (left) compared with the (b) The query-key matrix of the first ground truth (right).
diffusion transformer block.
Figure 6: In panel (a), we observe that the spatial-temporal dependencies (represented by the covariance matrix) have been well captured by diffusion transformers. We collect 1000 generated samples to calculate the dependencies. In panel (b), similar to Panel (b) in Figure 4, the learned query-key matrix demonstrates dominant scores between time embedding (bottom-right block).
CONCLUSION AND DISCUSSION
We have studied diffusion transformers for learning Gaussian process data. We have developed a score function approximation theory, by leveraging transformers to unroll gradient descent algorithm. Further, we have established sample complexities of diffusion transformer and discussed the influence of spatial-temporal dependencies on learning efficiency. While Gaussian process data enjoys mathematical simplicity and is relatively preliminary, our theoretical insights and experimental findings can provide invaluable intuition to analyze and design sequential data modeling using diffusion processes. An interesting future direction is consider generic dynamic models. We expect broad and positive societal impact on advancing diffusion models for sequential data synthesis, forecasting and editing, including video and audio snippets.
10

---

Published as a conference paper at ICLR 2025
8 ACKNOWLEDGMENTS
Mengdi Wang acknowledges support by NSF grants DMS-1953686, IIS IIS-2107344, RMMI-1653433 and ONR grant 1006977.
REFERENCES
Juan Miguel Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecast- ing with structured state space models. arXiv preprint arXiv:2208.09330 2022.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34: 19981-19999,, 2021.
Anjishnu Banerjee, David B Dunson, and Surya T Tokdar. Efficient gaussian process regression for large datasets. Biometrika, 100(1):75-89, 2013.
Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang. Gang Yue, Yue Cao, Hang Su, and Jun Zhu. One transformer fits all distributions in multi-modal diffusion at scale. In International Conference Machine Learning, pages 1692-1717. PMLR, 2023.
Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising markov models. arXiv preprint arXiv:2211.0359) 2022.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. arXiv preprint wwx.iv2330.00666 2023.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear conver- genee bounds for diffusion models via stochastic localization. In Proceedings of the International Conference on Learning Representations, 2024.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, Varun Jampani, and Robin Rombach. Stable video diffusion: Scaling latent video diffusion models to large datasets, 2023. URL
https://arxive org/ abs//////////////////// / http./ooo/////.1...111111 2311.
Viacheslav Borovitskiy, Iskander Azangulov, Alexander Terenin, Peter Mostowsky, Marc Deisenroth, and Nicolas Durrande. Matérn gaussian processes on graphs. In International Conference on Artificial Intelligence and Statistics, pages 2593-2601. PMLR, 2021.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231 -37, 2015.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnad Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 15:28266-2827777 2022.
Clément Canonne. short note on an inequality between kl and tv, 2023.
Francesco Paolo Casale, Adrian V Dalca, Luca Saglietti, Jennifer Listgarten, and Nicolo Fusi. Gaussian process prior variational autoencoders, 2018. URL
http://xxxvv.o.g//bbaa//a /bbs//
Stanley Chan. Tutorial on diffusion models for imaging and vision. arXiv preprint arXiv:2403. 18103, 2024.
Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Nonparametric regression on low- dimensional manifolds using deep relu networks: Function approximation and statistical recovery. Information and Inference: A Journal of the IMA, 11(4):1203-122 1253, 2022a.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In Proceedings of the International Conference Machine Learning, pages 4672-4712, PMLR, 2023a.

---

Published as a conference paper at ICLR 2025
Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Ap- plications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.0777,, 2024.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209. 1115,, 2022b.
Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode provably fast, arXiv preprint arXiv:2305.11799 2023b.
Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degradaeee beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning, pages 4462-4484. PMLR, 2023c.
Yifan Chen, Bamdad Hosseini, Houman Owhadi, and Andrew M Stuart. Solving and learning nonlinear pdes with gaussian processes. Journal of Computational Physics, 047:1116600 2021.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems, 2(4):303-314 314,, 1989.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
Philippe Delsarte and Genin. Spectral properties of finite toeplitz matrices, In Mathematical Theory of Networks and Systems: Proceedings the MTNS-83 International Symposium Beer Sheva, Israel, June 20-24, 1983, pages 194-213, Springer, 2005.
Volker L Deringer, Albert P Bartók, Noam Bersttein, David M Wilkins, Michele Ceriotti, and Gábor Csányi Gaussian process regression for materials and molecules. Chemical Reviews, 121(16): 10073-10141, 2021.
Zihan Ding and Chi Jin. Consistency models as rich and efficient policy class for reinforcement learning, arXiv preprint arXiv:2309.1698; 2023.
Vincent Fortuin, Dmitry Baranchuk. Gunnar Rätsch. and Stephan Mandt. Gp-vae: Deep probabilistic time series imputation. In International conference on artificial intelligence and statistics, pages 1651-1661 PMLR, 2020.
Hengyu Fu, Zhuoran Yang, Mengdi Wang, and Minshuo Chen. Unveil conditional diffusion models with classifier-free guidance: A sharp statistical theory. arXiv preprint arXhv:2403.11968 2024.
Yuan Gao, Haokun Chen, Xiang Wang, Zhicai Wang, Xue Wang, Jinyang Gao, and Bolin Ding. Diffsformer: A diffusion transformer on stock factor augmentation, 2024. URL https:// // arxiv. abs 2402 06656.
Zhiye Guo, Jian Liu, Yanli Wang, Mengrui Chen, Duolin Wang, Dong Xu, and Jianlin Cheng. Diffusion models in bioinformatics: A wave of deep learning revolution in action. arXiv preprint arXiv:2302.10007, 2023.
Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint arXiv:2312.0666) 2023.
Philippe Hstruc, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. IDQL: Implicit Q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Philip Hartman and Aurel Wintner, On the spectra of toeplitz's matrices. American Journal of Mathematics, 72(2):59-366, 1950.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 63:6850-6851, 2020.
12

---

Published as a conference paper at ICLR 2025
Aapo Hyvärinen and Peter Dayan, Estimation of non-oooml statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Yuling Jiao, Lican Kang, Huazhen Lin, Jin Liu, and Heng Zuo. Latent schrödinger bridge diffusion model generative learning. arXiv preprint arXiv:2404. 13309, 2024.
Vladislav Kargin. Spectrum of random toeplitz matrices with band structure. 2009.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. arXiv preprint arXiv:2206.06227, 2022a..
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. arXiv preprint arXiv:2209. 1a 2022b.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.0925/ 10255,,,0..0,1,,, 2023.
Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably, arXiv preprint arXiv:2403.03852 2024a.
Zihao Li, Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Yinyu Ye, Minshuo Chen, and Mengdi Wang. Diffusion model for data-driven black-box optimization. arXiv preprint arXiv:2403. 132/1, 2024b.
Chenghao Liu, Enming Liang, and Minghua Chen. Characterizing resnet's 's universal approximation capability, 2024a. URL https::/p//.c wwwpopppeevveeeeeeeeeeeeeeeeeeee ew net. forum? ddppmmmmpoddcCoooooocococc o 2hv...........
Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, and Jianfeng Gao. Sora: review on background, technology, limitations, and opportunities of large vision models. arXiv preprint arXiv:2402.1717, 2024b.
Song Mei. U-nets belief propagation: Efficient classificatton,, denoising, and diffusion in generative hicrarchical models. arXiv preprint arXiv:2404.18444 2024.
Song Mei and Yuchen Wu. Deep networks denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309. /420, 2023.
Vishal Monga, Yuelong Li, and Yonina C Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2):18-44, 2021.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. arXiv preprint arXiv:2303.0186. 2023.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301 200.77331..0777, 2023.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
Giuseppe Da Prato and Luciano Tubaro. Wick powers in stochastic pdes: introduction. 2007.
URL https::///// /
/tt////pi. /api....... semant conmannicccholar org/ Corpporrrpsi :I: 55493217.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted Intervention-MICCA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer, 2015.
Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(00):69-106, 2004.
13

---

Published as a conference paper at ICLR 2025
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous- time discrete diffusion models. arXiv preprint arXiv:2211.167556 2022.
Rong Tang and Yun Yang. Adaptivity of diffusion models to manifold structures. In International Conference on Artificial Intelligence and Statistics, pages 1648-1656. PMLR, 2024.
Wenpin Tang and Hanyang Zhao. Score-based diffusion models via stochastic differential equations-a technical tutorial. arXiv preprint 20XXi:24002.04887. 2024.
Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-based dif-
fusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems, 34:24804-2481 2021.
Muhang Tian, Bernie Chen, Allan Guo, Shiyi Jiang, and Anru R Zhang. Fast and reliable generation of ehr time series via diffusion models. arXiv preprint 1rXiv:2310.15299 2023.
Paolo Tilli. A note the spectral distribution of toeplitz matrices. Linear and Multilinear Algebra, 457--9) 1998.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need. Advances Neural Information Processing Systems, 30, 2017.
Pascal Vincent. A connection between matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Yaohui Wang, Xinyuan Chen, Xin Ma, Shangchen Zhou, Ziqi Huang, Yi Wang. Ceyuan Yang, Yinan He, Jiashuo Yu, Peiqing Yang, Yuwei Guo, Tianxing Wu, Chenyang Si, Yuming Jiang, Cunjian Chen, Chen Change Loy, Bo Dai, Dahua Lin, Yu Qiao, and Ziwei Liu. Lavie: High-quality video generation with cascaded latent diffusion models, 2023. URL https: /ttp///aaivv..//ai///a/eeaeaawva org/ abs/ 2309. 15103.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for muchine learning, volume 2. MIT press Cambridge, MA, 2006.
Junde Wu, Wei Ji, Huazhu Fu, Min Xu, Yueming Jin, and Yanwu Xu. Medsegdiff-v2. Diffusion-barse medical image segmentation with transformer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 6030-6038, 2024.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint XXix:2203.009222 2022.
Dmitry Yarotsky. Optimal approximation of continuous functions by very deep relu networks. In Conference on learning theory, pages 639-649. PMLR, 2018.
Chenshuang Zhang, Chaoning Zhang, Sheng Zheng, Mengchun Zhang, Maryam Qamar, Sung-Ho Bae, and In So Kweon. A survey on audio diffusion models: Text to speech synthesis and enhancement in generative ai. arXiv preprint 2x3 2, 2023.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang. Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey, arXiv preprint arXiv:2311.01223 2023.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL EXPERIMENTS
We provide additional numerical results to support and validate our theory: 1) We present comparisons between diffusion transformers and vanilla diffusion models with UNet, and 2) We vary the size of the diffusion transformers and demonstrate the training dynamic of gradually capturing spatial-tempont dependencies.
Comparison of DiT to diffusion with UNet To conduct a comparison between DiT and diffusion
with UNet, we choose one instance of our Gaussian process data with d = 8 and N 128. The covariance function is Gaussian kernel with v = 2 and = 64. We collect n = 10000 sequences for training a DiT with 12 transformer blocks and UNet-based diffusion model with 4 down/up sampling procedures. Each of the down/up sampling in UNet consists of 3 residual and convolution layers that the DiT and UNet have approximately the same model size.
Each model is sufficiently trained for 400 epochs, when the training error has converged. In the testing stage, we collect 10000 samples generated separately from each model and estimate the spatial-temporal dependencies for comparison. As shown in the first column of Figure 7, both DiT and UNet-based diffusion model capture the decay pattern in the temporal correlation. However, DiT exhibits much better learning result, matching the ground truth. The temporal correlation of the samples generated by UNet-based diffusion model presents a "piecewise" pattern, not as smooth as the ground truth. We conjecture that it is caused by the size of the filter in convolution layers that prevents UNets from learning complete temporal correlation.
Moreover, DiT exhibits significant strength in learning the spatial correlation. As shown in the remaining columns of Figure 7, DiT successfully captures the spatial correlation between two tokens sufficiently separated, i.c., the correlation is rather weak. In contrast, UNet-based diffusion struggles in learning spatial correlation. We find clear inconsistent patterns of the estimated spatial correlation in the third row, not to mention that the pattern deviates from the ground truth. This result not only demonstrates the surprising learning and generalization capabilities of DiT in sequential data, but also indicates some advantage of DiT over UNet-based diffusion models.
Performance of DiT with varying network size and sequence length We study the influence of the sequence length N and the number of transformer blocks I. on the performance. We choose the sequence length N and the number of transformer blocks L follows:
1) we fix = 16, while change N {2 25, 2",
2) we fix N = 128, while change L in {1,2, 16).
We adopt a uniform Gaussian process setting of == and t= 64 as considered in the experiments shown in Figure 7. After training, we independently generate 10000 sequences for performance evaluation. The metric is the relative error e of the estimated sample covariance matrix to its ground truth (see a definition in Appendix F), which is plotted in Figure 8. As shown in the left panel of the figure, the relative error increases mildly as the sequence length increases, which aligns with our v VN-dependenn in Theorem 2, demonstrating DiT's strength in handling long sequences. The right panel of the figure shows that the performance of DiT improves the number of transformer blocks increases, yet at a marginally diminishing speed when is sufficiently large. Using& 8 transformer blocks suffices for efficient learning in this case, which approximately verifies our construction of transformer with (log N) blocks.
Training dynamic of DiT To further examine how DiT captures spatial-temporal dependencies, we present the training dynamic of DIT. As shown in Figure 9, the spatial correlation is captured relatively accurately at an early stage of the training (after the first few epochs), while the temporal correlation is learned slower (after epoch 50). We observe the gradual change in the temporal pattern to match eventually the ground truth. In particular, large temporal dependencies are casier to learn, while weak temporal dependencies require more epochs to learn.
15

---

Published as a conference paper at ICLR 2025
Turh Iespami condatiio
Spatial
Spatial
Spatial
46:01N
4.4-22H
Gpopperl ocnderkon of DrI cempies
Spetial cocrelction (10, Hi1
Spatial cocrelation (3), 301
Spatial cocrelation (30,
Tympurl Ieb Tiat
Special corteletim at (10, Il)
Spetial coctelation at (20, 301
Spotial corrllttin at (30, 501
pAdc
Figure 7: Visualization of ground truth spatial-temponts dependencies (the first row), spatial-temporal dependencies of the DiT-generated samples (the second row), and the UNet-generated samples (the third row). We visualize the temporal dependencies in the first column, while visualize the spatial dependencies of three pairs (i,j) E {(10, 11), (20, 30), (30, 50)} of tokens respectively, representing short-horizon, medium-horizon. and long- long-horizon spatial dependencies. For each pair (i.j), we compute the (empirical) covariance matrix between the i-th token and the j-token and visualize the matrix in normalized version Yi.j Σ, where we set the true spatial covariance E.
Σ[[F Ellp for a better comparison with
Relative ervor for varions sequeath lengths N
Relative ervor for varions depths
10-1
N=22
L-2
X-126
1-8
X-206
Relative error
Relative
10=
10
10a
10a
Number of training samples
Number training samples
(a) Relative error with varying N
(b) Relative error with varying L
Figure 8: Relative error with different sequence lengths and numbers of transformer blocks. We observe a mild error increase when increasing the length the sequence. Meanwhile, increasing number of blocks in transformer leads to initial gain of the performance. However, the performance gain saturated when the number of transformer blocks sufficiently large. These observations support theoretical results.
16

---

Published as a conference paper at ICLR 2025
Twined spoclts
Thined epodls
Trained
Figure 9: Spatial-temporal correlations of samples generated DiT trained with different number of epochs. We train DiT with 10000 samples and use Adam optimizer with a mini-batch size of 128. We visualize the temporal correlation in the first row, and in the second row, plot the estimated spatial correlation between the 40-th token and the 60-token using generated sequences.
B RELATED WORK
Our work establishes score approximation and distribution estimation theories of diffusion transformer with sequential data. Prior works focus on static data and provide sampling and learning guarantees of diffusion models. In particular, assuming access to relatively accurate estimated score function, Benton et al. (2022; 2023; 2024); Li al. (2024a; 2023); Chen et al. (2022b); Lee et al. (2022a:b); Chen et al. (2023c;b) show that the generated distribution of diffusion models stays close to the ground- truth distribution. Towards an end-to-end analysis, i.e., involving the score estimation procedure, Chen et al. (2023a): Oko et al. (2023); Li et al. (2024b); Mei and Wu (2023); Tang and Yang (2024); Jiao et al. (2024) all provide sample complexity bounds of diffusion models for various types of data, including manifold data and graphical models. Yet these results not directly applicable understanding how spatial-tempon dependencies are captured by diffusion transformers in sequential data.
For score approximation using transformers, we adopt an algorithm unrolling approach (Monga et al., 2021). In particular, view the score function as the last iterate of gradient descent algorithm and utilize transformers to implement gradient descent iterations. We are aware of Mei and Wu (2023); Mei (2024) showing U-Net performing algorithm unrolling in diffusion models.
On the empirical side, Diffusion Transformer (DiT) (Peebles and Xie, 2023) challenges the common choice of U-Net for image generation, providing state-of-the-arr performance. Moreover, diffusion transformer exhibits appealing scalability towards better generation qualities with larger model sizes. More recently, diffusion transformers are leveraged video generation (Gupta et al., 2023; Liu et al., 2024b), where the video data is patchified and rearranged into a long sequence. Besides, diffusion transformers are also used for other sequential data (Sun al., 2022; Austin et al., 2021; Campbell et al., 2022; Gao et al., 2024), such language, music and financial data.
C OMITTED PROOFS IN SECTION 3
Before diving into detailed proofs of our results (Theorem and Theorem 2), we list our assumptions for quick reference.
Data assumption. We consider Gaussian process data in Rd within interval [0, H). which defined in Section 2. The covariance function of the underlying Gaussian process verifies Assumption I. For technical convenience, we also assume the mean function pu.)) the Gaussian process can be efficiently represented by neural networks in Assumption 2.
Transformer architecture. We denote a transformer architecture by T(D, L, M, B, R), where D. L, M, B, R hyperparameters defining the size of the transformer. In Theorem 1, will choose these hyperparameters depending on the desired approximation accuracy. In Theorem we will further choose these hyperparameters depending on the training sample size n.
17

---

Published as a conference paper at ICLR 2025
C.1 DERIVING SCORE FUNCTION FOR GAUSSIAN PROCESS DATA
By the definition of Gaussian process, we know that the stacking vector [ , xx] E RRN follows a Gaussian distribution N(N. I & E). Along the forward process, we progressively add Gaussian noise to the initial data distribution and hence, for any Vr, we have
Pc(vi) =
(2moo) dNN/2 exp
lIve - voll
2
(A)
(-2 u) (T 1)) dvo.
exp
(2T) 12/2/N/2 det
(B)
Note that (A) is the Gaussian transition kernel corresponding to the forward process and (B) is the clean data density function. It is clear that Pe(v:) again a Gaussian distribution. By completing the squares and some algebraic manipulation, we have
pr(v.) x exp (-) (6)
+ ( Ev- -anu).
As sanity check, Pt is now a Gaussian density function of N(ae of of @ Σ), matching the marginal distribution of the forward process. Therefore, the score function is
V Poggr
(a T E +o (v. asp))
C.2 PROOF OF LEMMA 1
To prove the Lemma, first state standard result in convex optimization. Lemma 2 (Theorem 3.12 in Bubeck (2015)). Let f be 3-smooth and a-strongly convex on R& and x" the glohal minimizer. Then gradient descent with 7 a+B satisfies
(k+1)
x*
<
[(())
x" x+1,
k= 0,0,
Here x(+1) (x)
fx((()) is the outcome in (I: + 1)-th th iteration of GD and B/a.
With the lemma above, the proof of Lemma is quite straightworward.
Proof of Lemma 1. Denote the truncated score function by
s(v.) (a ( Σ) +oo1 (v. a..).
Let's consider the following quadratic target function
(s) 2 (T & Σ) + oJ (v: -it) s.
And iterate (k)) by gradient descent on this target function with s(() Since E, is Amax (a (T 2 Σ) + of smoooh and Amin (a & Σ) - -tronnlyy convex, by the approxi- mate GD formula (4) and Lemma 1, for any 2 0. we have
g (v+) s(v:))1
Re
5 (8)
(v:) s(vi) ||| - exp
2(k
[|5(v.)ll2
Here R is the condition number of o
+ I Moreover, we have T 36 Σ)
[S(v)))l2 < (a ② E) + +021) -1 -12 | a,u412
= = (a [
Σ) + of l| acHll22
1 - dillla
18

---

Published as a conference paper at ICLR 2025
Thus, by taking the number of iterations K [ log(1/c)], we have
| (v.) ( -S(vs)
|| - 02.eel126 ||V
oi
Besides, denote +, of [ Σ) If and क, of (T Σ) o.I. The difference between the truth score function and the truncated score function is
(Vv,) s svv.)ll2= =
(a @ Σ) + + - (a & Σ) + oft) -) (ve anti)
(,,- "7 )v an/4)l2
< क, क; 1l Ii. olllll2
2, ' (
2,) p; 11, llv 11.1lv=00 anHll2
< 9, 11 0, I12 p. $ llve - a:14122
S 0. 11 4 12 live aa
Let's focus on bounding [b.
Without any assumptions on d natural bound is
/
Φ $ 1| S
D Tellr at 1=1|F
Σ iーd
r
(8)
by
Altogether, have
[(x))vv)) - V Vlogpi(wi)(1
<
V/ - aolllll
+
- |E
Σ r
of
V
|iーsi>d
Σ F
Σ
r
2 live
auplll2
(9)
of
|
In the last inequality, we invoke of S for any t 0. The proof is complete.
GD with Approximation Error In the process of using transformers to express GD, additional approximation error will be induced, i.e., the update formula becomes
((++1) = (A) L) (())
Here {(k) represents the approximation error. Then we have the following convergence analysis:
Lemma 3. Given > 0, if the approximation error in each step satisfies [|5( ||| S e, then after
K #c+1 log (1/c)] steps of GD on minimizing the truncated target function L. we have
(Kt + 1)e
El of
(v.) - (logp
r パー-ز
Σ
of [|:
ac14l22
Proof of Lemma 3. By similar deduction in the proof of Lemma 1, for any k 2 0, we have
11+1)
(k) 1/ Cs (++ + (A)
=
2
2
(i) 1/ (s
+ ||| (k)
-
2
K1
||(0)) -
+
Thus, we have
(k+1) -8
(Nt + 1)e
1
(1s(00)
(st + 1)e
Kt +
19

---

Published as a conference paper at ICLR 2025
Note that if there exists & K such that
(s(h) - -5|12 - (Re + 1)c/2, we have for any b k.
|s(hh) - $|12 I (Ne + 1)e/2, so ||s(K) $ S (Kt + 1) which finishes the proof. Now we
assume for any S K S K, |s(k) - 8112 2 (Ki + 1) 1)e/2. Then we have
|(()] 8
(ri + 1)E
1
((00))
(Ng + 1)6
Kiz
+
2
Rit
(((0)
S
+
2k:
S
exp
s (0)
(10)
+
Substituting k = K [ log
(1/e)] into the inequality, we have
(Ni + 1)e
((())
e <| S0, l|ve-au -
Combining the result with (8), we complete our proof.
C.3 PROOF OF COROLLARY
Let's first analyse the approximation error by truncating T. For simplicity, denote AT r r.
Lemma 4. Suppose Assumption holds. Then for any 0, by taking J ( log (og(((())) 'M1
1, we have ||A||| || c.
Proof of Lemma 4. By Assumption 1, we have f(m) 2 cm. According to the definition of AT, have
| 1
Σ
T(1,,
ドーン
25(b)
(2N 2k) exp
exp
e
2(ck)
(2N
2k) exp
N-1
20"k""
< 2N
Σ
exp
E
k=
20
2N
exp
dt
8
J-1
20¹
1)2d
NC
by taking J (2e log (NE/e2)) 1/00 The proof is complete.
+1 (1=0(((o(NN/e))))) we ensure that ||
Remark 1. From the proof of Lemma 4, we also observe that I is diagonally dominant, i.e..
Te 2 Ej [Ty. when we have > exp 2f(k)) 1. According to the proof above, a
sufficient condition for this to hold is c c- < 1, i.c., & s c. When truncating I' by any length J, r
remians diagonally dominant, therefore, I is positive semidefinite.
Now turn back to the proof of Corollary 1. By taking
( 20
log (NE((o(&(')) 1/07
(((o))(N(((o)))))))
J =
20

---

Published as a conference paper at ICLR 2025
ensure that | S off?
Then plugging the approximation error into (9) concludes proof.
Moreover, we could bound the operator norm of I in the same way. For any v R such that ||vl2 = 1, we have
Tv
Σ
1.j
N
14 +
2 2exp(-22'') 12)
k=
N-1
N- k
N
N
$ + 2
exp(- 22c2
2
++
k=
N-1
$ 1142 2
1ap(-2e'k1/0
k=1
$ <1+c
Thus, we have
|| 51+c 1+
(II)
We will apply this bound Lemma 12.
D OMITTED PROOFS IN SECTION
D.1 PROOF OF THEOREM 1
By the previous derivation, we know the truth score function is written as
V Voggp(vv,) (o [ Due to the fast decay of D, consider truncated score function
Σ + o (v. - azp1).
(12)
s s(v)) = - (a [ + o (ve - acp1
with Fi Tiy if Ji - < J and Tij 0 otherwise. Denoting r = r + AT (see the formal statement in Lemma 4), by appropriately choosing M = ((log(1/e))) we could guarantee
|aT|| S for any > 0. Now we decompose the L2 error into the following two items:
S VlogPellLa(P)) S S $ (L. +
+|
V Po (L))))))
Proposition
Lemma
Here (ss)))(P))) I [ls (v. is the squared L2 norm w.I.. a density function Pt of distribution P.. We provide the error analysis for the two items Proposition 1 and Lemma separately. Under these two supporting statements, the proof of Theorem I is quite straightforward.
Now we state Proposition and Lemma 5 with their detailed proof deferred to next sections. Proposition provides an approximation guarantee on the difference between 5 and in (2(P)), which is our main result throughout the proof of Theorem I.
Proposition Suppose Assumption holds. Given to (0, TJ. there exists transformer architec- ture š T(D, L, M, B, R2) such that with proper weight parameters, it yields an approximation & to the truncated score function s with
2² S
S - 51 5
La(P)
7 t 2 to
oi
The transformer architecture satisfies
= 9d + d, + de + 1, L (1o, (og(NN/ = (((elog(Nd [Elle 1))
B = Ollog(Ndd(eeot))))oeeo 2Nd( + | b))) RS = ((lggNNd/(oo.))))
21

---

Published as a conference paper at ICLR 2025
The proof is provided in Appendix D.2. Moreover, Lemma bounds the difference between the truncated score function s and the ground truth in (2()) distance.
Lemma S. Given any > 0 and t > 0, by choosing
J o (flog(N |E ((or))
holds that for any & 0,
V log Pelle Pegleelnss(p)) S
e 0e
The proof is provided in Appendix D.4.
Now back to the proof of Theorem 1, recall that we can decompose the score error as in (12). Thus, combining Proposition 1 and Lemma 5 and adjusting the constants to bound the error by c/0r, finish the proof of Theorem 1.
D.2 PROOF OF PROPOSITION
To prove Proposition which is the backbone of the proof.
first need uniform approximation theory of s bounded region,
Lemma 6. Given any radius Ro 2 1, error level 0 < E and to 0, there exists transformer
architecture T(D, L, M, B, R)) that gives rise to s satisfying
(((v)) s(vi)ll2 c, for any [|s(v))l12 S Roo 2 to.
The transformer architecture satisfies
D 9d + de + de 11, I= ( log( Ro Nd) M = 4J, ?r² + | (0)
B = (ogg(Rod//(oo
Ri = O(Roo
The proof is provided in Appendix D.3. Then we could upper bound the second moments of the truncated score function $ in Pt.
Lemma 7. For any 0 < 1, by setting
1/07
elog
we have the second moment of the truncated score function bounded by | (2(P)) S
.2Nd
Moreover, we need to derive uniform upper bound of the true score function for convenience of truncation arguments.
Lemma 8. Suppose Assumption holds, then with probability exxpppppppcc (CC), the range of truth score function can be bounded by [|Viogp((v>)] < o. 2 (N +8VNd). Here C' is an absolute constant.
The proof's of the lemmas are provided in Appendix D.4. Now we are ready to prove Proposition 1.
Proof of Proposition I. By Lemma 6, obtain transformers such that [[((v)) s(ve)ll2 for any [|s(v:)]l2 S Rog,1 and i. 2 ln- Moreover, have |[((v))l2 CR for some absolute
22

---

Published as a conference paper at ICLR 2025
constant C. Now can decompose the Lg error as
$llix(ra) Ev. [[((v) s(v:)ll
[[((v)) s(v.)1l2 {I[s(v))la 1)
[[S(v.) - s(v.)) 1 (1s((v))l22 Root!! "
Ev.
C c (Il((v))ll> Roo,
2 2E,
07
+ 2 [|5(vi)]21 (1|s(v))l2 Roo |1]
S <6² + 20
Pr [|s(v))]l2 > Ro
of
+ 21,v [|s(v))11
1/2
Pr [Ils(v))ll Roo 11 1/2
x 20 Rg Pr [Ils(v.)ll> Roo
=²
C2 [[5(v:) Pr [[s((v))l2> Roo, 200
< < oi
2 + |1 + C2,22Nd Pr [[|s(v)]]l2 > Roo
Here we invoke Lemma 25 in the second-to-last inequality and invoke Lemma 7 in the last inequality.
By Lemma 8, choosing Rg = C1 CiVNooo [اll CN for some absolute constant C1, we can bound Pr [[ls(v:)]l > R by
2²
Pr [I(s(v))1l2 > Roo 4(Ro + 11
thus bounding the La error by O(e//o?) By adjusting the constant, we can ensure that there exists transformer architecture T(D, L, M, B, R.) that gives rise to S with
D 9d + de do
L (me, log(Ro.Nd/(eor)))) O log(Nd/(<ort)))) B= ( 2 + | )))
M 0 ((llog(dN [E /(<ar.))))
R₂ o ((lo((Nd((ooto))))
such that
|s - s|| (1)))())) S e² 307
The proof is complete.
D.3 PROOF OF LEMMA 6
D.3.1 TRANSFORMER ARCHITECTURE
Following Figure 3, construct our targeted transformer architecture
fout feD o...o fGD fpre fin-
Encoder For simplicity, suppose the encoder converts the initial input into Y fin fix XxI]..,NN]] =
, E DxN satisfies
y: [x] e, σt (1), 0 1,0
where ((t) = In, an, of aJ E R with d, = Here 0, and 0g2 serve as the buffer space for storing the components necessary for expressing gradient descent algorithm.
23

---

Published as a conference paper at ICLR 2025
Transformer Components Besides the encoder and decoder, fpre represents multi-layer trans- formers that prepare the necessary components for the gradient descent block, such as the mean μa; for each input token X;- faD represents multi-layer transformers that approximately express one step of gradient descent, which is the key component of our network. We elaborate on the construction of these subnetworks in D.3.2.
Decoder Suppose the output tokens from past layers has produced score approximator in matrix shape, we design the decoder as follows:
fout = fuorm o fiiner:
where flinear DxN Rd extracts d N block from the input and flattens it into vector to align with the dimension of the score function. fuorm IRdN 7 RdN controls the output range of the
network by Re, which can be mathematically written as
8, if ||s||2 S Rr,
fnorm (s)
Rc |
S, otherwise.
We remark that such clipping strategy is also applied other theoretical works Oko et al. (2023) to better adapt the magnitude of the score function different diffusion time t. since the score function of a degenerate Gaussian distribution could blow up as + O. Morcover, such clipping layer can be easily expressed by finite layers of feed-forward networks if we take both and R, as the input. See Equation 5 in Lemma of Chen et al. (2022a) for an example of the construction.
Raw Transformer Network In the proof of Lemma 6, we mainly focus on the transformer blocks instead of the encoders and decoders. We denote the raw transformer network as
T(D, I, M, B) (f f (FFN,, Attn,) o ... (FFN, Attn;),
The input and out dimension D,
Attn; uses entrywise ReLU activation for i = 1...., number heads in each Attn bounded by M, the Frobenius norm of each weight matrix is bounded by B).
D.3.2 APPROXIMATE EACH TRANSFORMER SUBNETWORKS
For simplicity, for the subnetwork fpre. assume that the mean vectors μa; can be directly expressed by a constant number of transformer blocks:
Assumption There exists raw transformer fa E Traw(D, Lys My O(d) such that for any input token y; [x]_ T,,TT (t), OT 1, 0,22" we have
faly.) [x]_ eT. . (t)) uT 1, 0.1
Here Ly and M₁ all constants. This assumption is mild, given the approximation ability of transformers.
After applying fus we begin implement gradient descent algorithm using transformer blocks. Starting at ;
(0) O₂ for i E [N], the first iteration can be written as
(1) ; (0)
0,5(0) + (x; (1t41) --mi(x (x
71
Thus, we only need apply a multiplication module fmult (ar, (4:
Ro aep, that approximately re- alizes the product operation. We defer the detailed construction of the multiplication module to Appendix G.3, which utilizes Ologg( max; [| |M /fmmtt) transformer blocks to reach the accuracy Il (at; pa) acp4illso 6mui for any E [N].
Lemma (Construct the first step GD). Suppose the input is Y [yT, ........... vN E RDN where each token is
y; [xxT,, c. Φ (f), 4d: p 1, TTI E +
24

---

Published as a conference paper at ICLR 2025
Given error level and learning rate 7/, there exists transformer fGD,1 E Traw (D, M,,,M.B. B) such that
JGD,1 = x, ,0), Smu (Th, X; Smutt (Cr, (((,,,,))))))))))))))(()))) (.)) 0
Smate (ae, pa)) 0gd
Here fi0 = ||//llooo and the two multiplication modules satisfy [mmute p4)) 0.p4ill2 S</VN
and also mutt (M,, (mmtee(ks (art,((t)))))) 14)) "x(x; a,11i)12 S S/VNN for E [N] and the parame- ters of the networks satisty
D = 9d de + de + I, O(log (1xx |10 Nd/c)),
M ==1, B (d([|x|l + [|s|l=o)).
Proof of Lemma 9.
To prove the lemma, apply Corollary 3 and use the last 3d dimensions 03a such that
the buffer space to approximate the product operation fmult R + R
XN
e
eN
(b(1)
((1)
fmult( (Y)=
Old fmutt (at, m1)
Oad fmulc (at, NN)
O3d
O3d
where fimlt (at, (4:) = 0s + Cus with [|Euell2 S c/vN. After obtaining faults we use one
transformer block TB FFN Attn with the attention block being trivial, i.e., all weight parameters being zero so that Attn(y) = y. We then choose FFN(y) = y + W2ReLU(W iy) with
W1
Id
(dx(1d+d+ad) Odx(ad+d+4d)
La
(xx(3d+1)
<Rd)
Is
-a (3d+11)
and
O O(add.+d;))) +dtd La
W
-Id
(2dd)
0 (174+1)x(d)
0 ((dd+1)x(d)
We further apply another multiplication module fuult,, rescale the gradient computed by the transformer block by the learning rate 7/. We can check that the first-step GD is represented by
faD, fmult, TB fmult.
Size of Transformer Blocks for Approximating The First GD Iteration We summarize in the following table the resulting network size of transformer blocks for implementing the first GD iteration when initialized with s; = O.
Table I: Transformer size for approximating the first GD iteration
Input dimension
N with 9d + d, + d, +1
of blocks L
1 2Lmult with Lumutt Olog
Nd/e))
# of attention heads M
Parameter bound B3
Od[[[x
Here the parameter bound B is obtained by noting that the magnitude of weight parameters in the multiplication module is at most O O[|xx|]ac + ||s|| and there are at most O(d) nonzero weight parameters each weight matrix. The proof is complete.
The next lemma presents the construction of next K steps of GD.
25

---

Published as a conference paper at ICLR 2025
Lemma 10 (Construct next steps of GD). Suppose the input token is
Yi = | e (t).s, 0 0 0 fmule(as, u.) 0,
then there exists a transformer fGD,2 E Traw (D, M, B) that upproximately iterates s; by following the GD update formula, i.e.,
N Σ
fGD.2ly.) = ,c. (T)).
ss m
in oî + (x: asptt + +
0 0 02 (aui mo) , 0
Here ||il|2 S c/vN for i [N] Moreover, the parameters of the networks satisfy
D 9d + de + de +,, L = Ollog (I|x || Nd/e)), M = 4J, B = 0 El +))) ||||
+ |s|l
Proof of Lemma 10. It suffices to construct several transformer blocks to represent one gradient descent iteration with truncated matrix r, which takes the form
(k+1)
(k) 8
TJt Σ a E)s
(k) + oss! +xi, - acpi))
8
(k)
1tt
2 r
Es
(k)
(n,02s/ + nx(x, aspli)
(13)
(B)
(A)
To ease the presentation, we consider fixed time and drop the subscript We also drop the superscript (k). In the following, will use two transformer blocks to approximate (A) and (B) separately. The input to those transformer blocks is Y .y.....yy E DDN, where D is a larger dimension and will be specified shortly. For each column vector y:, it stores copies of s, and other relevant information. Recall each token y, is
yi x eT_ Φ (1), 8 0ī 0 0 Smult (or, (4)
0 E 1
Here we reserve 0g as the buffer space for the multiplication module (the additional hidden width).
Multiplication Module Before diving into approximating terms (A) and (B), we use multi- plication module consisting of series of transformer blocks to transform each column vector y; into
y: x [ TT)), famute (a², sss (o?, fuute ,,,00 ,,,0%/l
where fmult denotes an approximation to the entrywise multiplication realized by the multiplication module. We defer the detailed construction of the multiplication module to Appendix G.3, which utilizes O[log[]ss|| /Emule)) transformer blocks to reach the accuracy [].fmut - -a>slla emult-
Compared to the raw input before the multiplication module, we have an easy access to the useful quantities ass, a²s, and au,, which simplifies our next step.
The First Attention Block for Approximating (A) Here we will construct transformer block TB, = FFN, o Attn, for approximating (A). Despite the huge dimension matrix T, is Toeplitz due to the uniform time grid ( .. ). There are at most N different entries in T and the
(i entry only depends on the gap j j|. As a result, we denote Tm = 1ii - jl = m) and rewrite term (A) as
N N-1 N
(A) - Σ
Σ
2 Ym 11 {li - = m|Es,
1==0 1=1
26

---

Published as a conference paper at ICLR 2025
The display above suggests a construction of multi-head attention layer. Formally, for arbitrary value of m, we construct four attention heads with ReLU activation. By Assumption 1, the indicator m) can be realized by calculating the inner product e, ej of time embedding.
function f To see this, we observe
1 7(202 2 _(2r² (2r2 lle, -e,l()) -2-
1
e, e,
(2 j()).
Therefore, it holds that
1 - = m) = 11 c o, = 1 p
since f is strictly increasing. Directly approximating an indicator function using ReLU network can be difficult. Yet we note that i j can only take integer values. Therefore, we can slightly widen the decision band for the indicator function. Specifically, we denote a minimum gap A min;=1,....N N- ( 1) (1)).. Thus, we deduce
oTT E 12 1 (()) -
1 A. ( 1
1
}
1 ( - J| =
m m]
We can use four ReLU functions to approximate the right-hand side of the last display. In specific, construct trapezoid funetion follows,
(ee e,)
8
ReLU
c.ee eTe,-- -2
f
1
8
ReLU
e_e,--"+ p fm)) 1
A
2
A
8 ReLU
c.e, ,2
1
(m)
8 ReLU
e e, ,² 2 (m) 1
A
2
is straightforward to check that v = in a 4 width interval centered at r² (m). To this end, can use four attention heads to realize the function sb. In particular, for the first attention head, choose
A
(Q T diag (OoEaa: Id. Od, xde x -r² +
1 (m) + A A (csurccaa and
To +. ((d))d)((((())))))) 0 (dd+dd+d)xd (dx(2d+d+td))
0 (c
V
A Am Σ
(6d+1)
((d4+))((2d++
1(dd+1)xd
((d++1)x
It is not difficult to check that this attention head calculates the first ReLU function in vh. Analogously, for the second attention head, we choose
[oaxa: It, Od.xd, ((5d)x(5)) + 2 P + 1 A. ( )) and
(Q") K diag
-²=-V'
for realizing the second ReLU function. The third and fourth attention heads have the following parameters,
(Q²) T K 3 = diag ([ossa Id,, Od, ((od) ((d)) -r2 2 2 ((() A, (23)) (A.0a)xr
8
V3_ -v
(oosea It., xdexde ((d)) -r² 2 + 2. (m)
1 (A.00axxaa)) and
(Q*) T K diag
4
V4 - vl.
By summing up the output of the four attention heads, we derive the output of such an attention layer. For the i-th patch, the output is
x 7 Φ (1), Si, Smmttt (a² ( Smule (0²,
7m (n (eTT,))T))))))))))))))>))))))) Simutt (a² 83 ())) Smute (a,, us) 1, 1,47
27

---

Published as a conference paper at ICLR 2025
Note that this output calculates for fixed value of m. In order summing over different values of m, we utilize 4J attention heads. Here 4J attention heads are enough, since Fış = 0 - 2 We have
(y:) ee .6 (1), S;, (a ), fmult (o+ s,),
N-1
Σ Σ
7m im ej,E Smule(a? 83) maute ue)) 3d
For the feedforward layer FFN,, W O, W2 0, b, 0 and b such that the output of the first attention block is
(B: = x e
ee, Φ (t), Si, (a ; (mmlit(² (x²
N-1
Σ (e, e,) ( m=
(a s8] famli p4) O
The Second Transformer Block for Approximating (B) Similar to the first block, the goal here is to construct 7B, = FFN2 o Attna for realizing (B). This is much casier than approximating (A), we only need the feed forward layer while set the attention layer trivial. In particular, choose Q K and V being all zero matrices so as to maintain the input to the attention layer. For the feedforward layer FFN2, we choose W
Ia
Od, xd.
Od,
La
Odxd
I Ia
La
Ls
0 (3d+1)))
Ia
0d.xdo
Odexds
la
Odxd
Ia
I
0xx
Ia
Ia
W1
-Id
R (sd)xD
I,
-aa
Ia
IIa
where the missing values all zero. Then we have W Attng TTB (y:) as
s; -N11 Σ j=1 7m (e E fmule (a sy) + fmult (o², (x, fmult (a, v4))
Ej=1
2m= (e ej E. f muntc (a² 8j) ) Smult (o, s:) (xi fmult (ar, mk))
Smult (a² (8))
S;
fmultt (a² 22.55 S;)
fmnte (a s;)
Smaate (o²,
mm=0
N-1
Ym (e, e, fmule (a²,
-=1 7m (e, e, E. fmule (a?,s,)
m=0
suffices to choose b by = and W₂ equal to
Oldtd, (x(((()))
Ia
La
W₂-
Ia
I
R Dx(8d)
Ia
La
Ia
I
0(ad+1)x(84))
where missing values are all zero. Using the fact that ReLU(x) - ReLU(-=)) = we have
W2 ReLU(W, Attn2 o (y:())))
N-1
02 OT TH S.
Σ
Σ
Yon im (e,eee)) (aui - mmmnt( (x) - -mutt (ar, pr)) ))
m=0j=1
N-1 N
Tim (e e, Smutt (a² ;8)),, - -mmileaa (8)), -Smute (,,, T
m==0 1=1
28

---

Published as a conference paper at ICLR 2025
Therefore, by concatenating the two transformer blocks, have
TB2 o TBi(y.)
x x,
(1),
N N
Σ
I'm (ee e fmilt (a² (a²,s) (o -x - fmmltea p4 )),
-xx
mh 0 jm
8, 0g 02 Sumit (a, p4) T 0gd
Lastly, by implementing another multiplication module fmult that scales the gradient by the learning rate Tre we obtain
fmalt TB2 TBi(y.)
N-1
Smult
A1:
Σ
Σ
I'm (ee ej) Smatt (a? ,s]) Smute (on (x] Smutte (1,,,,,,) v))))
m=0j=1
s, OJ 0 fmult (as, p.) 1, 0
It is straightforward to implement a feedforward layer to sum up S; with the gradient increment in fmult TB2 TB Clearly, the feedforward layer can be realized by a transformer block with trivial attention layer similar to TB₁ with much simpler weight matrices. As a result, we have implemented one gradient descent iteration, where the output replaces s; in the initial input vector. We denote faD,2 our network implementation of one GD iteration.
We can repeat the multiplication module followed by the two transformer blocks approximate the gradient descent for K iterations as needed. Note that, we don't need to calculate f p4) after the first step of GD, which leads to the construction fGD.2.
Bounding Approximation Error Examining the output of our approximation for one gradient descent iteration, it is not difficult to observe that the approximation error determined by the error in the multiplication module. For the k-th iteration, denote the approximation realized by the transformer as
8 (k+1) (.()) + fmult Thi Σ
Σ 7' (eTT,e fmule fmmit(a?
m==
Sumle (0² 20,, (.(()) (x7 Smaut (ar, m;))
We compare the output with the exact GD update with last iteration at (t,)), which
N-1
(k+1) -s(A)
Te Σ
Σ
mm (e, e,)= 25(( d (x] aqp)).
j=1
For any index E fl...., ......N]. N). we bound the difference
N-1
(k++)
ΣΣ
E(a²s!!) (a
8
Cault + T
+ 4/ 0 20 (k) Smalt (a2 ((()) || + me llape - fmaute (o, p3)ll2
(86)
mullt + 1/. EE12 N / + 21/, de where in inequality (i), the fact that Ym 7m (e,;" e,) 1, and in incquality (ii), plug in the
approximation error of multiplication module. By setting Cault = n | = 333/2 1/2 we
ensure that (o(k+) -((++1) Se/VN.
29

---

Published as a conference paper at ICLR 2025
Size of Transformer Blocks for Approximating One GD Iteration Similar to the proof of Lemma 9, we summarize the network size for implementing one GD iteration (except the first iteration).
Table 2: Transformer size for approximating one GD iteration (except the first iteration)
Input dimension
D x N with D = 9d + d. + d. +1
of blocks L
Lmuis with Liuut = O | [ Nd/e))
# of attention heads M Parameter bound B
4. with J the covariance truncation length
O (d(( Σ o0 + +r)) |
To see the parameter bound B, we observe that the parameters in the constructed transformer are bounded by max (1, 2 |>|| )²). Here, || maxy |Ea denotes the maximum magnitude of entries. Since there are at most (dd nonzero weights in each weight matrix, we have the norm of the parameters bounded by
(dd(s-1 |> + +r)))
|x + ||||l=)) (d()EE) + +)²)) ||x| + ||s|loo).
The proof is complete.
D.3.3 FORMAL PROOF OF LEMMA
Recall that the score function is written
s(v,) = (a & +o (v. - a.4))
Thus, [[s(ve)]l2 & 7 we have [|v 0,1ll2 - of x Σ) + o [s((.))l2 S CR where
= (T @ Σ) + o 2
< (T 8 Σ)|, + |
| |E +1
$ /N( + |
The last inequality follows from Lemma 4. Thus, the infinity norm of v, can be bounded by
[[velloo < | arpll + ||a/ا|| I &Roo, + 14. Moreover, by the proof lemma 5, we can also bound the norm of the truncated score function by
[|S(v.)]] S [[s(v:) - s(v:)ll2 + [|s(v.)]l2
S 4 |r IV, - aa/4l2 + ||s(v.)]lz
0 | (CRoo,1 + + + Roo
Thus, according to Lemma 4, by choosing O((elog(Ne)])) /o.))) we can ensure that [|s(v:)]] < 2R
Now we construct the transformers as follows:
5(x)= fnorm Jíinear o SGD o o.n, ofGD, o Jn o fi-
(K-1)xJeD22
Here four extracts the (d + de + de + 1)-th to (2d + de + d:)-th rows of output as score approximator, and we choose the clipping range as R = 2
Moreover, let
funear fGD.2 FGD fGD.2 fGD.1 fp fin to be the score approximator without clipping, then under the condition that [|s(vo)l2 Roo have [}((v.) - s(v.))l, & (vv)) - s(v.))lz-
Now let's bound the difference between šu (vt) and s. By Lemma 9 and Lemma 10, So(v)) expresses the output after K steps of GD with error level VN s/vN- Here multiply VN because
30

---

Published as a conference paper at ICLR 2025
considering the entire score function instead of each patch. Then by Lemma 2 and following the proof of Lemma 3, we have
[[o((v)) - =(vi)ll2 - (Ki 1)E
2K
+ texp
||((v))ll2
Ait +
(Ni + +1e
2K
+ +exp
Ro
Fig +
Thus, by replacing the error level by 2e 2Co + 1) < 2</(xe + 1) for sufficiently large constant C' and K - log( (o,e)) we have || - s(v,)llp I c for any 2 tg. Moreover, we remark that each patch of the score approximator lies in the euclidean ball with radius
(*)
(ve) S - 2 - 1 + 4Ro0,
throughout the network. Here s (k) (v:) represents the i-th patch of the score approximator after k GD blocks.
The Overall Size of Transformer Architecture for Approximating The Score Function We combine the network sizes in Table and Table 2, which gives rise to a characterization of the overall network size for approximating the score function of Gaussian process data. The result is summarized in the following table.
Table Overall transformer network size for approximating the score function
Input dimension
D N with D 9d + de + d, +1
of blocks L
Lμ + (2 + Lmultt KK = O (1i,, log No
# of attention heads M
4J = =0 ( (elog(NN || /010)) 1/0
Parameter bound B
log (o () Noooo (1² + |8
Output range R₁
2Ro0,"
Note that to obtain the bound on L and M, we substitute the choice of J and K into the network sizes in Tables and 2. The parameter bound is obtained by explicitly evaluating the magnitude ||s| of the score function in Table 2. The proof is complete.
D.4 OMITTED PROOFS OF LEMMAS ABOUT THE SCORE FUNCTION
Proof of Lemma 5. Suppose V, P. = [a Σ) +ofl)) Denote 4. too and . = af² @ Σ) + I. Then we have T, = 3, + af & Σ). Moreover, by Lemma 4, choosing J = O(log(NN N |E (((ot))) ensures that || S E | F of Thus, we have
||
V Vlogp:ll2,(2)) Ey,
[[((ve) V lo
-
|1( [- r (v a:11
(0.1 r. 1/2
tr ( r, r
④ : - r
< ( r. [r
5 - r.''l
< (Φ, - -
( - r
r 1
| (Ar
T ② Σ) +0
||
[af(T
||
& >>اه> [| & m) +
at ||
||
of
31

---

Published as a conference paper at ICLR 2025
In the last inequality, we invoke || I e| oî. The proof is complete.
Proof of Lemma 7. Suppose V1 P. = (a of +) +) Denote 4 of + and . = [T & Σ) + I.I. Then we have
| Ev,
|| v+)
V Vlogpo(ve)(22
Ev.
[
&, 1/2
'T, -1,1)
|
F
X 6 (Ar
+ 6 F
(1 12 a
[A Σ)[lF +
vNd
VNd क 1|
af || P |Σ
VNd
+
Nd
of
of
VNd
(c+ NN
of
2Nd
of
Here we invoke |T S | of by Lemma 4 in the second last inequality.
Proof of Lemma 8, Note that if ~ P. we can write v, = act4 1/2 z, where o, - aaTTE=oo and (0, I). Thus, the score function can be written
V llggs(Vv.) = Φ. -v. ox144) +> 1/2 2.
Thus, we have (v,)
z =, which quadratic form of the standard Gaussian. By in Lemma 24, we have
taking gzz) = V log (v:)
P V lo.):(v:))l B[/ log 2 8V Var([/Voogp.(v:)] 2exp (-C20).
Since we have
Ejll log/Vlopp. (v: (x I Ndo,² and Var([[VVog | 1| $ Ndo,
have |/ log so (N with probability 1 2exp (-C20).
D.5 RESULTS ON TRANSFORMERS WITH SOFTMAX ACTIVATION
In this section, we showcase the capability of softmax transformers in unrolling gradient descent algorithm. In this section, we consider the quadratic kernel function, i.e.,
(h, h2)) exp (- ||c, ,, /t)
In the following lemma, we reproduce our results in Lemma 10 with softmax activation. Moreover, due to the exponential scaling strategy of softmax, we show that one head of attention layer is enough to express the GD with the target function without truncation on T.
Lemma 11 (Unroll GD in softmax transformers). Suppose the input token is
y; [x]__ c, Φ (t).s, 0 0 fimate (as, p4) 1, 0dd
32

---

Published as a conference paper at ICLR 2025
then there exists softmax transformer sootsof Traw (D, Olog (||x|| || Nd/e)), 1, O(N2 + d + | + s)soomaax) that approximately iterates by following the GD update formula, i.c.,
x; 0, T (t), sf --1 ? Σ (a E oJ) (x, acpp;) + Ci. soDtaotmmx(()) (y:)
0 0 02 (au p4) 1, 0 1:
Here ||<|l| S Se/VN for E [N]. Compared with Lemma 10, we add an additional time embedding i at the end of the input for technical convenience.
Proof of Lemma 11. The proof is similar to that of Lemma 10 but with different construction on the attention layer. We also decompose the gradient descent iteration
(k+1)
(k) =s, T)t
(K) + ofss!" T E),,
+xi, acti))
5
(k)
TJ
af Es
(2)
(noiss (k) + nx(xi aspi))
(14)
(A)
By the proof of Lemma 10, part (B) only utilizes feed-forward networks, thus no changes need to be made. We will elaborate on approximating (A) using softmax attention layers. To ease the presentation, we consider fixed time and drop the subscript We also drop the superscript (k). Recall each token is
yi x 6 Φ TT(TTT()))))))))(
(1).s, 0T OJ, 0T fmut (ot, 14)) 1, 12 T E 1d +d++1
Here reserve 0f as the buffer space for the multiplication module (the additional hidden width). Before diving into approximating terms (A), we follow the proof of Lemma 10 to use a multiplication module consisting of a series of transformer blocks transform each column vector y, into
y1 y== [xx] ee, o6 (t), fmult (a [,, fo s. ),
Oo fmult (at, pu) 1, 0,
where fmult denotes an approximation to the entrywise multiplication realized by the multiplication module. We defer the detailed construction of the multiplication module to Appendix G.3, which utilizes (log((1 /c m))))) transformer blocks to reach the accuracy Umuut(a2, Here we could still utilize the multiplication module after changing the activation function attention layer to softmax because its construction only relies on feed-forward networks, the attention layers trivial throughout the network in multiplication module. See Appendix G.3 for more discussions. Now we begin to approximate (A).
a²s S mmult
po
can simply set
Comparison with ReLU activation In softmax transformers, we leverage the exponential scaling mechanism of the softmax activation to directly construct the entire kernel matrix I without truncation, so only one head of attention layer needed. On the contrary, in ReLU transformers, we only construct the main diagonals of I with multiple attention heads. See Lemma 10 for more details.
Approximate (A) Here we will construct a transformer block TB FFN; Attn, for approxi- mating (A). Our construction depends on the following fact:
N
(A) = na?
exp(- - lle; - -
1=1
naa2
exp((2r2 - 20, {)))/{{)Es
1=1
11a² D, - e /l, ,2r -2e, (e]Te)/()) EEs,.... .... Esv],
33

---

Published as a conference paper at ICLR 2025
where Dj Ej=1 exp( - 2e, ey){() is the normalizing constant in the softmax activation.
Thus, we can use one attention head to first realize the unnormalized version:
(A') a² softmax - 2c, e1)/e, ,2r - 20, ex)/ [Es..... Esn]
Use Attn to approximate (A') In particular, for the first attention head, we choose
(Q) K diag (Oaxd, 28 LId. Od, xde O (6d) 2r ((a)) )
and
O (1d+t,
(d)
V
(dx(2dtd=td)
E
(Bx++)+1)
0 ((
0 (44+1))x
(4d+1)x(66+1))
Thus, we have
Attni = x o x (t),,, St, fumli (a²,,)) Sumlt(o?,
N
D
exp((2r2 - 20,e;)/EEE (a ()), famile p4) 1, 12
j=1
To multiply our constructed value by the normalizing constant Dj, will construct a feed-forward layer FFN, such that the output of the first attention block is
TB =
x eT. Φ (t), fmutt (a², (,,,s]) fuu ss
DT exp( 2e,e;)/E)E fa .s)) fmmlt (ats pz) D,, 0 3d-21
(15)
Here D, is an approximation to D,. Now we show the approximation strategy to approximate D,.
Use FFN to approximate D, Note that we can rewrite D,
N-1
D. 1+ [1
Σ
k})g(().
Where o(k) cxp(-fffi f()) /E) is the correlation function. Let D, (m) = 1+ Σ 11)
1(i & (k) to be an approximation of D,, we know by the proof of Lemma 4, for any i E [N].
N-1
(-
2c 2cmm² m²
D(1) - (i)
Σ k=m+1
g(k) s "
l
Thus, to derive an error bound, we only need to be in the order of O((llgg11//0) (1)
For each k. we have
(1 > k + 1) + 1(i N (k))(k)
(())(R - k) ReLU(i k 1) ReLU(N k+1 -1) ReLU(N - k: 1))
Thus, the entire function D( can be expressed by 4m = o ((Elogg (/) ReLU neurons the feed-forward network. Specifically, we take FFN(y)
holds for any integer
W WeRelU(WiN)) W, and choose W1 as
D-3--1 D-3-11
-1
3d-1
2
3d-1
W1
0
D-3-11
-m
1
0
3d-1
I.. (2m++))D
0
0-3d-1
0
3d-1
-1
N
0
D-33-1
3d-1
34

---

Published as a conference paper at ICLR 2025
Then we have W1 Attn (y:)
i - 2
m N i
N
We choose b1 = 0, b2 (0 3d 3al- Denote
Wg - (o(1),9(2) - x(1),g(3) ((). ((m) - gm( - 1), (m(m]
and let
W = 0 3d
Wg g W9 (3d)x (2m
check that this exactly express D(m) and store it in the correpsonding location according to the equation (15). Taking TB = FFN o Attn, we reach our target in equation (15). After that, we could apply another multiplication module fmult to multiply the gradient component by D;, which gives rise to
Smult TB = x, ,, 15 (1), 8i, Samit(a ), Smult (o²,s
N Dr,D,'
exp((2r² 22,0;)/EfE fmall (a² (a², (a².s))) }
funll
fmult (a D, 0
Then, we can completely follow the proof of Lemma 10 to constuct the gradient and scale the gradient by the learning rate The and the approximation error the gradient only adds an additional term induced by approximating D;, which has been well controlled by setting m O ( (elo (1/c))
to reach error level of We remark that D, can be approximated for only once throughout the transformer networks and we can reuse each GD block.
Size of Transformer Architecture with Softmax Activation We list here the size of transformer architecture equipped with the softmax activation. Readers may draw a quick comparison with the size of ReLU transformer in Table 3.
Table 4: Size of transformer with Softmax activation for approximating score function (with Gaussian covariance function)
Input dimension
D x N with D 9d + de + de +
of blocks L
2+ Limult
# of attention heads M Parameter bound B
ONN + d
ΣΣ )²)
Output range Rt
2Ro0.
It is observed that using Softmax activation leads to reduced number of transformer blocks. The reason behind is that the Gaussian covariance function can be approximately represented by one attention block. The proof is complete.
With Lemma and following the proof in Appendix D while omitting the part of truncating i.e. setting J = N, we could similarly construct a transfoorrssssrmerss as
Seoftmaxx Sout SGD,softmaxx
JGD,softmax o fin.
such that [anftmax loggelllgg()) S o, for any 2 ln-
35

---

Published as a conference paper at ICLR 2025
E OMITTED PROOFS IN SECTION 5
E.1 PROOF OF THEOREM 2
Throughout this section, assume Assumption holds, and there exists absolute constant C as in the statement of Theorem 2 such that C-1 & X(E) 2 C and | S C, which means
| = O(1), 1 = (1), | = O(vd) and || S V.d.
Let's first analyse the effect of carly-stopping time to on the accuracy of distribution estimation. The following lemma presents the distance between the truth distribution P and the early-stopped distribution Pto W22 distance.
Lemma 12. Suppose P NM, I & )), taking the early stopping time to o(1), we have
( P..)) & (1 ato )² 1 + Co.Noo
Here C Amax (T
Σ) VI Σ 8 & / (11).
This means that W2(PP & Co We defer the proof to Appendix E.1.1.
Then state our theory objective function as
estimation error. Our estimator is chosen minimize the
Σ
(v(v)) - V log v. (/
6 = arg min FrEF
for each time step t j0, T). The next proposition provides guarantee the score estimation error, which we will transfer into the distribution estimation error later.
Proposition 2. By taking ndv in Theorem have the expected score estimation error bounded by
Ep[((3)] & (T log(L1to)) log (kandN/to) 411// Alu
n
The proof is provided in Appendix E.2. Although our assumption does not ensure the Novikov's condition to hold for sure, according Chen al. (2022b), long as we have bounded second moment for the score estimation eгror and finite KL divergence w.r.t. the standard Gaussian, we could still adopt Girsanov's Theorem and bound the KL divergence between the two distribution. We restate the lemma as follows:
Lemma 13 (Proposition D.1 1.1 in Oko et al. (2023), see also Theorem in Chen et al. (2022b)). Let Po be probability distribution, and let Y = (Yt)eejo.r) and Y' (Y7) be(oo)) be two stochastic processes that satisfy the following SDEs:
dW,, Yo ~ Po dY'' s' (Y/, t)dt + dW,, Yg ~ Po We further define the distributions of Y, and Y by Pt and pt. Suppose that px(x) [|(s s') (1)) s
dY, s sYY,,
(16)
for any t E [0, 2j. Then we have
(PT < /-1 (x( (ss - s')( 1) dxdt.
Now we are ready to prove Theorem 2.
Proof of Theorem 2. Now transfer the score estimation error to TV-distance bound using Lemma 13. Note that under our assumptions and for any s F, we have
I (v.)) (vs) - pv.) & /. po(vi)
|Ive aspall2
C' dv, S of
1
+ of
of
36

---

Published as a conference paper at ICLR 2025
Thus, the condition (16) holds for to 1 S T, which means that we could apply Girsanov's theorem in this time range. Remember that the backward process is written as
V : dt + dW with X N(O,I).
We denote the distribution of X PF_ In real setting, we replace Vlogp, by its score estimator S, which gives rise to the following backward process:
x (X; dt + dW with X N(O,I).
We denote the generated distribution of X as Pr_.. Besides, we consider the truth backward process the inverse process of the forward one, which is defined as
dx
V loggpT-.XX dt + dW, with X Pr.
We denote the distribution of X by P/- then we have P. P. for any 1 S T.
Since X and X are obtained through the same backward SDE but with different initial distribu- tions, by Data Processing Inequality and Pinsker's Inequality (see e.g., Lemma Canonne (2023)),
we have
TV(Pt, Pt.) TV(P) PTo)
$ (L((P% (1P))
(K(.P)
S (KL,P) exp(-T).
Thus, we could decompose the TV bound into
Pto, S TV PT. + TV(P: Pro)
$ exp(--T) //11 J. pe(v.) ()(v),y,t) - V log (v))l (17)
Thus, by taking expectation over the dataset 'D and invoking Jensen inequality, we have Ep TV P(o) & EDTV(PL Pi:)) EDTV(P, Pr..)
x Sexp(-T) (T (1)g(1/))))3( log (1andN(1))) 4+1/u Nd
n By taking T = Ollog n) and combining the results in Lemma 12, we have completed our proof.
E.1.1
PROOF OF LEMMA 12
Proof.
Since P (H, I & Σ) and P. NNaau, af @ E o,I), by the formula of the W2
distance between two multivariate Gaussian distributions, we have
(P, Pto) = Il - a. 12 + ( x 1) - (a & t
Suppose
& Σ DP, where E RdNx an orthogonal matrix and D
diag(>1, ..... daN) is diagonal matrix with A; 2 0. Then have
(T 1/ (afr Σ + 1/2
p 1 P - (
11/22 (o D + 1/2
Nd
of(2; - 1)2
(VTe + Vatd: o7)²
of2 i=1 of
Nd
< NdC.,o
37

---

Published as a conference paper at ICLR 2025
By taking = tu and plugging the inequality above into the expression of (P. (P. Pto), we complete our proof.
E.2 PROOF OF PROPOSITION
E.2.1 ADDITIONAL NOTATIONS
For any score estimator S,, we denote its population loss as:
((.)) = /- dt Eve. ([)((v)) - V Plogp.
and we define the empirical loss as
((S.) _ Σ dt vo (v(v,) + Vi VgVo
of
Here, D D-
(v(0)
i.i.d samples from true distribution P(= Po). When taken expectation
over the choice of samples, we have Ep Ep((())] = {(s)) + according to Vincent (2011) for any score estimator S. Here C' is constant independent with S, Our score estimator & is chosen to
minimize the objective function as
S arg min a.E
Σ n i=1
E velvo (1) [s((v)) - Vlogpo(v. 1 v
for each time step t E [0,T].
E.2.2 PROOF PROPOSITION 2
The proof follows that Theorem 1. in Fu et al. (2024) by neglecting the conditional information. Specifically, we replace Lemma D.1 in Fu et al. (2024) by Lemma 14, which provides an 2. (og(1/to))(E3 + 1) uniform upper bound on the magnitude of the empirical loss function. Lemma 14 (Counterpart of Lemma D.1 in Fu et al. (2024)). Then for any estimator š
T(D, L, M, B, R, /NNo," and data point Vo. have single-point score loss bounded by
h(8. vo) = ['²
dt at..E, (c(v:)) + Vt VyVo 60
S 2Nd(T + log(1/to)) + 1).
(18)
of
The proof provided in Appendix E.2.3. Besides, to implement covering number techniques, we also introduce a truncated loss truune as in Fu et al. (2024) and bound its difference with the truth loss with small error. We define the truncated loss
| Vi a
Frunn
St
1 TI
dt vilggo
1
( 11100
s Ro
o²
1o
The follow Lemma aligns with equation (D.12) in Fu et al. (2024). Lemma 15 (Truncation error of the truncated loss function). There exists a constant Cn such that for [
any 1, by choosing Ro /Nootr tr(E) CR log(Nd/e)) we have for any & E
E O(log(Nd/e)VNNd)
Evg L dt B B (v(v))+
V. avVo of
{ [llllll Z Ro}
S (T log(L/to))
Moreover, based the truncated loss function, have the following results on bounding the difference between the losses of two score estimators that are close to each other, which enables us to apply the results in Appendix E.3.
Lemma 16. Given the truncation radius Ro 0. Suppose g() E F such that (e()) (v) ((2)
(v) < for any llvll2 S Ro + a C N.CNNdlog(dN// and 2 lo where C is
an absolute constant. Then we have (arune((1))) - Jarune (3(2))
( < 2c(T +lo((1/40) (VNd(R, (VNdll +Clog(dNee) + 2Ro + 1).
38

---

Published as a conference paper at ICLR 2025
Besides, combining the results of Lemmas 23 and 16, direcly have the following results bounding the covering number of the truncated loss function class, which aligns with Lemma D.3 in Fu et al. (2024).
Lemma 17 (Counterpart of Lemma D.3 in Fu et al. (2024)). We consider the truncated loss function class defined as
S(Ro) h (s, IR RSEFF
(19)
Here htrune (s, v) h(s, v)1 fllv|12 I Ro)- Then the log-covering number of this loss truncated function class with output range lying in the Euclidean ball with radius Ro can be bounded by
(Og:)( S 8D2M (
L 1oo LFAF + log B2
(20)
66
where R2 satisfies R2 - (r + Cau) VN + Ro N.CCNNllog((N/es)) and es satisfies
= 2 log(1 (dd Cllo(dNNess)) 2Rg + +1)
(21)
Note that if we take (ndN for some constant C. we have log E5 The proofs all the supporting lemmas above are provided in Appendix E.2.3.
(log(ndNN 1)).
With the lemmas and statements above, completely follow the main proof of Theorem 4.1 of Fu et al. (2024) to prove Proposition 2.
Proof of Proposition 2. First we set dy = in Fu al. (2024) for our unconditioned setting. Then by choosing the covering accuracy 8 in Lemma 17 and taking the truncation range as in Lemma 15, i.e. Ro = OOloo(((o)))))) (Nd/es) we reproduce (D.17) in Fu et al. (2024) as
ED [(S) S 2 inf /' Ev, (llp(vv)-VVVogpvv.))))) sEF 20
(l.g
(T + log(L/too)
+
log + 2(T + log(1/to))es + 78.
n
< 2(T + log(1/to))e
+ 2Nd(T + (og(1/lo)) +1) .DD2 . lo Ly log
24
C6
+ 2(T + log(1 (//1// 14ca(T + log(1/14)) (d, + C'lo(dN//ss) + 2no + 1).
In the inequality, we invoke (20) and substititue d by its expression w.r.t. according to (21).
Choosing the both the approximation error and the covering accuracy as Cs 1/n and plugging the corresponding parameters about the size of the transformer classes according to Theorem gives rise to
Ep Ep(((S) $ (T + log(L/to)) log (kandNVt 4+1/v Ndd
The proof is complete.
E.2.3 PROOFS OF OTHER SUPPORTING LEMMAS FOR PROPOSITION
Proof of Lemma 14. Notice that when Ve ~ (.(. v()),, we have v, - ov = 0,2 where z ~ (O, INa) is a standard Gaussian variable. Therefore,
12
V2 a
E 1
(v(v))
< 2F
( 2 22
of
v.lvo
2Nd(RR +
2 2R 2Nd/0²
of
39

---

Published as a conference paper at ICLR 2025
Now can take integral over t [tu, T and obtain that:
dt-EE v.lv8 Bv(v.)
v of
S 22Nd(12+11) (''
et <2Nd((+og((/))))
e'dt
Here, we use the fact that
J/ e'dl
e¹ - = log(e" -1) -log(e"o - 1) s + log(1/1t)
The proof is complete.
To prove Lemma 15, we first need to bound the range of the data with high probability.
Lemma 18 (Range of the data). Given 5 > 0, with probability 1 2n exp( -66) over the dataset, Here C' is the absolute constant C2/ v3 3 in Lemma 24.
Proof of Lemma 18. This statement is directly related to the concentration of high-dimensional Gaussian distribution. For a random variable v ~ (O, Σn) where Σg N Na is the covariance matrix. Then, we use the polynomial concentration lemma (Lemma 24) and let g(-) = 112 be the 2-degree polynomial applied on the Gaussian. Then, we have:
Nd
E[g(v) = tr (E), B[g(v)²1= Σ
E0) +
(Eo)) + (E))),) < <
Therefore, we apply Lemma 24 and conclude that with probability at least 1 exp( (where is the C2 in Lemma 24), we have:
|| -| S / Var(Ilu13) < v E
For our case, Eg = < Σ, and therefore we can conclude that: with probability at least 1 2n exp(-C8):
ElE S Notrt(() | I 112 tr() tr(z)= 36
holds for Vi E [n]. Finally, we replace 5 with 5/v3 and it comes to our conclusion.
Now we are ready to prove Lemma 15.
Proof of Lemma 15. This conclusion can be made by combining the two lemmas above. By using Cauchy-Schward inequality and Lemma 14, we have
EEv (' dt vive Sc(v.) +
avvo (6)
2
1 1 Ro}
< 2Nd(R + 1)(T + log(1/to)))
1 2 Rg]
< 2Nd(R + 1)(T + (og(())))))
/2exp (-COR log(Nd/e)/2)
Let Cp = 2/C, then we have:
/'
T
(2)
2
E
dt
E
v.lv
(9)
Se(v.) +
of
(1voll2 122 2 Ro)
eel + +lg((1/1))),
which comes conclusion.
40

---

Published as a conference paper at ICLR 2025
Proof of Lemma 16. For any datapoint Vo such that |/voll Rg and diffusion time t 2 to, we have
Bvaivo
81 ( (v=)
Vi
a
E
velv!")
St (v+)
+
Vt
a
o7
0t
BrEN(On)
$t (arvo 012)
2 +
8
( (acVo +00
+ 2
T1
51 12
BRENion
of (((
sa:v (
(or(5/")
(arvo + + +2x)
S o, Ex- (1ll(1(())1 (z)(2)ll
where (2) (s (1) $/ (a-vo oz2) and Φ2(z) oo(s) (a.vo
alavvo+000 + 2z. By the
upper bound of s we know that
(z)l11 I 2R, vNdo,", and 1 S 2R, N. 2 2
Thus, they both have sub-linear growth with respect to || Now El (1022)))) by
decompose
Ex. (l) [z)(2] woo (62())l1 fllavvo + oæll2 s RBH
(0. (10 (loovo + t R1H
< Re(R, v Nd o o, H2))
+ 4R, v Noo, E, (R, (R.vNd+ Nd | fllavvo + +,zl12 2 R1}
In the inequality. we invoke [z)]2 < c and ||z||2 S s (Ro + R1) for llavv + o,2112 S Rq.
By Cauchy inequality, can bound A by
1/2
A Eg (R.Vdd
|z|11
Pr Po + oz2 R/11/2
V (2R3 2) Pr [|/||1 2 ( or Rn)].
Since z is standard Gaussian, we can set Ri = Ro + +o, C' so that A < e/(4vNd), Altogether, we have
for some absolute constant
Ez Φ (z) (z) s 2e (R.
1) CVNdloood//)))
Plugging the inequality into the expression of (fruno(((1) ?arune (3(2)) we have
T
Perune (3(1) Prunn (2(2)) <
oo -: R.V.N or + 1) CVVNllogggaN)//)
10
T
oi 22
Ri CdVVN++22H2 CD C)
$ 2e + +og((1/t))) (VNNd(L ( Clog(dNV/e) 2Ro +1)
The proof is complete.
E.3 COVERING NUMBER OF THE MULTI-LAYER TRANSFORMERS
The network apply satisfies the following form:
f = fi-10...
where the total layer number = 2L (which consists of L feed-forward layers and Transformer layers) and fr f2.... fi are either attention layers feed-forward networks whose input and output lies in DDN Denote B (R) = (X R DxN | F < R) a d-dimensional ball with radius R. Assume there exists sequence of radius Ro, R1, ......y Rí (which will be determined later)
41

---

Published as a conference paper at ICLR 2025
such that f; B(R + ((R)) holds for Vi = 0,.,....44 Then, need to compute the covering number with respect to the loc norm of the function space constructed by f ((R)) - ((R)) with the form above.
Notice that for two such functions f= oi.......0. fr and f'= fi f1-1 .... ff, the Loc norm of their difference can be upper bounded by the following lemma:
Lemma 19. For functions (ʃi RR - RdxN N )æEly and (/] RdxN + RdxN hEll we denote their composition := fi .1 and f == f/o
where R RdxN
For any mntrix-to-matrix function 9. we denote its || || norm as:
| F, sup (Xg 1|1
X
and its Lipschitz continuity
llg(X) g(X')llF
Lip(g) sup x
|X-I'
which is simple extension from the Lipschitz continuity with respect to l2 norm of vectors. Then, it holds that:
II - f' S
Σ 11
II Lip(S;).
j=++
Proof. For any == 2,3, 2....... and x E Ro), it holds that:
[lʃ ... f1(x) f/ o ... fi(x)) |1,00
< 11fk fk-1 ....o o fX( f& o fk-1 o .... f((X)
+ Ilf o fh-1 ..... fi( - fa fk-1o... (1
$ 11% - fal + (fx) -11f-1 .... fi fx-1 .... fillF,00.
After taking maximum over x E B(Ro), we conclude that [ o .... fi o.. Fl S Il - F + Lip(fx) f1 ... fi fhe .... F, By using the method of induction, we can easily derive our conclusion.
Now, in order to compute the covering number of the function space constructed by functions f with the form f ft fi-1 .... fi where f. F,, we firstly need to bound the Lipschitz constants for the function classes F;. Also, we need to estimate the covering number of each F.
Lemma 20. For the function space of feed-forward network
J=
FFN DxN + DNN Y Y + W2 ReLU(W,Y + +
W l. Welle: 11l12, b2ll2 < B, B
then all functions in the class jere are (1 Lipssiich The covering number can be bounded as:
1 (R (R+ VN)
// JFFN
|1 < 4D log
Proof. For FPP with Y- ReLU Y +b2 b+ T by the additivity of Lipschitz constant and the fact that ReLU is 1-Lisschitz continuous, we can conclude that the Lipschitz constant of function f is no larger than 1+ |+||W2ll2
11ll2 S 1+ W W2llF 1+c² which means f is (1 + -Lpsspcciitz continuous. On the other hand, for two functions f.g E FFEX where /(Y) Y + WN ReL.U(W, Y b,1 b, and (YY = Y + W Rel.U(W
b111 1 Then for vY E B(R). we have:
IIf(Y) g(Y)|l < VN. [lb - b' la + IW; (ReLU(W,X Y + bb, ReLu + b)1 ll
+ (W W2) ReLU(W, Y b 1
VN. |lb: bíll2 F/ (VN) b2 balla + R||W Willp)
+ + W/llF - WVll=-((/|| W ||| + / (2ll2)
VV |lb; bill BVN |b> bbbll' BR |W, W B(R+ VN). [W2 - W2llr.
42

---

Published as a conference paper at ICLR 2025
Therefore, it holds that:
g I V N |b 1ahh-b/ji
BBW -
Walln,
which leads to the upper bound of covering number:
M(6:FFHM (1e,,)))) S N(8//VNN (12)) N.0/A////BVVN 112)
-IF)
where = (vE |/vl22 M) E |M < B). It well-known that for > 0, we have (e; M1, (122) S (3B/s)² Ni; ||l) - (3B/e) D To sum up, we finally conclude that
1
12B (R VN)
log (6;
|| (IE,,c) I 4D2 log
The proof is complete. Lemma 21. For the function space of attention network
Juhinn Attn : DxN + DR DDN
Y Y v a ((Q"))
|V IQ F, K B for Vm E [M], Y E B(R)
then all functions in the class Jutte are (1 + BMVN + 2B MR N) Lipschitz. The covering number can be bounded as:
6B2M/NN3/2
log Jähta |l (1p,,)) I D2M log
8
Here, the metric F supYeB(n) [|f(Y))l.
Proof. Notice that for any two functions f1, fz and Y, Y', we have:
means that the Lipschitz constant of V VIYT
T is no larger than V I2 vN R/N|VV" ||2 Q IF RVN. Here, we use the fact that both ReLU and softmax activation function over vectors is I-Lipschitz under norm and that over matrices
1-Lipschitz under Frobenius norm. See Appendix G.2. Also, (Q" KmY is 2 Q T Km ||| RV RR/--Lipschitz under Frobenius norm when Y E (().. Therefore, in the function class Jutta, all functions included are (1 BM + 2B M R - continuous, which comes to our conclusion. For the covering number, given any two functions f.g E JAttn with their corresponding parameter sets {VT, Q K E [M] and (Vm
Q KT
[M]], then:
II((Y) , <
Σ
IVV
V a (Y Q" KT I|
M
+
V (o (Y Q Y (YTQu Q K ||
m=1
M
M IV VTlle RVN VN+ + Σ
BRVN Y ( - Y
m=1
m=1
M
[V - VT RN Σ
322
|B -B
43

---

Published as a conference paper at ICLR 2025
Here, BT "T and B Qr KT Then, | Il. [ || S B² holds for Vm E [M].
leads to the following upper bound of the covering number of Futta
JAhtnn 1l II N66/2MRN; M,,1 (Ie) II N66/2BMR 3/2 M22 -1e)))))
where matrix M₁ {V E RDXD sum up, we conclude that:
| S B} and M [B B
B B²). To
6BMRN
6B 6B²
log.N (6;; Jöttn (1F,0o) S DD2M log
+ D2 M log
5
which comes to conclusion.
Now, we are ready to combine these results together.
Lemma 22. Consider the multi-layer transformers class Traw (D, L, M, B) without encoders and decoders. Then the log-covering number with input range bounded by Ro (||l| S RR) can be bounded by
12R ML 22L3/22
lo;NV(;; Traw, Ra, ll 111,,000) S 4 (I2 Ly + log
Here LF 1 - (BMvN& + MREN) and AF (1+ 1²
Proof. For the score network class == Traw (D, L, M, B), we have:
F= f2L-10.... o fi f F; for Vi [2L]], each F, is either J=FN or Fltin).
Denote Li == 1= (BMvN R N) as the upper bound of Lipschitz constant for both 5 and Futee, Notice that for any input Y E DxN such that ||Y[|= Ro. then the output of Sced-forward network f E FFFN holds
I S || B(B VN) BVN (1 B 2BVN) (Ro V1).
The output of attention network f E Juhtta holds that
l| S | | + MB VN (1 MBVN)
Here we use the fact that || & VN holds for all probability matrix E DxN, Denote A5 = (1 + B2 + 2B v (1 + MB/N) > as the signal amplifier each layer, then we can set the sequence of radius as R; = (Ax)* Ro with Rg > 1 and fs B( + B(R;). According to Lemma 19, we have:
N(6; F. 11 S I
N(8/2LL 22- F.. (1F,,)),
which leads
log. (o;; F, || (11,00)
24
<
log N28/2LL -
1=1
12B RR + VN)
2LI2 21-4
6B 1 N3/2. .2L2
<
4D
V D2M log
24R, B² MLN3/2
< 8
[L'lo + log
8
The proof is complete.
Moreover, if we consider the entire transformers architecture with the encoder and decoder layers, directly have the following results:
44

---

Published as a conference paper at ICLR 2025
Lemma 23 (Covering number of transformers). Consider the entire transformer architecture F T(D, L, M, B, R) which the encoder that takes Vi E R Nd as the input and embeds it with the time embedding and the diffusion time-step embedding (tt), where ||e||p and || Caitt for some absolute constant Caitt >> 0. Then the log-covering numher with initial input range V S Ro bounded by
Jo; Ro, ll. 111 s 8DA M. (12 log L Ar log
24R B
RDXN with || - ||vel12 rVN Cair /N, and our decoder (defined Appendix D.3.I) 1) only extracts part of the entries from the output and impose clipping function with range being I on it, both of which lead to no increase the covering number.
Here R1 Ro + VN(r Can).. This because the encoder maps V:
F EXPERIMENTAL DETAILS IN SECTION
Patch Embedding Modification To implement simulation experiments Gaussian Process data, we slightly adapt the original DiT designed for image generation (Peebles and Xie, 2023). In the original DiT, pre-trained VAE encoder deployed convert image in the training dataset to feature representations and patch embedding is applied by splitring the image into subblocks(patches) and flattening and projecting each subblock to a feature. But our numeric experiments, we dropped the VAE encoder and split the data at time dimension, with each time step being a patch, and project the patch into feature of higher dimension.
Kernel Estimation Method In experiment, each sample generated by DiT in is denoted as = [.....,, SN] E R xd where s, = S[i] E TRd denotes the data patch at i-th time index. In our experiment we take N 128 and d 8 and with n E {103, 3200, 10ah 32000, 10 samples total. To evaluate the quality of generated data, we calculate the mean of data patch at i-th time step and covariance matrix between data patch at i-th and j-th time step as follows and compare them with the theoretical ones.
St k=1
ΣNJ
1 Σ S& (.A( TT))T k=1
Here jh E IRd is the empirical mean of the i-th patch and Σ., is the empirical covariance matrix between the i-th and the j-th patch. With Es,, we could estimate the empirical kernel value (i,
as
(i. argmin Σ. aΣ
where E dR is the true covariance matrix. After running each pair of (i, N,, we can get the empirical kernel = ()),,))) NxN
Relative Error In the experiments where we compare sample efficiency across different training data size n and kernel setting v, E, the metric a relative error of the estimated sample covariance matrix to its ground-truth:
E
r
&
cruth
r &
Σ
Here (f & D Druth is the estimated covariance matrix through the same method but under an equal amount of training data, instead of generated data. This relative enror eliminates the influence of different scaling & Σ and the concentration error caused by finite samples.
45

---

Published as a conference paper at ICLR 2025
Query-Key Matrices and Value Matrices From Figure 10, the query-key matrix Q K emphasize the time embedding part of the input across different layers. Interestingly, Figure 11 shows that the weight emphasis the value matrix V is on the data patch part instead. This observation aligns with our theoretical construction of the score function using the DiT structure. Specifically, the matrix (determined by the query-key matrix) captures the kernel T, representing temporal dependency, while the value matrix determines the covariance Σ, representing spatial dependency.
that the weights of
Layer
Layer
Layer 3
10
10
20
30
10
20
30
10
20
30
10
Layer 4
Layer
Layer 6
0
10
10
10
20
20
20
30
30
30
20
Layer 7
Laver
Layer
10
20
30
30
30
10
20
30
10
20
Layer 10
Layer 11
Laver 12
10
10
20
30
330
10
20
30
10
Figure 10: Query Key matrices in different transformer attention layers( 1-12),
Attention Score Different Attention Layers To further demonstrate our theory, we visualize the attention score Y Q KY, E N averaged over 10s data points at different backward diffusion times E {0, 50, 100, 200, 400, 800, 1000) to observe if gradually unveils the kernel I as the backward process progresses. According Figure 12, in the initial few layers (1-2), observe the early stages of kermel construction. the subsequent layers (3-12), the attention score matrix increasingly resembles the kemel,, becoming clearer the backward diffusion process advances.
46

---

Published as a conference paper at ICLR 2025
Layer
Laver
Laver 3
20
30
10
20
30
10
20
30
10
20
30
Layer 4
Laver
Laver 6
01
10
10
20
20
30
30
10
20
30
7
Laver
0
10
10
20
20
20
30
30
30
10
20
30
Layer 10
Layer
Laver 12
20
30
10
20
30
10
20
30
10 20
30
Figure 11: Value weight matrices in different transformer attention layers( 12))
Generated Motions of 2D Balls We collect latent sample from the diffusion transformer and map it to the original 2D space through the pretrained decoder of the VAE, forming a 240-frame video. We present 36 frames in Figure 13 from the video, which are chosen uniformly out of the 240 frames. As shown in the figure, the motion of the ball shows great time consistency and accurately captures the bouncing- back mechanism as expected.
47

---

Published as a conference paper at ICLR 2025
Layer
Layer 2
Layer
Layer
Layer
Layer
Layer 7
Layer 8
Layer
Layer 10
Layer 11
Layer 12
Figure 12: Attention score matrices in different attention layers and at different steps of the backward process.
Figure 13: Consecutive frames video generated by a trained diffusion transformer with a 2D VAE.
48

---

Published as a conference paper at ICLR 2025
G SUPPORTING TECHNICAL RESULTS
G.1 GAUSSIAN LEMMAS
In this subsection, we will introduce several Lemmas control the deviation of random variables which polynomially depend on some Gaussian random variables. We will use slightly modified version of Lemma 30 from Damian et al. (2022).
Lemma 24. Let g be polynomial of degree and ~ (O, La). Then there exists an absolute positive constant C₂ depending only on p such that for any 6 1,
P [[(x)) ) E[g())] 2 (VVar(g(z)) - exp (-C,62/)).
The next Lemma orginates from Theorem 4.3, Prato and Tubaro (2007).
Lemma 25, For any E N and f (,3 Ia)) be degree l polynomial, for any g 2 2, we have
Egnry [f(2) S Cg.e (Ezny [f())²) 8/2
where we use Cg.e to denote some universal constant that only depends on g.
G.2 LIPSCHITZ CONTINUITY OF ACTIVATION FUNCTIONS
Lemma (I-Lipsehitz continuity softmax function). For the softmax function TRd Rd is I-Lipschitz continuous under l2 norm.
Proof. For any I E Rd, denote [xx as the Jacobian matrix softmax function at x. Then, the Lipschitz continuity of (-) under norm can be upper bound by supzeRa [=](=)] By caleulation, we have
((z) =
Diag (a(z)) - (z))r(z))
Let (xx) = (p1.P2..... .........PP)) .P)) which is probability vector, then
d 1t((1 (1 - + + (p:p))² S
| S
Pi
=1
Therefore suPPER4 [I < 1, which comes to conclusion. A direct extension is that the column- wise softmax over matrices is I-Lipschitz continuous under ||p norm.
G.3 BASICS ON RESNET UNIVERSAL. APPROXIMATION THEORY
In this section, briefly introduce universal approximation theory of ResNet. An L-layer ResNet R(x) : Rd > Rd can be defined as
R(x) B FFN, FFN,-1 FFN, A(x).
(22)
Here A Rdd + RRd and B Rd' + Rdd are two linear transformations, and FFN; Rdd - Rdd are basic residual blocks defined (y) y W2, ReLU(W b2,s) + b with Wi, (d, do. d' L, the set of ResNet functions from gd to d with L layers, d' neurons each identity layer, maximum width W = max4e(2] (L) {d:} and nonzero weights S. Moreover, the Frobenius norm of the weight matrices W j.i and the Euclidean norm the bias vectors bj, are uniformly bounded by C > 0.
Rd W2, E IRd xd. b1,4 E Rd and b2. E Rd. We denote by RW
The next lemma shows that
approximate product operation with ResNet. Lemma 27 (Proposition 12 in Liu al. (2024a)). Let x, ResNet R (2, 1, Oloo(((/o)))) (llog(B/c). O(B) such that
B, B] with B 2 Then there exists
|R(,, v) zyl < e, x, y E [- B, B)
(23)
holds.
The construction process follows the standard techniques proposed by Yarotsky (2018), which first use feed-forward networks to approximate the square function (x) ss x² and transfer it into product operator x x(x, y) (Jaq y) fx 4))/4 xy.
Furthermore, for IRd and TR, we can approxiamtely construct the following mapping: f(xx .v]
yx.
(24)
49

---

Published as a conference paper at ICLR 2025
Corollary 2. Given 0, there exists ResNet fmult
R. + 1, 1,d, 4d,
((l))( dO (Brllog(BB))))))))) B/c)), (Bd)) such that fmult (y, x) = yx + e, where ||e S c.
Proof. By Lemma 27, we know there exists ResNet R R
R satisfies
Re 0,0.1444444..000////00000000/000
and
[R y) - zyl S c, x,y [-
Then let's consider the following two steps of mapping:
A([x] T y] = [1 y, 22, .. Id. Vl
B[ i, y. 22, y, F. y) [R(....) R(FI, ... R(xa, v)l
Here we could choose A be a (d 1) x 2d matrix
1
0
0
0
0
0
1
0
0
O
0
A
0
0
1 0 1
0
0
0
and take B parallelization of homogeneous networks R. Let fmult B A, we know fmuit RV(d+11 d,,,,,d,d,, 3d, 4d, Olog((B))) B/e)), (O(lo)(B/e)) (d)), and it approximately realizes the mapping (24) with Lo approximation error being e for the first entries and error for the last entry.
More generally, if consider the input to be [x 3y, z) and want to construct (FFN-only) transformers that approximately maps the input to [y .9.z] W] we have the following results:
Corollary 3. Suppose the input to be Y = y2.. yal E PxN with y; = [x O Wi, 2, where x2 E [-B, Bjd, wA E [, B) and z, Rd. Given any 0, there exists a (FFN-only) transformers
fmait FFN o FFN FFN
with L (log(B/c) layers that approximately multiplies each component x; with the weight w,, which keeping other dimensions the same. This can be formally written as
fmult (w1,x1)
fmultt [N)
O2d
Ozd
Emiit (Y) =
where [] xx) wixx 0o0 S e.
W1
WN N
Z
The inner dimension of the FFNs most Sd. Moreover, the number of nonzero coefficients in each weight matrices or bias vectors is at most Odd) and the norm of the matrices and bias are all bounded by ((B)).
Here require the buffer variables O2a in the input because Corollary 2 needs 3d neurons in each dimension to store the calculation results that are necessary for constructing the product function. Thus, add O2d that [x 00Tal E 3R for construction convenience.
Proof of Corollary 3. By Corollary 2, there exists a ResNet fmult
[Rd+1 + nd such that
fmult (y, x) B o FFN, FFN
FFN, A(x, y)
with fmult (y. x) yx + e, where || ||Ellgg I e. Since the FFNs in the transformers are position-wise, only need to consider the mapping of one token while others completely the same. For each i, suppose FFN, : 3d3d - 3d satisfies
FFN v Wa ReLU(W1,.v 1.v + b2,e) b
50

---

Published as a conference paper at ICLR 2025
with W1. e Rd.33d W2, E Radxd $1s E R&d and 2 E RR Here d, for all sis L. For the first FFN, let
x O2d w
x O2d w
x
FFN
W2.1 2.1 Od xd.
ReLU
w
b2,1
b
22d
Ou.
z
z
z
PFN,
A(x, y)
For
Sis
suppose the input is yv" w. TTTTT =[vT, w. TJT
with v E 1g and E 12
FFN
v
w
w
v
xd
2,
RcLU
([ x
v
w
+
12,6
bi, 0.2
FFN;(v)
z.
z
Here d' = d2 + 1. For the final layer, let
B
Oaxd'
Bv Odd W
- 13
Oaxed
FFN 7++1
w
+
W2,4+1 ReLU
Lzd
Oaxd
w
Z
Isd
Odxd'
z
Here
diag(1 1d, O2a)
diag(1 [a, O2d)
Lzd 0 ((2d+4))
Iad
W 2,
(22++d))xc
4(2d+d))xx )xd
0 d2
Thus, we have finished constructing the (FFN-only) transformers we want by taking
fuult
FFN 1+1 FFN,.1
FFN1,
and the hidden dimension of the FFNs is at most &d. Morcover, by the definition of the original FFN; and the new FFW, the number of nonzero coefficients in each weight matrix bias vector is most (d), and the norm of the matrices and bias are all bounded by (dd). The proof is complete. L
G.4 ASYMPTOTIC RESULTS ON THE SPECTRUM OF TOEPLITZ MATRICES
In this section, provide some existing results on the spectrum of Toeplitz matrix when both its size and bandwidth to infinity. For Toeplitz matrix T = N its (i, j)-th component Tıf aji-jl where M is known as
only depends on its distance diagonal. When k > M, we have ag the bandwidth. Denote Å1, A2........ Av to be the eigenvalues of T, with multiplicities counted and let μN
5x, be the empirical distribution of the spectrum. In the asymptotic case where M.N > co. we focus on the behavior of μN. Denote Fx (x) to be the cumulative distribution function of JNN, Kargin (2009) proposes the result that (x)(2) converges the standard Gaussian distribution when the Toeplitz matrix follows that Ea, 0, Ea 1 supA.N |E| MM
C. N - c (0, 1).
Do and most importantly, the band-to-size ratio N + 0. While the band-to-size ratio M the spectrum distribution (x) converges statistical works such as Hartman and Wintner (1 (1950): Tilli (1998); Delsarte and Genin (2005) study the spectrum of generalized Toeplitz matrices by using Fourier expansion.
non-Gaussian distribution V Some other
In our case, the condition number Kt, will keep in a constant range if the sequence {ax} introduced above sharply decays, while grow N goes larger if {ax) slowly decays. In the former case, can treat Ki, as constant, which does not affect our analysis. In the latter case, since we have a natural upper bound of
Kito S 10t Amax X Σ) S 25'
by taking ln = -1/3 in Theorem 2, we can still obtain a m 1/6 convergeence rate both W2 distance and TV distance.
51


Reviews:
[{"review_id": "uCEXAF1gkm", "reviewer": "Reviewer_XNMz", "summary": "This paper explores transformers as denoisers within diffusion models for generating spatiotemporal data, motivated by recent advancements in video foundation models, such as SORAH.\nThe approach begins by assuming a Gaussian process as the underlying data distribution and simplifies it to a scale-varying Gaussian process where the covariance changes only in scale over time. In this framework, the score function is related to the inverse of a modified covariance matrix. Since directly calculating this inverse is computationally prohibitive for high-dimensional latent vectors (dimension = d * N), the paper introduces a convex quadratic relaxation solvable via gradient descent. The authors then design a transformer architecture inspired by unrolled gradient descent steps. They also derive sample complexity bounds for the transformer’s ability to learn the covariance based on the smoothness and decay properties of the kernel matrix.\nThe paper includes basic experiments with synthetic Gaussian process data and a simple ball-motion scenario to validate the proposed method.", "questions": "- In the quadratic formulation in Equation (3), the score function estimation remains non-separable over timesteps (i.e., across x1,t, ..., xN,t). Would further decomposition allow for a distributed solution with minimal inter-timestep communication?\n\n- Prior work, such as [Sahiner et al., 2022], has already shown that attention layers capture correlations and reflect low-rank priors, so it’s unsurprising that attention proves effective for correlated and stationary data.\n\nSahiner, A., Ergen, T., Ozturkler, B., Pauly, J., Mardani, M., & Pilanci, M. (2022, June). Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In the International Conference on Machine Learning (pp. 19050-19088). PMLR.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper addresses a timely and relevant problem, aiming to understand why spatiotemporal transformers effectively learn distributions in spatiotemporal data.\n\n- The writing is clear, making the methodology and findings accessible.", "weaknesses": "- The goal of the paper lacks clarity. If the aim is to demonstrate the expressiveness of transformers for correlated or spatiotemporal data, this doesn’t necessarily relate to diffusion models and could be examined independently in a more simplified setting.\n\n-The assumption of a Gaussian process with scale-wise stationarity is a strong simplification, limiting the applicability to realistic distributions. Consequently, the results offer limited insights into transformer design for broader, more complex distributions.\n\n-The experiments are weak and unconvincing. To substantiate the theory, more realistic data should be used, and comparisons with alternative architectures, such as UNet (3D or 2+1D), known for capturing spatiotemporal correlations, would strengthen the argument.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0NMOdEUPfi", "reviewer": "Reviewer_fqSV", "summary": "This paper scrutinises diffusion transformers under a simplified setting where the input time series is assumed to follow a Gaussian process. The paper carefully derives the analytical forms of the dynamics and provides score approximation as well as data-distribution estimation guarantees in the simplified setting of interest. The authors cast this paper as a first step towards building more sophisticated theories.", "questions": "Please see weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I really like the framing in this paper. While sounding like toy settings, GPs can be quite general and could be a good starting point of analysis. I especially appreciate the analytical forms, for example in the score approximator, and how the covariance function relates to positional embeddings. The paper is also quite well written. I have gone through parts of the proofs and they made sense - although I have not exhaustively verified every step.", "weaknesses": "- Can we have a section in the appendix where all the assumptions and simplifications are listed? This would really ease the conclusions that the reader can draw from this paper.\n- Ln. 397: What's a \"properly transformer\" architecture? \n- Could the authors give more elaborations on what motivates the choices of $\\epsilon$ and $T$ in Thm. 2? \n- What about empirically testing the learning on GPLVMs? Wouldn't this be a natural fit? \n- There are some duplicate looking sections in the appendix, for example relating to the size of the transformer blocks. Can't we gather those in a big table instead of repeating them after each relevant section?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iLZW7ZuaRu", "reviewer": "Reviewer_zotT", "summary": "This paper provides a theoretical study on the limits of the learning capabilities of transformer-based diffusion models on learning spatial-temporal data. The authors achieve this by reformulating the task of learning the score function of Gaussian process data as a gradient descent algorithm. This algorithm can then be unrolled using a transformer archtitecture where the attention layers can be studied to evaluate how spatial-temporal dependencies are learned. \nThe results are thoroughly motivated and dervied and validated using a few numerical examples. \n \nI believe that the theoretical results could have an impact in terms of understanding the learning of temporal dependencies but currently the practical results are not extensive and the take away messages are not clear enough", "questions": "- The learned temporal dependencies presented in Figure 5 look very similar to what you get when you just plot the attention between two uncorrelated matrices with positional embedding (l. 300). What is the difference between these maps and the ones you present, e.g., in Fig. 5", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- studying gaussian processes to asses spatio-temporal dependencies is inherently reasonable and could be used as an important theoretical framework to study neural architectures.\n\n- The method mathematically is well grounded and derived\n\n- High-quality figures visualizing key components of the contribution", "weaknesses": "- While the results convincingly show that transformers can learn spatial-temporal dependencies, the experiments only show examples of properly learned tasks. Given that theoretical boundaries were derived, I would have loved more experiments investigating these limits. For example something like Figure 1 but with a degradation for higher length. The same goes for Section 6.2. What I am missing is a clear connection between model complexity and performance.\n\n- Lack of comparison to other methods: I think it would be helpful to see how typical diffusion models, including Unet-based architectures, perform on the task of learning Gaussian processes. Clear results could help in shifting the design towards transformer-based diffusion architectures.\n\n- Generalization to real-world problems: There are no experiments showing how well the insights generated by the experiments translate to real-world problems. You present the ball experiment but even there you have to consider how well your experiments align with the assumption of using Gaussian process data (cmp. l. 504-506). Real-world problems will often not follow Gaussian process dynamics\n\n- It would be nice if you could state some take-home messages for the reading. Ideally some general information such as design choices for the model architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ggUl0SKfvs", "reviewer": "Reviewer_FvaY", "summary": "In this paper, authors explore the reasons for Diffusion Transformer (DiT)’s capability of capturing spatial-temporal dependencies within the sequential data. They theoretically establish score approximation and distribution estimation guarantee of DiT for learning Gaussian process data. Specifically, they replace the score function in diffusion models with gradient descent process defined by truncated covariance, which are further implemented with transformer models. Authors conduct some experiments to prove the effectiveness of their theory under their assumptions.", "questions": "1. The theorems proposed in the paper are based on the authors’ assumptions, then:\n1.a) Does it mean that the theory is only valid under these assumptions, or does the model only work on the data satisfying these assumptions? \n1.b) How to ensure that the assumptions can represent the case in real sequential data? \n1.c) How do authors derive these assumptions? E.g., in Assumption 1, why the definition of the covariance function has that formula, are there any theoretical basis for it?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This paper studies the principle behind the Diffusion Transformer and provides a theoretical explanation of its capability of modeling spatial-temporal correlations within data, which is the very first trial in related fields. \n2. Authors propose a novel score function approximation theory as well as the sample complexity bound. All the lemmas and theorems are well defined with sufficient mathematical proofs. \n3. The correctness of the proposed theory has been evaluated on both synthesized and real data through experiments.", "weaknesses": "1. Although the overall writing is relatively clear, there’re some suggestions:\n1.a) The object of this paper is sequential data, which has the axis of time (real world time, $h$). while in the diffusion process, there’s another axis of diffusion step ($t$), and authors also name it as ‘time’. To avoid confusion, authors should use some other descriptions to make it clearer. \n1.b) The symbols should be consistent. E.g., the definitions of $v_t$, $\\mathcal{D}$ and $v_0^t$, the use of subscripts of data $\\mathbf{x}$ are confusing.\n1.c) The font size of Figures could be larger, e.g. Fig.2, otherwise the equations and words are difficult to recognize under normal physical paper size.   \n2. Although the effectiveness of the theory has been proved on synthesized Gaussian process data and semi-synthetic video data. More real data with complex contents should also be tested on.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)