# Diffusion Transformer Captures Spatial-Temporal Dependencies: A Theory for
Gaussian Process Data

Hengyu Fu Stanford University. Email<EMAIL_ADDRESS>Zehao Dou Yale
University. Email<EMAIL_ADDRESS>Jiawei Guo Northwestern University.
Email<EMAIL_ADDRESS>Mengdi Wang Princeton University. Email:
<EMAIL_ADDRESS>Minshuo Chen44footnotemark: 4

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

##  1 Introduction

Diffusion models have emerged as a powerful new technology for generative AI,
which is 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 et al., 2021;
Tian et al., 2023), reinforcement learning and control (Pearce et al., 2023;
Hansen-Estruch et al., 2023; Zhu et 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 a neural network. In order to train the
score neural network, diffusion models utilize a 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 as a score neural
network (Peebles and Xie, 2023; Wu et al., 2024; Bao et 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 in high-fidelity video generation strengthens
the seemingly unlimited potentiality of diffusion models (Gupta et al., 2023;
Liu et al., 2024b). These models advocate transformers over the traditional
U-Net for parameterizing the score function. A 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 the time. These dependencies resonate well
with the self-attention mechanism in transformers for capturing token-wise
correlation, suggesting benefits for learning with sequential video data. We
illustrate diffusion transformer accurately learning spatial-temporal
dependencies in Figure 1.

Figure 1: Diffusion transformer learns spatial-temporal dependencies. The
diffusion transformer is trained with data sampled from a 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 D 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.

Despite the empirical success, there lacks a rigorous understanding of
diffusion transformers for sequential data modeling. Different from static
data, such as images, sequential data consists of a series of dependent data
frames. The data sequence can be extensively long. For instance, a one-minute
video would contain over 1440 image frames. Therefore, naïvely handling the
sequential data by stacking all the frames as if independent will lead to a
large dimensionality dependence and the spatial-temporal correlation further
complicates matters. As a result, existing results of diffusion models for
static data can provide highly limited insights on the following fundamental
questions:

Can diffusion transformers efficiently capture spatial-temporal dependencies
in sequential data?
If yes, how do spatial-temporal dependencies affect the learning efficiency?

We answer these questions firmly by studying for the first time of 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-dimensional, 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 can encode a 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 a fluid. Gaussian process is also a powerful
statistical tool for regression, classification and forecasting problems
(Banerjee et al., 2013; Deringer et 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 so 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 as follows.

$\bullet$ We propose a novel score function approximation scheme for Gaussian
process data, which represents the score function by a gradient descent
algorithm (Lemma 1). Then we construct a transformer architecture to unroll
the gradient descent algorithm in Theorem 1. We particularly highlight that
the attention layer effectively captures the spatial-temporal dependencies.
Meanwhile, the decay pattern of those dependencies influences the
approximation efficiency.

$\bullet$ 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/\sqrt{n}$, where $n$ is the sample size. It also depends on the decay of
temporal dependencies and long-horizon dependencies introduces large
dependency on the sequence length.

$\bullet$ We provide numerical results to support our theory by showing the
learning performance under various settings. More interestingly, we
demonstrate that a well-trained diffusion transformer reproduces the ground-
truth spatial-temporal dependencies accurately within an attention layer,
emphasizing the applicability of our theoretical insights.

##### 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 a relatively accurate estimated score function,
Benton et al. (2022, 2023, 2024); Li et 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 are not directly
applicable to understanding how spatial-temporal 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, we view the score function as
the last iterate of a 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-art performance. Moreover, diffusion transformer exhibits appealing
scalability towards better generation qualities with larger model sizes. More
recently, diffusion transformers are leveraged in video generation (Gupta et
al., 2023; Liu et al., 2024b), where the video data is patchified and
rearranged into a long sequence. Besides, transformers are also used in
discrete diffusion models (Sun et al., 2022; Austin et al., 2021).

##### Notation

We use bold letters to denote vectors and matrices. For a vector $\mathbf{v}$,
we denote its Euclidean norm as $\left\lVert\mathbf{v}\right\rVert_{2}$. For a
matrix $\mathbf{A}$, we denote its operator, Frobenius norm as
$\left\lVert\mathbf{A}\right\rVert_{2}$ and
$\left\lVert\mathbf{A}\right\rVert_{\rm F}$, respectively. Moreover, we denote
$\left\lVert\mathbf{A}\right\rVert_{\infty}=\max_{i,j}{\left|\mathbf{A}\right|}$.
We denote the condition number of a positive definite matrix $\mathbf{A}$ by
$\kappa(\mathbf{A})=\lambda_{\max}(\mathbf{A})/\lambda_{\min}(\mathbf{A})$,
where $\lambda_{\max}$ and $\lambda_{\min}$ denote the maximum and minimum
eigenvalues. We denote $f\lesssim g$ when there exists a constant $C>0$ with
$f\leq Cg$.

##  2 Gaussian Process and Diffusion Transformer

In this section, we formalize our data modeling and sampling problem with
Gaussian process data. Meanwhile, we briefly introduce diffusion processes and
transformer architectures.

##### Gaussian Process

We denote $\\{\mathbf{X}_{h}\\}_{h\in[0,H]}$ as a continuous-time Gaussian
process in the time interval $[0,H]$. The process lives in the $d$-dimensional
Euclidean space, i.e., $\mathbf{X}_{h}\in\mathbb{R}^{d}$ for any $h\in[0,H]$.
A defining property of Gaussian process is that for any finite collection of
time indexes $h_{1},\dots,h_{N}$ with $N\in\mathbb{N}^{+}$, the joint
distribution of $\\{\mathbf{X}_{h_{1}},\dots,\mathbf{X}_{h_{N}}\\}$ is still
Gaussian. As a particular example, for a fixed time $h$, the marginal
distribution of $\mathbf{X}_{h}$ is Gaussian.

To fully describe a Gaussian process, we need the concept of mean and
covariance functions. Roughly speaking, a mean function
$\bm{\mu}(h)=\mathbb{E}[\mathbf{X}_{h}]$ characterizes the expected evolution
trend of the process. A covariance function
$\bm{\Gamma}(h_{1},h_{2})=\mathbb{E}[(\mathbf{X}_{h_{1}}-\bm{\mu}(h_{1}))(\mathbf{X}_{h_{2}}-\bm{\mu}(h_{2}))^{\top}]$
captures the correlation between two time indexes in the process. Note that,
when $h_{1}=h_{2}=h$, the covariance function computes the covariance matrix
of $\mathbf{X}_{h}$. 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 (strictly) stationary Gaussian processes,
that is, the covariance function only depends on the gap between time indexes.
Accordingly, we reparameterize the covariance function as
$\bm{\Gamma}(h_{1},h_{2})=\gamma(h_{1},h_{2})\bm{\Sigma}$, where
$\gamma(\cdot,\cdot)$ is a scalar-output function and
$\bm{\Sigma}=\operatorname{\rm Cov}[\mathbf{X}_{h}]$ (identical for any $h$).

##### Sequential Data Sampled from Gaussian Process

In real-world scenarios, an underlying continuous-time process is often
perceived by a 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 to the human eye, which cannot distinguish
between individual frames as if the video is continuous. Following the same
spirit, we denote $h_{1},\dots,h_{N}$ for a sufficiently large
$N\in\mathbb{N}^{+}$ as a uniform grid on the time interval $[0,H]$ and form a
discrete sequence $\\{\mathbf{X}_{h_{1}},\dots,\mathbf{X}_{h_{N}}\\}$ observed
at those times from an underlying Gaussian process. By the definition of
Gaussian process, if we stack $\mathbf{X}_{h_{1}},\dots,\mathbf{X}_{h_{N}}$
consecutively as a vector in $\mathbb{R}^{dN}$, it follows a Gaussian
distribution. The mean is
$\bm{\mu}=[\bm{\mu}_{1}^{\top},\dots,\bm{\mu}_{N}^{\top}]^{\top}$ with
$\bm{\mu}_{i}=\bm{\mu}(h_{i})$ and the covariance is a block-wise matrix
represented as $\bm{\Gamma}\otimes\bm{\Sigma}$, where

$\displaystyle\bm{\Gamma}=\begin{bmatrix}\gamma(h_{1},h_{1})&\cdots&\gamma(h_{1},h_{N})\\\
\vdots&\ddots&\vdots\\\
\gamma(h_{N},h_{1})&\cdots&\gamma(h_{N},h_{N})\end{bmatrix}\quad\text{and}\quad\bm{\Gamma}\otimes\bm{\Sigma}=\begin{bmatrix}\bm{\Gamma}_{11}\bm{\Sigma}&\cdots&\bm{\Gamma}_{1N}\bm{\Sigma}\\\
\vdots&\ddots&\vdots\\\
\bm{\Gamma}_{N1}\bm{\Sigma}&\cdots&\bm{\Gamma}_{NN}\bm{\Sigma}\end{bmatrix}.$

Here, $\otimes$ is the matrix Kronecker product. Notably, $\bm{\Gamma}$
captures the temporal dependency across time and $\bm{\Sigma}$ captures the
spatial dependency of entries in $\mathbf{X}_{h}$. Since $h_{1},\dots,h_{N}$
form a uniform grid, $\bm{\Gamma}$ is a symmetric Toeplitz matrix with entries
taking at most $N$ different values.

In a learning setting, we collect $n$ i.i.d. realizations of the discrete
sequence $\\{\mathbf{X}_{h_{1}},\dots,\mathbf{X}_{h_{N}}\\}$, denoted as
$\mathcal{D}=\\{\mathbf{x}_{1,j},\dots,\mathbf{x}_{N,j}\\}_{j=1}^{n}$. We aim
to use a diffusion transformer for learning and generating new samples
mimicking the distribution of the discrete sequence. The subtlety here is that
naïvely learning the joint distribution of
$\\{\mathbf{X}_{h_{1}},\dots,\mathbf{X}_{h_{N}}\\}$ is subject to a large
dimension factor $N$, heavily exaggerating the problem dimension and
jeopardizing the learning efficiency. Fortunately, we 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 a continuous-time perspective for a 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 a forward and a backward process for training and
sample generation, respectively:

$\displaystyle{\rm d}\mathbf{X}_{t}$
$\displaystyle=-\frac{1}{2}\mathbf{X}_{t}{\rm d}t+{\rm
d}\mathbf{W}_{t},~{}\text{for}~{}t\in[0,T]~{}\text{and}~{}\mathbf{X}_{0}\sim
P_{0},$ (Forward) $\displaystyle{\rm d}\mathbf{X}_{t}^{\leftarrow}$
$\displaystyle=\left[\frac{1}{2}\mathbf{X}_{t}^{\leftarrow}+\nabla\log
p_{T-t}(\mathbf{X}_{t}^{\leftarrow})\right]{\rm d}t+{\rm
d}\overline{\mathbf{W}}_{t},~{}\text{for}~{}t\in[0,T]~{}\text{and}~{}\mathbf{X}_{0}^{\leftarrow}\sim{\sf
N}(\mathbf{0},\mathbf{I}).$ (Backward)

Here, $\mathbf{W}_{t}$ and $\overline{\mathbf{W}}_{t}$ are independent Wiener
processes and $T$ is a finite terminal time. The initial distribution of
$\mathbf{X}_{0}$ is $P_{0}$, which is also the clean data distribution, and we
denote $P_{t}$ as the marginal distribution of $\mathbf{X}_{t}$. Since we are
corrupting $P_{0}$ by Gaussian noise, $P_{t}$ has a density function $p_{t}$.
Thus, $\nabla\log p_{t}$ in the backward process is recognized as the score
function, which is unknown and requires learning. Typically, the score
function will be parameterized by a neural network and trained by optimizing a
loss function, which we will introduce in Section 5. When generating new
samples, we simulate a discretized version of the backward process using the
learned score function.

For Gaussian process data, we interpret $\mathbf{X}_{t}$ as a vector in
$\mathbb{R}^{dN}$ 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 3 for a detailed discussion on the structure of
$\nabla\log p_{t}$.

##### Transformer Architecture

A transformer comprises a series of blocks and each block encompasses a multi-
head attention layer and a feedforward layer. Let
$\mathbf{Y}=[\mathbf{y}_{1},\dots,\mathbf{y}_{N}]\in\mathbb{R}^{D\times N}$ be
the (column) stacking matrix of $N$ patches. In a transformer block, multi-
head attention computes

$\displaystyle\textstyle{\tt
Attn}(\mathbf{Y})=\mathbf{Y}+\sum_{m=1}^{M}\mathbf{V}^{m}\mathbf{Y}\cdot\sigma\left((\mathbf{Q}^{m}\mathbf{Y})^{\top}\mathbf{K}^{m}\mathbf{Y}\right),$
(1)

where $\mathbf{V}^{m},\mathbf{Q}^{m}$, and $\mathbf{K}^{m}$ are weight
matrices of corresponding sizes in the $m$-th attention head, and $\sigma$ is
an activation function. The attention layer is followed by a feedforward
layer, which computes

$\displaystyle{\tt FFN}(\mathbf{Y})=\mathbf{Y}+\mathbf{W}_{1}\cdot{\rm
ReLU}(\mathbf{W}_{2}\mathbf{Y}+\mathbf{b}_{2}\mathbf{1}^{\top})+\mathbf{b}_{1}\mathbf{1}^{\top}.$

Here, $\mathbf{W}_{1},\mathbf{W}_{2}$ are weight matrices, $\mathbf{b}_{1}$
and $\mathbf{b}_{2}$ are offset vectors, $\mathbf{1}$ denotes a vector of
ones, and the ReLU activation function is applied entrywise. For our study,
the raw input to a transformer is $N$ patches of $d$-dimensional vectors and
time $t$ in the backward process. We refer to ${\mathcal{T}}(D,L,M,B,R)$ as a
transformer architecture defined by

$\displaystyle{\mathcal{T}}(D,L,M,B,R)=\\{f$ $\displaystyle:f=f_{\rm
out}\circ({\tt FFN}_{L}\circ{\tt Attn}_{L})\circ\dots\circ({\tt
FFN}_{1}\circ{\tt Attn}_{1})\circ f_{\rm in},$ $\displaystyle\quad\text{${\tt
Attn}_{i}$ uses entrywise ReLU activation for $i=1,\dots,L$},$
$\displaystyle\quad\text{number of heads in each~{}{\tt Attn}~{}is bounded by
$M$},$ $\displaystyle\quad\text{the Frobenius norm of each weight matrix is
bounded by $B$},$ $\displaystyle\quad\text{the output range $\left\lVert
f\right\rVert_{2}$ is bounded by $R$}\\}.$

See Figure 2 for an illustration and explanation of $f_{\rm in}$ and $f_{\rm
out}$. For attention layers, we consider ReLU activation for technical
convenience and postpone a discussion with softmax activation to Appendix B.5.

Figure 2: Transformer architecture. Here $f_{\rm in}$ is a linear layer to
lift input patch to $\mathbb{R}^{D}$, which appends the input raw data with
time embedding and other useful information. After passing through $L$
transformer blocks, $f_{\rm out}$ projects each patch into the data original
dimension $\mathbb{R}^{d}$ and clip the output range by $R$. We allow the
output range to be time $t$ dependent (denoted as $R_{t}$).

##  3 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, we study how can transformers
effectively represent the score function of our Gaussian process data. We are
particularly interested in understanding how can transformers capture the
spatial-temporal dependencies via multi-head attention.

###  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 time $t\in(0,T]$, given
$\mathbf{v}_{t}=[\mathbf{x}_{1,t}^{\top},\dots,\mathbf{x}_{N,t}^{\top}]^{\top}$
in the backward process, the score function is

$\displaystyle\nabla\log
p_{t}(\mathbf{v}_{t})=(\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I})^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\quad\text{for}\quad\alpha_{t}=e^{-t/2}~{}\text{and}~{}\sigma_{t}=\sqrt{1-e^{-t}},$
(2)

where $\bm{\mu}$ is the mean vector. We defer the derivation to Appendix A.1.
Observe that the score function for the $i$-th patch depends on the evolution
of all patches, reflecting temporal dependencies in sequential data. Moreover,
the influence of each patch is determined by the covariance function
$\bm{\Gamma}_{ij}=\gamma(h_{i},h_{j})$. In the extreme case of
$\bm{\Gamma}_{ij}=\mathds{1}\\{i=j\\}$, i.e., there is no correlation between
patches, the score function reduces to
$(\alpha_{t}^{2}\mathbf{I}\otimes\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I})^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})$,
which isolates patches and reproduces the score function of the Gaussian
distribution ${\sf N}(\bm{\mu}_{i},\bm{\Sigma})$ for each patch.

Although the score function assumes the closed-form expression in (2),
effectively representing it may suffer from difficulties. In fact, there are
dependencies among $N$ patches and the correlation is encoded by the inverse
covariance matrix
$(\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I})^{-1}$.
Directly representing it using a transformer is subject to high complexity and
a 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 time $t\in(0,T]$. It is straightforward to check by the
first-order optimality that $\nabla\log p_{t}(\mathbf{v}_{t})$ is the
minimizer of the following quadratic objective function,

$\displaystyle\nabla\log
p_{t}(\mathbf{v}_{t})=\mathop{\mathrm{argmin}}_{\mathbf{s}\in\mathbb{R}^{dN}}~{}\mathcal{L}_{t}(\mathbf{s}):=\frac{1}{2}\mathbf{s}^{\top}\left(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)\mathbf{s}-(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})^{\top}\mathbf{s}.$
(3)

Examining (3) reveals that the objective function is strongly convex as long
as $t>0$. Moreover importantly, the formulation (3) is free of matrix inverse
and the optimal solution can be found by a gradient descent algorithm.
Specifically, in the $k$-th iteration, gradient descent for $\mathcal{L}_{t}$
computes

$\displaystyle\mathbf{s}^{(k+1)}=\mathbf{s}^{(k)}-\eta_{t}\cdot\left(\left(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)\mathbf{s}^{(k)}-(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\right),$
(4)

where $\eta_{t}$ 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
$\bm{\Gamma}\otimes\bm{\Sigma}$, the gradient descent algorithm converges
exponentially fast for approximating the ground-truth score function.
Moreover, we could substitute $\bm{\Gamma}$ in (4) by any of its approximation
$\bar{\bm{\Gamma}}$. We quantify the representation error after sufficient
gradient descent iterations.

######  Lemma 1 (Gradient Descent Iterate Approximates the Score Function).

For an arbitrarily fixed time $t\in[0,T]$ and $\mathbf{v}_{t}$, given an error
tolerance $\epsilon>0$ and any integer $J<N$, if $\bar{\bm{\Gamma}}$ with
$\bar{\bm{\Gamma}}_{ij}=\bm{\Gamma}_{ij}\mathds{1}\\{|i-j|<J\\}$ is positive
semidefinite, then running gradient descent in (4) with a suitable step size
$\eta_{t}$ for
$K=\mathcal{O}{\left(\kappa_{t}\log{\left(1/{\epsilon}\right)}\right)}$
iterations gives rise to

$\displaystyle\left\lVert\mathbf{s}^{(K)}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}\leq\underbrace{\frac{1}{\sigma_{t}^{2}}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\epsilon}_{\mathcal{E}_{1}:\text{GD
representation
error}}+\underbrace{\frac{\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}}{\sigma_{t}^{4}}\sqrt{\sum_{{\left|i-j\right|}\geq
J}\bm{\Gamma}_{ij}^{2}}}_{\mathcal{E}_{2}:\text{correlation truncation
error}},$

where
$\kappa_{t}=\kappa{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}$
is the condition number of
$\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$.

The proof is deferred to Appendix A.2. To interpret the lemma, we first
consider $J=N$, which implies $\mathcal{E}_{2}=0$ and $\mathcal{E}_{1}$
recovers the typical convergence guarantee of gradient descent algorithm: for
a 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 we aim
to model in the score function. In particular, instead of working with (4), we
consider a surrogate gradient descent iteration driven by replacing
$\bm{\Gamma}$ by $\bar{\bm{\Gamma}}$, which neglects correlations beyond
length $J$. We discuss sufficient conditions for ensuring $\bar{\bm{\Gamma}}$
being positive semidefinite after Assumption 1. Introducing
$\bar{\bm{\Gamma}}$ incurs the truncation error $\mathcal{E}_{2}$, whose
magnitude depends on the pattern of temporal dependencies. Apparently, with
long-horizon dependencies, the truncation error $\mathcal{E}_{2}$ 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 $\mathcal{E}_{2}$

In many dynamical systems, temporal dependencies decay rather fast as a
function of the time gap. The decay pattern may lead to a well controlled
truncation error $\mathcal{E}_{2}$. As the decay is determined by the
covariance function in a Gaussian process, we impose the following assumption.

######  Assumption 1.

Let $\mathbf{e}_{1},\dots,\mathbf{e}_{N}\in\mathbb{R}^{d_{e}}$ be
$d_{e}$-dimensional time embedding of $h_{1},\dots,h_{N}$ with
$\left\lVert\mathbf{e}_{i}\right\rVert_{2}=r$ for each $i$ such that there
exists a positive and strictly increasing function $f$ with
$\left\lVert\mathbf{e}_{i}-\mathbf{e}_{j}\right\rVert_{2}=f(|i-j|)\geq
c{\left|i-j\right|}$ for an absolute constant $c>0$. Further, the covariance
function $\gamma$ satisfies

$\displaystyle\gamma(h_{i},h_{j})=\exp\left(-\left\lVert\mathbf{e}_{i}-\mathbf{e}_{j}\right\rVert_{2}^{\nu}/\ell\right)\quad\text{for}\quad\nu\in[1,2]\quad\text{and}\quad\ell>0.$

Firstly, Assumption 1 says that the time embedding (a.k.a. position embedding)
preserves the gap between real times. For example, a common time embedding
used in sequential data modeling (Vaswani et al., 2017) is sinusoidal
transformations, where
$\mathbf{e}_{i}=[r\sin(2i\pi/C),r\cos(2i\pi/C)]^{\top}\in\mathbb{R}^{2}$ for a
positive radius $r$ and a constant $C>0$. We can check that
$\left\lVert\mathbf{e}_{i}-\mathbf{e}_{j}\right\rVert_{2}=2r\sin({\left|i-j\right|}\pi/C)\geq
4r{\left|i-j\right|}/C$ is positive and approximately linearly increasing for
a sufficiently large $C$.

Secondly, the covariance function decays exponentially fast and the speed is
controlled by the exponent $\nu$ and the bandwidth $\ell$. Large $\ell$ or
small $\nu$ indicates that the correlation between different times decays
relatively slowly. Thus, the sequential data has some long-horizon
dependencies. The range of $\nu$ includes the well-known quadratic-exponential
(Gaussian) covariance function ($\nu=2$). Moreover, when $\nu=1$, the
covariance function coincides with the correlation in a Brownian motion.
Varying $\nu\in[1,2]$ can capture abundant temporal dependency patterns.

Besides, under Assumption 1, we can provide a sufficient condition for
ensuring $\bar{\bm{\Gamma}}$ being positive semidefinite for any $J$.
Specifically, when $\ell\leq c^{\nu}$, $\bm{\Gamma}$ is symmetric diagonally
dominant, i.e., the diagonal entry has a larger magnitude than the sum of the
magnitudes of off diagonal entries. This implies that for any truncation
length $J$, $\bar{\bm{\Gamma}}$ is always positive semidefinite. See Remark 1
in Appendix A.3 for a formal justification. Apparently, requiring $\ell\leq
c^{\nu}$ is not necessary and as long as the covariance function decays
sufficiently fast, $\bar{\bm{\Gamma}}$ can be positive semidefinite. We now
establish the following corollary to demonstrate a reasonable choice of $J$.

######  Corollary 1 (Correlation Truncation with Decay).

Suppose Assumption 1 holds with $\ell\leq c^{\nu}$. For any
$\epsilon<\lambda_{\min}(\bm{\Gamma})$, under the setup of Lemma 1, by setting
$J=\mathcal{O}{\left((\ell\log(N/(\epsilon\sigma_{t})))^{1/\nu}\right)}$, it
holds that

$\displaystyle\left\lVert\mathbf{s}^{(K)}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}\leq
2\sigma_{t}^{-2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\epsilon.$

The proof is deferred to Appendix A.3, where we keep the $\ell\leq c^{\nu}$
condition for technical convenience. This should not be considered
restrictive, as our theory holds as long as the truncation of $\bm{\Gamma}$ at
length $J$ is positive semidefinite. Corollary 1 shows that the truncation
error $\mathcal{E}_{2}$ can be controlled at the same order of
$\mathcal{E}_{1}$. The truncation length $J$ is logarithmically dependent on
the full length of the sequence, indicating that we can focus on relatively
short-horizon dependencies proportional to the bandwidth $\ell$. This
observation is the key to promote the representation and learning efficiency
of diffusion transformer.

##  4 Approximation Theory of Score Function using Transformers

This section devotes to establishing a transformer approximation theory of the
score function. Different 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 $\ell\leq c^{\nu}$. Given any $t_{0}\in(0,T]$
and a small $\epsilon<\lambda_{\min}(\bm{\Gamma})$, there exists a transformer
architecture ${\mathcal{T}}(D,L,M,B,R_{t})$ such that with proper weight
parameters, it yields an approximation $\widetilde{\mathbf{s}}$ to the score
function $\nabla\log p_{t}$ with

$\displaystyle\int\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}p_{t}(\mathbf{v}_{t}){\rm
d}\mathbf{v}_{t}\leq\sigma_{t}^{-2}\epsilon\quad\text{for all}\quad
t\in[t_{0},T].$

The transformer architecture satisfies

$\displaystyle
D=\mathcal{O}(d+d_{e}),~{}~{}L=\mathcal{O}{\left(\kappa_{t_{0}}\log^{2}(Nd/(\epsilon\sigma_{t_{0}}))\right)},~{}~{}M=\mathcal{O}{\left({\left(\ell\log(Nd\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/(\epsilon\sigma_{t_{0}}))\right)}^{1/\nu}\right)},$ $\displaystyle\hskip
10.00002ptB=\mathcal{O}{\left(\log(Nd/(\epsilon\sigma_{t_{0}}))\sigma^{-2}_{t_{0}}Nd(r^{2}+\left\lVert\bm{\Sigma}\right\rVert_{\infty})\right)},~{}~{}R_{t}=\mathcal{O}{\left({\log(Nd/(\epsilon\sigma_{t_{0}}))\sqrt{Nd}}/{\sigma_{t}}\right)}.$

The proof is deferred to Appendix B.1. Here we observe that the approximation
error depends on $\sigma_{t}$, 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 is very different from existing works. Our results also hold for
softmax activated transformer architectures, which is discussed in Appendix
B.5. We remark on other interpretations.

Figure 3: Construction of score function approximation using a transformer. By
rewriting the score function as the optimizer of a quadratic objective
function, we use gradient descent algorithm to approximate the optimizer. We
allow correlation truncation to 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 is realized by a multiplication module
followed by two transformer blocks. In the first transformer block, its
attention layer calculates correlation $\bm{\Gamma}_{ij}$ utilizing time
embedding. The second transformer block calculates the linear offset
$-\eta_{t}\sigma_{t}^{2}\mathbf{s}+\eta_{t}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})$
in (4).

##### A Glimpse of Transformer Architecture

The constructed transformer architecture is demonstrated in Figure 3. To
achieve an approximation, the transformer unrolls the gradient descent
algorithm. For realizing a single step, a multiplication module calculates
time dependent rescaling parameters, e.g., $\alpha_{t}$ and $\sigma_{t}$. Then
two transformer blocks implement the iteration in (4). As can 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 the condition number $\kappa_{t_{0}}$.
For sharply decaying covariance functions, $\kappa_{t_{0}}$ is a constant.
However, for slowly decaying covariance functions, $\kappa_{t_{0}}$ can be
large, indicating the fundamental difficulty of capturing long-horizon
dependencies. While explicitly bounding $\kappa_{t_{0}}$ for a finite $N$ goes
beyond the current technical limit, we discuss its asymptotic behavior in
Appendix E.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 $\mathbf{Q}$ and $\mathbf{K}$
matrices, so that it finds the correlation between a pair of time embedding
$\mathbf{e}_{i}$ and $\mathbf{e}_{j}$. Specifically, it calculates the inner
product $\mathbf{e}_{i}^{\top}\mathbf{e}_{j}$ for approximating the
correlation coefficient $\gamma(h_{i},h_{j})$. 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.

##  5 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 to estimate
the score function. Conceptually, we can use a quadratic loss,

$\displaystyle\textstyle\mathop{\mathrm{argmin}}_{\mathbf{s}\in{\mathcal{T}}}\int_{0}^{T}\mathbb{E}_{\mathbf{v}_{t}\sim
P_{t}}\left\lVert\mathbf{s}(\mathbf{v}_{t},t)-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}{\rm d}t,$ (5)

However, this loss function is not directly implementable due to $\nabla\log
p_{t}$ being unknown and numerical instability when $t$ approaches zero.
Therefore, we consider the following loss for practical implementation,

$\displaystyle\textstyle\mathop{\mathrm{argmin}}_{\mathbf{s}\in{\mathcal{T}}}\int_{t_{0}}^{T}\mathbb{E}_{\mathbf{v}_{0}}\mathbb{E}_{\mathbf{v}_{t}\sim{\sf
N}(\alpha_{t}\mathbf{v}_{0},\sigma_{t}^{2}\mathbf{I})}\left\lVert\mathbf{s}(\mathbf{v}_{t},t)-\frac{\alpha_{t}\mathbf{v}_{0}-\mathbf{v}_{t}}{\sigma_{t}^{2}}\right\rVert_{2}^{2}{\rm
d}t.$ (6)

Here, $t_{0}$ is an early-stopping time and
$\frac{\alpha_{t}\mathbf{v}_{0}-\mathbf{v}_{t}}{\sigma_{t}^{2}}$ substitutes
the unknown score function $\nabla\log p_{t}$. The equivalence of (6) to (5)
is established in the seminal works (Vincent, 2011; Hyvärinen and Dayan,
2005). When given the collected data set
$\mathcal{D}=\\{\mathbf{x}_{1,j},\dots,\mathbf{x}_{N,j}\\}_{j=1}^{n}$, we
replace the population expectation $\mathbb{E}_{\mathbf{v}_{0}}$ in (6) by a
sample empirical distribution. We denote
$\mathbf{v}^{i}_{0}=[\mathbf{x}_{i,1}^{\top},\dots,\mathbf{x}_{i,N}^{\top}]^{\top}$
as the stacking vector of a data sequence. Then the estimated score function
$\widehat{\mathbf{s}}$ can be written as an empirical risk minimizer,

$\displaystyle\textstyle\widehat{\mathbf{s}}\in\mathop{\mathrm{argmin}}_{\mathbf{s}\in{\mathcal{T}}}\frac{1}{n}\sum_{i=1}^{n}\int_{t_{0}}^{T}\mathbb{E}_{\mathbf{v}_{t}\sim{\sf
N}(\alpha_{t}\mathbf{v}_{0}^{i},\sigma_{t}^{2}\mathbf{I})}\left\lVert\mathbf{s}(\mathbf{v}_{t},t)-\frac{\alpha_{t}\mathbf{v}_{0}^{i}-\mathbf{v}_{t}}{\sigma_{t}^{2}}\right\rVert_{2}^{2}{\rm
d}t.$

To generate new samples, diffusion transformer uses $\widehat{\mathbf{s}}_{n}$
in the backward process. Correspondingly, we denote the distribution learned
by such a diffusion transformer as $\widehat{P}$. We bound the divergence of
$\widehat{P}$ to our ground-truth data distribution $P_{0}$ in the following
theorem.

######  Theorem 2 (Sample Complexity of Diffusion Transformer).

Suppose Assumption 1 holds with $\ell\leq c^{\nu}$. Assume there exists a
constant $C>1$ such that
$C^{-1}\leq\lambda_{\min}(\bm{\Sigma})\leq\lambda_{\max}(\bm{\Sigma})\leq C$
and $\left\lVert\bm{\mu}\right\rVert_{\infty}\leq C$. We choose the
transformer architecture ${\mathcal{T}}(D,L,M,B,R_{t})$ as in Theorem 1 with
$\epsilon=1/n$. By setting the terminal time $T=\log(n)$, considering
$\widehat{P}$ generated by the empirical risk minimizer
$\widehat{\mathbf{s}}$, we have

$\displaystyle\mathbb{E}_{\mathcal{D}}\left[{\rm
TV}(\widetilde{P}_{0},\widehat{P})\right]\lesssim\sqrt{\frac{\ell^{1/\nu}\kappa_{t_{0}}^{2}Nd^{3}}{n}}\cdot\log^{\frac{5\nu+1}{2\nu}}{\left(\kappa_{t_{0}}ndNt_{0}^{-1}\right)},$
(7)

where $\widetilde{P}_{0}$ is a perturbed data distribution satisfying
$W_{2}(P_{0},\widetilde{P}_{0})\lesssim\ell\sqrt{t_{0}Nd}$.

The proof is deferred to Appendix C. We remark that the emergence of
$\widetilde{P}_{0}$ owes to the early-stopping in score estimation, which is
obtained by evolving $P_{0}$ along the forward process for time $t_{0}$. The
optimal choice on $t_{0}$ depends on the condition number $\kappa_{t_{0}}$.
For fast decay covariance functions, $\kappa_{t_{0}}=\mathcal{O}(1)$ and we
can choose $t_{0}=1/n$ and the generalization error is in the order of
$\widetilde{\mathcal{O}}(\sqrt{\ell Nd^{3}/n})$. In this case, we have a
relatively weak dependence on $N$, demonstrating the efficiency of diffusion
transformers in sequential data modeling. On the other hand, when
$\kappa_{t_{0}}$ is large, i.e., with the presence of long-horizon
dependencies, the learning efficiency suffers from heavier dependence on $N$.
We experiment with various decaying patterns in Section 6 and show the
corresponding performance.

##  6 Numerical Results

In this section, we conduct experiments on diffusion transformer for learning
synthetic Gaussian process data. We test various aspects of our theory,
including learning efficiency and the capture of spatial-temporal
dependencies.

##### Experiment Setup

We consider synthetic Gaussian process data with $d=8$ and $N=128$. We stick
to the stationary setting, where the covariance matrix is generated by
$\bm{\Sigma}=\mathbf{A}^{\top}\mathbf{A}\in\mathbb{R}^{8\times 8}$ for a
Gaussian random matrix $\mathbf{A}\in\mathbb{R}^{8\times 8}$. We set the mean
vector $\bm{\mu}=\mathbf{0}$ for simplicity. The covariance function is chosen
as in Assumption 1 with
$\gamma(h_{i},h_{j})=\exp\left(-|h_{i}-h_{j}|^{\nu}/\ell\right).$ We set
$\ell$ and $\nu$ as hyperparameters and report their influences on the
learning. We generate $n\in\\{1000,3200,10000,32000,100000\\}$ sequences from
the Gaussian process as our training data in different settings. The diffusion
transformer 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 D.

##### Influence of Covariance Function Decay on Capturing Spatial-Temporal
Dependencies

We study the influence of covariance functions on the performance of diffusion
transformer. We vary the exponent $\nu$ and bandwidth $\ell$ as follows: 1) we
keep $\ell=64$, while $\nu$ varies in $\\{1,2^{1/4},2^{1/2},2^{3/4},2\\}$; 2)
we keep $\nu=2$, while $\ell$ varies in $\\{$4,16,64,256,1024$\\}$. After
training, we generate 1000 samples at each time step for performance
evaluation. The metric is a relative error $\epsilon$ of the estimated sample
covariance matrix to its ground-truth (see a definition in Appendix D), which
is reported in Figure LABEL:fig:error-n-v(a).

##### 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 time
embedding and a data vector written as
$[\mathbf{z}_{i,t}^{\top},\mathbf{e}_{i}^{\top}]^{\top}\in\mathbb{R}^{32}$,
where $\mathbf{e}_{i}\in\mathbb{R}^{16}$ is the time embedding and
$\mathbf{z}_{i,t}\in\mathbb{R}^{16}$ is the output of the patchify module in
DiT, obtained by a linear transformation on raw data. We plot the heat map of
query and key matrices in the $5$th transformer block in Figure
LABEL:fig:error-n-v(b). Plots for other transformer blocks are provided in
Appendix D. As can be seen, the interaction between time embedding is dominant
(bottom-right block), which aligns with our approximation theory for
constructing the transformer architecture.

##### Backward Diffusion Process Unveils the Kernel Matrix in Attention Scores

To further understand how DiT captures the temporal dependencies, i.e., matrix
$\bm{\Gamma}$ for Gaussian process data, we plot the evolution of score
matrices $(\mathbf{Q}\mathbf{Y})^{\top}\mathbf{K}\mathbf{Y}$ 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 $\bm{\Gamma}$. 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 truth
temporal dependencies. In further subsequent layers, we observe that the
learned pattern of temporal dependencies is kept. We refer readers to Figure 8
in Appendix D for a complete plot of score matrices in each attention 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, we observe that the temporal dependencies are well captured
starting from the 3rd layer.

##  7 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 a 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 to 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.

## References

  * Alcaraz and Strodthoff (2022) Juan Miguel Lopez Alcaraz and Nils Strodthoff.  Diffusion-based time series imputation and forecasting with structured state space models.  _arXiv preprint arXiv:2208.09399_ , 2022.
  * Austin et al. (2021) 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:17981–17993, 2021.
  * Banerjee et al. (2013) Anjishnu Banerjee, David B Dunson, and Surya T Tokdar.  Efficient gaussian process regression for large datasets.  _Biometrika_ , 100(1):75–89, 2013.
  * Bao et al. (2023) 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 on Machine Learning_ , pages 1692–1717. PMLR, 2023.
  * Benton et al. (2022) Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet.  From denoising diffusions to denoising markov models.  _arXiv preprint arXiv:2211.03595_ , 2022.
  * Benton et al. (2023) Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis.  Linear convergence bounds for diffusion models via stochastic localization.  _arXiv preprint arXiv:2308.03686_ , 2023.
  * Benton et al. (2024) Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis.  Nearly $d$-linear convergence bounds for diffusion models via stochastic localization.  In _Proceedings of the International Conference on Learning Representations_ , 2024.
  * Borovitskiy et al. (2021) 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.
  * Bubeck (2015) Sébastien Bubeck.  Convex optimization: Algorithms and complexity.  _Foundations and Trends® in Machine Learning_ , 8(3-4):231–357, 2015.
  * Canonne (2023) Clément L. Canonne.  A short note on an inequality between kl and tv, 2023.
  * Chan (2024) Stanley H Chan.  Tutorial on diffusion models for imaging and vision.  _arXiv preprint arXiv:2403.18103_ , 2024.
  * Chen et al. (2022a) 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–1253, 2022a.
  * Chen et al. (2023a) 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 on Machine Learning_ , pages 4672–4712. PMLR, 2023a.
  * Chen et al. (2024) Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang.  An overview of diffusion models: Applications, guided generation, statistical rates and optimization.  _arXiv preprint arXiv:2404.07771_ , 2024.
  * Chen et al. (2022b) Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang.  Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.  _arXiv preprint arXiv:2209.11215_ , 2022b.
  * Chen et al. (2023b) Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim.  The probability flow ode is provably fast.  _arXiv preprint arXiv:2305.11798_ , 2023b.
  * Chen et al. (2023c) Sitan Chen, Giannis Daras, and Alex Dimakis.  Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers.  In _International Conference on Machine Learning_ , pages 4462–4484. PMLR, 2023c.
  * Chen et al. (2021) Yifan Chen, Bamdad Hosseini, Houman Owhadi, and Andrew M Stuart.  Solving and learning nonlinear pdes with gaussian processes.  _Journal of Computational Physics_ , 447:110668, 2021.
  * Cybenko (1989) George Cybenko.  Approximation by superpositions of a sigmoidal function.  _Mathematics of Control, Signals and Systems_ , 2(4):303–314, 1989.
  * Damian et al. (2022) 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.
  * Delsarte and Genin (2005) Philippe Delsarte and Y Genin.  Spectral properties of finite toeplitz matrices.  In _Mathematical Theory of Networks and Systems: Proceedings of the MTNS-83 International Symposium Beer Sheva, Israel, June 20–24, 1983_ , pages 194–213. Springer, 2005.
  * Deringer et al. (2021) Volker L Deringer, Albert P Bartók, Noam Bernstein, David M Wilkins, Michele Ceriotti, and Gábor Csányi.  Gaussian process regression for materials and molecules.  _Chemical Reviews_ , 121(16):10073–10141, 2021.
  * Ding and Jin (2023) Zihan Ding and Chi Jin.  Consistency models as a rich and efficient policy class for reinforcement learning.  _arXiv preprint arXiv:2309.16984_ , 2023.
  * Fu et al. (2024) Hengyu Fu, Zhuoran Yang, Mengdi Wang, and Minshuo Chen.  Unveil conditional diffusion models with classifier-free guidance: A sharp statistical theory.  _arXiv preprint arXiv:2403.11968_ , 2024.
  * Guo et al. (2023) Zhiye Guo, Jian Liu, Yanli Wang, Mengrui Chen, Duolin Wang, Dong Xu, and Jianlin Cheng.  Diffusion models in bioinformatics: A new wave of deep learning revolution in action.  _arXiv preprint arXiv:2302.10907_ , 2023.
  * Gupta et al. (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.06662_ , 2023.
  * Hansen-Estruch et al. (2023) Philippe Hansen-Estruch, 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.
  * Hartman and Wintner (1950) Philip Hartman and Aurel Wintner.  On the spectra of toeplitz’s matrices.  _American Journal of Mathematics_ , 72(2):359–366, 1950.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Hyvärinen and Dayan (2005) Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(4), 2005.
  * Jiao et al. (2024) Yuling Jiao, Lican Kang, Huazhen Lin, Jin Liu, and Heng Zuo.  Latent schrödinger bridge diffusion model for generative learning.  _arXiv preprint arXiv:2404.13309_ , 2024.
  * Kargin (2009) Vladislav Kargin.  Spectrum of random toeplitz matrices with band structure.  2009\.
  * Lee et al. (2022a) Holden Lee, Jianfeng Lu, and Yixin Tan.  Convergence for score-based generative modeling with polynomial complexity.  _arXiv preprint arXiv:2206.06227_ , 2022a.
  * Lee et al. (2022b) Holden Lee, Jianfeng Lu, and Yixin Tan.  Convergence of score-based generative modeling for general data distributions.  _arXiv preprint arXiv:2209.12381_ , 2022b.
  * Li et al. (2023) Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi.  Towards faster non-asymptotic convergence for diffusion-based generative models.  _arXiv preprint arXiv:2306.09251_ , 2023.
  * Li et al. (2024a) 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.
  * Li et al. (2024b) 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.13219_ , 2024b.
  * Liu et al. (2024a) Chenghao Liu, Enming Liang, and Minghua Chen.  Characterizing resnet’s universal approximation capability, 2024a.  URL https://openreview.net/forum?id=PCTqol2hvy.
  * Liu et al. (2024b) Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, and Jianfeng Gao.  Sora: A review on background, technology, limitations, and opportunities of large vision models.  _arXiv preprint arXiv:2402.17177_ , 2024b.
  * Mei (2024) Song Mei.  U-nets as belief propagation: Efficient classification, denoising, and diffusion in generative hierarchical models.  _arXiv preprint arXiv:2404.18444_ , 2024.
  * Mei and Wu (2023) Song Mei and Yuchen Wu.  Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models.  _arXiv preprint arXiv:2309.11420_ , 2023.
  * Monga et al. (2021) 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.
  * Oko et al. (2023) Kazusato Oko, Shunta Akiyama, and Taiji Suzuki.  Diffusion models are minimax optimal distribution estimators.  _arXiv preprint arXiv:2303.01861_ , 2023.
  * Pearce et al. (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.10677_ , 2023.
  * Peebles and Xie (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.
  * Prato and Tubaro (2007) Giuseppe Da Prato and Luciano Tubaro.  Wick powers in stochastic pdes: an introduction.  2007\.  URL https://api.semanticscholar.org/CorpusID:55493217.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18_ , pages 234–241. Springer, 2015.
  * Seeger (2004) Matthias Seeger.  Gaussian processes for machine learning.  _International journal of neural systems_ , 14(02):69–106, 2004.
  * Song et al. (2020) Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_ , 2020.
  * Song and Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Sun et al. (2022) Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai.  Score-based continuous-time discrete diffusion models.  _arXiv preprint arXiv:2211.16750_ , 2022.
  * Tang and Yang (2024) 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.
  * Tang and Zhao (2024) Wenpin Tang and Hanyang Zhao.  Score-based diffusion models via stochastic differential equations–a technical tutorial.  _arXiv preprint arXiv:2402.07487_ , 2024.
  * Tashiro et al. (2021) Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon.  CSDI: Conditional score-based diffusion models for probabilistic time series imputation.  _Advances in Neural Information Processing Systems_ , 34:24804–24816, 2021.
  * Tian et al. (2023) 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 arXiv:2310.15290_ , 2023.
  * Tilli (1998) Paolo Tilli.  A note on the spectral distribution of toeplitz matrices.  _Linear and Multilinear Algebra_ , 45(2-3):147–159, 1998.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Vincent (2011) Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011.
  * Williams and Rasmussen (2006) Christopher KI Williams and Carl Edward Rasmussen.  _Gaussian processes for machine learning_ , volume 2.  MIT press Cambridge, MA, 2006.
  * Wu et al. (2024) Junde Wu, Wei Ji, Huazhu Fu, Min Xu, Yueming Jin, and Yanwu Xu.  Medsegdiff-v2: Diffusion-based medical image segmentation with transformer.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 38, pages 6030–6038, 2024.
  * Xu et al. (2022) Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang.  Geodiff: A geometric diffusion model for molecular conformation generation.  _arXiv preprint arXiv:2203.02923_ , 2022.
  * Yarotsky (2018) Dmitry Yarotsky.  Optimal approximation of continuous functions by very deep relu networks.  In _Conference on learning theory_ , pages 639–649. PMLR, 2018.
  * Zhang et al. (2023) 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 arXiv:2303.13336_ , 2, 2023.
  * Zhu et al. (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.

##  Appendix A Omitted Proofs in Section 3

###  A.1 Deriving Score Function for Gaussian Process Data

By the definition of Gaussian process, we know that the stacking vector
$[\mathbf{x}_{1}^{\top},\dots,\mathbf{x}_{N}^{\top}]^{\top}\in\mathbb{R}^{dN}$
follows a Gaussian distribution ${\sf
N}(\bm{\mu},\bm{\Gamma}\otimes\bm{\Sigma})$. Along the forward process, we
progressively add Gaussian noise to the initial data distribution and hence,
for any $\mathbf{v}_{t}$, we have

$\displaystyle p_{t}(\mathbf{v}_{t})$
$\displaystyle=\int\underbrace{\frac{1}{(2\pi\sigma_{t}^{2})^{dN/2}}\exp\left(-\frac{1}{2\sigma_{t}^{2}}\left\lVert\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}\right\rVert_{2}^{2}\right)}_{(A)}$
$\displaystyle\quad\cdot\underbrace{\frac{1}{(2\pi)^{dN/2}\sqrt{\det(\bm{\Gamma}\otimes\bm{\Sigma})}}\exp\left(-\frac{1}{2}(\mathbf{v}_{0}-\bm{\mu})^{\top}(\bm{\Gamma}\otimes\bm{\Sigma})^{-1}(\mathbf{v}_{0}-\bm{\mu})\right)}_{(B)}{\rm
d}\mathbf{v}_{0}.$

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
$p_{t}(\mathbf{v}_{t})$ is again a Gaussian distribution. By completing the
squares and some algebraic manipulation, we have

$\displaystyle
p_{t}(\mathbf{v}_{t})\propto\int\exp\left(-\frac{1}{2}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})^{\top}(\sigma_{t}^{2}\mathbf{I}+\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma})^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\right).$

As a sanity check, $p_{t}$ is now a Gaussian density function of ${\sf
N}(\alpha_{t}\bm{\mu},\sigma_{t}^{2}\mathbf{I}+\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma})$,
matching the marginal distribution of the forward process. Therefore, the
score function is

$\displaystyle\nabla\log
p_{t}(\mathbf{v}_{t})=\left(\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I}\right)^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu}).$

###  A.2 Proof of Lemma 1

To prove the Lemma, we first state a standard result in convex optimization.

######  Lemma 2 (Theorem 3.12 in Bubeck [2015]).

Let $f$ be $\beta$-smooth and $\alpha$-strongly convex on $\mathbb{R}^{d}$ and
$\mathbf{x}^{\star}$ is the global minimizer. Then gradient descent with
$\eta=\frac{2}{\alpha+\beta}$ satisfies

$\displaystyle\left\lVert\mathbf{x}^{(k+1)}-\mathbf{x}^{\star}\right\rVert_{2}\leq{\left(\frac{\kappa-1}{\kappa+1}\right)}\left\lVert\mathbf{x}^{(k)}-\mathbf{x}^{\star}\right\rVert_{2},~{}~{}k=0,1,\dots.$

Here $\mathbf{x}^{(k+1)}=\mathbf{x}^{(k)}-\eta\nabla f(\mathbf{x}^{(k)})$ is
the outcome in $(k+1)-$th iteration of GD and $\kappa=\beta/\alpha$.

With the lemma above, the proof of Lemma 1 is quite straightworward.

###### Proof of Lemma 1.

Denote the truncated score function by

$\bar{\mathbf{s}}(\mathbf{v}_{t})={\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu}).$

Let’s consider the following quadratic target function

$\displaystyle\bar{\mathcal{L}}_{t}(\mathbf{s})=\frac{1}{2}\mathbf{s}^{\top}{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}\mathbf{s}-(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})^{\top}\mathbf{s},$

And we iterate $\mathbf{s}^{(k)}$ by gradient descent on this target function
with $\mathbf{s}^{(0)}=\mathbf{0}$. Since $\bar{\mathcal{L}}_{t}$ is
$\lambda_{\max}{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}$-smooth
and
$\lambda_{\min}{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}$-strongly
convex, by the approximate GD formula (4) and Lemma 1, for any $k\geq 0$, we
have

$\displaystyle\left\lVert\mathbf{s}^{(k+1)}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq{\left(\frac{\bar{\kappa}_{t}-1}{\bar{\kappa}_{t}+1}\right)}\left\lVert\mathbf{s}^{(k)}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\exp{\left(-\frac{2(k+1)}{\bar{\kappa}_{t}+1}\right)}\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}.$

Here $\bar{\kappa}_{t}$ is the condition number of
$\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$.
Moreover, we have

$\displaystyle\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle={\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})$
$\displaystyle\leq\left\lVert{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}^{-1}\right\rVert_{2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle=\lambda_{\min}^{-1}{\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle\leq\frac{\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}}{\sigma_{t}^{2}}$

Thus, by taking the number of iterations
$K=\lceil\frac{\bar{\kappa}_{t}+1}{2}\log(1/\epsilon)\rceil$, we have

$\left\lVert\mathbf{s}^{(K)}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\frac{\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\epsilon}{\sigma_{t}^{2}}.$

Besides, denote
$\bm{\Phi}_{t}=\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$
and
$\bar{\bm{\Phi}}_{t}=\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$.
The difference between the truth score function and the truncated score
function is

$\displaystyle\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})-\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle=\left\lVert{\left({\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}^{-1}-{\left(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)}^{-1}\right)}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\right\rVert_{2}$
$\displaystyle=\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-\bm{\Phi}_{t}^{-1}\right)}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bar{\bm{\Phi}}_{t}^{-1}-\bm{\Phi}_{t}^{-1}\right\rVert_{2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle=\left\lVert\bm{\Phi}_{t}^{-1}{\left(\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right)}\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bm{\Phi}_{t}^{-1}\right\rVert_{2}\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{2}\left\lVert\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right\rVert_{2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle\leq\sigma_{t}^{-4}\left\lVert\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right\rVert_{2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}.$

Let’s focus on bounding
$\left\lVert\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right\rVert_{2}$. Without any
assumptions on $\bm{\Phi}_{t}$, a natural bound is

$\displaystyle\left\lVert\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bm{\Phi}_{t}-\bar{\bm{\Phi}}_{t}\right\rVert_{\rm
F}=\alpha^{2}_{t}\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}\sqrt{\sum_{{\left|i-j\right|}>J}\bm{\Gamma}_{ij}^{2}}.$ (8)

Altogether, we have

$\displaystyle\left\lVert\mathbf{s}^{(K)}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq\frac{\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\epsilon}{\sigma_{t}^{2}}+\sigma_{t}^{-4}\alpha^{2}_{t}\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}\sqrt{\sum_{{\left|i-j\right|}>J}\bm{\Gamma}_{ij}^{2}}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}$
$\displaystyle\leq{\left(\epsilon+\frac{\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}}{\sigma_{t}^{2}}\sqrt{\sum_{{\left|i-j\right|}>J}\bm{\Gamma}_{ij}^{2}}\right)}\sigma_{t}^{-2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}.$
(9)

In the last inequality, we invoke $\alpha_{t}^{2}\leq 1$ for any $t\geq 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

$\displaystyle\mathbf{s}^{(k+1)}=\mathbf{s}^{(k)}-\eta\nabla\mathcal{L}_{t}(\mathbf{s}^{(k+1)})+\bm{\xi}^{(k)}.$

Here $\bm{\xi}^{(k)}$ represents the approximation error. Then we have the
following convergence analysis:

######  Lemma 3.

Given $\epsilon>0$, if the approximation error in each step satisfies
$\left\lVert\bm{\xi}^{(k)}\right\rVert_{2}\leq\epsilon$, then after
$K=\lceil\frac{\kappa_{t}+1}{2}\log\left({1}/{\epsilon}\right)\rceil$ steps of
GD on minimizing the truncated target function $\bar{\mathcal{L}}_{t}$, we
have

$\displaystyle\left\lVert\mathbf{s}^{(K)}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq\frac{(\kappa_{t}+1)\epsilon}{2}+{\left(\epsilon+\frac{\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}}{\sigma_{t}^{2}}\sqrt{\sum_{{\left|i-j\right|}>J}\bm{\Gamma}_{ij}^{2}}\right)}\sigma_{t}^{-2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}.$

###### Proof of Lemma 3.

By a similar deduction in the proof of Lemma 1, for any $k\geq 0$, we have

$\displaystyle\left\lVert\mathbf{s}^{(k+1)}-\bar{\mathbf{s}}\right\rVert_{2}$
$\displaystyle=\left\lVert\mathbf{s}^{(k)}-\eta\nabla\mathcal{L}_{t}(\mathbf{s}^{(k+1)})+\bm{\xi}^{(k)}-\bar{\mathbf{s}}\right\rVert_{2}$
$\displaystyle\leq\left\lVert\mathbf{s}^{(k)}-\eta\nabla\mathcal{L}_{t}(\mathbf{s}^{(k+1)})-\bar{\mathbf{s}}\right\rVert_{2}+\left\lVert\bm{\xi}^{(k)}\right\rVert_{2}$
$\displaystyle\leq{\left(\frac{\kappa_{t}-1}{\kappa_{t}+1}\right)}\left\lVert\mathbf{s}^{(k)}-\bar{\mathbf{s}}\right\rVert_{2}+\epsilon.$

Thus, we have

$\displaystyle\left\lVert\mathbf{s}^{(k+1)}-\bar{\mathbf{s}}\right\rVert_{2}-\frac{(\kappa_{t}+1)\epsilon}{2}\leq{\left(\frac{\kappa_{t}-1}{\kappa_{t}+1}\right)}{\left(\left\lVert\mathbf{s}^{(k)}-\bar{\mathbf{s}}\right\rVert_{2}-\frac{(\kappa_{t}+1)\epsilon}{2}\right)}.$

Note that if there exists $k\leq K$ such that
$\left\lVert\mathbf{s}^{(k)}-\bar{\mathbf{s}}\right\rVert_{2}\leq{(\kappa_{t}+1)\epsilon}/{2}$,
we have for any $k_{1}\geq k$,
$\left\lVert\mathbf{s}^{(k_{1})}-\bar{\mathbf{s}}\right\rVert_{2}\leq{(\kappa_{t}+1)\epsilon}/{2}$,
so
$\left\lVert\mathbf{s}^{(K)}-\bar{\mathbf{s}}\right\rVert_{2}\leq{(\kappa_{t}+1)\epsilon}/{2}$,
which finishes the proof. Now we assume for any $0\leq k\leq K$,
$\left\lVert\mathbf{s}^{(K)}-\bar{\mathbf{s}}\right\rVert_{2}\geq{(\kappa_{t}+1)\epsilon}/{2}$.
Then we have

$\displaystyle\left\lVert\mathbf{s}^{(k)}-\bar{\mathbf{s}}\right\rVert_{2}-\frac{(\kappa_{t}+1)\epsilon}{2}$
$\displaystyle\leq{\left(\frac{\kappa_{t}-1}{\kappa_{t}+1}\right)}^{k}{\left(\left\lVert\mathbf{s}^{(0)}-\bar{\mathbf{s}}\right\rVert_{2}-\frac{(\kappa_{t}+1)\epsilon}{2}\right)}$
$\displaystyle\leq{\left(\frac{\kappa_{t}-1}{\kappa_{t}+1}\right)}^{k}\left\lVert\mathbf{s}^{(0)}-\bar{\mathbf{s}}\right\rVert_{2}$
$\displaystyle\leq\exp\left(-\frac{2k}{\kappa_{t}+1}\right)\left\lVert\mathbf{s}^{(0)}-\bar{\mathbf{s}}\right\rVert_{2}.$
(10)

Substituting
$k=K=\lceil\frac{\kappa_{t}+1}{2}\log\left(1/{\epsilon}\right)\rceil$ into the
inequality, we have

$\displaystyle\left\lVert\mathbf{s}^{(K)}-\bar{\mathbf{s}}\right\rVert_{2}-\frac{(\kappa_{t}+1)\epsilon}{2}\leq\epsilon\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\sigma_{t}^{-2}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\epsilon.$

Combining the result with (8), we complete our proof. ∎

###  A.3 Proof of Corollary 1

Let’s first analyse the approximation error by truncating $\bm{\Gamma}$. For
simplicity, we denote $\Delta\bm{\Gamma}=\bm{\Gamma}-\bar{\bm{\Gamma}}$.

######  Lemma 4.

Suppose Assumption 1 holds. Then for any $\epsilon>0$, by taking
$J=\left\lceil{\left(\frac{\ell}{2}\log(N\ell/\epsilon^{2})\right)}^{1/\nu}\right\rceil+1$,
we have $\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm F}\leq\epsilon$.

###### Proof of Lemma 4.

By Assumption 1, we have $f(m)\geq cm$. According to the definition of
$\Delta\bm{\Gamma}$, we have

$\displaystyle\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm F}^{2}$
$\displaystyle=\sum_{{\left|i-j\right|}\geq J}\gamma(h_{i},h_{j})^{2}$
$\displaystyle=\sum_{k=J}^{N-1}(2N-2k)\exp{\left(-\frac{2f(k)^{\nu}}{\ell}\right)}$
$\displaystyle\leq\sum_{k=J}^{N-1}(2N-2k)\exp{\left(-\frac{2(ck)^{\nu}}{\ell}\right)}$
$\displaystyle\leq
2N\sum_{k=J}^{N-1}\exp{\left(-\frac{2c^{\nu}k^{\nu}}{\ell}\right)}$
$\displaystyle\leq
2N\int_{J-1}^{\infty}\exp{\left(-\frac{2c^{\nu}t^{\nu}}{\ell}\right)}{\rm d}t$
$\displaystyle\leq
c^{-\nu}N\ell\exp{\left(-\frac{2c^{\nu}(J-1)^{\nu}}{\ell}\right)}.$

by taking
$J=\left\lceil{\left(\frac{\ell}{2c^{\nu}}\log{\left(N\ell/\epsilon^{2}c^{\nu}\right)}\right)}^{1/\nu}\right\rceil+1=\mathcal{O}(\ell\log(N/\epsilon)^{1/\nu})$,
we ensure that $\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}^{2}\leq\epsilon^{2}$. The proof is complete. ∎

######  Remark 1.

From the proof of Lemma 4, we also observe that $\bm{\Gamma}$ is diagonally
dominant, i.e., $\bm{\Gamma}_{ii}\geq\sum_{j\neq i}|\bm{\Gamma}_{ij}|$, when
we have $2\sum_{k=1}^{N-1}\exp{\left(-\frac{2f(k)^{\nu}}{\ell}\right)}\leq 1$.
According to the proof above, a sufficient condition for this to hold is
$c^{-\nu}\ell\leq 1$, i.e., $\ell\leq c^{\nu}$. When truncating $\bm{\Gamma}$
by any length $J$, $\bar{\bm{\Gamma}}$ remians diagonally dominant, therefore,
$\bar{\bm{\Gamma}}$ is positive semidefinite.

Now we turn back to the proof of Corollary 1. By taking

$J=\left\lceil{\left(\frac{\ell}{2c^{\nu}}\log{\left(N\ell/(\sigma_{t}^{4}\epsilon^{2}c^{\nu})\right)}\right)}^{1/\nu}\right\rceil=\mathcal{O}{\left({\left(\ell\log(N/(\epsilon\sigma_{t}))\right)}^{1/\nu}\right)},$

we ensure that
$\left\lVert\Delta\bm{\Gamma}\right\rVert_{F}^{2}\leq\sigma_{t}^{4}\epsilon^{2}$.
Then plugging the approximation error into (9) concludes our proof.

Moreover, we could bound the operator norm of $\bm{\Gamma}$ in the same way.
For any $\mathbf{v}\in\mathbb{R}^{N}$ such that
$\left\lVert\mathbf{v}\right\rVert_{2}=1$, we have

$\displaystyle{\left|\mathbf{v}^{\top}\bm{\Gamma}\mathbf{v}\right|}$
$\displaystyle=\sum_{i,j}\bm{\Gamma}_{i,j}v_{i}v_{j}$
$\displaystyle=1+\sum_{k=1}^{N-1}2\exp(-2c^{\nu}k^{\nu}/\ell)\sum_{i=1}^{N-k}v_{i}v_{i+k}$
$\displaystyle\leq
1+2\sum_{k=1}^{N-1}\exp(-2c^{\nu}k^{\nu}/\ell)\sqrt{\sum_{i=1}^{N-k}v_{i}^{2}\sum_{i=1}^{N-k}v_{i+k}^{2}}$
$\displaystyle\leq 1+2\sum_{k=1}^{N-1}\exp(-2c^{\nu}k^{\nu}/\ell)$
$\displaystyle\leq 1+c^{-\nu}\ell.$

Thus, we have

$\displaystyle\left\lVert\bm{\Gamma}\right\rVert_{2}\leq
1+c^{-\nu}\ell\lesssim 1+\ell.$ (11)

We will apply this bound in Lemma 12.

##  Appendix B Omitted Proofs in Section 4

###  B.1 Proof of Theorem 1

By the previous derivation, we know the truth score function is written as

$\displaystyle\nabla\log
p_{t}(\mathbf{v}_{t})=\left(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu}).$
(12)

Due to the fast decay of $\bm{\Gamma}$, we consider a truncated score function

$\displaystyle\bar{\mathbf{s}}(\mathbf{v}_{t})=\left(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right)^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})$

with $\bar{\bm{\Gamma}}_{ij}=\bm{\Gamma}_{ij}$ if ${\left|i-j\right|}<J$ and
$\bar{\bm{\Gamma}}_{ij}=0$ otherwise. Denoting
$\bm{\Gamma}=\bar{\bm{\Gamma}}+\Delta\bm{\Gamma}$ (see the formal statement in
Lemma 4), by appropriately choosing $M=\mathcal{O}((\log(1/\epsilon))^{1/v})$,
we could guarantee $\left\lVert\Delta\bm{\Gamma}\right\rVert\leq\epsilon$ for
any $\epsilon>0$. Now we decompose the $L_{2}$ error into the following two
items:

$\displaystyle\left\lVert\widetilde{\mathbf{s}}-\nabla\log
p_{t}\right\rVert_{L_{2}(P_{t})}\leq\underbrace{\left\lVert\widetilde{\mathbf{s}}-\bar{\mathbf{s}}\right\rVert_{L_{2}(P_{t})}}_{\text{Proposition
\ref{prop::transformers approximate
shat}}}+\underbrace{\left\lVert\bar{\mathbf{s}}-\nabla\log
p_{t}\right\rVert_{L_{2}(P_{t})}}_{\text{Lemma \ref{lemma::trunc Gamma L2}}}.$

Here
$\left\lVert\mathbf{s}\right\rVert_{L_{2}(P_{t})}^{2}=\int\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}^{2}p_{t}(\mathbf{v}_{t}){\rm
d}t$ is the squared $L_{2}$ norm w.r.t. a density function $p_{t}$ of
distribution $P_{t}$. We provide the error analysis for the two items in
Proposition 1 and Lemma 5 separately. Under these two supporting statements,
the proof of Theorem 1 is quite straightforward.

Now we state Proposition 1 and Lemma 5 with their detailed proof deferred to
next sections. Proposition 1 provides an approximation guarantee on the
difference between $\widetilde{\mathbf{s}}$ and $\bar{\mathbf{s}}$ in
$L_{2}(P_{t})$, which is our main result throughout the proof of Theorem 1.

######  Proposition 1.

Suppose Assumption 1 holds. Given $t_{0}\in(0,T]$, there exists a transformer
architecture $\widetilde{\mathbf{s}}\in{\mathcal{T}}(D,L,M,B,R_{t})$ such that
with proper weight parameters, it yields an approximation
$\widetilde{\mathbf{s}}$ to the truncated score function $\bar{\mathbf{s}}$
with

$\displaystyle\left\lVert\widetilde{\mathbf{s}}-\bar{\mathbf{s}}\right\rVert_{L_{2}(P_{t})}^{2}\leq\frac{\epsilon^{2}}{\sigma_{t}^{2}},~{}~{}~{}t\geq
t_{0}.$

The transformer architecture satisfies

$\displaystyle
D=9d+d_{t}+d_{e}+1,L=\mathcal{O}{\left(\kappa_{t_{0}}\log(Nd/(\epsilon\sigma_{t_{0}}))^{2}\right)},M=\mathcal{O}{\left({\left(\ell\log(Nd\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/(\epsilon\sigma_{t_{0}}))\right)}^{1/\nu}\right)},$ $\displaystyle\hskip
20.00003ptB=\mathcal{O}{\left(\log(Nd/(\epsilon\sigma_{t_{0}}))\sigma^{-2}_{t_{0}}Nd(r^{2}+\left\lVert\bm{\Sigma}\right\rVert_{\infty})\right)},R_{t}=\mathcal{O}{\left({\log(Nd/(\epsilon\sigma_{t_{0}}))\sqrt{Nd}}/{\sigma_{t}}\right)}.$

The proof is provided in Appendix B.2.

Moreover, Lemma 5 bounds the difference between the truncated score function
$\bar{\mathbf{s}}$ and the ground truth in $L_{2}(P_{t})$ distance.

######  Lemma 5.

Given any $\epsilon>0$ and $t>0$, by choosing

$J=\mathcal{O}{\left({\left(\ell\log(N\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/(\epsilon\sigma_{t}))\right)}^{1/\nu}\right)},$

it holds that for any $t>0$,

$\displaystyle\left\lVert\bar{\mathbf{s}}-\nabla\log
p_{t}\right\rVert_{L_{2}(P_{t})}\leq\frac{\epsilon}{\sigma_{t}}.$

The proof is provided in Appendix B.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 $\epsilon/\sigma_{t}$, we finish the proof of
Theorem 1.

###  B.2 Proof of Proposition 1

To prove Proposition 1, we first need a uniform approximation theory of
$\bar{\mathbf{s}}$ on a bounded region, which is the backbone of the proof.

######  Lemma 6.

Given any radius $R_{0}\geq 1$, error level $0\leq\epsilon<1$ and $t_{0}>0$,
there exists a transformer architecture ${\mathcal{T}}(D,L,M,B,R)$ that gives
rise to $\widetilde{\mathbf{s}}$ satisfying

$\displaystyle\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\epsilon,~{}~{}\text{for
any}\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
R_{0}\sigma_{t}^{-1},~{}~{}t\geq t_{0}.$

The transformer architecture satisfies

$\displaystyle
D=9d+d_{t}+d_{e}+1,~{}~{}L=\mathcal{O}{\left(\kappa_{t_{0}}\log(R_{0}Nd/(\epsilon\sigma_{t_{0}}))^{2}\right)},~{}~{}M=4J,$
$\displaystyle
B=\mathcal{O}{\left(\log(R_{0}Nd/(\epsilon\sigma_{t_{0}}))\sqrt{N}R_{0}d\sigma^{-2}_{t_{0}}(r^{2}+\left\lVert\bm{\Sigma}\right\rVert_{\infty})\right)},~{}~{}R_{t}=\mathcal{O}{\left(R_{0}\sigma_{t}^{-1}\right)}.$

The proof is provided in Appendix B.3. Then we could upper bound the second
moments of the truncated score function $\bar{\mathbf{s}}$ in $P_{t}$.

######  Lemma 7.

For any $0\leq\epsilon<1$, by setting

$J>\left\lceil{\left(\frac{1}{2}\ell\log(N\ell\left\lVert\bm{\Sigma}\right\rVert^{2}_{\rm
F}/\sigma_{t}^{2})\right)}^{1/\nu}\right\rceil,$

we have the second moment of the truncated score function bounded by
$\left\lVert\bar{\mathbf{s}}\right\rVert_{L_{2}(P_{t})}\leq\frac{\sqrt{2Nd}}{\sigma_{t}}$.

Moreover, we need to derive a uniform upper bound of the true score function
for convenience of truncation arguments.

######  Lemma 8.

Suppose Assumption 1 holds, then with probability $1-2\exp(-C\delta)$, the
range of truth score function can be bounded by $\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}\leq\sigma_{t}^{-2}{\left(N+\delta\sqrt{Nd}\right)}$.
Here $C$ is an absolute constant.

The proofs of the lemmas are provided in Appendix B.4. Now we are ready to
prove Proposition 1.

###### Proof of Proposition 1.

By Lemma 6, we obtain a transformers such that
$\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq\epsilon$
for any $\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
R_{0}\sigma_{t}^{-1}$ and $t\geq t_{0}$. Moreover, we have
$\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq
CR_{0}\sigma^{-1}_{t}$ for some absolute constant $C$. Now we can decompose
the $L_{2}$ error as

$\displaystyle\left\lVert\widetilde{\mathbf{s}}-\bar{\mathbf{s}}\right\rVert_{L_{2}(P_{t})}^{2}$
$\displaystyle=\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert^{2}_{2}\right]}$
$\displaystyle\leq\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert^{2}_{2}{\mathbf{1}\left\\{\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
R_{0}\sigma_{t}^{-1}\right\\}}\right]}$
$\displaystyle\quad+\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert^{2}_{2}{\mathbf{1}\left\\{\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right\\}}\right]}$
$\displaystyle\leq\epsilon^{2}+2\mathbb{E}_{\mathbf{v}_{t}}{\left[\frac{C^{2}R_{0}^{2}}{\sigma_{t}^{2}}{\mathbf{1}\left\\{\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right\\}}\right]}$
$\displaystyle\quad+2\mathbb{E}_{\mathbf{v}_{t}}{\left[{\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}^{2}}{\mathbf{1}\left\\{\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right\\}}\right]}$
$\displaystyle\leq\epsilon^{2}+\frac{2C^{2}R_{0}^{2}}{\sigma^{2}_{t}}\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}$
$\displaystyle\quad+2\mathbb{E}_{\mathbf{v}_{t}}{\left[{\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}^{4}}\right]}^{1/2}\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}^{1/2}$
$\displaystyle\leq\epsilon^{2}+\frac{1}{\sigma_{t}^{2}}2C^{2}R_{0}^{2}\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}$
$\displaystyle\quad+C_{2,4}\mathbb{E}_{\mathbf{v}_{t}}{\left[{\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}^{2}}\right]}\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}^{1/2}$
$\displaystyle\leq\epsilon^{2}+\frac{1}{\sigma_{t}^{2}}\left({2C^{2}(R_{0}+\left\lVert\bm{\mu}\right\rVert_{2}^{2})}+C_{2,4}2Nd\right)\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}^{1/2}.$

Here we invoke Lemma 25 in the second-to-last inequality and invoke Lemma 7 in
the last inequality.

By Lemma 8, choosing
$R_{0}=C_{1}\sqrt{Nd}\log(\left\lVert\bm{\mu}\right\rVert_{2}CNd\epsilon^{-1}\sigma_{t_{0}}^{-1})$
for some absolute constant $C_{1}$, we can bound
$\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}$
by

$\displaystyle\Pr{\left[\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}>R_{0}\sigma_{t}^{-1}\right]}\leq\frac{\epsilon^{2}}{4(R_{0}+\left\lVert\bm{\mu}\right\rVert^{2}_{2})C^{2}Nd},$

thus bounding the $L_{2}$ error by
$\mathcal{O}({\epsilon^{2}}/{\sigma_{t}^{2}})$. By adjusting the constant, we
can ensure that there exists a transformer architecture
${\mathcal{T}}(D,L,M,B,R_{t})$ that gives rise to $\widetilde{\mathbf{s}}$
with

$\displaystyle D=9d+d_{t}+d_{e}+1,\quad
L=\mathcal{O}{\left(\kappa_{t_{0}}\log(R_{0}Nd/(\epsilon\sigma_{t_{0}}))^{2}\right)}=\mathcal{O}{\left(\kappa_{t_{0}}\log(Nd/(\epsilon\sigma_{t_{0}}))^{2}\right)},$
$\displaystyle
M=\mathcal{O}{\left({\left(\ell\log(dN\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/(\epsilon\sigma_{t_{0}}))\right)}^{1/\nu}\right)},\quad
B=\mathcal{O}{\left(\log(Nd/(\epsilon\sigma_{t_{0}}))\sigma^{-2}_{t_{0}}Nd(r^{2}+\left\lVert\bm{\Sigma}\right\rVert_{\infty})\right)},$
$\displaystyle\hskip
100.00015ptR_{t}=\mathcal{O}{\left(\log(Nd/(\epsilon\sigma_{t_{0}})){\sqrt{Nd}}/{\sigma_{t}}\right)}$

such that

$\left\lVert\widetilde{\mathbf{s}}-\bar{\mathbf{s}}\right\rVert_{L_{2}(P_{t})}^{2}\leq\epsilon^{2}/\sigma_{t}^{2}.$

The proof is complete. ∎

###  B.3 Proof of Lemma 6

####  B.3.1 Transformer Architecture

Following Figure 3, we construct our targeted transformer architecture as

$\displaystyle f=f_{\rm out}\circ f_{\rm GD}\circ\dots\circ f_{\rm GD}\circ
f_{\rm pre}\circ f_{\rm in}.$

##### Encoder

For simplicity, we suppose the encoder converts the initial input into
$\mathbf{Y}=f_{\rm
in}([\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{N}])=[\mathbf{y}_{1}^{\top},\dots,\mathbf{y}_{N}^{\top}]\in\mathbb{R}^{D\times
N}$ satisfies

$\displaystyle\mathbf{y}_{i}=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{0}_{5d}^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top},$

where
$\bm{\phi}(t)=[\eta_{t},\alpha_{t},\sigma^{2}_{t},\alpha_{t}^{2}]^{\top}\in\mathbb{R}^{d_{t}}$
with $d_{t}=4$. Here $\mathbf{0}_{5d}^{\top}$ and $\mathbf{0}_{3d}^{\top}$
serve as the buffer space for storing the components necessary for expressing
gradient descent algorithm.

##### Transformer Components

Besides the encoder and decoder, $f_{\rm pre}$ represents a multi-layer
transformers that prepare the necessary components for the gradient descent
block, such as the mean $\bm{\mu}_{i}$ for each input token $\mathbf{x}_{i}$.
$f_{\rm GD}$ represents a 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 B.3.2.

##### Decoder

Suppose the output tokens from past layers has produced a score approximator
in matrix shape, we design the decoder as follows:

$\displaystyle f_{\rm out}=f_{\rm norm}\circ f_{\rm linear},$

where $f_{\rm linear}:\mathbb{R}^{D\times N}\rightarrow\mathbb{R}^{dN}$
extracts a $d\times N$ block from the input and flattens it into a vector to
align with the dimension of the score function. $f_{\rm
norm}:\mathbb{R}^{dN}\rightarrow\mathbb{R}^{dN}$ controls the output range of
the network by $R_{t}$, which can be mathematically written as

$\displaystyle f_{\rm norm}(\mathbf{s})=\left\\{\begin{aligned}
\mathbf{s},~{}~{}\text{if}~{}~{}\left\lVert\mathbf{s}\right\rVert_{2}\leq
R_{t},\\\
\frac{R_{t}}{\left\lVert\mathbf{s}\right\rVert_{2}}\mathbf{s},~{}~{}\text{otherwise.}\end{aligned}\right.$

We remark that such clipping strategy is also applied in other theoretical
works Oko et al. [2023] to better adapt to the magnitude of the score function
at different diffusion time $t$, since the score function of a degenerate
Gaussian distribution could blow up as $t\rightarrow 0$. Moreover, such
clipping layer can be easily expressed by finite layers of feed-forward
networks if we take both $\mathbf{s}$ and $R_{t}$ as the input. See Equation 5
in Lemma 2 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

$\displaystyle{\mathcal{T}}(D,L,M,B)=\\{f$ $\displaystyle:f=({\tt
FFN}_{L}\circ{\tt Attn}_{L})\circ\dots\circ({\tt FFN}_{1}\circ{\tt
Attn}_{1}),$ $\displaystyle\quad\text{The input and out dimension is $D$},$
$\displaystyle\quad\text{${\tt Attn}_{i}$ uses entrywise ReLU activation for
$i=1,\dots,L$},$ $\displaystyle\quad\text{number of heads in each~{}{\tt
Attn}~{}is bounded by $M$},$ $\displaystyle\quad\text{the Frobenius norm of
each weight matrix is bounded by $B$}\\}.$

####  B.3.2 Approximate Each Transformer Subnetworks

For simplicity, for the subnetwork $f_{\rm pre}$, we assume the mean vectors
$\bm{\mu}_{i}$ can be directly expressed by constant number of transformer
blocks:

######  Assumption 2 (Construct $\bm{\mu}$).

There exists a raw transformer $f_{\mu}\in{\mathcal{T}}_{{\rm
raw}}(D,L_{\mu},M_{\mu},\mathcal{O}(d))$ such that for any input token
$\mathbf{y}_{i}=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{0}_{5d}^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top}$,
we have

$\displaystyle
f_{\mu}(\mathbf{y}_{i})=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{0}_{4d}^{\top},\bm{\mu}_{i}^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top}.$

Here $L_{\mu}$ and $M_{\mu}$ are all constants. After this block, we begin to
express gradient descent method using transformers. Starting at
$\mathbf{s}_{i}^{(0)}=\mathbf{0}_{d}$ for $i\in[N]$, the first iteration can
be written as

$\displaystyle\mathbf{s}_{i}^{(1)}=\mathbf{s}_{i}^{(0)}-\eta_{t}{\left(\sum_{j=1}^{N}(\alpha_{t}^{2}\bar{\bm{\Gamma}}_{ij}\bm{\Sigma})\mathbf{s}_{j}^{(0)}+\sigma_{t}^{2}\mathbf{s}_{i}^{(0)}-(\mathbf{x}_{i}-\alpha_{t}\bm{\mu}_{i})\right)}=\eta_{t}{\left(\mathbf{x}_{i}-\alpha_{t}\bm{\mu}_{i}\right)}.$

Thus, we only need apply a multiplication module $f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})\approx\alpha_{t}\bm{\mu}_{i}$ that
approximately realizes the product operation. We defer the detailed
construction of the multiplication module to Appendix E.3, which utilizes
$\mathcal{O}(\log(\max_{i}\left\lVert\bm{\mu}_{i}\right\rVert_{2}/\epsilon_{\text{mult}}))$
transformer blocks to reach the accuracy $\left\lVert f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})-\alpha_{t}\bm{\mu}_{i}\right\rVert_{\infty}\leq\epsilon_{\text{mult}}$
for any $i\in[N]$.

######  Lemma 9 (Construct the first step of GD).

Suppose the input is
$\mathbf{Y}=[\mathbf{y}_{1}^{\top},\dots,\mathbf{y}_{N}^{\top}]\in\mathbb{R}^{D\times
N}$, where each token is

$\displaystyle\mathbf{y}_{i}=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{0}_{4d}^{\top},\bm{\mu}_{i}^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top}\in\mathbb{R}^{9d+d_{e}+d_{t}+1}.$

Given error level $\epsilon<1$ and learning rate $\eta_{t}>0$, there exists a
transformer $f_{\rm GD,1}\in{\mathcal{T}}_{\rm raw}(D,L,M,B)$ such that

$\displaystyle f_{{\rm
GD},1}{\left(\mathbf{y}_{i}\right)}=\Bigg{[}\mathbf{x}_{i}^{\top},$
$\displaystyle\bm{e}_{i}^{\top},\bm{\phi}^{\top}(t),f_{\rm
mult}(\eta_{t},\mathbf{x}_{i}-f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i}))^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},$
$\displaystyle f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\Bigg{]}^{\top}.$

Here $\mu_{0}=\left\lVert\bm{\mu}\right\rVert_{\infty}$, and the two
multiplication modules satisfy $\left\lVert f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top}-\alpha_{t}\bm{\mu}_{i}\right\rVert_{2}\leq\epsilon/\sqrt{N}$
and also $\left\lVert f_{\rm mult}(\eta_{t},\mathbf{x}_{i}-f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i}))-\eta_{t}(\mathbf{x}_{i}-\alpha_{t}\bm{\mu}_{i})\right\rVert_{2}\leq\epsilon/\sqrt{N}$
for $i\in[N]$, and the parameters of the networks satisfy

$\displaystyle
D=9d+d_{e}+d_{t}+1,~{}~{}L=\mathcal{O}(\log{\left(\left\lVert\mathbf{x}\right\rVert_{\infty}\left\lVert\mathbf{s}\right\rVert_{\infty}Nd/\epsilon\right)}),$
$\displaystyle\hskip
50.00008ptM=1,~{}~{}B=\mathcal{O}{\left(d(\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty})\right)}.$

###### Proof of Lemma 9.

To prove the lemma, we apply Corollary 3 and use the last $3d$ dimensions
$\mathbf{0}_{3d}$ as the buffer space to approximate the product operation
$f_{\rm mult}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ such that

$\displaystyle f_{\rm
mult}(\mathbf{Y})=\begin{bmatrix}\mathbf{x}_{1}&\cdots&\mathbf{x}_{N}\\\
\mathbf{e}_{1}&\cdots&\mathbf{e}_{N}\\\ \bm{\phi}(t)&\cdots&\bm{\phi}(t)\\\
\mathbf{0}_{4d}&\cdots&\mathbf{0}_{4d}\\\ f_{\rm
mult}(\alpha_{t},\bm{\mu}_{1})&\cdots&f_{\rm mult}(\alpha_{t},\bm{\mu}_{N})\\\
1&\cdots&1\\\ \mathbf{0}_{3d}&\cdots&\mathbf{0}_{3d}\end{bmatrix},$

where $f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})=\alpha_{t}\bm{\mu}_{i}+\bm{\epsilon}_{\mu,i}$
with $\left\lVert\bm{\epsilon}_{\mu,i}\right\rVert_{2}\leq\epsilon/\sqrt{N}$.
After $f_{\rm mult}$, we could simply add one layer of transformers
$\mathcal{TB}={\tt FFN}\circ{\tt Attn}$ with the attention block being
trivial. We just take ${\tt FFN}(\mathbf{y})=\mathbf{y}+\mathbf{W}_{2}{\rm
ReLU}(\mathbf{W}_{1}\mathbf{y})$ with

$\displaystyle\mathbf{W}_{1}=\begin{bmatrix}\mathbf{I}_{d}&\mathbf{0}_{d\times(4d+d_{e}+d_{t})}&-\mathbf{I}_{d}&\mathbf{0}_{d\times(3d+1)}\\\
-\mathbf{I}_{d}&\mathbf{0}_{d\times(4d+d_{e}+d_{t})}&\mathbf{I}_{d}&\mathbf{0}_{d\times(3d+1)}\\\
\end{bmatrix}\in\mathbb{R}^{(2d)\times D}$

and

$\displaystyle\mathbf{W}_{2}=\begin{bmatrix}\mathbf{0}_{(d+d_{e}+d_{t})\times(d)}&\mathbf{0}_{(d+d_{e}+d_{t})\times(d)}\\\
\mathbf{I}_{d}&-\mathbf{I}_{d}\\\
\mathbf{0}_{(7d+1)\times(d)}&\mathbf{0}_{(7d+1)\times(d)}\end{bmatrix}\in\mathbb{R}^{D\times(2d)}.$

After this transformer block, we apply another multiplication module $f_{\rm
mult,1}$ to scale the gradient by the learning rate $\eta_{t}$. We have
explicitly constructed the first-step GD as $f_{\rm GD,1}=f_{\rm
mult,1}\circ\mathcal{TB}\circ f_{\rm mult}$.

##### Size of Transformer Blocks for Approximating One GD Iteration

We summarize our construction by characterizing the size of the architecture:

  * •

The input to the transformer is of dimension $D\times N$ with
$D=9d+d_{e}+d_{t}+1$.

  * •

The number of transformer blocks is bounded by $L=1+2L_{\rm mult}$. By Lemma
3, the number of layers of the multiplication module $L_{\rm
mult}=\mathcal{O}(\log{\left(\left\lVert\mathbf{x}\right\rVert_{\infty}\left\lVert\mathbf{s}\right\rVert_{\infty}Nd/\epsilon\right)})$.

  * •

The number of attention heads is $M=1$.

  * •

The magnitude of the parameter in the multiplication module is at most
$\mathcal{O}(\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty})$.
Since there are at most $\mathcal{O}(d)$ nonzero weights in each weight
matrix, we have the norm of parameter bounded by
$\mathcal{O}{\left(d(\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty})\right)}.$

The proof is complete. ∎ The next lemma presents the construction of next
$K-1$ steps of GD.

######  Lemma 10 (Construct next steps of GD).

Suppose the input token is

$\displaystyle\mathbf{y}_{i}=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top},$

then there exists a transformer $f_{\rm GD,2}\in{\mathcal{T}}_{\rm
raw}(D,L,M,B)$ that approximately iterates $\mathbf{s}_{i}$ by following the
GD update formula, i.e.,

$\displaystyle f_{{\rm GD},2}{\left(\mathbf{y}_{i}\right)}$
$\displaystyle=\Bigg{[}\mathbf{x}_{i}^{\top},\bm{e}_{i}^{\top},\bm{\phi}^{\top}(t),\quad\mathbf{s}_{i}^{\top}-\eta{\left(\sum_{j=1}^{N}(\alpha_{t}^{2}\bar{\bm{\Gamma}}_{ij}\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I})\mathbf{s}_{j}^{\top}-(\mathbf{x}_{i}-\alpha_{t}\bm{\mu}_{i})^{\top}\right)}+\bm{\epsilon}_{i},$
$\displaystyle\hskip
72.26999pt\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\Bigg{]}^{\top}.$

Here
$\left\lVert\bm{\epsilon}^{\prime}_{i}\right\rVert_{2}\leq\epsilon/\sqrt{N}$
for $i\in[N]$. Moreover, the parameters of the networks satisfy

$\displaystyle
D=9d+d_{e}+d_{t}+1,~{}~{}L=\mathcal{O}(\log{\left(\left\lVert\mathbf{x}\right\rVert_{\infty}\left\lVert\mathbf{s}\right\rVert_{\infty}Nd/\epsilon\right)}),$
$\displaystyle\hskip
50.00008ptM=4J,~{}~{}B=\mathcal{O}{\left(d(\left\lVert\bm{\Sigma}\right\rVert_{\infty}+r^{2}))+\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty}\right)}.$

###### Proof of Lemma 10.

It suffices to construct several transformer blocks to represent one gradient
descent iteration with truncated matrix $\bar{\bm{\Gamma}}$, which takes the
form

$\displaystyle\mathbf{s}_{i}^{(k+1)}$
$\displaystyle=\mathbf{s}_{i}^{(k)}-\eta_{t}\left[\sum_{j=1}^{N}(\alpha_{t}^{2}\bar{\bm{\Gamma}}_{ij}\bm{\Sigma})\mathbf{s}_{j}^{(k)}+\sigma^{2}_{t}\mathbf{s}_{i}^{(k)}-(\mathbf{x}_{i,t}-\alpha_{t}\bm{\mu}_{i})\right]$
$\displaystyle=\mathbf{s}_{i}^{(k)}-\underbrace{\eta_{t}\sum_{j=1}^{N}\alpha_{t}^{2}\bar{\bm{\Gamma}}_{ij}\bm{\Sigma}\mathbf{s}_{j}^{(k)}}_{(A)}+\underbrace{\left(-\eta_{t}\sigma_{t}^{2}\mathbf{s}_{i}^{(k)}+\eta_{t}(\mathbf{x}_{i,t}-\alpha_{t}\bm{\mu}_{i})\right)}_{(B)}.$
(13)

To ease the presentation, we consider a fixed time $t$ and drop the subscript
$t$. We also drop the superscript $(k)$. In the following, we will use two
transformer blocks to approximate $(A)$ and $(B)$ separately. The input to
those transformer blocks is
$\mathbf{Y}=[\mathbf{y}_{1},\dots,\mathbf{y}_{N}]\in\mathbb{R}^{D\times N}$,
where $D$ is a larger dimension and will be specified shortly. For each column
vector $\mathbf{y}_{i}$, it stores copies of $\mathbf{s}_{i}$ and other
relevant information. Recall each token $\mathbf{y}_{i}$ is

$\displaystyle\mathbf{y}_{i}=[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}]^{\top}\in\mathbb{R}^{9d+d_{e}+d_{t}+1}.$

Here we reserve $\mathbf{0}_{3d}^{\top}$ as the buffer space for the
multiplication module (the additional hidden width).

##### Multiplication Module

Before diving into approximating terms $(A)$ and $(B)$, we use a
multiplication module consisting of a series of transformer blocks to
transform each column vector $\mathbf{y}_{i}$ into

$\displaystyle\mathbf{y}_{i}=\left[\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i}^{\top},f_{\rm
mult}(\alpha^{2},\mathbf{s}_{i}^{\top}),f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top}),\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\right]^{\top},$

where $f_{\rm mult}$ denotes an approximation to the entrywise multiplication
realized by the multiplication module. We defer the detailed construction of
the multiplication module to Appendix E.3, which utilizes
$\mathcal{O}(\log(\left\lVert\mathbf{s}\right\rVert_{\infty}/\epsilon_{\rm
mult}))$ transformer blocks to reach the accuracy $\left\lVert f_{\rm
mult}(\alpha^{2},\mathbf{s})-\alpha^{2}\mathbf{s}\right\rVert_{\infty}\leq\epsilon_{\rm
mult}$.

Compared to the raw input before the multiplication module, we have an easy
access to the useful quantities $\alpha^{2}\mathbf{s}_{i}$,
$\sigma^{2}\mathbf{s}_{i}$ and $\alpha\bm{\mu}_{i}$, which simplifies our next
step.

##### The First Attention Block for Approximating $(A)$

Here we will construct a transformer block ${\mathcal{T}}\mathcal{B}_{1}={\tt
FFN}_{1}\circ{\tt Attn}_{1}$ for approximating $(A)$. Despite the huge
dimension of matrix $\bar{\bm{\Gamma}}$, it is Toeplitz due to the uniform
time grid $\\{h_{1},\dots,h_{N}\\}$. There are at most $N$ different entries
in $\bar{\bm{\Gamma}}$ and the $(i,j)$-th entry only depends on the gap
$|i-j|$. As a result, we denote
$\gamma_{m}=\bar{\bm{\Gamma}}_{ij}\mathds{1}\\{|i-j|=m\\}$ and rewrite term
$(A)$ as

$\displaystyle(A)=\eta_{t}\sum_{m=0}^{N-1}\sum_{j=1}^{N}\alpha^{2}\gamma_{m}\mathds{1}\\{|i-j|=m\\}\bm{\Sigma}\mathbf{s}_{j}.$

The display above suggests a construction of a multi-head attention layer.
Formally, for an arbitrary value of $m$, we construct four attention heads
with ReLU activation. By Assumption 1, the indicator function
$\mathds{1}\\{|i-j|=m\\}$ can be realized by calculating the inner product
$\mathbf{e}_{i}^{\top}\mathbf{e}_{j}$ of time embedding. To see this, we
observe

$\displaystyle\mathbf{e}_{i}^{\top}\mathbf{e}_{j}=\frac{1}{2}(2r^{2}-\left\lVert\mathbf{e}_{i}-\mathbf{e}_{j}\right\rVert_{2}^{2})=\frac{1}{2}(2r^{2}-f^{2}(|i-j|)).$

Therefore, it holds that

$\displaystyle\mathds{1}\\{|i-j|=m\\}=\mathds{1}\left\\{\mathbf{e}_{i}^{\top}\mathbf{e}_{j}=r^{2}-\frac{1}{2}f^{2}(m)\right\\},$

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
$\Delta=\min_{i=1,\dots,N-1}\\{f^{2}(i+1)-f^{2}(i)\\}$. Thus, we deduce

$\displaystyle\mathds{1}\\{|i-j|=m\\}=\mathds{1}\left\\{\mathbf{e}_{i}^{\top}\mathbf{e}_{j}\in\left[r^{2}-\frac{1}{2}f^{2}(m)-\frac{1}{4}\Delta,r^{2}-\frac{1}{2}f^{2}(m)+\frac{1}{4}\Delta\right]\right\\}.$

We can use four ReLU functions to approximate the right-hand side of the last
display. In specific, we construct a trapezoid function as follows,

$\displaystyle\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})$
$\displaystyle=\frac{8}{\Delta}{\rm
ReLU}\left(\mathbf{e}_{i}^{\top}\mathbf{e}_{j}-r^{2}+\frac{1}{2}f^{2}(m)+\frac{1}{4}\Delta\right)-\frac{8}{\Delta}{\rm
ReLU}\left(\mathbf{e}_{i}^{\top}\mathbf{e}_{j}-r^{2}+\frac{1}{2}f^{2}(m)+\frac{1}{8}\Delta\right)$
$\displaystyle\quad-\frac{8}{\Delta}{\rm
ReLU}\left(\mathbf{e}_{i}^{\top}\mathbf{e}_{j}-r^{2}+\frac{1}{2}f^{2}(m)-\frac{1}{8}\Delta\right)+\frac{8}{\Delta}{\rm
ReLU}\left(\mathbf{e}_{i}^{\top}\mathbf{e}_{j}-r^{2}+\frac{1}{2}f^{2}(m)-\frac{1}{4}\Delta\right).$

It is straightforward to check that $\psi=1$ in a $\frac{4}{\Delta}$-width
interval centered at $r^{2}-\frac{1}{2}f^{2}(m)$. To this end, we can use four
attention heads to realize the function $\psi$. In particular, for the first
attention head, we choose

$\displaystyle(\mathbf{Q}^{1})^{\top}\mathbf{K}^{1}$ $\displaystyle={\rm
diag}\left(\left[\mathbf{0}_{d\times
d},\mathbf{I}_{d_{e}},\mathbf{0}_{d_{t}\times
d_{t}},\mathbf{0}_{(5d)\times(5d)},-r^{2}+\frac{1}{2}f^{2}(m)+\frac{1}{4}\Delta,\mathbf{0}_{(3d)\times(3d)}\right]\right)~{}\text{and}$
$\displaystyle\mathbf{V}^{1}$
$\displaystyle=\begin{bmatrix}\mathbf{0}_{(4d+d_{e}+d_{t})\times(2d+d_{e}+d_{t})}&\mathbf{0}_{(4d+d_{e}+d_{t})\times
d}&\mathbf{0}_{(4d+d_{e}+d_{t})\times(6d+1)}\\\
\mathbf{0}_{d\times(2d+d_{e}+d_{t})}&\frac{8}{\Delta}\gamma_{m}\bm{\Sigma}&\mathbf{0}_{d\times(6d+1)}\\\
\mathbf{0}_{(4d+1)\times(2d+d_{e}+d_{t})}&\mathbf{0}_{(4d+1)\times
d}&\mathbf{0}_{(4d+1)\times(6d+1)}\end{bmatrix}.$

It is not difficult to check that this attention head calculates the first
ReLU function in $\psi$. Analogously, for the second attention head, we choose

$\displaystyle(\mathbf{Q}^{2})^{\top}\mathbf{K}^{2}$ $\displaystyle={\rm
diag}\left(\left[\mathbf{0}_{d\times
d},\mathbf{I}_{d_{e}},\mathbf{0}_{d_{t}\times
d_{t}},\mathbf{0}_{(5d)\times(5d)},-r^{2}+\frac{1}{2}f^{2}(m)+\frac{1}{8}\Delta,\mathbf{0}_{(3d)\times(3d)}\right]\right)~{}\text{and}$
$\displaystyle\mathbf{V}^{2}$ $\displaystyle=-\mathbf{V}^{1}$

for realizing the second ReLU function. The third and fourth attention heads
have the following parameters,

$\displaystyle(\mathbf{Q}^{3})^{\top}\mathbf{K}^{3}$ $\displaystyle={\rm
diag}\left(\left[\mathbf{0}_{d\times
d},\mathbf{I}_{d_{e}},\mathbf{0}_{d_{t}\times
d_{t}},\mathbf{0}_{(5d)\times(5d)},-r^{2}+\frac{1}{2}f^{2}(m)-\frac{1}{8}\Delta,\mathbf{0}_{(3d)\times(3d)}\right]\right),$
$\displaystyle\mathbf{V}^{3}$ $\displaystyle=-\mathbf{V}^{1}$
$\displaystyle(\mathbf{Q}^{4})^{\top}\mathbf{K}^{4}$ $\displaystyle={\rm
diag}\left(\left[\mathbf{0}_{d\times
d},\mathbf{I}_{d_{e}},\mathbf{0}_{d_{t}\times
d_{t}},\mathbf{0}_{(5d)\times(5d)},-r^{2}+\frac{1}{2}f^{2}(m)-\frac{1}{4}\Delta,\mathbf{0}_{(3d)\times(3d)}\right]\right),~{}\text{and}$
$\displaystyle\mathbf{V}^{4}$ $\displaystyle=\mathbf{V}^{1}.$

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

$\displaystyle\bigg{[}\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i},f_{\rm
mult}(\alpha^{2},\mathbf{s}_{i}^{\top}),f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top}),$
$\displaystyle\qquad\qquad\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top}),f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\bigg{]}^{\top}.$

Note that this output calculates for a fixed value of $m$. In order to summing
over different values of $m$, we utilize $4J$ attention heads. Here $4J$
attention heads are enough, since $\bar{\bm{\Gamma}}_{ij}=0$ if $|i-j|\geq J$.
We have

$\displaystyle{\tt Attn}_{1}(\mathbf{y}_{i})$
$\displaystyle=\bigg{[}\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i},f_{\rm
mult}(\alpha^{2},\mathbf{s}_{i}^{\top}),f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top}),$
$\displaystyle\qquad\qquad\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top}),f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\bigg{]}^{\top}.$

For the feedforward layer ${\tt FFN}_{1}$, we set $\mathbf{W}_{1}=\mathbf{0}$,
$\mathbf{W}_{2}=\mathbf{0}$, $\mathbf{b}_{1}=\mathbf{0}$ and
$\mathbf{b}_{2}=\mathbf{0}$ such that the output of the first attention block
is

$\displaystyle{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i})$
$\displaystyle=\bigg{[}\mathbf{x}_{i}^{\top},\mathbf{e}_{i}^{\top},\bm{\phi}^{\top}(t),\mathbf{s}_{i},f_{\rm
mult}(\alpha^{2},\mathbf{s}_{i}^{\top}),f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top}),$
$\displaystyle\qquad\qquad\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top}),f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\bigg{]}^{\top}.$

##### The Second Transformer Block for Approximating $(B)$

Similar to the first block, the goal here is to construct
${\mathcal{T}}\mathcal{B}_{2}={\tt FFN}_{2}\circ{\tt Attn}_{2}$ for realizing
$(B)$. This is much easier than approximating $(A)$, we only need the feed
forward layer while set the attention layer trivial. In particular, we choose
$\mathbf{Q},\mathbf{K}$ and $\mathbf{V}$ being all zero matrices so as to
maintain the input to the attention layer. For the feedforward layer ${\tt
FFN}_{2}$, we choose $\mathbf{W}_{1}$ as

$\displaystyle\mathbf{W}_{1}=\begin{bmatrix}-\mathbf{I}_{d}&\mathbf{0}_{d_{e}\times
d_{e}}&\mathbf{0}_{d_{t}\times d_{t}}&\mathbf{0}_{d\times
d}&\mathbf{0}_{d\times
d}&\mathbf{I}_{d}&\mathbf{I}_{d}&\mathbf{I}_{d}&\mathbf{0}_{d\times(3d+1)}\\\
\mathbf{I}_{d}&\mathbf{0}_{d_{e}\times d_{e}}&\mathbf{0}_{d_{t}\times
d_{t}}&\mathbf{0}_{d\times d}&\mathbf{0}_{d\times
d}&-\mathbf{I}_{d}&-\mathbf{I}_{d}&-\mathbf{I}_{d}&\mathbf{0}_{d\times(3d+1)}\\\
&&&&\mathbf{I}_{d}&&&&\\\ &&&&-\mathbf{I}_{d}&&&&\\\ &&&&&\mathbf{I}_{d}&&&\\\
&&&&&-\mathbf{I}_{d}&&&\\\ &&&&&&\mathbf{I}_{d}&&\\\
&&&&&&-\mathbf{I}_{d}&&\end{bmatrix}\in\mathbb{R}^{(8d)\times D},$

where the missing values are all zero. Then we have $\mathbf{W}_{1}\cdot{\tt
Attn}_{2}\circ{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i})$ as

$\displaystyle\begin{bmatrix}\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j})+f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i})-(\mathbf{x}_{i}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}))\\\
-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j})-f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i})+(\mathbf{x}_{i}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}))\\\ f_{\rm mult}(\alpha^{2},\mathbf{s}_{i})\\\
-f_{\rm mult}(\alpha^{2},\mathbf{s}_{i})\\\ f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i})\\\ -f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i})\\\
\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j})\\\
-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j})\end{bmatrix}.$

It suffices to choose $\mathbf{b}_{1}=\mathbf{b}_{2}=\mathbf{0}$ and
$\mathbf{W}_{2}$ equal to

$\displaystyle\mathbf{W}_{2}=\left[\begin{array}[]{c}\mathbf{0}_{(d+d_{e}+d_{t})\times(8d)}\\\
\hline\cr\begin{array}[]{cccccccc}-\mathbf{I}_{d}&\mathbf{I}_{d}&&&&&&\\\
&&-\mathbf{I}_{d}&\mathbf{I}_{d}&&&&\\\
&&&&-\mathbf{I}_{d}&\mathbf{I}_{d}&&\\\
&&&&&&-\mathbf{I}_{d}&\mathbf{I}_{d}\end{array}\\\
\hline\cr\mathbf{0}_{(4d+1)\times(8d)}\end{array}\right]\in\mathbb{R}^{D\times(8d)},$

where missing values are all zero. Using the fact that ${\rm ReLU}(x)-{\rm
ReLU}(-x)=x$, we have

$\displaystyle\quad\mathbf{W}_{2}\cdot{\rm ReLU}(\mathbf{W}_{1}\cdot{\tt
Attn}_{2}\circ{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i}))$
$\displaystyle=\Bigg{[}\mathbf{0}_{d}^{\top},\mathbf{0}_{d_{e}}^{\top},\mathbf{0}_{d_{t}}^{\top},-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top})-f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top})+(\mathbf{x}_{i}^{\top}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}^{\top})),$
$\displaystyle\qquad-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top}),-f_{\rm
mult}(\alpha^{2},\mathbf{s}_{i}^{\top}),-f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top}),\mathbf{0}_{d}^{\top},0,\mathbf{0}_{3d}^{\top}\Bigg{]}^{\top}.$

Therefore, by concatenating the two transformer blocks, we have

$\displaystyle{\mathcal{T}}\mathcal{B}_{2}\circ{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i})=\Bigg{[}\mathbf{x}_{i}^{\top},\bm{e}_{i}^{\top},\bm{\phi}^{\top}(t),$
$\displaystyle\quad\mathbf{s}_{i}-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top})-f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top})+(\mathbf{x}_{i}^{\top}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}^{\top})),$ $\displaystyle\hskip
144.54pt\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\Bigg{]}^{\top}.$

Finally, by implementing another multiplication module $f_{\rm mult}$ that
scales the gradient by the learning rate $\eta_{t}$, we realize

$\displaystyle f_{\rm GD,2}$ $\displaystyle=f_{\rm
mult}\circ{\mathcal{T}}\mathcal{B}_{2}\circ{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i})=\Bigg{[}\mathbf{x}_{i}^{\top},\bm{e}_{i}^{\top},\bm{\phi}^{\top}(t),$
$\displaystyle\quad f_{\rm
mult}\bigg{(}\eta_{t},\mathbf{s}_{i}-\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top})-f_{\rm
mult}(\sigma^{2},\mathbf{s}_{i}^{\top})+(\mathbf{x}_{i}^{\top}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}^{\top}))\bigg{)},$ $\displaystyle\hskip
144.54pt\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},\mathbf{0}_{d}^{\top},f_{\rm
mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\mathbf{0}_{3d}^{\top}\Bigg{]}^{\top}.$

we have implemented one gradient descent iteration, where the output replaces
$\mathbf{s}_{i}$ in the initial input vector. Consequently, we can repeat the
multiplication module followed by the three transformer blocks to approximate
the gradient descent for $K$ iterations as needed. Note that, we don’t need to
calculate $f_{\rm mult}(\alpha_{t},\bm{\mu}_{i})^{\top}$ after the first step
of GD, which leads to the construction of $f_{\rm GD,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 is determined by
the error in the multiplication module. For the $k$-th iteration, we denote
the approximation realized by the transformer as

$\displaystyle\widehat{\mathbf{s}}_{i}^{(k+1)}=\widehat{\mathbf{s}}_{i}^{(k)}-f_{\rm
mult}\bigg{(}\eta_{t},\sum_{m=0}^{N-1}$
$\displaystyle\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot
f_{\rm mult}(\alpha^{2},\widehat{\mathbf{s}}_{j}^{(k)})$ $\displaystyle-f_{\rm
mult}(\sigma^{2},\widehat{\mathbf{s}}_{i}^{(k)})+(\mathbf{x}_{i}^{\top}-f_{\rm
mult}(\alpha,\bm{\mu}_{i}))\bigg{)}.$

We compare the output with the exact GD update with last iteration at
$\widehat{\mathbf{s}}_{i}^{(k)}$, which is

$\displaystyle\widehat{\mathbf{s}}_{i,\star}^{(k+1)}=\widehat{\mathbf{s}}_{i}^{(k)}-\eta_{t}\sum_{m=0}^{N-1}\sum_{j=1}^{N}\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\bm{\Sigma}\cdot\alpha^{2}\widehat{\mathbf{s}}_{j}^{(k)}-\sigma^{2}\widehat{\mathbf{s}}_{i}^{(k)}+(\mathbf{x}_{i}^{\top}-\alpha\bm{\mu}_{i}).$

For any index $i\in\\{1,\dots,N\\}$, we bound the difference

$\displaystyle\left\lVert\widehat{\mathbf{s}}_{i,\star}^{(k+1)}-\widehat{\mathbf{s}}_{i}^{(k+1)}\right\rVert_{2}$
$\displaystyle\overset{(i)}{\leq}\epsilon_{\rm
mult}+\eta_{t}\sum_{m=0}^{N-1}\sum_{j=1}^{N}\left\lVert\bm{\Sigma}(\alpha^{2}\mathbf{s}_{j}^{(k)}-f_{\rm
mult}(\alpha^{2},\widehat{\mathbf{s}}_{j}^{(k)})\right\rVert_{2}$
$\displaystyle\quad+\eta_{t}\left\lVert\sigma^{2}\mathbf{s}_{i}^{(k)}-f_{\rm
mult}(\sigma^{2},\widehat{\mathbf{s}}_{i}^{(k)})\right\rVert_{2}+\eta_{t}\left\lVert\alpha\bm{\mu}_{i}-f_{\rm
mult}(\alpha,\bm{\mu}_{i})\right\rVert_{2}$
$\displaystyle\overset{(ii)}{\leq}\epsilon_{\rm
mult}+\eta_{t}\left\lVert\bm{\Sigma}\right\rVert_{2}N\sqrt{d}\epsilon_{\rm
mult}+2\eta_{t}\sqrt{d}\epsilon_{\rm mult},$

where in inequality $(i)$, we use the fact that
$\gamma_{m}\psi(\mathbf{e}_{i}^{\top}\mathbf{e}_{j})\leq 1$, and in inequality
$(ii)$, we plug in the approximation error of multiplication module. By
setting $\epsilon_{\rm
mult}=\frac{1}{6}\eta_{t}^{-1}\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}^{-1}N^{-3/2}d^{-1/2}\epsilon$, we ensure that
$\left\lVert\widehat{\mathbf{s}}_{i,\star}^{(k+1)}-\widehat{\mathbf{s}}_{i}^{(k+1)}\right\rVert_{2}\leq\epsilon/\sqrt{N}$.

##### Size of Transformer Blocks for Approximating One GD Iteration

We summarize our construction by characterizing the size of the architecture:

  * •

The input to the transformer is of dimension $D\times N$ with
$D=9d+d_{e}+d_{t}+1$.

  * •

The number of transformer blocks is bounded by $L=2+L_{\rm mult}$. The number
of attention heads in each block is bounded by $M=4J$. According to Corollary
3, the number of layers of the multiplication module $L_{\rm
mult}=\mathcal{O}(\log{\left(\left\lVert\mathbf{x}\right\rVert_{\infty}\left\lVert\mathbf{s}\right\rVert_{\infty}Nd/\epsilon\right)})$.

  * •

By the construction above, the parameters in the transformer are bounded by
$\max\\{1,\frac{8}{\Delta}\eta\left\lVert\bm{\Sigma}\right\rVert_{\infty},r^{2}\\}$.
Here,
$\left\lVert\bm{\Sigma}\right\rVert_{\infty}=\max_{ij}|\bm{\Sigma}_{ij}|$
denotes the largest magnitude of entries. Since there are at most
$\mathcal{O}(d)$ nonzero weights in each weight matrix, we have the norm of
parameter bounded by

$\mathcal{O}{\left(d(\Delta^{-1}\left\lVert\bm{\Sigma}\right\rVert_{\infty}+r^{2}))+\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty}\right)}=\mathcal{O}{\left(d(\left\lVert\bm{\Sigma}\right\rVert_{\infty}+r^{2}))+\left\lVert\mathbf{x}\right\rVert_{\infty}+\left\lVert\mathbf{s}\right\rVert_{\infty}\right)}.$

The proof is complete. ∎

####  B.3.3 Formal Proof of Lemma 6

Recall that the score function is written as

$\displaystyle\mathbf{s}(\mathbf{v}_{t})=(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu}).$

Thus, if $\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
R_{0}\sigma_{t}^{-1}$, we have
$\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}\leq\left\lVert\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}\right\rVert_{2}\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
CR_{0}\sigma_{t}^{-1}$, where

$\displaystyle C$
$\displaystyle=\left\lVert(\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})\right\rVert_{2}$
$\displaystyle\leq\left\lVert\alpha_{t}^{2}(\bm{\Gamma}\otimes\bm{\Sigma})\right\rVert_{2}+\left\lVert\sigma_{t}^{2}\mathbf{I}\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bm{\Gamma}\right\rVert_{\rm
F}\left\lVert\bm{\Sigma}\right\rVert_{2}+1$
$\displaystyle\leq\sqrt{N(\ell+1)}\left\lVert\bm{\Sigma}\right\rVert_{2}+1.$

The last inequality follows from Lemma 4. Thus, the infinity norm of
$\mathbf{v}_{t}$ can be bounded by
$\left\lVert\mathbf{v}_{t}\right\rVert_{\infty}\leq\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{\infty}+\left\lVert\alpha_{t}\bm{\mu}\right\rVert_{\infty}\leq
CR_{0}\sigma_{t}^{-1}+\mu_{0}$. Moreover, by the proof of lemma 5, we can also
bound the norm of the truncated score function by

$\displaystyle\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert$
$\displaystyle\leq\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})-\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}+\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq\sigma_{t}^{-4}\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}\left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}+\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq\sigma_{t}^{-4}\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}(CR_{0}\sigma_{t}^{-1}+\mu_{0})+R_{0}\sigma_{t}^{-1}.$

Thus, according to Lemma 4, by choosing
$J=\mathcal{O}({\left(\ell\log(N\ell\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/\sigma_{t})\right)}^{1/\nu})$, we can ensure that
$\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert\leq
2R_{0}\sigma_{t}^{-1}$.

Now we construct the transformers as follows:

$\displaystyle\widetilde{\mathbf{s}}(\mathbf{x})=f_{\rm norm}\circ f_{\rm
linear}\circ\underbrace{f_{\rm GD,2}\circ f_{\rm GD,2}\circ\cdots\circ f_{\rm
GD,2}}_{(K-1)\times f_{\rm GD,2}}\circ f_{\rm GD,1}\circ f_{\mu}\circ f_{\rm
in}.$

Here $f_{\rm out}$ extracts the $(d+d_{e}+d_{t}+1)$-th to
$(2d+d_{e}+d_{t})$-th rows of the output as the score approximator, and we
choose the clipping range as $R=2R_{0}\sigma_{t}^{-1}$.

Moreover, let

$\widetilde{\mathbf{s}}_{0}=f_{\rm linear}\circ f_{\rm GD,2}\circ f_{\rm
GD,2}\circ\cdots\circ f_{\rm GD,2}\circ f_{\rm GD,1}\circ f_{\mu}\circ f_{\rm
in}$

to be the score approximator without clipping, then under the condition that
$\left\lVert\mathbf{s}(\mathbf{v}_{t})\right\rVert_{2}\leq
R_{0}\sigma_{t}^{-1}$, we have
$\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\left\lVert\widetilde{\mathbf{s}}(\mathbf{v}_{t})_{0}-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}$.

Now let’s bound the difference between
$\widetilde{\mathbf{s}}_{0}(\mathbf{v}_{t})$ and $\bar{\mathbf{s}}$. By Lemma
9 and Lemma 10, $\widetilde{\mathbf{s}}_{0}(\mathbf{v}_{t})$ expresses the
output after $K$ steps of GD with error level
$\sqrt{N}\cdot\epsilon/\sqrt{N}=\epsilon$. Here we multiply $\sqrt{N}$ because
we are considering the entire score function instead of each patch. Then by
Lemma 2 and following the proof of Lemma 3, we have

$\displaystyle\left\lVert\widetilde{\mathbf{s}}_{0}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle\leq\frac{(\kappa_{t}+1)\epsilon}{2}+\exp{\left(-\frac{2K}{\kappa_{t}+1}\right)}\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}$
$\displaystyle=\frac{(\kappa_{t}+1)\epsilon}{2}+\exp{\left(-\frac{2K}{\kappa_{t}+1}\right)}2R_{0}\sigma_{t}^{-1}.$

Thus, by replacing the error level $\epsilon$ by
$2\epsilon/(C\sigma_{t_{0}}^{-2}+1)\leq 2\epsilon/(\kappa_{t}+1)$ for
sufficiently large constant $C$ and
$K=\lceil\frac{\kappa_{t_{0}}+1}{2}\log(4R_{0}/(\sigma_{t}\epsilon))\rceil$,
we have
$\left\lVert\widetilde{\mathbf{s}}_{0}(\mathbf{v}_{t})-\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq\epsilon$
for any $t\geq t_{0}$. Moreover, we remark that each patch of the score
approximator lies in the euclidean ball with radius

$\left\lVert\widetilde{\mathbf{s}}^{(k)}_{i}(\mathbf{v}_{t})\right\rVert_{2}\leq\epsilon+2\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})\right\rVert_{2}\leq
1+4R_{0}\sigma_{t}^{-1}$

throughout the network. Here
$\widetilde{\mathbf{s}}^{(k)}_{i}(\mathbf{v}_{t})$ represents the $i$-th patch
of the score approximator after $k$ GD blocks.

##### Size of Transformer Architecture

We summarize our construction by characterizing the size of the architecture
of the entire network:

  * •

The input to the transformer is of dimension $D\times N$ with
$D=9d+d_{e}+d_{t}+1$.

  * •

The number of transformer blocks

$L=L_{\mu}+(2+L_{\rm
mult})K=\mathcal{O}(\kappa_{t_{0}}\log(R_{0}Nd/\epsilon)^{2}).$

  * •

The number of attention heads in each block is bounded by

$M=4J=\mathcal{O}{\left({\left(\ell\log(N\ell\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/\sigma_{t_{0}})\right)}^{1/v}\right)}.$

  * •

By the construction above and Lemma 9, 10 and our upper bound on the input in
each transformer block, we conclude that the Frobenius norm of each parameter
is bounded by

$B=\mathcal{O}{\left(\log(R_{0}Nd/(\epsilon\sigma_{t_{0}}))\sqrt{N}R_{0}\sigma^{-2}_{t_{0}}(r^{2}+\left\lVert\bm{\Sigma}\right\rVert_{\infty})\right)}.$

  * •

The output range is $R_{t}=2R_{0}\sigma_{t}^{-1}$.

The proof is complete.

###  B.4 Omitted Proofs of Lemmas about the Score Function

###### Proof of Lemma 5.

Suppose $\mathbf{v}_{t}\sim
P_{t}=\mathcal{N}(\alpha_{t}\bm{\mu},\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})$.
Denote
$\bm{\Phi}_{t}=\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$
and
$\bar{\bm{\Phi}}_{t}=\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$.
Then we have
${\bm{\Gamma}}_{t}=\bar{\bm{\Phi}}_{t}+\alpha_{t}^{2}(\Delta\bm{\Gamma}\otimes\bm{\Sigma})$.
Moreover, by Lemma 4, choosing
$J=\mathcal{O}((\log(N\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}/(\epsilon\sigma_{t})))$ ensures that
$\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}\leq\epsilon\left\lVert\bm{\Sigma}\right\rVert^{-2}_{\rm F}\sigma_{t}^{2}$.
Thus, we have

$\displaystyle\left\lVert\bar{\mathbf{s}}-\nabla\log
p_{t}\right\rVert_{L_{2}(P_{t})}^{2}$
$\displaystyle=\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}\right]}$
$\displaystyle=\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})\right\rVert_{2}^{2}\right]}$
$\displaystyle=\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}{\bm{\Gamma}}_{t}^{1/2}\right\rVert_{\rm
F}^{2}$
$\displaystyle=\mathop{\mathrm{tr}}{\left({\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}{\bm{\Gamma}}_{t}{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}\right)}$
$\displaystyle\leq\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}{\bm{\Gamma}}_{t}\right\rVert_{\rm
F}\left\lVert\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right\rVert_{\rm
F}$
$\displaystyle\leq\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}{\bm{\Gamma}}_{t}\right\rVert_{\rm
F}\left\lVert{\left(\bar{\bm{\Phi}}_{t}^{-1}-{\bm{\Gamma}}_{t}^{-1}\right)}{\bm{\Gamma}}_{t}\right\rVert_{\rm
F}\left\lVert{\bm{\Gamma}}^{-1}_{t}\right\rVert_{2}$
$\displaystyle=\left\lVert\alpha_{t}^{2}\bar{\bm{\Phi}}_{t}^{-1}(\Delta\bm{\Gamma}\otimes\bm{\Sigma})\right\rVert^{2}_{\rm
F}\left\lVert(\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})^{-1}\right\rVert_{2}$
$\displaystyle\leq\left\lVert(\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})^{-1}\right\rVert_{2}^{2}\left\lVert\alpha_{t}^{2}(\Delta\bm{\Gamma}\otimes\bm{\Sigma})\right\rVert^{2}_{\rm
F}\left\lVert(\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})^{-1}\right\rVert_{2}$
$\displaystyle\leq\alpha_{t}^{4}\left\lVert\Delta\bm{\Gamma}\right\rVert^{2}_{\rm
F}\left\lVert\bm{\Sigma}\right\rVert^{2}_{\rm
F}\left\lVert\sigma_{t}^{2}\mathbf{I}\right\rVert^{-3}_{2}$
$\displaystyle\leq\frac{\epsilon^{2}}{\sigma_{t}^{2}}.$

In the last inequality, we invoke
$\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}\leq\epsilon\left\lVert\bm{\Sigma}\right\rVert^{-2}_{\rm F}\sigma_{t}^{2}$.
The proof is complete. ∎

###### Proof of Lemma 7.

Suppose $\mathbf{v}_{t}\sim
P_{t}=\mathcal{N}(\alpha_{t}\bm{\mu},\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I})$.
Denote
$\bm{\Phi}_{t}=\alpha_{t}^{2}({\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$
and
$\bar{\bm{\Phi}}_{t}=\alpha_{t}^{2}(\bar{\bm{\Gamma}}\otimes\bm{\Sigma})+\sigma_{t}^{2}\mathbf{I}$.
Then we have

$\displaystyle\left\lVert\bar{\mathbf{s}}\right\rVert^{2}_{L_{2}(P_{t})}$
$\displaystyle=\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\bar{\mathbf{s}}(\mathbf{v}_{t})-\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}\right]}$
$\displaystyle=\mathbb{E}_{\mathbf{v}_{t}}{\left[\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{2}^{2}\right]}$
$\displaystyle=\left\lVert\bar{\bm{\Phi}}_{t}^{-1}{\bm{\Gamma}}_{t}^{1/2}\right\rVert_{\rm
F}^{2}$
$\displaystyle=\mathop{\mathrm{tr}}{\left(\bar{\bm{\Phi}}_{t}^{-1}{\bm{\Gamma}}_{t}{\bar{\bm{\Phi}}_{t}^{-1}}\right)}$
$\displaystyle\leq\left\lVert{\bar{\bm{\Phi}}_{t}^{-1}}{\bm{\Gamma}}_{t}\right\rVert_{\rm
F}\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{\rm F}$
$\displaystyle\leq\left\lVert{\bar{\bm{\Phi}}_{t}^{-1}}\alpha_{t}^{2}(\Delta\bm{\Gamma}\otimes\bm{\Sigma})+\mathbf{I}\right\rVert_{\rm
F}\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{\rm F}$
$\displaystyle\leq{\left(\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{2}\left\lVert{}\alpha_{t}^{2}(\Delta\bm{\Gamma}\otimes\bm{\Sigma})\right\rVert_{\rm
F}+\sqrt{Nd}\right)}\sqrt{Nd}\left\lVert\bar{\bm{\Phi}}_{t}^{-1}\right\rVert_{2}$
$\displaystyle\leq{\left(\frac{\alpha_{t}^{2}\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}\left\lVert\bm{\Sigma}\right\rVert_{\rm
F}}{\sigma_{t}^{2}}+\sqrt{Nd}\right)}\frac{\sqrt{Nd}}{\sigma_{t}^{2}}$
$\displaystyle\leq{\left(\epsilon+\sqrt{Nd}\right)}\frac{\sqrt{Nd}}{\sigma_{t}^{2}}$
$\displaystyle\leq\frac{2Nd}{\sigma_{t}^{2}}.$

Here we invoke $\left\lVert\Delta\bm{\Gamma}\right\rVert_{\rm
F}\leq\left\lVert\bm{\Sigma}\right\rVert_{\rm F}^{-1}\sigma_{t}^{2}$ by Lemma
4 in the second last inequality. ∎

###### Proof of Lemma 8.

Note that if $\mathbf{v}_{t}\sim P_{t}$ we can write
$\mathbf{v}_{t}=\alpha_{t}\bm{\mu}+\bm{\Phi}_{t}^{1/2}\mathbf{z}$, where
$\bm{\Phi}_{t}=\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}^{2}\mathbf{I}$
and $\mathbf{z}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. Thus, the score
function can be written as

$\nabla\log
p_{t}(\mathbf{v}_{t})=\bm{\Phi}_{t}^{-1}(\mathbf{v}_{t}-\alpha_{t}\bm{\mu})=\bm{\Phi}_{t}^{-1/2}\mathbf{z}.$

Thus, we have $\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}=\mathbf{z}^{\top}\bm{\Phi}_{t}^{-1}\mathbf{z}$,
which is a quadratic form of the standard Gaussian. By taking
$g(\mathbf{z})=\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}$ in Lemma 24, we have

$\displaystyle\mathbb{P}\left[|\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}-\mathbb{E}[\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}]|\geq\delta\sqrt{\operatorname{Var}(\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2})}\right]\leq
2\exp\left(-C_{2}\delta\right).$

Since we have

$\mathbb{E}[\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}]=\mathop{\mathrm{tr}}{\left(\bm{\Phi}_{t}^{-1}\right)}\leq
Nd\sigma_{t}^{-2}~{}\text{ and }~{}\operatorname{Var}(\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2})\leq\left\lVert\bm{\Phi}_{t}^{-1}\right\rVert^{2}_{\rm
F}\leq Nd\sigma_{t}^{-4},$

we have $\left\lVert\nabla\log
p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}\leq\sigma_{t}^{-2}{\left(N+\delta\sqrt{Nd}\right)}$
with probability $1-2\exp\left(-C_{2}\delta\right)$. ∎

###  B.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.,