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

## Abstract

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, ; Song et al., ; Ho et al., ; Zhang et al., ), sequential data modeling (Alcaraz and Strodthoff, ; Tashiro et al., ; Tian et al., ), reinforcement learning and control (Pearce et al., ; Hansen-Estruch et al., ; Zhu et al., ; Ding and Jin, ), as well as computational biology (Xu et al., ; Guo et al., ). 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., ). Afterwards, a few works also demonstrate the capability of using transformers as a score neural network (Peebles and Xie, ; Wu et al., ; Bao et al., ). Throughout the paper, we adopt the terminology in (Peebles and Xie, ) 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., ; Liu et al., ). 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: 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.
Refer to caption: x1.png

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, ; Williams and Rasmussen, ). 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., ; Deringer et al., ; Chen et al., ; Borovitskiy et al., ).

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. (, , ); Li et al. (, ); Chen et al. (); Lee et al. (, ); Chen et al. (, ) 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. (); Oko et al. (); Li et al. (); Mei and Wu (); Tang and Yang (); Jiao et al. () 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., ). 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 (); Mei () showing U-Net performing algorithm unrolling in diffusion models.

On the empirical side, Diffusion Transformer (DiT) (Peebles and Xie, ) 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., ; Liu et al., ), where the video data is patchified and rearranged into a long sequence. Besides, transformers are also used in discrete diffusion models (Sun et al., ; Austin et al., ).

##### 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, ). 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., ; Tang and Zhao, ; Chan, ). 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: 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}$).
Refer to caption: x2.png

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

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

##### 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 ℰ 2 subscript ℰ 2 \mathcal{E}_{2} caligraphic_E start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT

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

###### Assumption 1 .

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

## 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, ; Yarotsky, ), 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., ; Chen et al., ; Tang and Yang, ). 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: 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).
Refer to caption: x3.png

###### Theorem 1 ( Score Approximation by Transformers ) .

##### 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, ; Hyvärinen and Dayan, ). 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.

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

## 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, ) 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: 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.
Refer to caption: x4.png

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

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

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

###### Proof of Lemma 1 .

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

###### Lemma 3 .

###### Proof of Lemma 3 .

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

###### Lemma 4 .

###### Proof of Lemma 4 .

###### Remark 1 .

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

###### Proposition 1 .

###### Lemma 5 .

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

###### Lemma 6 .

###### Lemma 7 .

###### Lemma 8 .

###### Proof of Proposition 1 .

### 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. [] 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. [] 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}$.

###### Assumption 2 (Construct 𝝁 𝝁 \bm{\mu} bold_italic_μ ) .

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

###### Proof of Lemma 9 .

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

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

###### Proof of Lemma 10 .

##### 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 ) 𝐴 (A) ( italic_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 ) 𝐵 (B) ( italic_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)$.
∎

###### Proof of Lemma 5 .

###### Proof of Lemma 7 .

###### Proof of Lemma 8 .

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

$$ $\displaystyle\gamma(h_{i},h_{j})=\exp\left(-\left\lVert\mathbf{e}_{i}-\mathbf{ e}_{j}\right\rVert_{2}^{2}/\ell\right).$ $$

In the following lemma, we reproduce our results in Lemma 10 with softmax activation. Moreover, due to the exponential scaling strategy of softmax, we show that one head of attention layer is enough to express the GD with the target function without truncation on $\bm{\Gamma}$.

###### Lemma 11 (Unroll GD in softmax transformers) .

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-1}^{\top},i]^{\top},$ $$

then there exists a softmax transformer $f_{{\rm GD,softmax}}\in{\mathcal{T}}_{\rm raw}(D,\mathcal{O}(\log{\left(\left
\lVert\mathbf{x}\right\rVert_{\infty}\left\lVert\mathbf{s}\right\rVert_{\infty
}Nd/\epsilon\right)}),$ $1,\mathcal{O}(N^{2}+d+\left\lVert\bm{\Sigma}\right\rVert_{F}+r^{2});{\rm
softmax})$ that approximately iterates $\mathbf{s}_{i}$ by following the GD update formula, i.e.,

$$ $\displaystyle f_{{\rm GD,softmax}}{\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}{\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-1}^{\top},i\Bigg{]}^{\top}.$ $$

Here $\left\lVert\bm{\epsilon}^{\prime}_{i}\right\rVert_{2}\leq\epsilon/\sqrt{N}$ for $i\in[N]$. Compared with Lemma 10, we add an additional time embedding $i$ at the end of the input for technical convenience.

###### Proof of Lemma 11 .

The proof is similar to that of Lemma 10 but with different construction on the attention layer. We also decompose the gradient descent iteration as

$$ $\displaystyle\mathbf{s}_{i}^{(k+1)}$ $\displaystyle=\mathbf{s}_{i}^{(k)}-\eta_{t}\left[\sum_{j=1}^{N}(\alpha_{t}^{2} {\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}{\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)}.$ (14) $$

By the proof of Lemma 10, part $(B)$ only utilizes feed-forward networks, thus no changes need to be made. We will elaborate on approximating $(A)$ using softmax attention layers. To ease the presentation, we consider a fixed time $t$ and drop the subscript $t$. We also drop the superscript $(k)$. 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-1}^{\top},i]^{\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). Before diving into approximating terms $(A)$, we follow the proof of Lemma 10 to use a multiplication module consisting of a series of transformer blocks 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-1}^{\top},i \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(1/\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}$. Here we could still utilize the multiplication module after changing the activation function in attention layer to softmax because its construction only relies on feed-forward networks, so we can simply set the attention layers trivial throughout the network in multiplication module. See Appendix E.3 for more discussions. Now we begin to approximate $(A)$.

###### Lemma 11 (Unroll GD in softmax transformers) .

###### Proof of Lemma 11 .

##### Comparison with ReLU activation

In softmax transformers, we leverage the exponential scaling mechanism of the softmax activation to directly construct the entire kernel matrix $\Gamma$ without truncation, so only one head of attention layer is needed. On the contrary, in ReLU transformers, we only construct the main diagonals of $\bm{\Gamma}$ with multiple attention heads. See Lemma 10 for more details.

##### Approximate ( A ) 𝐴 (A) ( italic_A )

Here we will construct a transformer block ${\mathcal{T}}\mathcal{B}_{1}={\tt FFN}_{1}\circ{\tt Attn}_{1}$ for approximating $(A)$. Our construction depends on the following fact:

$$ $\displaystyle(A)$ $\displaystyle=\eta\alpha^{2}\sum_{j=1}^{N}\exp(-\left\lVert\mathbf{e}_{i}- \mathbf{e}_{j}\right\rVert_{2}^{2}/\ell)\bm{\Sigma}\mathbf{s}_{j}$ $\displaystyle=\eta\alpha^{2}\sum_{j=1}^{N}\exp((2r^{2}-2\mathbf{e}^{\top}_{i} \mathbf{e}_{j})/\ell)\bm{\Sigma}\mathbf{s}_{j}$ $\displaystyle={\eta\alpha^{2}D_{i}}\cdot{\rm softmax}{\left([(2r^{2}-2\mathbf{ e}^{\top}_{i}\mathbf{e}_{1})/\ell,\cdots,(2r^{2}-2\mathbf{e}^{\top}_{i}\mathbf {e}_{N})/\ell]\right)}\cdot{\left[\bm{\Sigma}\mathbf{s}_{1},\dots,\bm{\Sigma} \mathbf{s}_{N}\right]},$ $$

where $D_{i}=\sum_{j=1}^{N}\exp((2r^{2}-2\mathbf{e}^{\top}_{i}\mathbf{e}_{j})/\ell)$ is the normalizing constant in the softmax activation.
Thus, we can use one attention head to first realize the unnormalized version:

$$ $\displaystyle(A^{\prime})=\alpha^{2}\cdot{\rm softmax}{\left([(2r^{2}-2\mathbf {e}^{\top}_{i}\mathbf{e}_{1})/\ell,\cdots,(2r^{2}-2\mathbf{e}^{\top}_{i} \mathbf{e}_{N})/\ell]\right)}\cdot{\left[\bm{\Sigma}\mathbf{s}_{1},\dots,\bm{ \Sigma}\mathbf{s}_{N}\right]}.$ $$

##### Use Attn to approximate ( A ′ ) superscript 𝐴 ′ (A^{\prime}) ( italic_A start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT )

In particular, for the first attention head, we choose

$$ $\displaystyle(\mathbf{Q})^{\top}\mathbf{K}$ $\displaystyle={\rm diag}\left(\left[\mathbf{0}_{d\times d},-{2}{\ell^{-1}} \mathbf{I}_{d_{e}},\mathbf{0}_{d_{t}\times d_{t}},\mathbf{0}_{(5d)\times(5d)}, {2r^{2}}{{\ell}^{-1}},\mathbf{0}_{(3d)\times(3d)}\right]\right)~{}\text{and}$ $\displaystyle\mathbf{V}$ $\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})}&\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}.$ $$

Thus, 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 D_{i}^{-1}\sum_{j=1}^{N}\exp((2r^{2}-2\mathbf{e}_{i} \mathbf{e}_{j})/\ell)\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-1}^{\top} ,i\bigg{]}^{\top}.$ $$

To multiply our constructed value by the normalizing constant $D_{i}$, we will construct a feed-forward layer ${\tt FFN}_{1}$ such that the output of the first attention block is

$$ $\displaystyle{\mathcal{T}}\mathcal{B}(\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 D_{i}^{-1}\sum_{j=1}^{N}\exp((2r^{2}-2\mathbf{e}_{i}\mathbf{e}_{ j})/\ell)\bm{\Sigma}\cdot f_{\rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top}),f_{ \rm mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1,\widehat{D_{i}},\mathbf{0}_{3d-2}^ {\top},i\bigg{]}^{\top}.$ (15) $$

Here $\widehat{D_{i}}$ is an approximation to $D_{i}$. Now we show the approximation strategy to approximate $D_{i}$.

##### Use FFN to approximate D i subscript 𝐷 𝑖 D_{i} italic_D start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT

Note that we can rewrite $D_{i}$ as

$$ $D_{i}=1+\sum_{k=1}^{N-1}(\mathbf{1}\{i\geq k+1\}+\mathbf{1}\{i\leq N-k\})g(k).$ $$

Where $g(k)=\exp(-f(k)^{\nu}/\ell)$ is the correlation function. Let $\widehat{D_{i}}^{(m)}=1+\sum_{k=1}^{m}(\mathbf{1}\{i\geq k+1\}+\mathbf{1}\{i
\leq N-k\})g(k)$ to be an approximation of $D_{i}$, we know by the proof of Lemma 4, for any $i\in[N]$,

$$ $\displaystyle\left|D(i)-D_{m}(i)\right|\leq 2\sum_{k=m+1}^{N-1}g(k)\leq c^{- \nu}\ell\exp{\left(-\frac{2c^{\nu}m^{\nu}}{\ell}\right)}.$ $$

Thus, to derive an $\epsilon$ error bound, we only need $m$ to be in the order of $m=\mathcal{O}{\left({\left(\ell\log\left(1/\epsilon\right)\right)}^{1/\nu}
\right)}$. For each $k$, we have

$$ $\displaystyle\qquad(\mathbf{1}\{i\geq k+1\}+\mathbf{1}\{i\leq N-k\})g(k)$ $\displaystyle=g(k){\left({\rm ReLU}(i-k)-{\rm ReLU}(i-k-1)+{\rm ReLU}(N-k+1-i) -{\rm ReLU}(N-k-i)\right)}$ $$

holds for any integer $i$ . Thus, the entire function $\widehat{D}_{i}^{(m)}$ can be expressed by $4m=\mathcal{O}{\left({\left(\ell\log\left(1/\epsilon\right)\right)}^{1/\nu}
\right)}$ ReLU neurons in the feed-forward network. Specifically, we take ${\tt FFN}(\mathbf{y})=\mathbf{W}_{2}{\rm RelU}(\mathbf{W}_{1}\mathbf{y})+
\mathbf{W}_{1}$ and choose $\mathbf{W}_{1}$ as

$$ $\displaystyle\mathbf{W}_{1}=\begin{bmatrix}\mathbf{0}_{D-3d-1}^{\top}&-1& \mathbf{0}_{3d-1}^{\top}&1\\ \mathbf{0}_{D-3d-1}^{\top}&-2&\mathbf{0}_{3d-1}^{\top}&1\\ \vdots&\vdots&\vdots&\vdots\\ \mathbf{0}_{D-3d-1}^{\top}&-m-1&\mathbf{0}_{3d-1}^{\top}&1\\ \mathbf{0}_{D-3d-1}^{\top}&n&\mathbf{0}_{3d-1}^{\top}&-1\\ \vdots&\vdots&\vdots&\vdots\\ \mathbf{0}_{D-3d-1}^{\top}&N-m&\mathbf{0}_{3d-1}^{\top}&-1\\ \end{bmatrix}\in\mathbb{R}^{(2m+2)\times D}.$ $$

Then we have $\mathbf{W}_{1}\cdot{\tt Attn}\circ{\mathcal{T}}\mathcal{B}_{1}(\mathbf{y}_{i})$ as

$$ $\displaystyle\begin{bmatrix}i-1\\ i-2\\ i-m-1\\ N-i\\ N-i-1\\ N-i-m\end{bmatrix}.$ $$

We choose $\mathbf{b}_{1}=\mathbf{0}$, $\mathbf{b}_{2}=[\mathbf{0}_{D-3d-1}^{\top},1,\mathbf{0}_{3d}^{\top}]$. Denote

$$ $\mathbf{w}_{g}=[g(1),g(2)-g(1),g(3)-g(2),\cdots,g(m)-g(m-1),-g(m)]$ $$

and let

$$ $\displaystyle\mathbf{W}_{2}=\begin{bmatrix}\mathbf{0}_{(D-3d-1)\times(2m+2)}\\ \mathbf{w}_{g}~{}~{}\mathbf{w}_{g}\\ \mathbf{0}_{(3d)\times(2m+2)}\end{bmatrix},$ $$

we can check that this exactly express $\widehat{D}^{(m)}_{i}$ and store it in the correpsonding location according to the equation (15). Taking ${\mathcal{T}}\mathcal{B}={\tt FFN}\circ{\tt Attn}$, we reach our target in equation (15).
After that, we could apply another multiplication module $f_{\rm mult}$ to multiply the gradient component by $\widehat{D_{i}}$, which gives rise to

$$ $\displaystyle f_{\rm mult}\circ{\mathcal{T}}\mathcal{B}$ $\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 f_{\rm mult}{\left(\widehat{D_{i}},D_{i}^{-1}\sum_{j =1}^{N}\exp((2r^{2}-2\mathbf{e}_{i}\mathbf{e}_{j})/\ell)\bm{\Sigma}\cdot f_{ \rm mult}(\alpha^{2},\mathbf{s}_{j}^{\top})\right)},$ $\displaystyle\hskip 50.00008ptf_{\rm mult}(\alpha_{t},\bm{\mu}_{i})^{\top},1, \widehat{D_{i}},\mathbf{0}_{3d-2}^{\top},i\bigg{]}^{\top}.$ $$

Then, we can completely follow the proof of Lemma 10 to constuct the gradient and scale the gradient by the learning rate $\eta_{t}$, and the approximation error of the gradient only adds an additional term induced by approximating $D_{i}$, which has been well controlled by setting $m=\mathcal{O}{\left({\left(\ell\log\left(1/\epsilon\right)\right)}^{1/\nu}
\right)}$ to reach a error level of $\epsilon$. We remark that $\widehat{D_{i}}$ can be approximated for only once throughout the transformer networks and we can reuse it in each GD block.

##### Size of Transformer Architecture

We summarize our construction by characterizing the size of the architecture of $f_{\rm GD,softmax}$:

- •
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=2+L_{\rm mult}=\mathcal{O}(\log(\left\lVert\mathbf{x}\right\rVert_{\infty}
\left\lVert\mathbf{s}\right\rVert_{\infty}Nd/\epsilon)^{2}).$
- •
Only one attention head is needed, i.e., $M=1$.
- •
By the construction above, we conclude that the norm of each parameter is bounded by $\mathcal{O}(N^{2}+d+\left\lVert\bm{\Sigma}\right\rVert_{F}+r^{2}))$.

The proof is complete.
∎
With Lemma 11 and following the proof in Appendix B while omitting the part of truncating $\bm{\Gamma}$, i.e., setting $J=N$, we could similarly construct a softmax-transformers as

$$ $\displaystyle\widetilde{\mathbf{s}}_{\rm softmax}=f_{\rm out}\circ f_{\rm GD, softmax}\circ\dots\circ f_{\rm GD,softmax}\circ f_{\mu}\circ f_{\rm in}.$ $$

such that $\left\lVert\widetilde{\mathbf{s}}_{\rm softmax}-\nabla\log p_{t}\right\rVert_{
L_{2}(P_{t})}\leq\sigma_{t}^{-2}\epsilon$ for any $t\geq t_{0}$.

## Appendix C Omitted Proofs in Section 5

### C.1 Proof of Theorem 2

Throughout this section, we assume Assumption 1 holds, and there exists an absolute constant $C$ as in the statement of Theorem 2 such that $C^{-1}\leq\lambda(\bm{\Sigma})\leq C$ and $\left\lVert\bm{\mu}\right\rVert_{\infty}\leq C$, which means

$$ $\left\lVert\bm{\Sigma}\right\rVert_{\infty}=\mathcal{O}(1),~{}~{}\left\lVert \bm{\Sigma}\right\rVert_{\rm op}=\mathcal{O}(1),~{}~{}\left\lVert\bm{\Sigma} \right\rVert_{\rm F}=\mathcal{O}(\sqrt{d})~{}~{}\text{and}~{}~{}\left\lVert\bm {\mu}\right\rVert_{2}\leq\sqrt{Nd}.$ $$

Let’s first analyse the effect of early-stopping time $t_{0}$ on the accuracy of distribution estimation. The following lemma presents the distance between the truth distribution $P$ and the early-stopped distribution $P_{t_{0}}$ in $W_{2}$ distance.

###### Lemma 12 .

Suppose $P\sim\mathcal{N}(\bm{\mu},\bm{\Gamma}\otimes\bm{\Sigma})$, taking the early stopping time $t_{0}=o(1)$, we have

$$ $\displaystyle W_{2}^{2}{\left(P,P_{t_{0}}\right)}\leq(1-\alpha_{t_{0}})^{2} \left\lVert\bm{\mu}\right\rVert^{2}_{2}+C_{\nu}^{2}Nd\sigma_{t_{0}}^{2}.$ $$

Here $C_{\nu}=\lambda_{\max}(\bm{\Gamma}\otimes\bm{\Sigma})\vee 1\leq 1+\left\lVert
\bm{\Sigma}\right\rVert_{\rm op}\ell\lesssim 1+\ell$ by (11).

This means that $W_{2}(P,P_{t_{0}})\lesssim C_{\nu}\sqrt{t_{0}Nd}.$ We defer the proof to Appendix C.1.1.

Then we state our theory on score estimation error. Our score estimator $\widehat{\mathbf{s}}_{t}$ is chosen to minimize the objective function as

$$ $\widehat{\mathbf{s}}_{t}=\arg\min_{s_{t}\in\mathcal{F}}\frac{1}{n}\sum_{i=1}^{ n}\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\|\mathbf{s}_{t}( \mathbf{v}_{t})-\nabla\log p_{t}(\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)})\right \|_{2}^{2}.$ $$

for each time step $t\in[0,T]$.
The next proposition provides guarantee on the score estimation error, which we will transfer into the distribution estimation error later.

###### Proposition 2 .

By taking $\epsilon=\frac{1}{ndN}$ in Theorem 1, we have the expected score estimation error bounded by

$$ $\displaystyle\mathbb{E}_{\mathcal{D}}{\left[\ell(\widehat{\mathbf{s}})\right]} \lesssim(T+\log(1/t_{0}))\log{\left(\kappa_{t}ndN/t_{0}\right)}^{4+1/\nu}\cdot \frac{\ell^{1/\nu}\kappa_{t_{0}}Nd^{3}}{n}.$ $$

The proof is provided in Appendix C.2.
Although our assumption does not ensure the Novikov’s condition to hold for sure, according to Chen et al. [], as long as we have bounded second moment for the score estimation error and finite KL divergence w.r.t. the standard Gaussian, we could still adopt Girsanov’s Theorem and bound the KL divergence between the two distribution.
We restate the lemma as follows:

###### Lemma 13 (Proposition D.1 in Oko et al. [ 2023 ] , see also Theorem 2 in Chen et al. [ 2022b ] ) .

Let $p_{0}$ be a probability distribution, and let $Y={\left\{Y_{t}\right\}}_{t\in[0,T]}$ and $Y^{\prime}={\left\{Y^{\prime}_{t}\right\}}_{t\in[0,T]}$ be two stochastic processes that satisfy the following SDEs:

$$ $\displaystyle{\rm d}Y_{t}$ $\displaystyle=s(Y_{t},t){\rm d}t+{\rm d}W_{t},~{}~{}~{}Y_{0}\sim p_{0}$ $\displaystyle{\rm d}Y^{\prime}_{t}$ $\displaystyle=s^{\prime}(Y^{\prime}_{t},t){\rm d}t+{\rm d}W_{t},~{}~{}~{}Y^{ \prime}_{0}\sim p_{0}$ $$

We further define the distributions of $Y_{t}$ and $Y^{\prime}_{t}$ by $p_{t}$ and $p_{t}$. Suppose that

$$ $\displaystyle\int_{\mathbf{x}}p_{t}(\mathbf{x})\left\lVert(\mathbf{s}-\mathbf{ s}^{\prime})(\mathbf{x},t)\right\rVert_{2}^{2}{\rm d}\mathbf{v}_{t}\leq C$ (16) $$

for any $t\in[0,T]$. Then we have

$$ $\displaystyle\text{KL}\left(p_{T}\|p^{\prime}_{T}\right)\leq\int_{0}^{T}\frac{ 1}{2}\int_{\mathbf{x}}p_{t}(\mathbf{x})\left\lVert(\mathbf{s}-\mathbf{s}^{ \prime})(\mathbf{x},t)\right\rVert_{2}^{2}{\rm d}\mathbf{x}{\rm d}t.$ $$

Now we are ready to prove Theorem 2.

###### Proof of Theorem 2 .

Now we transfer the score estimation error to a TV-distance bound using Lemma 13. Note that under our assumptions and for any $\mathbf{s}\in\mathcal{F}$, we have

$$ $\displaystyle\int_{\mathbf{v}_{t}}p_{t}(\mathbf{v}_{t})\left\lVert\mathbf{s}_{ t}(\mathbf{v}_{t})-\nabla\log p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}{\rm d} \mathbf{v}_{t}$ $\displaystyle\lesssim\int_{\mathbf{v}_{t}}p_{t}(\mathbf{v}_{t}){\left(\frac{ \left\lVert\mathbf{v}_{t}-\alpha_{t}\bm{\mu}\right\rVert_{2}^{2}}{\sigma^{4}_{ t}}+\frac{C}{\sigma_{t}^{2}}\right)}{\rm d}\mathbf{v}_{t}\lesssim\frac{1}{ \sigma^{2}_{t}}.$ $$

Thus, the condition (16) holds for $t_{0}\leq t\leq T$, which means that we could apply Girsanov’s theorem in this time range. Remember that the backward process is written as

$$ $\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{with}~{}~{}\mathbf{X}_{0}^{\leftarrow}\sim{\sf N}(\mathbf{0}, \mathbf{I}).$ $$

We denote the distribution of $\mathbf{X}_{t}^{\leftarrow}$ as $P^{\leftarrow}_{T-t}$. In real setting, we replace $\nabla\log p_{t}$ by its score estimator $\widehat{\mathbf{s}}_{t}$, which gives rise to the following backward process:

$$ $\displaystyle{\rm d}\widehat{\mathbf{X}}_{t}^{\leftarrow}$ $\displaystyle=\left[\frac{1}{2}\widehat{\mathbf{X}}_{t}^{\leftarrow}+\widehat{ \mathbf{s}}_{T-t}(\widehat{\mathbf{X}}_{t}^{\leftarrow})\right]{\rm d}t+{\rm d }\overline{\mathbf{W}}_{t},~{}~{}\text{with}~{}~{}\mathbf{X}_{0}^{\leftarrow} \sim{\sf N}(\mathbf{0},\mathbf{I}).$ $$

We denote the generated distribution of $\widehat{\mathbf{X}}_{t}^{\leftarrow}$ as $\widehat{P}_{T-t}$. Besides, we consider the truth backward process as the inverse process of the forward one, which is defined as

$$ ${\rm d}\mathbf{X}_{t}^{\prime\leftarrow}=\left[\frac{1}{2}\mathbf{X}_{t}^{ \prime\leftarrow}+\nabla\log p_{T-t}(\mathbf{X}_{t}^{\prime\leftarrow})\right] {\rm d}t+{\rm d}\overline{\mathbf{W}}_{t}\quad\text{with}\quad\mathbf{X}_{0}^{ \prime\leftarrow}\sim P_{T}.$ $$

We denote the distribution of $\mathbf{X}_{t}^{\prime\leftarrow}$ by $P^{\prime}_{T-t}$, then we have $P^{\prime}_{t}\sim P_{t}$ for any $t\leq T$.

Since $\mathbf{X}^{\prime\leftarrow}$ and $\mathbf{X}^{\leftarrow}$ are obtained through the same backward SDE but with different initial distributions, by Data Processing Inequality and Pinsker’s Inequality (see e.g., Lemma 2 in Canonne []), we have

$$ $\displaystyle{\rm TV}(P_{t_{0}},P^{\leftarrow}_{t_{0}})$ $\displaystyle={\rm TV}(P^{\prime}_{t_{0}},P^{\leftarrow}_{t_{0}})$ $\displaystyle\lesssim\sqrt{{\rm KL}(P_{t_{0}}^{\prime}||P^{\leftarrow}_{t_{0}})}$ $\displaystyle\lesssim\sqrt{{\rm KL}(P_{T}||{\sf N}(\mathbf{0},\mathbf{I}))}$ $\displaystyle\lesssim\sqrt{{\rm KL}(P||{\sf N}(\mathbf{0},\mathbf{I}))}\exp(-T).$ $$

Thus, we could decompose the TV bound into

$$ $\displaystyle{\rm TV}(P_{t_{0}},\widehat{P}_{t_{0}})$ $\displaystyle\lesssim{\rm TV}(P_{t_{0}},P^{\leftarrow}_{t_{0}})+{\rm TV}(P^{ \leftarrow}_{t_{0}},\widehat{P}_{t_{0}})$ $\displaystyle\lesssim\exp(-T)+\sqrt{\int_{t_{0}}^{T}\frac{1}{2}\int_{\mathbf{v }_{t}}p_{t}(\mathbf{v}_{t})\left\lVert\widehat{\mathbf{s}}(\mathbf{v}_{t}, \mathbf{y},t)-\nabla\log p_{t}(\mathbf{v}_{t})\right\rVert_{2}^{2}{\rm d} \mathbf{v}_{t}{\rm d}t}.$ (17) $$

Thus, by taking expectation over the dataset $\mathcal{D}$ and invoking Jensen inequality, we have

$$ $\displaystyle\mathbb{E}_{\mathcal{D}}{\left[{\rm TV}(P_{t_{0}},\widehat{P}_{t_ {0}})\right]}$ $\displaystyle\lesssim\mathbb{E}_{\mathcal{D}}{\rm TV}(P_{t_{0}},P^{\leftarrow} _{t_{0}})+\mathbb{E}_{\mathcal{D}}{\rm TV}(P^{\leftarrow}_{t_{0}},\widehat{P}_ {t_{0}})$ $\displaystyle\lesssim\exp(-T)+\sqrt{(T+\log(1/t_{0}))\kappa_{t}^{2}\ell^{1/\nu }\log{\left(\kappa_{t}ndNt_{0}^{-1}\right)}^{4+1/\nu}\frac{Nd^{3}}{n}}.$ $$

By taking $T=\mathcal{O}(\log n)$ and combining the results in Lemma 12, we have completed our proof.
∎

###### Lemma 12 .

###### Proposition 2 .

###### Lemma 13 (Proposition D.1 in Oko et al. [ 2023 ] , see also
Theorem 2 in Chen et al. [ 2022b ] ) .

###### Proof of Theorem 2 .

#### C.1.1 Proof of Lemma 12

###### Proof.

Since $P\sim\mathcal{N}(\bm{\mu},\bm{\Gamma}\otimes\bm{\Sigma})$ and $P_{t}\sim\mathcal{N}(\alpha_{t}\bm{\mu},\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{
\Sigma}+\sigma_{t}\bm{I})$, by the formula of the $W_{2}$ distance between two multivariate Gaussian distributions, we have

$$ $\displaystyle W_{2}^{2}{\left(P,P_{t_{0}}\right)}=\left\lVert\bm{\mu}-\alpha_{ t}\bm{\mu}\right\rVert^{2}_{2}+\left\lVert(\bm{\Gamma}\otimes\bm{\Sigma})^{1/2 }-(\alpha_{t}^{2}\bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}\bm{I})^{1/2}\right \rVert^{2}_{\rm F}.$ $$

Suppose $\bm{\Gamma}\otimes\bm{\Sigma}=\bm{P}^{\top}\bm{D}\bm{P}$, where $P\in\mathbb{R}^{dN\times dN}$ is an orthogonal matrix and $\bm{D}={\rm diag}(\lambda_{1},\dots,\lambda_{dN})$ is a diagonal matrix with $\lambda_{i}\geq 0$. Then we have

$$ $\displaystyle\left\lVert(\bm{\Gamma}\otimes\bm{\Sigma})^{1/2}-(\alpha_{t}^{2} \bm{\Gamma}\otimes\bm{\Sigma}+\sigma_{t}\bm{I})^{1/2}\right\rVert^{2}_{\rm F}$ $\displaystyle=\left\lVert\bm{P}^{\top}\bm{D}^{1/2}\bm{P}^{\top}-\bm{P}^{\top}( \alpha_{t}^{2}\bm{D}+\sigma_{t}^{2}\bm{I})^{1/2}\bm{P}\right\rVert^{2}_{\rm F}$ $\displaystyle=\left\lVert\bm{D}^{1/2}-(\alpha_{t}^{2}\bm{D}+\sigma_{t}^{2}\bm{ I})^{1/2}\right\rVert^{2}_{\rm F}$ $\displaystyle=\sum_{i=1}^{Nd}\frac{\sigma_{t}^{4}(\lambda_{i}-1)^{2}}{(\sqrt{ \lambda_{t}}+\sqrt{\alpha_{t}^{2}\lambda_{i}+\sigma_{t}^{2}})^{2}}$ $\displaystyle\leq\sum_{i=1}^{Nd}\frac{\sigma_{t}^{4}C_{\nu}^{2}}{\sigma_{t}^{2}}$ $\displaystyle\leq NdC_{\nu}\sigma_{t}^{2}.$ $$

By taking $t=t_{0}$ and plugging the inequality above into the expression of $W_{2}^{2}{\left(P,P_{t_{0}}\right)}$, we complete our proof.
∎

###### Proof.

### C.2 Proof of Proposition 2

#### C.2.1 Additional Notations

For any score estimator $\widehat{\mathbf{s}}_{t}$, we denote its population loss $\ell$ as:

$$ $\ell(\widehat{\mathbf{s}}_{t}):=\int_{t_{0}}^{T}{\rm d}t\cdot\mathbb{E}_{ \mathbf{v}_{t}\sim p_{t}}\left\|\widehat{\mathbf{s}}_{t}(\mathbf{v}_{t})- \nabla\log p_{t}(\mathbf{v}_{t})\right\|_{2}^{2}$ $$

and we define the empirical loss $\widehat{\ell}$ as

$$ $\displaystyle\widehat{\ell}(\widehat{\mathbf{s}}_{t})$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\int_{t_{0}}^{T}{\rm d}t\cdot\mathbb{E} _{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\|\widehat{\mathbf{s}}_{t}( \mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}^{(i)}}{\sigma_{t }^{2}}\right\|_{2}^{2}.$ $$

Here, $\mathcal{D}=\left\{\mathbf{v}_{0}^{(i)}\right\}_{i\in[n]}$ are $n$ i.i.d samples from true distribution $P(=P_{0})$. When taken expectation over the choice of samples, we have $\mathbb{E}_{\mathcal{D}}[\widehat{\ell}(\widehat{\mathbf{s}}_{t})]=\ell(
\widehat{\mathbf{s}}_{t})+C$ according to Vincent [] for any score estimator $\widehat{\mathbf{s}}$. Here $C$ is a constant independent with $\widehat{\mathbf{s}}_{t}$. Our score estimator $\widehat{\mathbf{s}}_{t}$ is chosen to minimize the objective function as

$$ $\widehat{\mathbf{s}}_{t}=\arg\min_{s_{t}\in\mathcal{F}}\frac{1}{n}\sum_{i=1}^{ n}\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\|s_{t}(\mathbf{v}_{ t})-\nabla\log p_{t}(\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)})\right\|_{2}^{2}.$ $$

for each time step $t\in[0,T]$.

#### C.2.2 Proof of Proposition 2

The proof follows that of Theorem 4.1 in Fu et al. [] by neglecting the conditional information. Specifically, we replace Lemma D.1 in Fu et al. [] by Lemma 14, which provides an $2Nd(T+\log(1/t_{0}))(R_{s}^{2}+1)$ uniform upper bound on the magnitude of the empirical loss function.

###### Lemma 14 (Counterpart of Lemma D.1 in Fu et al. [ 2024 ] ) .

Then for any score estimator $\widehat{\mathbf{s}}\in{\mathcal{T}}(D,L,M,B,R_{s}\sqrt{Nd}\sigma_{t}^{-1})$ and data point $\mathbf{v}_{0}$, we have single-point score loss bounded by

$$ $\displaystyle h(\widehat{\mathbf{s}},\mathbf{v}_{0}):=\int_{t_{0}}^{T}{\rm d}t \cdot\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}}\left\|\widehat{\mathbf{s}}_ {t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}}{\sigma_{t}^ {2}}\right\|_{2}^{2}\leq 2Nd(T+\log(1/t_{0}))(R_{s}^{2}+1).$ (18) $$

The proof is provided in Appendix C.2.3.
Besides, to implement covering number techniques, we also introduce a truncated loss $\widehat{\ell}^{\rm trunc}$ as in Fu et al. [] and bound its difference with the truth loss with small error. We define the truncated loss as

$$ $\displaystyle\widehat{\ell}^{\rm trunc}(\widehat{\mathbf{s}}_{t})$ $\displaystyle:=\frac{1}{n}\sum_{i=1}^{n}\int_{t_{0}}^{T}{\rm d}t\cdot\mathbb{E }_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}{\left[\left\|\widehat{\mathbf{s}}_{ t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}^{(i)}}{\sigma _{t}^{2}}\right\|_{2}^{2}{\mathbf{1}\left\{\left\lVert v_{0}^{(i)}\right\rVert _{2}\leq R_{0}\right\}}\right]}.$ $$

The follow Lemma aligns with equation (D.12) in Fu et al. [].

###### Lemma 15 (Truncation error of the truncated loss function) .

There exists a constant $C_{R}$ such that for any $\epsilon<1$, by choosing $R_{0}=\sqrt{N\gamma_{0}\mathop{\mathrm{tr}}(\bm{\Sigma})+C_{R}\log(Nd/\epsilon
)\left\lVert\bm{\Gamma}\right\rVert_{\mathrm{F}}\left\lVert\bm{\Sigma}\right
\rVert_{\mathrm{F}}}=\mathcal{O}(\log(Nd/\epsilon)\sqrt{Nd})$, we have for any $\widehat{s}\in\mathcal{F}$,

$$ $\displaystyle{\left|\mathbb{E}_{\mathbf{v}_{0}}{\left[\int_{t_{0}}^{T}{\rm d}t \cdot\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\|\widehat{ \mathbf{s}}_{t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}} {\sigma_{t}^{2}}\right\|_{2}^{2}{\mathbf{1}\left\{\left\lVert\mathbf{v}_{0} \right\rVert\geq R_{0}\right\}}\right]}\right|}\lesssim(T+\log(1/t_{0}))\epsilon.$ $$

Moreover, based on the truncated loss function, we have the following results on bounding the difference between the losses of two score estimators that are close to each other, which enables us to apply the results in Appendix C.3.

###### Lemma 16 .

Given the truncation radius $R_{0}>0$. Suppose $\widehat{\mathbf{s}}^{(1)},\widehat{\mathbf{s}}^{(2)}\in\mathcal{F}$
such that $\left\lVert\widehat{\mathbf{s}}^{(1)}_{t}(\mathbf{v})-\widehat{\mathbf{s}}^{(2
)}_{t}(\mathbf{v})\right\rVert_{2}\leq\epsilon$ for any $\left\lVert\mathbf{v}\right\rVert_{2}\leq R_{0}+\sigma_{t}C\sqrt{Nd}\log(dN/\epsilon)$ and $t\geq t_{0}$, where $C$ is an absolute constant. Then we have

$$ ${\left|\widehat{\ell}^{\rm trunc}(\widehat{\mathbf{s}}^{(1)})-\widehat{\ell}^{ \rm trunc}(\widehat{\mathbf{s}}^{(2)})\right|}\leq 2\epsilon{(T+\log(1/t_{0})) {\left(\sqrt{Nd}(R_{s}+C\log(dN/\epsilon))+2R_{0}+1\right)}}.$ $$

Besides, combining the results of Lemmas 23 and 16, we direcly have the following results on bounding the covering number of the truncated loss function class, which aligns with Lemma D.3 in Fu et al. [].

###### Lemma 17 (Counterpart of Lemma D.3 in Fu et al. [ 2024 ] ) .

We consider the truncated loss function class defined as

$$ ${\mathcal{S}}(R_{0})={\left\{h^{\rm trunc}(\mathbf{s},\cdot):\mathbb{R}^{d} \rightarrow\mathbb{R}\bigg{|}\mathbf{s}\in\mathcal{F}\right\}}.$ (19) $$

Here $h^{\rm trunc}(\mathbf{s},\mathbf{v})=h(\mathbf{s},\mathbf{v}){\mathbf{1}\left
\{\left\lVert\mathbf{v}\right\rVert_{2}\leq R_{0}\right\}}$. Then the log-covering number of this loss truncated function class with output range lying in the Euclidean ball with radius $R_{0}$ can be bounded by

$$ $\displaystyle\log\mathcal{N}(\delta;{\mathcal{S}}(R_{0}),\left\lVert\cdot \right\rVert_{\infty})\leq 8D^{2}M\cdot\left(L^{2}\log L_{\mathcal{F}}A_{ \mathcal{F}}+\log\frac{24R_{2}B^{2}MLN^{3/2}}{\epsilon_{\delta}}\right)$ (20) $$

where $R_{2}$ satisfies $R_{2}\leq(r+C_{\rm diff})\sqrt{N}+R_{0}+\sigma_{t}C\sqrt{Nd}\log(dN/\epsilon_{
\delta})$, and $\epsilon_{\delta}$ satisfies

$$ $\displaystyle\delta=2\epsilon_{\delta}{(T+\log(1/t_{0})){\left(\sqrt{Nd}(R_{s} +C\log(dN/\epsilon_{\delta}))+2R_{0}+1\right)}}.$ (21) $$

Note that if we take $\delta=(ndNt_{0}^{-1})^{-C}$ for some constant $C$, we have $\log\epsilon_{\delta}=\mathcal{O}(\log(ndNt_{0}^{-1}))$. The proofs of all the supporting lemmas above are provided in Appendix C.2.3.

With the lemmas and statements above, we can completely follow the main proof of Theorem 4.1 of Fu et al. [] to prove Proposition 2.

###### Proof of Proposition 2 .

First we set $d_{y}=0$ in Fu et al. [] for our unconditioned setting. Then by choosing the covering accuracy $\delta$ in Lemma 17 and taking the truncation range as in Lemma 15, i.e., $R_{0}=\mathcal{O}(\log(Nd/\epsilon_{\delta})\sqrt{Nd})$, we reproduce (D.17) in Fu et al. [] as

$$ $\displaystyle\mathbb{E}_{\mathcal{D}}{\left[\ell(\widehat{\mathbf{s}})\right]}$ $\displaystyle\leq 2\inf_{\mathbf{s}\in\mathcal{F}}\int_{t_{0}}^{T}\mathbb{E}_{ \mathbf{v}_{t}}\left\lVert\mathbf{s}(\mathbf{v}_{t})-\nabla\log p_{t}(\mathbf{ v}_{t})\right\rVert_{2}^{2}{\rm d}t$ $\displaystyle\quad+\frac{(T+\log(1/t_{0}))}{n}\cdot\log\mathcal{N}+2(T+\log(1/ t_{0}))\epsilon_{\delta}+7\delta.$ $\displaystyle\leq 2(T+\log(1/t_{0}))\epsilon$ $\displaystyle\quad+\frac{2Nd(T+\log(1/t_{0}))(R_{s}^{2}+1)}{n}\cdot 8D^{2}M \cdot\left(L^{2}\log L_{\mathcal{F}}A_{\mathcal{F}}+\log\frac{24R_{2}B^{2}MLN^ {3/2}}{\epsilon_{\delta}}\right)$ $\displaystyle\quad+2(T+\log(1/t_{0}))\epsilon_{\delta}+14\epsilon_{\delta}{(T+ \log(1/t_{0})){\left(\sqrt{Nd}(R_{s}+C\log(dN/\epsilon_{\delta}))+2R_{0}+1 \right)}}.$ $$

In the inequality, we invoke (20) and substititue $\delta$ by its expression w.r.t. $\epsilon_{\delta}$ according to (21).
Choosing the both the approximation error and the covering accuracy as $\epsilon=\epsilon_{\delta}=1/n$ and plugging the corresponding parameters about the size of the transformer classes according to Theorem 1 gives rise to

$$ $\displaystyle\mathbb{E}_{\mathcal{D}}{\left[\ell(\widehat{\mathbf{s}})\right]} \lesssim(T+\log(1/t_{0}))\kappa_{t}^{2}\ell^{1/\nu}\log{\left(\kappa_{t}ndNt_{ 0}^{-1}\right)}^{4+1/\nu}\frac{Nd^{3}}{n}.$ $$

The proof is complete.
∎

###### Lemma 14 (Counterpart of Lemma D.1 in Fu et al. [ 2024 ] ) .

###### Lemma 15 (Truncation error of the truncated loss function) .

###### Lemma 16 .

###### Lemma 17 (Counterpart of Lemma D.3 in Fu et al. [ 2024 ] ) .

###### Proof of Proposition 2 .

#### C.2.3 Proofs of Other Supporting Lemmas for Proposition 2

###### Proof of Lemma 14 .

Notice that when $\mathbf{v}_{t}\sim p_{t}(\cdot\mid\mathbf{v}_{0}^{(i)})$, we have $\mathbf{v}_{t}-\alpha\mathbf{v}_{0}^{(i)}=\sigma_{t}z$ where $z\sim\mathcal{N}(0,\bm{I}_{Nd})$ is a standard Gaussian variable. Therefore,

$$ $\displaystyle\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\| \widehat{\mathbf{s}}_{t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t} \mathbf{v}_{0}^{(i)}}{\sigma_{t}^{2}}\right\|_{2}^{2}$ $\displaystyle\leq 2\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\| \widehat{\mathbf{s}}_{t}(\mathbf{v}_{t})\right\|_{2}^{2}+2\mathbb{E}_{\mathbf{ z}\sim\mathcal{N}(0,\bm{I})}\|\mathbf{z}/\sigma_{t}\|^{2}$ $\displaystyle\leq 2R_{s}^{2}Nd/\sigma_{t}^{2}+2Nd/\sigma_{t}^{2}=\frac{2Nd(R_{ s}^{2}+1)}{\sigma_{t}^{2}}.$ $$

Now we can take integral over $t\in[t_{0},T]$ and obtain that:

$$ $\int_{t_{0}}^{T}{\rm d}t\cdot\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i) }}\left\|\widehat{\mathbf{s}}_{t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_ {t}\mathbf{v}_{0}^{(i)}}{\sigma_{t}^{2}}\right\|_{2}^{2}\leq 2Nd(R_{s}^{2}+1) \cdot\int_{t_{0}}^{T}\frac{e^{t}{\rm d}t}{e^{t}-1}\leq 2Nd(T+\log(1/t_{0}))(R_ {s}^{2}+1).$ $$

Here, we use the fact that

$$ $\int_{t_{0}}^{T}\frac{e^{t}{\rm d}t}{e^{t}-1}=\log(e^{T}-1)-\log(e^{t_{0}}-1) \leq T+\log(1/t_{0}).$ $$

The proof is complete.
∎

To prove Lemma 15, we first need to bound the range of the data with high probability.

###### Lemma 18 (Range of the data) .

Given $\delta>0$, with probability $1-2n\exp(-C\delta)$ over the dataset,
Here $C$ is the absolute constant $C_{2}/\sqrt{3}$ in Lemma 24.

###### Proof of Lemma 18 .

This statement is directly related to the concentration of high-dimensional Gaussian distribution. For a random variable $\mathbf{v}\sim\mathcal{N}(\mathbf{0},\bm{\Sigma}_{0})$ where $\bm{\Sigma}_{0}\in\mathbb{R}^{Nd\times Nd}$ is the covariance matrix. Then, we use the polynomial concentration lemma (Lemma 24) and let $g(\cdot)=\|\cdot\|_{2}^{2}$ be the 2-degree polynomial applied on the Gaussian. Then, we have:

$$ $\mathbb{E}[g(v)]=\mathrm{tr}(\bm{\Sigma}_{0}),~{}~{}\mathbb{E}[g(v)^{2}]=3\sum _{i=1}^{Nd}\left(\bm{\Sigma}_{0}\right)_{ii}^{2}+2\sum_{i<j}\left(\left(\bm{ \Sigma}_{0}\right)_{ij}+\left(\bm{\Sigma}_{0}\right)_{ji}\right)^{2}\leq 3 \left\|\bm{\Sigma}_{0}\right\|_{\mathrm{F}}^{2}.$ $$

Therefore, we apply Lemma 24 and conclude that with probability at least $1-2\exp(-C\delta)$ (where $C$ is the $C_{2}$ in Lemma 24), we have:

$$ $\left|\|v\|_{2}^{2}-\mathbb{E}[\|v\|_{2}^{2}]\right|\leq\delta\sqrt{\mathrm{ Var}(\|v\|_{2}^{2})}\leq\sqrt{3}\delta\cdot\left\|\bm{\Sigma}_{0}\right\|_{ \mathrm{F}}.$ $$

For our case, $\bm{\Sigma}_{0}=\bm{\Gamma}\otimes\bm{\Sigma}$, and therefore we can conclude that: with probability at least $1-2n\exp(-C\delta)$:

$$ $\|\mathbf{v}_{0}^{(i)}\|_{2}^{2}\leq\mathrm{tr}(\bm{\Gamma})\cdot\mathrm{tr}( \bm{\Sigma})+\sqrt{3}\delta\|\bm{\Gamma}\|_{\mathrm{F}}\|\bm{\Sigma}\|_{ \mathrm{F}}\leq N\gamma_{0}\mathrm{tr}(\bm{\Sigma})+\sqrt{3}\delta\|\bm{\Gamma }\|_{\mathrm{F}}\|\bm{\Sigma}\|_{\mathrm{F}}.$ $$

holds for $\forall i\in[n]$. Finally, we replace $\delta$ with $\delta/\sqrt{3}$ and it comes to our conclusion.
∎

Now we are ready to prove Lemma 15.

###### Proof of Lemma 15 .

This conclusion can be made by combining the two lemmas above. By using Cauchy-Schwarz inequality and Lemma 14, we have

$$ $\displaystyle~{}~{}~{}{\left|\mathbb{E}_{\mathbf{v}_{0}}{\left[\int_{t_{0}}^{T }{\rm d}t\cdot\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\| \widehat{\mathbf{s}}_{t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t} \mathbf{v}_{0}^{(i)}}{\sigma_{t}^{2}}\right\|_{2}^{2}{\mathbf{1}\left\{\left \lVert\mathbf{v}_{0}\right\rVert\geq R_{0}\right\}}\right]}\right|}$ $\displaystyle\leq 2Nd(R_{s}^{2}+1)(T+\log(1/t_{0}))\cdot\sqrt{\mathbb{P}[\left \lVert\mathbf{v}_{0}\right\rVert^{2}\geq R_{0}^{2}]}$ $\displaystyle\leq 2Nd(R_{s}^{2}+1)(T+\log(1/t_{0}))\cdot\sqrt{2}\exp\left(-CC_ {R}\log(Nd/\varepsilon)/2\right).$ $$

Let $C_{R}=2/C$, then we have:

$$ ${\left|\mathbb{E}_{\mathbf{v}_{0}}{\left[\int_{t_{0}}^{T}{\rm d}t\cdot\mathbb{ E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}^{(i)}}\left\|\widehat{\mathbf{s}}_{t}( \mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}^{(i)}}{\sigma_{t }^{2}}\right\|_{2}^{2}{\mathbf{1}\left\{\left\lVert\mathbf{v}_{0}\right\rVert \geq R_{0}\right\}}\right]}\right|}\lesssim\varepsilon(T+\log(1/t_{0})),$ $$

which comes to our conclusion.
∎

###### Proof of Lemma 16 .

For any datapoint $\mathbf{v}_{0}$ such that $\left\lVert\mathbf{v}_{0}\right\rVert\leq R_{0}$ and diffusion time $t\geq t_{0}$, we have

$$ $\displaystyle\quad{\left|\mathbb{E}_{\mathbf{v}_{t}\mid\mathbf{v}_{0}}\left\| \widehat{\mathbf{s}}^{(1)}_{t}(\mathbf{v}_{t})+\frac{\mathbf{v}_{t}-\alpha_{t} \mathbf{v}_{0}}{\sigma_{t}^{2}}\right\|_{2}^{2}-\mathbb{E}_{\mathbf{v}_{t}\mid \mathbf{v}_{0}^{(i)}}\left\|\widehat{\mathbf{s}}^{(2)}_{t}(\mathbf{v}_{t})+ \frac{\mathbf{v}_{t}-\alpha_{t}\mathbf{v}_{0}}{\sigma_{t}}\right\|_{2}^{2} \right|}$ $\displaystyle={\left|\mathbb{E}_{\mathbf{z}\in\mathcal{N}(\mathbf{0},\bm{I})}{ \left[\left\|\widehat{\mathbf{s}}^{(1)}_{t}(\alpha_{t}\mathbf{v}_{0}+\sigma_{t }\mathbf{z})+\frac{\mathbf{z}}{\sigma_{t}}\right\|_{2}^{2}-\left\|\widehat{ \mathbf{s}}^{(2)}_{t}(\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z})+\frac{ \mathbf{z}}{\sigma_{t}}\right\|_{2}^{2}\right]}\right|}$ $\displaystyle={\left|\mathbb{E}_{\mathbf{z}\in\mathcal{N}(\mathbf{0},\bm{I})}{ \left[\sigma_{t}^{-1}{\left((\widehat{\mathbf{s}}^{(1)}_{t}-\widehat{\mathbf{s }}^{(2)}_{t})(\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z})\right)}^{\top}{ \left(\sigma_{t}(\widehat{\mathbf{s}}^{(1)}_{t}+\widehat{\mathbf{s}}^{(2)}_{t} )(\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z})+2\mathbf{z}\right)}\right]} \right|}$ $\displaystyle\leq\sigma_{t}^{-1}\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{ 0},\bm{I})}{\left[\left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left \lVert\bm{\phi}_{2}(\mathbf{z})\right\rVert_{2}\right]},$ $$

where $\bm{\phi}_{1}(\mathbf{z})=(\widehat{\mathbf{s}}^{(1)}_{t}-\widehat{\mathbf{s}}
^{(2)}_{t})(\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z})$ and $\bm{\phi}_{2}(\mathbf{z})=\sigma_{t}(\widehat{\mathbf{s}}^{(1)}_{t}+\widehat{
\mathbf{s}}^{(2)}_{t})(\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z})+2\bm{z}$. By the upper bound of $s^{(i)}$, we know that

$$ $\displaystyle\left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}$ $\displaystyle\leq 2R_{s}\sqrt{Nd}\sigma_{t}^{-1},~{}~{}\text{and}~{}~{}\left \lVert\bm{\phi}_{2}(\mathbf{z})\right\rVert_{2}\leq 2R_{s}\sqrt{Nd}+2\left \lVert\mathbf{z}\right\rVert_{2}.$ $$

Thus, they both have sub-linear growth w.r.t. $\left\lVert\mathbf{z}\right\rVert_{2}$. Now we decompose $\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{0},\bm{I})}{\left[\left\lVert\bm
{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left\lVert\bm{\phi}_{2}(\mathbf{z})
\right\rVert_{2}\right]}$ by

$$ $\displaystyle\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{0},\bm{I})}{\left[ \left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left\lVert\bm{\phi}_{2}( \mathbf{z})\right\rVert_{2}\right]}$ $\displaystyle=\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{0},\bm{I})}{\left[ \left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left\lVert\bm{\phi}_{2}( \mathbf{z})\right\rVert_{2}{\mathbf{1}\left\{\left\lVert\alpha_{t}\mathbf{v}_{ 0}+\sigma_{t}\mathbf{z}\right\rVert_{2}\leq R_{1}\right\}}\right]}$ $\displaystyle\quad+\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{0},\bm{I})}{ \left[\left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left\lVert\bm{\phi} _{2}(\mathbf{z})\right\rVert_{2}{\mathbf{1}\left\{\left\lVert\alpha_{t}\mathbf {v}_{0}+\sigma_{t}\mathbf{z}\right\rVert_{2}\geq R_{1}\right\}}\right]}$ $\displaystyle\leq 2\epsilon(R_{s}\sqrt{Nd}+\sigma_{t}^{-1}(R_{0}+R_{1}))$ $\displaystyle\quad+4R_{s}\sqrt{Nd}\sigma_{t}^{-1}\underbrace{\mathbb{E}_{ \mathbf{z}}{\left[(R_{s}\sqrt{Nd}+\left\lVert\mathbf{z}\right\rVert_{2}){ \mathbf{1}\left\{\left\lVert\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z} \right\rVert_{2}\geq R_{1}\right\}}\right]}}_{A}.$ $$

In the inequality, we invoke $\left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\leq\epsilon$ and $\left\lVert\mathbf{z}\right\rVert_{2}\leq\sigma_{t}^{-1}(R_{0}+R_{1})$ for $\left\lVert\alpha_{t}\mathbf{v}_{0}+\sigma_{t}\mathbf{z}\right\rVert_{2}\leq R
_{1}$. By Cauchy inequality, we can bound $A$ by

$$ $\displaystyle A$ $\displaystyle\leq\mathbb{E}_{\mathbf{z}}{\left[(R_{s}\sqrt{Nd}+\left\lVert \mathbf{z}\right\rVert_{2})^{2}\right]}^{1/2}\Pr{\left[\left\lVert\alpha_{t} \mathbf{v}_{0}+\sigma_{t}\mathbf{z}\right\rVert_{2}\geq R_{1}\right]}^{1/2}$ $\displaystyle\leq\sqrt{(2R_{s}^{2}+2)Nd\Pr{\left[\left\lVert\mathbf{z}\right \rVert_{2}\geq\sigma_{t}^{-1}(R_{1}-\alpha_{t}R_{0})\right]}}.$ $$

Since $\mathbf{z}$ is standard Gaussian, we can set $R_{1}=R_{0}+\sigma_{t}C\sqrt{Nd}\log(dN/\epsilon)$ for some absolute constant $C$ so that $A\leq\epsilon/(4\sqrt{Nd})$. Altogether, we have

$$ $\displaystyle\mathbb{E}_{\mathbf{z}\sim\mathcal{N}(\mathbf{0},\bm{I})}{\left[ \left\lVert\bm{\phi}_{1}(\mathbf{z})\right\rVert_{2}\left\lVert\bm{\phi}_{2}( \mathbf{z})\right\rVert_{2}\right]}\leq 2\epsilon{\left(R_{s}\sqrt{Nd}+\sigma_ {t}^{-1}(2R_{0}+1)+C\sqrt{Nd}\log(dN/\epsilon)\right)}.$ $$

Plugging the inequality into the expression of ${\left|\widehat{\ell}^{\rm trunc}(\widehat{\mathbf{s}}^{(1)})-\widehat{\ell}^{
\rm trunc}(\widehat{\mathbf{s}}^{(2)})\right|}$, we have

$$ $\displaystyle{\left|\widehat{\ell}^{\rm trunc}(\widehat{\mathbf{s}}^{(1)})- \widehat{\ell}^{\rm trunc}(\widehat{\mathbf{s}}^{(2)})\right|}$ $\displaystyle\leq\int_{t_{0}}^{T}\sigma_{t}^{-1}2\epsilon{\left(R_{s}\sqrt{Nd} +\sigma_{t}^{-1}(2R_{0}+1)+C\sqrt{Nd}\log(dN/\epsilon)\right)}$ $\displaystyle\leq\int_{t_{0}}^{T}\sigma_{t}^{-2}2\epsilon{\left(R_{s}\sqrt{Nd} +(2R_{0}+1)+C\sqrt{Nd}\log(dN/\epsilon)\right)}$ $\displaystyle\leq 2\epsilon{(T+\log(1/t_{0})){\left(\sqrt{Nd}(R_{s}+C\log(dN/ \epsilon))+2R_{0}+1\right)}}.$ $$

The proof is complete.
∎

###### Proof of Lemma 14 .

###### Lemma 18 (Range of the data) .

###### Proof of Lemma 18 .

###### Proof of Lemma 15 .

###### Proof of Lemma 16 .

### C.3 Covering Number of the Multi-layer Transformers

The score network we apply satisfies the following form:

$$ $f=f_{l}\circ f_{l-1}\circ\ldots\circ f_{1}$ $$

where the total layer number $l=2L$ (which consists of $L$ feed-forward layers and $L$ Transformer layers) and $f_{1},f_{2},\ldots,f_{l}$ are either attention layers or feed-forward networks whose input and output lies in $\mathbb{R}^{D\times N}$. Denote $\mathcal{B}(R)=\{\mathbf{X}\in\mathbb{R}^{D\times N}:~{}\|\mathbf{X}\|_{
\mathrm{F}}\leq R\}$ is a $d$-dimensional ball with radius $R$. Assume there exists a sequence of radius $R_{0},R_{1},\ldots,R_{l}>0$ (which will be determined later) such that $f_{i}:\mathcal{B}(R_{i-1})\rightarrow\mathcal{B}(R_{i})$ holds for $\forall i=0,1,\ldots,l$. Then, we need to compute the covering number with respect to the $l_{\infty}$ norm of the function space constructed by $f:\mathcal{B}(R_{0})\rightarrow\mathcal{B}(R_{l})$ with the form above.

Notice that for two such functions $f=f_{l}\circ f_{l-1}\circ\ldots\circ f_{1}$ and $f^{\prime}=f_{l}^{\prime}\circ f_{l-1}^{\prime}\circ\ldots\circ f_{1}^{\prime}$, the $l_{\infty}$ norm of their difference can be upper bounded by the following lemma:

###### Lemma 19 .

For functions $\{f_{i}:\mathbb{R}^{d\times N}\rightarrow\mathbb{R}^{d\times N}\}_{i\in[l]}$ and $\{f_{i}^{\prime}:\mathbb{R}^{d\times N}\rightarrow\mathbb{R}^{d\times N}\}_{i
\in[l]}$, we denote their composition as $f:=f_{l}\circ f_{l-1}\circ\ldots\circ f_{1}$ and $f^{\prime}:=f_{l}^{\prime}\circ f_{l-1}^{\prime}\circ\ldots\circ f_{1}^{\prime}$, where $f,f^{\prime}:\mathbb{R}^{d\times N}\rightarrow\mathbb{R}^{d\times N}$. For any matrix-to-matrix function $g$, we denote its $\|\cdot\|_{\mathrm{F},\infty}$ norm as:

$$ $\|g\|_{\mathrm{F},\infty}:=\sup_{\mathbf{X}}\|g(\mathbf{X})\|_{\mathrm{F}}$ $$

and its Lipschitz continuity as

$$ $\mathrm{Lip}(g):=\sup_{\mathbf{X},\mathbf{X}^{\prime}}\frac{\|g(\mathbf{X})-g( \mathbf{X}^{\prime})\|_{\mathrm{F}}}{\|\mathbf{X}-\mathbf{X}^{\prime}\|_{ \mathrm{F}}},$ $$

which is a simple extension from the Lipschitz continuity with respect to $l_{2}$ norm of vectors. Then, it holds that:

$$ $\|f-f^{\prime}\|_{\mathrm{F},\infty}\leq\sum_{i=1}^{l}\|f_{i}-f_{i}^{\prime}\| _{\mathrm{F},\infty}\cdot\prod_{j=i+1}^{l}\mathrm{Lip}(f_{j}).$ $$

###### Proof.

For any $k=2,3,\ldots,l$ and $x\in\mathcal{B}(R_{0})$, it holds that:

$$ $\displaystyle~{}~{}~{}~{}~{}\|f_{k}\circ\ldots f_{1}(x)-f_{k}^{\prime}\circ \ldots f_{1}^{\prime}(x)\|_{\mathrm{F},\infty}$ $\displaystyle\leq\|f_{k}^{\prime}\circ f_{k-1}^{\prime}\circ\ldots\circ f_{1}^ {\prime}(\mathbf{X})-f_{k}\circ f_{k-1}^{\prime}\circ\ldots\circ f_{1}^{\prime }(\mathbf{X})\|_{\mathrm{F},\infty}$ $\displaystyle~{}~{}~{}~{}~{}+\|f_{k}\circ f_{k-1}^{\prime}\circ\ldots\circ f_{ 1}^{\prime}(\mathbf{X})-f_{k}\circ f_{k-1}\circ\ldots\circ f_{1}(\mathbf{X})\| _{\mathrm{F},\infty}$ $\displaystyle\leq\|f_{k}^{\prime}-f_{k}\|_{\mathrm{F},\infty}+\mathrm{Lip}(f_{ k})\cdot\|f_{k-1}^{\prime}\circ\ldots\circ f_{1}^{\prime}-f_{k-1}\circ\ldots \circ f_{1}\|_{\mathrm{F},\infty}.$ $$

After taking maximum over $x\in\mathcal{B}(R_{0})$, we conclude that

$$ $\|f_{k}\circ\ldots f_{1}-f_{k}^{\prime}\circ\ldots f_{1}^{\prime}\|_{\mathrm{F },\infty}\leq\|f_{k}-f_{k}^{\prime}\|_{\mathrm{F},\infty}+\mathrm{Lip}(f_{k}) \cdot\|f_{k-1}\circ\ldots f_{1}-f_{k-1}^{\prime}\circ\ldots f_{1}^{\prime}\|_{ \mathrm{F},\infty}.$ $$

By using the method of induction, we can easily derive our conclusion.
∎

Now, in order to compute the covering number of the function space constructed by functions $f$ with the form $f=f_{l}\circ f_{l-1}\circ\ldots\circ f_{1}$ where $f_{i}\in\mathcal{F}_{i}$, we firstly need to bound the Lipschitz constants for the function classes $\mathcal{F}_{i}$. Also, we need to estimate the covering number of each $\mathcal{F}_{i}$.

###### Lemma 20 .

For the function space of feed-forward network

$$ $\displaystyle\mathcal{F}^{{\tt FFN}}=\Big{\{}$ $\displaystyle{\tt FFN}:\mathbb{R}^{D\times N}\rightarrow\mathbb{R}^{D\times N} ,~{}\mathbf{Y}\mapsto\mathbf{Y}+\mathbf{W}_{2}\cdot\mathrm{ReLU}(\mathbf{W}_{1 }\mathbf{Y}+\mathbf{b}_{2}\mathbf{1}^{\top})+\mathbf{b}_{1}\mathbf{1}^{\top}:$ $\displaystyle~{}~{}~{}~{}\|\mathbf{W}_{1}\|_{\mathrm{F}},\|\mathbf{W}_{2}\|_{ \mathrm{F}},\|\mathbf{b}_{1}\|_{2},\|\mathbf{b}_{2}\|_{2}<B,\mathbf{Y}\in \mathcal{B}(R)\Big{\}},$ $$

then all functions in the class $\mathcal{F}^{{\tt FFN}}$ are $(1+B^{2})$-Lipschitz. The covering number can be bounded as:

$$ $\log\mathcal{N}(\delta;\mathcal{F}^{{\tt FFN}},\|\cdot\|_{\mathrm{F},\infty}) \leq 4D^{2}\log\frac{12B^{2}(R+\sqrt{N})}{\delta}.$ $$

###### Proof.

For $\forall f\in\mathcal{F}^{{\tt FFN}}$ with $f(\mathbf{Y})=\mathbf{Y}\mapsto\mathbf{Y}+\mathbf{W}_{2}\cdot\mathrm{ReLU}(
\mathbf{W}_{1}Y+\mathbf{b}_{2}\mathbf{1}^{\top})+\mathbf{b}_{1}\mathbf{1}^{\top}$, by the additivity of Lipschitz constant and the fact that $\mathrm{ReLU}$ is 1-Lipschitz continuous, we can conclude that the Lipschitz constant of function $f$ is no larger than $1+\|\mathbf{W}_{2}\|_{2}\cdot\|\mathbf{W}_{1}\|_{2}\leq 1+\|\mathbf{W}_{1}\|_{
\mathrm{F}}\|\mathbf{W}_{2}\|_{\mathrm{F}}<1+c^{2}$, which means $f$ is $(1+c^{2})$-Lipschitz continuous. On the other hand, for two functions $f,g\in\mathcal{F}^{{\tt FFN}}$ where $f(\mathbf{Y})=\mathbf{Y}+\mathbf{W}_{2}\cdot\mathrm{ReLU}(\mathbf{W}_{1}
\mathbf{Y}+\mathbf{b}_{2}\mathbf{1}^{\top})+\mathbf{b}_{1}\mathbf{1}^{\top}$ and $g(\mathbf{Y})=\mathbf{Y}+\mathbf{W}_{2}^{\prime}\cdot\mathrm{ReLU}(\mathbf{W}_
{1}^{\prime}\mathbf{Y}+\mathbf{b}_{2}^{\prime}\mathbf{1}^{\top})+\mathbf{b}_{1
}^{\prime}\mathbf{1}^{\top}$. Then for $\forall\mathbf{Y}\in\mathcal{B}(R)$, we have:

$$ $\displaystyle\|f(\mathbf{Y})-g(\mathbf{Y})\|_{\mathrm{F}}\leq\sqrt{N}\cdot\| \mathbf{b}_{1}-\mathbf{b}_{1}^{\prime}\|_{2}+\|\mathbf{W}_{2}^{\prime}\cdot \left(\mathrm{ReLU}(\mathbf{W}_{1}\mathbf{Y}+\mathbf{b}_{2}\mathbf{1}^{\top})- \mathrm{ReLU}(\mathbf{W}_{1}^{\prime}\mathbf{Y}+\mathbf{b}_{2}^{\prime}\mathbf {1}^{\top})\right)\|_{\mathrm{F}}$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}+\|(\mathbf{W}_{2}-\mathbf {W}_{2}^{\prime})\cdot\mathrm{ReLU}(\mathbf{W}_{1}\mathbf{Y}+\mathbf{b}_{2} \mathbf{1}^{\top})\|_{\mathrm{F}}$ $\displaystyle~{}~{}~{}~{}~{}\leq\sqrt{N}\cdot\|\mathbf{b}_{1}-\mathbf{b}_{1}^{ \prime}\|_{2}+\|\mathbf{W}_{2}\|_{\mathrm{F}}\cdot\left(\sqrt{N}\|\mathbf{b}_{ 2}-\mathbf{b}_{2}^{\prime}\|_{2}+R\|\mathbf{W}_{1}-\mathbf{W}_{1}^{\prime}\|_{ \mathrm{F}}\right)$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}+\|\mathbf{W}_{2}-\mathbf{W}_{2}^{ \prime}\|_{\mathrm{F}}\cdot(R\|\mathbf{W}_{1}\|_{\mathrm{F}}+\sqrt{N}\|\mathbf {b}_{2}\|_{2})$ $\displaystyle~{}~{}~{}~{}~{}\leq\sqrt{N}\cdot\|\mathbf{b}_{1}-\mathbf{b}_{1}^{ \prime}\|_{2}+B\sqrt{N}\cdot\|\mathbf{b}_{2}-\mathbf{b}_{2}^{\prime}\|_{2}+BR \cdot\|\mathbf{W}_{1}-\mathbf{W}_{1}^{\prime}\|_{\mathrm{F}}+B(R+\sqrt{N}) \cdot\|\mathbf{W}_{2}-\mathbf{W}_{2}^{\prime}\|_{\mathrm{F}}.$ $$

Therefore, it holds that:

$$ $\|f-g\|_{\mathrm{F},\infty}\leq\sqrt{N}\cdot\|\mathbf{b}_{1}-\mathbf{b}_{1}^{ \prime}\|_{2}+B\sqrt{N}\cdot\|\mathbf{b}_{2}-\mathbf{b}_{2}^{\prime}\|_{2}+BR \cdot\|\mathbf{W}_{1}-\mathbf{W}_{1}^{\prime}\|_{\mathrm{F}}+B(R+\sqrt{N}) \cdot\|\mathbf{W}_{2}-\mathbf{W}_{2}^{\prime}\|_{\mathrm{F}},$ $$

which leads to the upper bound of covering number:

$$ $\displaystyle\mathcal{N}(\delta;\mathcal{F}^{{\tt FFN}},\|\cdot\|_{\mathrm{F}, \infty})$ $\displaystyle\leq\mathcal{N}(\delta/4\sqrt{N};\mathcal{M}_{1},\|\cdot\|_{2}) \cdot N(\delta/4B\sqrt{N};\mathcal{M}_{1},\|\cdot\|_{2})\cdot$ $\displaystyle~{}~{}~{}~{}~{}\mathcal{N}(\delta/4BR;\mathcal{M}_{2},\|\cdot\|_{ \mathrm{F}})\cdot\mathcal{N}(\delta/4B(R+\sqrt{N});\mathcal{M}_{2},\|\cdot\|_{ \mathrm{F}})$ $$

where $\mathcal{M}_{1}=\{\mathbf{v}\in\mathbb{R}^{D}:~{}\|\mathbf{v}\|_{2}\leq B\},
\mathcal{M}_{2}=\{\mathbf{M}\in\mathbb{R}^{D\times D}:~{}\|\mathbf{M}\|_{
\mathrm{F}}<B\}$. It is well-known that for $\forall\varepsilon>0$, we have $\mathcal{N}(\varepsilon;\mathcal{M}_{1},\|\cdot\|_{2})\leq(3B/\varepsilon)^{D}$ and $\mathcal{N}(\varepsilon;\mathcal{M}_{2},\|\cdot\|_{\mathrm{F}})\leq(3B/
\varepsilon)^{D^{2}}$. To sum up, we finally conclude that

$$ $\log\mathcal{N}(\delta;\mathcal{F}^{{\tt FFN}},\|\cdot\|_{\mathrm{F},\infty}) \leq 4D^{2}\log\frac{12B^{2}(R+\sqrt{N})}{\delta}.$ $$

∎

The proof is complete.

###### Lemma 21 .

For the function space of attention network

$$ $\displaystyle\mathcal{F}^{{\tt Attn}}=\Big{\{}$ $\displaystyle{\tt Attn}:\mathbb{R}^{D\times N}\rightarrow\mathbb{R}^{D\times N },~{}\mathbf{Y}\mapsto\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):$ $\displaystyle~{}~{}~{}~{}\|\mathbf{V}^{m}\|_{\mathrm{F}},\|\mathbf{Q}^{m}\|_{ \mathrm{F}},\|\mathbf{K}^{m}\|_{\mathrm{F}}<B~{}~{}\text{for}~{}\forall m\in[M ],\mathbf{Y}\in\mathcal{B}(R)\Big{\}},$ $$

then all functions in the class $\mathcal{F}^{{\tt Attn}}$ are $(1+BM\sqrt{N}+2B^{3}MR^{2}N)$-Lipschitz. The covering number can be bounded as:

$$ $\log\mathcal{N}(\delta;\mathcal{F}^{{\tt Attn}},\|\cdot\|_{\mathrm{F},\infty}) \leq D^{2}M\cdot\log\frac{6B^{2}MR^{3}N^{3/2}}{\delta}.$ $$

Here, the metric $\|f\|_{\mathrm{F},\infty}:=\sup_{\mathbf{Y}\in\mathcal{B}(R)}\|f(\mathbf{Y})\|
_{\mathrm{F}}$.

###### Proof.

Notice that for any two functions $f_{1},f_{2}$ and $\mathbf{Y},\mathbf{Y}^{\prime}$, we have:

$$ $\|f_{1}(\mathbf{Y})f_{2}(\mathbf{Y})-f_{1}(\mathbf{Y}^{\prime})f_{2}(\mathbf{Y }^{\prime})\|_{\mathrm{F}}\leq\|f_{1}(\mathbf{Y}^{\prime})\|_{\mathrm{F}}\cdot \|f_{2}(\mathbf{Y})-f_{2}(\mathbf{Y}^{\prime})\|_{\mathrm{F}}+\|f_{2}(\mathbf{ Y})\|_{\mathrm{F}}\cdot\|f_{1}(\mathbf{Y})-f_{1}(\mathbf{Y}^{\prime})\|_{ \mathrm{F}}.$ $$

It means that the Lipschitz constant of $\mathbf{V}^{m}\mathbf{Y}\cdot\sigma\left((\mathbf{Q}^{m}\mathbf{Y})^{\top}
\mathbf{K}^{m}\mathbf{Y}\right)$ is no larger than $\|\mathbf{V}^{m}\|_{2}\cdot\sqrt{N}+R\sqrt{N}\|\mathbf{V}^{m}\|_{2}\cdot 2\|
\mathbf{Q}^{m\top}\mathbf{K}^{m}\|_{\mathrm{F}}\cdot R\sqrt{N}$. Here, we use the fact that both ReLU and softmax activation function over vectors is 1-Lipschitz under $l_{2}$ norm and that over matrices is 1-Lipschitz under Frobenius norm. See Appendix E.2. Also, $(\mathbf{Q}^{m}\mathbf{Y})^{\top}\mathbf{K}^{m}\mathbf{Y}$ is $2\|\mathbf{Q}^{m\top}\mathbf{K}^{m}\|_{\mathrm{F}}\cdot R\sqrt{N}$-Lipschitz under Frobenius norm when $\mathbf{Y}\in\mathcal{B}(R)$. Therefore, in the function class $\mathcal{F}^{{\tt Attn}}$, all functions included are $(1+BM\sqrt{N}+2B^{3}MR^{2}N)$-Lipschitz continuous, which comes to our conclusion. For the covering number, given any two functions $f,g\in\mathcal{F}^{{\tt Attn}}$ with their corresponding parameter sets $\{\mathbf{V}_{1}^{m},\mathbf{Q}_{1}^{m},\mathbf{K}_{1}^{m}:~{}m\in[M]\}$ and $\{\mathbf{V}_{2}^{m},\mathbf{Q}_{2}^{m},\mathbf{K}_{2}^{m}:~{}m\in[M]\}$, then:

$$ $\displaystyle\|f(\mathbf{Y})-g(\mathbf{Y})\|_{\mathrm{F}}$ $\displaystyle\leq\sum_{m=1}^{M}\left\|(\mathbf{V}_{1}^{m}-\mathbf{V}_{2}^{m}) \mathbf{Y}\cdot\sigma\left(\mathbf{Y}^{\top}\mathbf{Q}_{1}^{m\top}\mathbf{K}_{ 1}^{m}\mathbf{Y}\right)\right\|_{\mathrm{F}}$ $\displaystyle~{}~{}~{}~{}+\sum_{m=1}^{M}\left\|\mathbf{V}_{2}^{m}\mathbf{Y} \cdot\left(\sigma\left(\mathbf{Y}^{\top}\mathbf{Q}_{1}^{m\top}\mathbf{K}_{1}^{ m}\mathbf{Y}\right)-\sigma\left(\mathbf{Y}^{\top}\mathbf{Q}_{2}^{m\top}\mathbf {K}_{2}^{m}\mathbf{Y}\right)\right)\right\|_{\mathrm{F}}$ $\displaystyle\leq\sum_{m=1}^{M}\|\mathbf{V}_{1}^{m}-\mathbf{V}_{2}^{m}\|_{ \mathrm{F}}\cdot R\sqrt{N}\cdot\sqrt{N}+\sum_{m=1}^{M}BR\sqrt{N}\cdot\|\mathbf {Y}^{\top}(\mathbf{B}_{2}^{m}-\mathbf{B}_{1}^{m})\mathbf{Y}\|_{\mathrm{F}}$ $\displaystyle\leq\sum_{m=1}^{M}\|\mathbf{V}_{1}^{m}-\mathbf{V}_{2}^{m}\|_{ \mathrm{F}}\cdot RN+\sum_{m=1}^{M}BR^{3}N^{3/2}\cdot\|\mathbf{B}_{2}^{m}- \mathbf{B}_{1}^{m}\|_{\mathrm{F}}.$ $$

Here, $\mathbf{B}_{1}^{m}:=\mathbf{Q}_{1}^{m\top}\mathbf{K}_{1}^{m}$ and $\mathbf{B}_{2}^{m}:=\mathbf{Q}_{2}^{m\top}\mathbf{K}_{2}^{m}$. Then, $\|\mathbf{B}_{1}^{m}\|_{\mathrm{F}},\|\mathbf{B}_{2}^{m}\|_{\mathrm{F}}\leq B^
{2}$ holds for $\forall m\in[M]$. It leads to the following upper bound of the covering number of $\mathcal{F}^{{\tt Attn}}$.

$$ $\displaystyle\mathcal{N}(\delta;\mathcal{F}^{{\tt Attn}},\|\cdot\|_{\mathrm{F} ,\infty})\leq\prod_{m=1}^{M}\mathcal{N}(\delta/2MRN;\mathcal{M}_{1},\|\cdot\|_ {\mathrm{F}})\cdot\prod_{m=1}^{M}\mathcal{N}(\delta/2BMR^{3}N^{3/2};\mathcal{M }_{2},\|\cdot\|_{\mathrm{F}})$ $$

where matrix set $\mathcal{M}_{1}=\{V\in\mathbb{R}^{D\times D}:~{}\|\mathbf{V}\|_{\mathrm{F}}
\leq B\}$ and $\mathcal{M}_{2}=\{\mathbf{B}\in\mathbb{R}^{D\times D}:~{}\|\mathbf{B}\|_{
\mathrm{F}}\leq B^{2}\}$. To sum up, we conclude that:

$$ $\log\mathcal{N}(\delta;\mathcal{F}^{{\tt Attn}},\|\cdot\|_{\mathrm{F},\infty}) \leq D^{2}M\cdot\log\frac{6BMRN}{\delta}+D^{2}M\cdot\log\frac{6B^{2}MR^{3}N^{3 /2}}{\delta},$ $$

which comes to our conclusion.
∎

Now, we are ready to combine these results together.

###### Lemma 22 .

Consider the multi-layer transformers class ${\mathcal{T}}_{\rm raw}(D,L,M,B)$ without encoders and decoders.
Then the log-covering number with input range bounded by $R_{0}$ ($\left\lVert\mathbf{Y}\right\rVert_{\rm F}\leq R_{0}$) can be bounded by

$$ $\displaystyle\log\mathcal{N}(\delta;{\mathcal{T}}_{\rm raw},R_{0},\|\cdot\|_{ \mathrm{F},\infty})$ $\displaystyle\leq 4D^{2}M\cdot\left(L^{2}\log L_{\mathcal{F}}A_{\mathcal{F}}+ \log\frac{12R_{0}B^{2}MLN^{3/2}}{\delta}\right).$ $$

Here $L_{\mathcal{F}}=1+c^{2}\vee(BM\sqrt{N}+2B^{3}MR^{2}N)$ and $A_{\mathcal{F}}=(1+B^{2}+2B\sqrt{N})\vee(1+MB\sqrt{N})$.

###### Proof.

For the score network class $\mathcal{F}={\mathcal{T}}_{\rm raw}(D,L,M,B)$, we have:

$$ $\mathcal{F}=\{f_{2L}\circ f_{2L-1}\circ\ldots\circ f_{1}:~{}f_{i}\in\mathcal{F }_{i}~{}\text{for}~{}\forall i\in[2L],~{}\text{each }\mathcal{F}_{i}\text{ is either }\mathcal{F}^{{\tt FFN}}\text{ or }\mathcal{F}^{{\tt Attn}}\}.$ $$

Denote $L_{\mathcal{F}}:=1+B^{2}\vee(BM\sqrt{N}+2B^{3}MR^{2}N)$ as the upper bound of Lipschitz constant for both $\mathcal{F}^{{\tt FFN}}$ and $\mathcal{F}^{{\tt Attn}}$. Notice that for any input $Y\in\mathbb{R}^{D\times N}$ such that $\|Y\|_{\mathrm{F}}\leq R_{0}$, then the output of feed-forward network $f\in\mathcal{F}^{{\tt FFN}}$ holds

$$ $\|f(Y)\|_{\mathrm{F}}\leq\|Y\|_{\mathrm{F}}+B(B\|Y\|_{\mathrm{F}}+\sqrt{N})+B \sqrt{N}=(1+B^{2}+2B\sqrt{N})\cdot(R_{0}\vee 1).$ $$

The output of attention network $f\in\mathcal{F}^{{\tt Attn}}$ holds that

$$ $\|f(Y)\|_{\mathrm{F}}\leq\|Y\|_{\mathrm{F}}+MB\|Y\|_{\mathrm{F}}\cdot\sqrt{N} \leq(1+MB\sqrt{N})R_{0}.$ $$

Here we use the fact that $\|P\|_{\mathrm{F}}\leq\sqrt{N}$ holds for all probability matrix $P\in\mathbb{R}^{D\times N}$. Denote $A_{\mathcal{F}}:=(1+B^{2}+2B\sqrt{N})\vee(1+MB\sqrt{N})>1$ as the signal amplifier of each layer, then we can set the sequence of radius as $R_{i}=\left(A_{\mathcal{F}}\right)^{i}\cdot R_{0}$ with $R_{0}>1$ and $f_{i}:\mathcal{B}(R_{i-1})\rightarrow\mathcal{B}(R_{i})$. According to Lemma 19, we have:

$$ $\mathcal{N}(\delta;\mathcal{F},\|\cdot\|_{\mathrm{F},\infty})\leq\prod_{i=1}^{ 2L}\mathcal{N}(\delta/2LL_{\mathrm{F}}^{2L-i};\mathcal{F}_{i},\|\cdot\|_{ \mathrm{F},\infty}),$ $$

which leads to

$$ $\displaystyle\quad\log\mathcal{N}(\delta;\mathcal{F},\|\cdot\|_{\mathrm{F}, \infty})$ $\displaystyle\leq\sum_{i=1}^{2L}\log\mathcal{N}(\delta/2LL_{\mathrm{F}}^{l-i}; \mathcal{F}_{i},\|\cdot\|_{\mathrm{F},\infty})$ $\displaystyle\leq\sum_{i=1}^{2L}\left(4D^{2}\log\frac{12B^{2}(R_{i-1}+\sqrt{N} )\cdot 2LL_{\mathrm{F}}^{2L-i}}{\delta}~{}\vee~{}D^{2}M\cdot\log\frac{6B^{2}MR _{i-1}^{3}N^{3/2}\cdot 2LL_{\mathrm{F}}^{2L-i}}{\delta}\right)$ $\displaystyle\leq 8D^{2}M\cdot\left(L^{2}\log L_{\mathcal{F}}A_{\mathcal{F}}+ \log\frac{24R_{0}B^{2}MLN^{3/2}}{\delta}\right).$ $$

The proof is complete.
∎

Moreover, if we consider the entire transformers architecture with the encoder and decoder layers, we directly have the following results:

###### Lemma 23 (Covering number of transformers) .

Consider the entire transformer architecture $\mathcal{F}={\mathcal{T}}(D,L,M,B,R)$ in which the encoder that takes $\mathbf{v}_{t}\in\mathbb{R}^{Nd}$ as the input and embeds it with the time embedding $\mathbf{e}$ and the diffusion time-step embedding $\bm{\phi}(t)$, where $\left\lVert\mathbf{e}\right\rVert_{2}=r$ and $\left\lVert\bm{\phi}(t)\right\rVert_{2}\leq C_{\rm diff}$ for some absolute constant $C_{\rm diff}>0$. Then the log-covering number with initial input range $\left\lVert\mathbf{v}_{t}\right\rVert_{2}\leq R_{0}$ is bounded by

$$ $\displaystyle\log\mathcal{N}(\delta;\mathcal{F},R_{0},\|\cdot\|_{\mathrm{F}, \infty})$ $\displaystyle\leq 8D^{2}M\cdot\left(L^{2}\log L_{\mathcal{F}}A_{\mathcal{F}}+ \log\frac{24R_{1}B^{2}MLN^{3/2}}{\delta}\right).$ $$

Here $R_{1}=R_{0}+\sqrt{N}(r+C_{\rm diff})$. This is because the encoder maps $\mathbf{v}_{t}$ to $\mathbf{Y}\in\mathbb{R}^{D\times N}$ with $\left\lVert\mathbf{Y}\right\rVert_{\rm F}\leq\left\lVert\mathbf{v}_{t}\right
\rVert_{2}+r\sqrt{N}+C_{\rm diff}\sqrt{N}$, and our decoder (defined in Appendix B.3.1) only extracts part of the entries from the output and impose a clipping function with range being $R$ on it, both of which lead to no increase in the covering number.

###### Lemma 19 .

###### Proof.

###### Lemma 20 .

###### Proof.

###### Lemma 21 .

###### Proof.

###### Lemma 22 .

###### Proof.

###### Lemma 23 (Covering number of transformers) .

## Appendix D Experimental Details in Section 6

##### Patch Embedding Modification

To implement our simulation experiments on Gaussian Process data, we slightly adapt the original DiT designed for image generation [Peebles and Xie, ]. In the original DiT, a pre-trained VAE encoder is deployed to convert image in the training dataset to feature representations and patch embedding is applied by splitting the image into subblocks(patches) and flattening and projecting each subblock to a feature. But in our numeric experiments, we dropped the VAE encoder and split the data at time dimension, with each time step being a patch, and project the patch into feature of higher dimension.

##### Kernel Estimation Method

In our experiment, each sample generated by DiT in is denoted as $\mathbf{S}=[\mathbf{s}_{1},...,\mathbf{s}_{N}]\in\mathbb{R}^{N\times d}$, where $\mathbf{s}_{i}=\mathbf{S}[i]\in\mathbb{R}^{d}$ denotes the data patch at $i$-th time index. In our experiment we take $N=128$ and $d=8$ and with $n\in\{10^{3},3200,10^{4},32000,10^{5}\}$ samples in total. To evaluate the quality of generated data, we calculate the mean of data patch at $i$-th time step and covariance matrix between data patch at $i$-th and $j$-th time step as follows and compare them with the theoretical ones.

$$ $\displaystyle\widehat{\bm{\mu}}_{i}=\frac{1}{n}\sum_{k=1}^{n}\mathbf{S}_{k}[i] ,~{}~{}~{}~{}\widehat{\bm{\Sigma}}_{i,j}=\frac{1}{n}\sum_{k=1}^{n}(\mathbf{S}_ {k}[i]-\widehat{\bm{\mu}}_{i})(\mathbf{S}_{k}[j]-\widehat{\bm{\mu}}_{j})^{\top}.$ $$

Here $\widehat{\bm{\mu}}_{i}\in\mathbb{R}^{d}$ is the empirical mean of the $i$-th patch and $\widehat{\bm{\Sigma}}_{i,j}$ is the empirical covariance matrix between the $i$-th and the $j$-th patch. With $\widehat{\bm{\Sigma}}_{i,j}$, we could estimate the empirical kernel value $\widehat{\gamma}(i,j)$ as

$$ $\displaystyle\widehat{\gamma}(i,j)=\mathop{\mathrm{argmin}}_{\alpha}\left \lVert\widehat{\bm{\Sigma}}_{i,j}-\alpha\bm{\Sigma}\right\rVert_{\rm F}.$ $$

where $\bm{\Sigma}\in\mathbb{R}^{d\times d}$ is the true covariance matrix.
After running on each pair of $(i,j)\in\{1,2,...,N\}^{2}$, we can get the empirical kernel $\widehat{\bm{\Gamma}}=[\widehat{\gamma}(i,j)]_{ij}\in\mathbb{R}^{N\times N}$.

##### Relative Error

In the experiments where we compare sample efficiency across different training data size $n$ and kernel setting $\nu,\ell$, the metric is a relative error of the estimated sample covariance matrix to its ground-truth:

$$ $\displaystyle\epsilon=\frac{\left\|\widehat{\bm{\Gamma}}\otimes\widehat{\bm{ \Sigma}}-\bm{\Gamma}\otimes\bm{\Sigma}\right\|_{\rm F}^{2}}{\left\|(\widehat{ \bm{\Gamma}}\otimes\widehat{\bm{\Sigma}})_{\rm truth}-\bm{\Gamma}\otimes\bm{ \Sigma}\right\|_{\rm F}^{2}}.$ $$

Here $(\widehat{\bm{\Gamma}}\otimes\widehat{\bm{\Sigma}})_{\rm truth}$ is the estimated covariance matrix through the same method but under an equal amount of training data, instead of generated data. This relative error eliminates the influence of different scaling of $\bm{\Gamma}\otimes\bm{\Sigma}$ and the concentration error caused by finite samples.

##### Query-Key Matrices and Value Matrices

From Figure  6, we can see that the weights of the query-key matrix $\mathbf{Q}^{\top}\mathbf{K}$ emphasize the time embedding part of the input across different layers. Interestingly, Figure 7 shows that the weight emphasis on the value matrix $\mathbf{V}$ is on the data patch part instead. This observation aligns with our theoretical construction of the score function using the DiT structure. Specifically, the score matrix (determined by the query-key matrix) captures the kernel $\bm{\Gamma}$, representing temporal dependency, while the value matrix determines the covariance $\bm{\Sigma}$, representing spatial dependency.

Figure: Figure 6: Query-Key matrices in different transformer attention layers(1~12).
Refer to caption: x5.png

Figure: Figure 7: Value weight matrices $\mathbf{V}$ in different transformer attention layers(1~12).
Refer to caption: x6.png

##### Attention Score in Different Attention Layers

To further demonstrate our theory, we visualize the attention score $\mathbf{Y}_{t}^{\top}\mathbf{Q}^{\top}\mathbf{K}\mathbf{Y}_{t}\in\mathbb{R}^{N
\times N}$ averaged over $n=10^{5}$ data points at different backward diffusion times $t\in\{0,50,100,200,400,800,1000\}$ to observe if it gradually unveils the kernel $\bm{\Gamma}$ as the backward process progresses. According to Figure 8, in the initial few layers (1~2), we observe the early stages of kernel construction. In the subsequent layers (3~12), the attention score matrix increasingly resembles the kernel, becoming clearer as the backward diffusion process advances.

Figure: Figure 8: Attention score matrices in different attention layers and at different steps of the backward process.
Refer to caption: x7.png

## Appendix E Supporting Technical Results

### E.1 Gaussian Lemmas

In this subsection, we will introduce several Lemmas to control the deviation of random variables which polynomially depend on some Gaussian random variables.
We will use a slightly modified version of Lemma 30 from Damian et al. [].

###### Lemma 24 .

Let $g$ be a polynomial of degree $p$ and $x\sim\mathcal{N}(0,I_{d})$. Then there exists an absolute positive constant $C_{p}$ depending only on $p$ such that for any $\delta>1$,

$$ $\mathbb{P}\left[|g(x)-\mathbb{E}[g(x)]|\geq\delta\sqrt{\operatorname{Var}(g(x) )}\right]\leq 2\exp\left(-C_{p}\delta^{2/p}\right).$ $$

The next Lemma orginates from Theorem 4.3, Prato and Tubaro [].

###### Lemma 25 .

For any $\ell\in\mathbb{N}$ and $f\in L^{2}(\mathcal{N}(0,\bm{I}_{d}))$ to be a degree $\ell$ polynomial, for any $q\geq 2$, we have

$$ $\mathbb{E}_{z\sim\gamma}\left[f(z)^{q}\right]\leq C_{q,\ell}\left(\mathbb{E}_{ z\sim\gamma}\left[f(z)^{2}\right]\right)^{q/2}.$ $$

where we use $C_{q,\ell}$ to denote some universal constant that only depends on $q,\ell$.

###### Lemma 24 .

###### Lemma 25 .

### E.2 Lipschitz Continuity of Activation Functions

###### Lemma 26 (1-Lipschitz continuity of softmax function) .

For the softmax function $\sigma:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$, it is 1-Lipschitz continuous under $l_{2}$ norm.

###### Proof.

For any $x\in\mathbb{R}^{d}$, denote $J(x)$ as the Jacobian matrix of softmax function at $x$. Then, the Lipschitz continuity of $\sigma(\cdot)$ under $l_{2}$ norm can be upper bound by $\sup_{x\in\mathbb{R}^{d}}\|J(x)\|_{\mathrm{F}}$. By calculation, we have

$$ $J(x)=\mathrm{Diag}\left(\sigma(x)\right)-\sigma(x)\sigma(x)^{\top}.$ $$

Let $\sigma(x)=(p_{1},p_{2},\ldots,p_{d})^{\top}$ which is a probability vector, then

$$ $\|J(x)\|_{\mathrm{F}}^{2}\leq\sum_{i=1}^{d}p_{i}^{2}(1-p_{i})^{2}+\sum_{i\neq j }(p_{i}p_{j})^{2}\leq\left(\sum_{i=1}^{d}p_{i}^{2}\right)^{2}\leq\left(\sum_{i =1}^{d}p_{i}\right)^{4}=1.$ $$

Therefore $\sup_{x\in\mathbb{R}^{d}}\|J(x)\|_{\mathrm{F}}\leq 1$, which comes to our conclusion. A direct extension is that the column-wise softmax over matrices is 1-Lipschitz continuous under $\|\cdot\|_{\mathrm{F}}$ norm.
∎

###### Lemma 26 (1-Lipschitz continuity of softmax function) .

###### Proof.

### E.3 Basics on ResNet Universal Approximation Theory

In this section, we briefly introduce universal approximation theory of ResNet. An $L$-layer ResNet $\mathbf{R}(\mathbf{x}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{o}}$ can be defined as

$$ $\displaystyle\mathbf{R}(\mathbf{x})=\mathcal{B}\circ{\tt FFN}_{L}\circ{\tt FFN }_{L-1}\circ\cdots\circ{\tt FFN}_{1}\circ\mathcal{A}(\mathbf{x}).$ (22) $$

Here $\mathcal{A}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d^{\prime}}$ and $\mathcal{B}:\mathbb{R}^{d^{\prime}}\rightarrow\mathbb{R}^{d_{o}}$ are two linear transformations, and ${\tt FFN}_{i}:\mathbb{R}^{d^{\prime}}\rightarrow\mathbb{R}^{d^{\prime}}$ are basic residual blocks defined as
${\tt FFN}_{i}(\mathbf{y})=\mathbf{y}+\mathbf{W}_{2,i}\cdot{\rm ReLU}(\mathbf{W
}_{1,i}\mathbf{y}+\mathbf{b}_{2,i})+\mathbf{b}_{1,i}$ with $\mathbf{W}_{1,i}\in\mathbb{R}^{d_{i}\times d^{\prime}}$, $\mathbf{W}_{2,i}\in\mathbb{R}^{d^{\prime}\times d_{i}}$, $\mathbf{b}_{1,i}\in\mathbb{R}^{d^{\prime}}$ and $\mathbf{b}_{2,i}\in\mathbb{R}^{d_{i}}$. We denote by $\mathcal{RN}(d,d_{o},d^{\prime},W,L,S,C)$ the set of ResNet functions from $\mathbb{R}^{d}$ to $\mathbb{R}^{d_{o}}$ with $L$ layers, $d^{\prime}$ neurons in each identity layer, maximum width $W=\max_{i\in[L]}\{d_{i}\}$ and nonzero weights $S$. Moreover, the Frobenius norm of the weight matrices $\mathbf{W}_{j,i}$ and the Euclidean norm of the bias vectors $\bm{b}_{j,i}$ are uniformly bounded by $C>0$.

The next lemma shows that we can approximate product operation with ResNet.

###### Lemma 27 (Proposition 12 in Liu et al. [ 2024a ] ) .

Let $x,y\in[-B,B]$ with $B\geq 1$. Then there exists a ResNet
$R\in\mathcal{RN}(2,1,3,4,\mathcal{O}(\log(B/\epsilon),\mathcal{O}(\log(B/
\epsilon),\mathcal{O}(B))$ such that

$$ $\displaystyle{\left|R(x,y)-xy\right|}\leq\epsilon,~{}~{}~{}x,y\in[-B,B]^{d}$ (23) $$

holds.

The construction process follows the standard techniques proposed by Yarotsky [], which first use feed-forward networks to approximate the square function $f_{\rm sq}(x)\approx x^{2}$ and transfer it into a product operator $\times(x,y)=(f_{\rm sq}(x+y)-f_{\rm sq}(x-y))/4\approx xy$.

Furthermore, for $\mathbf{x}\in\mathbb{R}^{d}$ and $y\in\mathbb{R}$, we can approxiamtely construct the following mapping:

$$ $\displaystyle f([\mathbf{x}^{\top},y]^{\top})=y\mathbf{x}.$ (24) $$

###### Corollary 2 .

Given $\epsilon>0$, there exists a ResNet $f_{\rm mult}\in\mathcal{RN}(d+1,d,3d,4d,$ $\mathcal{O}(\log(B/\epsilon)),d\mathcal{O}(\log(B/\epsilon)),\mathcal{O}(Bd))$ such that
$f_{\rm mult}(y,\mathbf{x})=y\mathbf{x}+\bm{\epsilon}$, where $\left\lVert\bm{\epsilon}\right\rVert_{\infty}\leq\epsilon$.

###### Proof.

By Lemma 27, we know there exists a ResNet $R:\mathbb{R}^{2}\rightarrow\mathbb{R}$ satisfies

$$ $R\in\mathcal{RN}(2,1,3,4,\mathcal{O}(\log(B/\epsilon)),\mathcal{O}(\log(B/ \epsilon)),\mathcal{O}(1))$ $$

and

$$ $\displaystyle{\left|R(x,y)-xy\right|}\leq\epsilon,~{}~{}x,y\in[-B,B].$ $$

Then let’s consider the following two steps of mapping:

$$ $\displaystyle\mathcal{A}([\mathbf{x}^{\top},y]^{\top})$ $\displaystyle=[x_{1},y,x_{2},y,\cdots,x_{d},y]^{\top},$ $\displaystyle\mathcal{B}([x_{1},y,x_{2},y,\cdots,x_{d},y]^{\top})$ $\displaystyle=[R(x_{1},y),R(x_{1},y),\cdots,R(x_{d},y)]^{\top}.$ $$

Here we could choose $\mathcal{A}$ to be a $(d+1)\times 2d$ matrix

$$ $\displaystyle\mathcal{A}=\begin{bmatrix}1&0&\cdots&0&0\\ 0&0&\cdots&0&1\\ 0&1&\cdots&0&0\\ 0&0&\cdots&0&1\\ \vdots&&\ddots&&\vdots\\ 0&0&\cdots&1&0\\ 0&0&\cdots&0&1\end{bmatrix},$ $$

and take $\mathcal{B}$ as parallelization of $d$ homogeneous networks $R$. Let $f_{\rm mult}=\mathcal{B}\circ\mathcal{A}$, we know $f_{\rm mult}\in\mathcal{RN}(d+1,d,3d,4d,\mathcal{O}(\log(B/\epsilon)),d
\mathcal{O}(\log(B/\epsilon)),\mathcal{O}(d))$, and it approximately realizes the mapping (24) with $L_{\infty}$ approximation error being $\epsilon$ for the first $d$ entries and no error for the last entry.
∎

More generally, if we consider the input to be $[x^{\top},y,\mathbf{z}]^{\top}$ and want to construct a (FFN-only) transformers that approximately maps the input to $[yx^{\top},y,\mathbf{z}]^{\top}$, we have the following results:

###### Corollary 3 .

Suppose the input to be $\mathbf{Y}=[\mathbf{y}_{1},\mathbf{y}_{2},\dots,\mathbf{y}_{N}]\in\mathbb{R}^{
D\times N}$ with $\mathbf{y}_{i}=[\mathbf{x}^{\top}_{i},\mathbf{0}_{2d}^{\top},w_{i},\mathbf{z}_
{i}^{\top}]$, where $\mathbf{x}_{i}\in[-B,B]^{d}$, $w_{i}\in[-B,B]$ and $\mathbf{z}_{i}\in\mathbb{R}^{d_{z}}$. Given any $\epsilon>0,$ there exists a (FFN-only) transformers

$$ $\displaystyle\mathbf{f}_{\texttt{mult}}={\tt FFN}_{L}\circ{\tt FFN}_{L-1}\circ \cdots\circ{\tt FFN}_{1}$ $$

with $L=\mathcal{O}(\log(B/\epsilon))$ layers that approximately multiplies each component $\mathbf{x}_{i}$ with the weight $w_{i}$, which keeping other dimensions the same. This can be formally written as

$$ $\displaystyle\mathbf{f}_{\texttt{mult}}(\mathbf{Y})=\begin{bmatrix}f_{\rm mult }(w_{1},\mathbf{x}_{1})&\cdots&f_{\rm mult}(w_{N},\mathbf{x}_{N})\\ \mathbf{0}_{2d}&\cdots&\mathbf{0}_{2d}\\ w_{1}&\cdots&w_{N}\\ \mathbf{z}_{1}&\cdots&\mathbf{z}_{N}\end{bmatrix},~{}~{}\text{where}~{}~{} \left\lVert f_{\rm mult}(w_{i},\mathbf{x}_{i})-w_{i}\mathbf{x}_{i}\right\rVert _{\infty}\leq\epsilon.$ $$

The inner dimension of the FFNs is at most $8d$. Moreover, the number of nonzero coefficients in each weight matrices or bias vectors is at most $\mathcal{O}(d)$, and the norm of the matrices and bias are all bounded by $\mathcal{O}(Bd)$.

Here we require the buffer variables $\mathbf{0}_{2d}$ in the input because Corollary 2 needs $3d$ neurons in each dimension to store the calculation results that are necessary for constructing the product function. Thus, we add $\mathbf{0}_{2d}$ so that $[\mathbf{x}^{\top},\mathbf{0}_{2d}^{\top}]\in\mathbb{R}^{3d}$ for construction convenience.

###### Proof of Corollary 3 .

By Corollary 2, there exists a ResNet $f_{\rm mult}:\mathbb{R}^{d+1}\rightarrow\mathbb{R}^{d}$ such that

$$ $\displaystyle f_{\rm mult}(y,\mathbf{x})=\mathcal{B}\circ{\tt FFN}_{L}\circ{ \tt FFN}_{L-1}\circ\cdots\circ{\tt FFN}_{1}\circ\mathcal{A}(\mathbf{x},y)$ $$

with $f_{\rm mult}(y,\mathbf{x})=y\mathbf{x}+\bm{\epsilon}$, where $\left\lVert\bm{\epsilon}\right\rVert_{\infty}\leq\epsilon$. Since the FFNs in the transformers are position-wise, we only need to consider the mapping of one token while others are completely the same. For each $i$, suppose ${\tt FFN}_{i}:\mathbb{R}^{3d}\rightarrow\mathbb{R}^{3d}$ satisfies

$$ $\displaystyle{\tt FFN}_{i}(\mathbf{v})=\mathbf{v}+\mathbf{W}_{2,i}\cdot{\rm ReLU }(\mathbf{W}_{1,i}\mathbf{v}+\mathbf{b}_{2,i})+\mathbf{b}_{1,i}$ $$

with $\mathbf{W}_{1,i}\in\mathbb{R}^{d_{i}\times 3d}$, $\mathbf{W}_{2,i}\in\mathbb{R}^{3d\times d_{i}}$, $\mathbf{b}_{1,i}\in\mathbb{R}^{3d}$ and $\mathbf{b}_{2,i}\in\mathbb{R}^{d_{i}}$. Here $d_{i}\leq 4d$ for all $1\leq i\leq L$. For the first FFN, let

$$ $\displaystyle{\tt FFN}^{\prime}_{1}{\left(\begin{bmatrix}\mathbf{x}\\ \mathbf{0}_{2d}\\ w\\ \mathbf{z}\end{bmatrix}\right)}$ $\displaystyle=\begin{bmatrix}\mathbf{x}\\ \mathbf{0}_{2d}\\ w\\ \mathbf{z}\end{bmatrix}+\begin{bmatrix}\mathbf{W}_{2,1}\\ \mathbf{0}_{d_{z}\times d_{i}}\end{bmatrix}\cdot{\rm ReLU}\left({\left[\mathbf {W}_{1,1}\mathcal{A},\mathbf{0}_{d_{i}\times(d_{z}+2d+1)}\right]}\begin{ bmatrix}\mathbf{x}\\ w\\ \mathbf{0}_{2d}\\ \mathbf{z}\end{bmatrix}+\mathbf{b}_{2,1}\right)+\begin{bmatrix}\mathbf{b}_{1,1 }\\ \mathbf{0}_{d_{z}}\end{bmatrix}$ $\displaystyle=\begin{bmatrix}{\tt FFN}_{1}\circ\mathcal{A}(\mathbf{x},y)\\ w\\ \mathbf{z}\end{bmatrix}.$ $$

For $2\leq i\leq L$, suppose the input is $[\mathbf{v}^{\top},w,\mathbf{z}^{\top}]^{\top}=[\mathbf{v}^{\top},w,\mathbf{z}
^{\top}]^{\top}$ with $\mathbf{v}\in\mathbb{R}^{3d}$ and $\mathbf{z}\in\mathbb{R}^{d_{z}}$, let

$$ $\displaystyle{\tt FFN}^{\prime}_{i}{\left(\begin{bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}\right)}=\begin{bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}+\begin{bmatrix}\mathbf{W}_{2,i}\\ \mathbf{0}_{d^{\prime}_{z}\times d_{i}}\end{bmatrix}\cdot{\rm ReLU}\left({ \left[\mathbf{W}_{1,i},\mathbf{0}_{d_{i}\times d^{\prime}_{z}}\right]}\begin{ bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}+\mathbf{b}_{2,i}\right)+\begin{bmatrix}\mathbf{b}_{1,i }\\ \mathbf{0}_{d^{\prime}_{z}}\end{bmatrix}=\begin{bmatrix}{\tt FFN}_{i}(\mathbf{ v})\\ \mathbf{z}\end{bmatrix}.$ $$

Here $d^{\prime}_{z}=d_{z}+1$. For the final layer, let

$$ $\displaystyle{\tt FFN}^{\prime}_{L+1}{\left(\begin{bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}\right)}=\begin{bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}+\bm{W}_{2,L+1}\cdot{\rm ReLU}\left(\begin{bmatrix} \mathcal{B}&\mathbf{0}_{d\times d^{\prime}_{z}}\\ -\mathcal{B}&\mathbf{0}_{d\times d^{\prime}_{z}}\\ \mathbf{I}_{3d}&\mathbf{0}_{d\times d^{\prime}_{z}}\\ -\mathbf{I}_{3d}&\mathbf{0}_{d\times d^{\prime}_{z}}\end{bmatrix}\begin{ bmatrix}\mathbf{v}\\ w\\ \mathbf{z}\end{bmatrix}\right)=\begin{bmatrix}\mathcal{B}\mathbf{v}\\ \mathbf{0}_{2d}\\ w\\ \mathbf{z}\end{bmatrix}.$ $$

Here

$$ $\displaystyle\mathbf{W}_{2,L+1}=\begin{bmatrix}\text{diag}(\mathbf{1}_{d}, \mathbf{0}_{2d})&-\text{diag}(\mathbf{1}_{d},\mathbf{0}_{2d})&-\mathbf{I}_{3d} &\mathbf{I}_{3d}\\ \mathbf{0}_{(2d+d^{\prime}_{z})\times d}&\mathbf{0}_{(2d+d^{\prime}_{z})\times d }&\mathbf{0}_{(2d+d^{\prime}_{z})\times 3d}&\mathbf{0}_{(2d+d^{\prime}_{z}) \times 3d}\end{bmatrix}.$ $$

Thus, we have finished constructing the (FFN-only) transformers we want by taking

$$ $\mathbf{f}_{\texttt{mult}}={\tt FFN}^{\prime}_{L+1}\circ{\tt FFN}_{L-1}\circ \cdots\circ{\tt FFN}^{\prime}_{1},$ $$

and the hidden dimension of the FFNs is at most $8d$. Moreover, by the definition of the original ${\tt FFN}_{i}$ and the new ${\tt FFN}^{\prime}_{i}$, the number of nonzero coefficients in each weight matrix or bias vector is at most $\mathcal{O}(d)$, and the norm of the matrices and bias are all bounded by $\mathcal{O}(d)$. The proof is complete.
∎

###### Lemma 27 (Proposition 12 in Liu et al. [ 2024a ] ) .

###### Corollary 2 .

###### Proof.

###### Corollary 3 .

###### Proof of Corollary 3 .

### E.4 Asymptotic Results on the Spectrum of Toeplitz Matrices

In this section, we provide some existing results on the spectrum of Toeplitz matrix when both its size and bandwidth go to infinity. For a Toeplitz matrix $\mathbf{T}\in\mathbb{R}^{N\times N}$, its $(i,j)$-th component $\mathbf{T}_{ij}=a_{|i-j|}$ only depends on its distance to diagonal. When $k>M$, we have $a_{k}=0$ where $M$ is known as the bandwidth. Denote $\lambda_{1},\lambda_{2},\ldots,\lambda_{N}$ to be the eigenvalues of $\mathbf{T}$, with multiplicities counted and let $\mu_{N}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{\lambda_{i}}$ be the empirical distribution of the spectrum. In the asymptotic case where $M,N\rightarrow\infty$, we focus on the behavior of $\mu_{N}$. Denote $F_{N}(x)$ to be the cumulative distribution function of $\mu_{N}$, Kargin [] proposes the result that $F_{N}(x)$ converges to the standard Gaussian distribution when the Toeplitz matrix follows that $\mathbb{E}a_{k}=0,\mathbb{E}a_{k}^{2}=\frac{1}{M}$, $\sup_{k,N}\mathbb{E}|\sqrt{M}a_{k}|^{4}<C<\infty$ and most importantly, the band-to-size ratio $\frac{M}{N}\rightarrow 0$. While the band-to-size ratio $\frac{M}{N}\rightarrow c\in(0,1)$, the spectrum distribution $F_{N}(x)$ converges to some non-Gaussian distribution $\Psi_{c}$. Some other statistical works such as Hartman and Wintner [], Tilli [], Delsarte and Genin [] study the spectrum of generalized Toeplitz matrices by using Fourier expansion.

In our case, the condition number $\kappa_{t_{0}}$ will keep in a constant range if the sequence $\{a_{k}\}$ introduced above sharply decays, while grow as $N$ goes larger if $\{a_{k}\}$ slowly decays. In the former case, we can treat $\kappa_{t_{0}}$ as a constant, which does not affect our analysis. In the latter case, since we have a natural upper bound of

$$ $\kappa_{t_{0}}\leq\sigma_{t_{0}}^{-2}\lambda_{\max}(\bm{\Gamma}\otimes\bm{ \Sigma})\lesssim\ell\sigma_{t_{0}}^{-2}\lesssim t_{0}^{-1},$ $$

by taking $t_{0}=n^{-1/3}$ in Theorem 2, we can still obtain a $n^{-1/6}$-convergence rate in both $W_{2}$ distance and ${\rm TV}$ distance.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - In the quadratic formulation in Equation (3), the score function estimation remains non-separable over timesteps (i.e., across x1,t, ..., xN,t). Would further decomposition allow for a distributed solution with minimal inter-timestep communication?

- Prior work, such as [Sahiner et al., 2022], has already shown that attention layers capture correlations and reflect low-rank priors, so it’s unsurprising that attention proves effective for correlated and stationary data.

Sahiner, A., Ergen, T., Ozturkler, B., Pauly, J., Mardani, M., & Pilanci, M. (2022, June). Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In the International Conference on Machine Learning (pp. 19050-19088). PMLR.
soundness: 2
strengths: - The paper addresses a timely and relevant problem, aiming to understand why spatiotemporal transformers effectively learn distributions in spatiotemporal data.

- The writing is clear, making the methodology and findings accessible.
summary: This paper explores transformers as denoisers within diffusion models for generating spatiotemporal data, motivated by recent advancements in video foundation models, such as SORAH.
The approach begins by assuming a Gaussian process as the underlying data distribution and simplifies it to a scale-varying Gaussian process where the covariance changes only in scale over time. In this framework, the score function is related to the inverse of a modified covariance matrix. Since directly calculating this inverse is computationally prohibitive for high-dimensional latent vectors (dimension = d * N), the paper introduces a convex quadratic relaxation solvable via gradient descent. The authors then design a transformer architecture inspired by unrolled gradient descent steps. They also derive sample complexity bounds for the transformer’s ability to learn the covariance based on the smoothness and decay properties of the kernel matrix.
The paper includes basic experiments with synthetic Gaussian process data and a simple ball-motion scenario to validate the proposed method.
weaknesses: - The goal of the paper lacks clarity. If the aim is to demonstrate the expressiveness of transformers for correlated or spatiotemporal data, this doesn’t necessarily relate to diffusion models and could be examined independently in a more simplified setting.

-The assumption of a Gaussian process with scale-wise stationarity is a strong simplification, limiting the applicability to realistic distributions. Consequently, the results offer limited insights into transformer design for broader, more complex distributions.

-The experiments are weak and unconvincing. To substantiate the theory, more realistic data should be used, and comparisons with alternative architectures, such as UNet (3D or 2+1D), known for capturing spatiotemporal correlations, would strengthen the argument.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see weaknesses.
soundness: 3
strengths: I really like the framing in this paper. While sounding like toy settings, GPs can be quite general and could be a good starting point of analysis. I especially appreciate the analytical forms, for example in the score approximator, and how the covariance function relates to positional embeddings. The paper is also quite well written. I have gone through parts of the proofs and they made sense - although I have not exhaustively verified every step.
summary: This paper scrutinises diffusion transformers under a simplified setting where the input time series is assumed to follow a Gaussian process. The paper carefully derives the analytical forms of the dynamics and provides score approximation as well as data-distribution estimation guarantees in the simplified setting of interest. The authors cast this paper as a first step towards building more sophisticated theories.
weaknesses: - Can we have a section in the appendix where all the assumptions and simplifications are listed? This would really ease the conclusions that the reader can draw from this paper.
- Ln. 397: What's a "properly transformer" architecture? 
- Could the authors give more elaborations on what motivates the choices of $\epsilon$ and $T$ in Thm. 2? 
- What about empirically testing the learning on GPLVMs? Wouldn't this be a natural fit? 
- There are some duplicate looking sections in the appendix, for example relating to the size of the transformer blocks. Can't we gather those in a big table instead of repeating them after each relevant section?

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - The learned temporal dependencies presented in Figure 5 look very similar to what you get when you just plot the attention between two uncorrelated matrices with positional embedding (l. 300). What is the difference between these maps and the ones you present, e.g., in Fig. 5
soundness: 2
strengths: - studying gaussian processes to asses spatio-temporal dependencies is inherently reasonable and could be used as an important theoretical framework to study neural architectures.

- The method mathematically is well grounded and derived

- High-quality figures visualizing key components of the contribution
summary: This paper provides a theoretical study on the limits of the learning capabilities of transformer-based diffusion models on learning spatial-temporal data. The authors achieve this by reformulating the task of learning the score function of Gaussian process data as a gradient descent algorithm. This algorithm can then be unrolled using a transformer archtitecture where the attention layers can be studied to evaluate how spatial-temporal dependencies are learned. 
The results are thoroughly motivated and dervied and validated using a few numerical examples. 
 
I believe that the theoretical results could have an impact in terms of understanding the learning of temporal dependencies but currently the practical results are not extensive and the take away messages are not clear enough
weaknesses: - While the results convincingly show that transformers can learn spatial-temporal dependencies, the experiments only show examples of properly learned tasks. Given that theoretical boundaries were derived, I would have loved more experiments investigating these limits. For example something like Figure 1 but with a degradation for higher length. The same goes for Section 6.2. What I am missing is a clear connection between model complexity and performance.

- Lack of comparison to other methods: I think it would be helpful to see how typical diffusion models, including Unet-based architectures, perform on the task of learning Gaussian processes. Clear results could help in shifting the design towards transformer-based diffusion architectures.

- Generalization to real-world problems: There are no experiments showing how well the insights generated by the experiments translate to real-world problems. You present the ball experiment but even there you have to consider how well your experiments align with the assumption of using Gaussian process data (cmp. l. 504-506). Real-world problems will often not follow Gaussian process dynamics

- It would be nice if you could state some take-home messages for the reading. Ideally some general information such as design choices for the model architecture.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. The theorems proposed in the paper are based on the authors’ assumptions, then:
1.a) Does it mean that the theory is only valid under these assumptions, or does the model only work on the data satisfying these assumptions? 
1.b) How to ensure that the assumptions can represent the case in real sequential data? 
1.c) How do authors derive these assumptions? E.g., in Assumption 1, why the definition of the covariance function has that formula, are there any theoretical basis for it?
soundness: 3
strengths: 1. This paper studies the principle behind the Diffusion Transformer and provides a theoretical explanation of its capability of modeling spatial-temporal correlations within data, which is the very first trial in related fields. 
2. Authors propose a novel score function approximation theory as well as the sample complexity bound. All the lemmas and theorems are well defined with sufficient mathematical proofs. 
3. The correctness of the proposed theory has been evaluated on both synthesized and real data through experiments.
summary: In this paper, authors explore the reasons for Diffusion Transformer (DiT)’s capability of capturing spatial-temporal dependencies within the sequential data. They theoretically establish score approximation and distribution estimation guarantee of DiT for learning Gaussian process data. Specifically, they replace the score function in diffusion models with gradient descent process defined by truncated covariance, which are further implemented with transformer models. Authors conduct some experiments to prove the effectiveness of their theory under their assumptions.
weaknesses: 1. Although the overall writing is relatively clear, there’re some suggestions:
1.a) The object of this paper is sequential data, which has the axis of time (real world time, $h$). while in the diffusion process, there’s another axis of diffusion step ($t$), and authors also name it as ‘time’. To avoid confusion, authors should use some other descriptions to make it clearer. 
1.b) The symbols should be consistent. E.g., the definitions of $v_t$, $\mathcal{D}$ and $v_0^t$, the use of subscripts of data $\mathbf{x}$ are confusing.
1.c) The font size of Figures could be larger, e.g. Fig.2, otherwise the equations and words are difficult to recognize under normal physical paper size.   
2. Although the effectiveness of the theory has been proved on synthesized Gaussian process data and semi-synthetic video data. More real data with complex contents should also be tested on.

Meta-review: The paper develops a theoretical framework to explain how diffusion transformers (DiTs) capture spatial-temporal dependencies in sequential data, using Gaussian processes as a simplified model. It introduces a novel score function approximation and establishes sample complexity bounds, supported by mathematical proofs and experiments on synthetic Gaussian process data and semi-synthetic tasks. Strengths include a strong theoretical contribution to understanding transformers, well-founded mathematical derivations, and high-quality visualizations. Weaknesses include the reliance on Gaussian process assumptions, limited practical experiments on real-world data, and insufficient comparisons to alternative architectures like UNets. Missing elements include more diverse datasets and guidelines for architectural design beyond theoretical insights. The paper’s significant theoretical contributions and clear presentation justify acceptance, with suggestions to expand empirical validation and real-world applicability in future work.

Decision: Accept (Poster)