Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time

## Abstract

The computational complexity of the self-attention mechanism in popular transformer architectures poses significant challenges for training and inference, and becomes the bottleneck for long inputs. Is it possible to significantly reduce the quadratic time complexity of computing the gradients in multi-layer transformer models? This paper proves that a novel fast approximation method can calculate the gradients in almost linear time n 1 + o ⁢ ( 1 ) superscript 𝑛 1 𝑜 1 n^{1+o(1)} italic_n start_POSTSUPERSCRIPT 1 + italic_o ( 1 ) end_POSTSUPERSCRIPT where n 𝑛 n italic_n is the input sequence length, while it maintains a polynomially small approximation error 1 / poly ⁡ ( n ) 1 poly 𝑛 1/\operatorname{poly}(n) 1 / roman_poly ( italic_n ) across the entire model.
Our theory holds for general loss functions and when the multi-layer transformer model contains many practical sub-modules, such as residual connection, casual mask, and multi-head attention.
By improving the efficiency of gradient computation, we hope that this work will facilitate more effective training and deployment of long-context language models based on our theoretical results.

## 1 Introduction

Large Language Models (LLMs), such as ChatGPT [], GPT-4 [], Claude 3.5 [], Llama 3.1 [], and others, have demonstrated immense potential to enhance various aspects of our daily lives, e.g., conversation AI [], AI agent [, ], search AI [], AI assistant [, ] and many so on.
One of the most emergent abilities of LLMs is dealing with long-context information, a format that is crucial for recording material like academic papers, official reports, legal documents, and so on.
LLMs have proven adept at tackling long-context tasks, including Retrieval Augmented Generation (RAG) [, ], zero-shot summarization [, ], and maintaining very long-term conversations [, ], and so on.
This proficiency has necessitated the development of long-context modeling capabilities within LLMs.

The self-attention mechanism is crucial for the success of LLMs, since LLMs are mainly based on Transformer architecture whose key module is attention.
In attention computation, we will compute the attention score between each pair of tokens, which is the complexity bottleneck during long context training and inference.
In detail, we need to spend $O(n^{2}d)$ running time for each self-attention block, which is quadratic in $n$, where $n$ is the length of the context input and $d$ is the hidden feature dimension of the model.
For example, LLaMA 3.1 405B [], one of the cutting-edge LLMs, supports $n=$128k and $d=4096$, while taking $30.84$M GPU training hours, which underscores the need for more efficient training processes for such extensive context models. Given the extensive context lengths of LLMs, this quadratic time complexity results in critical challenges: ($i$) a marked decrease in training efficiency [, ]; and ($ii$) significant energy usage, which in turn contributes to higher carbon dioxide emissions [, ].

One seminal work [] showed that the self-attention inference can be approximated in almost linear time. However, this result is for the *inference* time (forward pass), but does not address the main challenge, which is the expensive computation in the *training* time (backward pass). In this work, we address this main challenge, by proving that the gradient computation in the back-propagation of self-attention can be approximated in almost linear time. This suggests we may be able to save the substantial resources required for training LLMs.

### 1.1 Key background

We first introduce some basic background, starting with defining the softmax function and the self-attention module.

###### Definition 1.1 (Softmax) .

Let $z\in\mathbb{R}^{n}$. We define
$\mathsf{Softmax}:\mathbb{R}^{n}\to\mathbb{R}^{n}$ satisfying

$$ $\displaystyle\mathsf{Softmax}(z):=\exp(z)/\langle\exp(z),{\bf 1}_{n}\rangle.$ $$

Here we apply $\exp$ to a vector entry-wise.

###### Definition 1.2 (Self-attention module) .

Let $X\in\mathbb{R}^{n\times d}$ denote the input sequence, where $n$ is the number of input tokens and $d$ is the hidden dimension size.
Let $W_{Q},W_{K},W_{V}\in\mathbb{R}^{d\times d}$ be the query, key and value weight matrix.
The self-attention function $\mathsf{Attn}(X)$ with weights is:

$$ $\displaystyle\mathsf{Attn}(X)=\mathsf{Softmax}(XW_{Q}W_{K}^{\top}X^{\top}/d) \cdot XW_{V}.$ $$

where $\mathsf{Softmax}$ is applied to each row of its input matrix.
The attention can be re-written as:

$$ $\displaystyle\mathsf{Attn}(X)=f(X)\cdot XW_{V},$ $$

where (1) $A:=\exp(XW_{Q}W_{K}^{\top}X^{\top}/d)\in\mathbb{R}^{n\times n}$ and $\exp$ is applied element-wise, (2) $D:=\operatorname{diag}(A{\bf 1}_{n})\in\mathbb{R}^{n\times n}$, and (3) $f(X):=D^{-1}A\in\mathbb{R}^{n\times n}$ is the attention matrix.

In contemporary LLMs, the architecture typically incorporates multiple layers of attention. Consequently, in order to design a fast training algorithm for the entire model, it is imperative to examine self-attention within the multi-layer transformer structure formally defined as follows.

###### Definition 1.3 (Multi-layer transformer) .

Let $m$ denote the number of transformer layers in the model. Let $X$ be the input sequence.
Let $g_{i}$ denote components other than self-attention in the $i$-th transformer layer, and assume its forward and backward computations can be run in time linear in its input sequence length.
Let $\mathsf{Attn}_{i}$ denote the self-attention module in the $i$-th transformer layer with weights $W_{Q_{i}},W_{K_{i}},W_{V_{i}}$ (see also Definition 1.2).
We define an $m$-layer transformer as

$$ $\displaystyle\mathsf{F}_{m}(X):=g_{m}\circ\mathsf{Attn}_{m}\circ g_{m-1}\circ \mathsf{Attn}_{m-1}\circ\dots\circ g_{1}\circ\mathsf{Attn}_{1}\circ g_{0}(X),$ $$

where $\circ$ denotes function composition.

In Definition 1.3, the $g_{i}$ includes the layer norm, MLP, residual connection, dropout, positional encoding, multi-head concatenation, and other operations. All forward and backward computations of these practical modules can be run in linear time with respect to $n$.
Thus, in this work, we mainly focus on the acceleration of self-attention module.
Specifically, as shown in Definition 1.2, the $n\times n$ attention matrix $f(X)$ dominates the computational complexity, introducing a quadratic bottleneck. In the exact computation case, if the attention matrix is full rank, no acceleration is possible. However, by compromising negligible accuracy, designing a fast sub-quadratic algorithm becomes feasible. Fortunately, by employing the polynomial kernel approximation method from AA [], we can approximate the attention matrix and achieve an almost linear time $n^{1+o(1)}$ algorithm, effectively breaking the quadratic bottleneck.

###### Definition 1.1 (Softmax) .

###### Definition 1.2 (Self-attention module) .

###### Definition 1.3 (Multi-layer transformer) .

### 1.2 Our contributions

We now state our main result as follows:

###### Theorem 1.4 (Main result, informal version of Theorem 4.2 ) .

Let $n$ be the number of tokens, and $d$ the hidden dimension size.
We assume $d=O(\log n)$ and each number in matrices can be written using $O(\log n)$ bits. Assume the number of layers $m=n^{o(1)}$.
There exists an algorithm (Algorithm 1) that can compute the gradient of multi-layer self-attention (see also Definition 1.3) in almost linear time $n^{1+o(1)}$, where the approximation error of the entire model can be bounded by $1/\operatorname{poly}(n)$.

Our assumption is mild when the context length $n$ is large, as the feature dimension $d$ is usually regarded as a constant, which is also used in AA []; similarly, the number of layers is usually much smaller than $n$ and regarded as a constant.
Our results indicate that large language models (LLMs) can be trained in almost linear time $n^{1+o(1)}$ and maintain a robust approximation guarantee, while the traditional way takes $\Omega(n^{2})$ time.
This advancement is realized through the application of polynomial kernel approximation [, ].
To be more specific, by leveraging the inherent sparsity within the dense attention matrix, we perform efficient low-rank approximation, thereby significantly accelerating the computation of the dense matrices.
Our framework is applicable to *general* loss functions, making it universally applicable.
Furthermore, our analysis holds when the multi-layer transformer model contains many practical sub-modules, such as residual connection, casual mask, and multi-head attention (Section 6).

Numerous studies, including FlashAttention [, , ], quantization techniques [, ], and sparsity approaches [, ], have empirically focused on accelerating attention mechanisms. However, theoretically, these methods are still constrained by quadratic time complexity. In this study, we introduce an innovative acceleration technique (Algorithm 1) that effectively overcomes this quadratic bottleneck, backed by solid theoretical foundations (Theorem 4.2).
Moreover, this new method is designed to be seamlessly integrated with existing approaches to further enhance their performance (see Section 6).

Our contributions are as follows:

- •
We introduce a fast computation method that allows the gradient of each self-attention layer to be approximated in almost linear time $n^{1+o(1)}$ with $1/\operatorname{poly}(n)$ error, where $n$ is the input sequence length, breaking the quadratic time complexity bottleneck (Theorem 4.1).
- •
We extend our single-layer results to module-wise gradient computation so that our Algorithm 1 approximates gradient computation in $m\cdot n^{1+o(1)}$ time for $m$-layer transformer. Importantly, the approximation of the gradient diverges from the exact gradient by an error of $1/\operatorname{poly}(n)$ across the entire model (Theorem 4.2).
- •
Additionally, our analysis holds for general loss functions and when the multi-layer transformer model contains residual connection, casual mask, and multi-head attention.
Our results can be applied to any gradient-based algorithm, e.g., training, full fine-tuning, prompt-tuning, and so on (Section 6).

###### Theorem 1.4 (Main result, informal version of Theorem 4.2 ) .

## 2 Related Work

#### Long-context modeling in LLMs.

As LLMs grow in size and capability, in-context learning (ICL) [, , , ] has become a preferred method for directing these models to perform a variety of tasks, as opposed to the resource-intensive process of fine-tuning. Nonetheless, research has indicated that longer prompts can impair LLMs performance due to the limitation on maximum sequence length during pre-training []. Consequently, extending the maximum sequence length during pre-training and fine-tuning stages is imperative. Enhancing training efficiency is crucial given the prevalent use of the Transformer architecture in LLMs, which incurs a quadratic computational cost relative to sequence length.
Addressing this challenge, some studies have explored continued fine-tuning of LLMs with extended context lengths [], while others have experimented with the interpolation and extrapolation capabilities of positional embedding []. [] handles long context by compressing the input tokens. However, these approaches have not fundamentally addressed the core issue: the quadratic computational cost associated with sequence length in the attention mechanism [, ]. In this study, we delve into accelerating the attention mechanism, thereby addressing the long-context modeling issue at its essence.

#### Attention acceleration.

Attention mechanism has faced criticism due to its quadratic time complexity with respect to context length, a concern exacerbated by the increasing length in modern large language models (LLMs) such as GPT-4 [], Claude 3.5 [], Llama 3.1 [, ], etc.
Nevertheless, this limitation can be circumvented by employing polynomial kernel approximation techniques [], which enable the derivation of a low-rank representation of the attention matrix. This innovation significantly accelerates both the training and inference processes of a single attention layer, achieving almost linear time complexity [, ], while our work supports both training and inference for any multi-layer transformer.
Furthermore, this approach can be extended to higher-order attention mechanisms, i.e., tensor attention, maintaining almost linear time complexity during both training and inference [, ].
Moreover, there are other theoretical approaches. For instance, LLS^+24d [] introduces the conv-basis method to accelerate attention computation. HJK^+ [] proposes a near-linear time algorithm under the assumptions of uniform softmax column norms and sparsity.

Roadmap. Our paper is organized as follows.
Section 3 provides essential conceptions
and key definitions
across the whole paper.
Section 4 presents our primary findings, where we articulate our novel algorithm that is capable of calculating gradients across the entire model in almost linear time.
In Section 5, we explain the techniques we employ, including low-rank approximation,
techniques for accelerating the computation of gradients,
and an analysis of the approximation error.
Section 6 provides various extensions of our algorithm.
Lastly, we conclude this paper in Section 7.

## 3 Preliminary

#### Notations.

For any positive integer $n$, we use $[n]$ to denote set $\{1,2,\cdots,n\}$.
For two vectors $x\in\mathbb{R}^{n}$ and $y\in\mathbb{R}^{n}$, we use $\langle x,y\rangle$ to denote the inner product between $x,y$. Namely, $\langle x,y\rangle=\sum_{i=1}^{n}x_{i}y_{i}$.
We use $e_{i}$ to denote a vector where only $i$-th coordinate is $1$, and other entries are 0 0.
For each $a,b\in\mathbb{R}^{n}$, we use $a\odot b\in\mathbb{R}^{n}$ to denote the Hardamard product, i.e. the $i$-th entry of $(a\odot b)$ is $a_{i}b_{i}$ for all $i\in[n]$.
We use ${\bf 1}_{n}$ to denote a length-$n$ vector where all the entries are ones.
We use $\|A\|_{\infty}$ to denote the $\ell_{\infty}$ norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e., $\|A\|_{\infty}:=\max_{i\in[n],j\in[d]}|A_{i,j}|$.
We use $\operatorname{poly}(n)$ to denote some polynomial in $n$.

### 3.1 Loss function

The loss function is the optimization objective in the training of LLMs, and we define it as follows.

###### Definition 3.1 (Loss function $L(X)$ ) .

For some input matrix $X\in\mathbb{R}^{n\times d}$, we define the one-unit loss function $\ell(X)_{j,k}:\mathbb{R}^{n\times d}\rightarrow\mathbb{R}$, for any $j\in[n],k\in[d]$, and assume differentiability.
Furthermore, we define the overall loss function $L(X)$, such that

$$ $\displaystyle L(X)=\sum_{j=1}^{n}\sum_{k=1}^{d}\ell(X)_{j,k}$ $$

###### Remark 3.2 .

Typically, the most widely used loss function in the LLM training procedure is the cross-entropy loss function, which can also be viewed as a summation of one unit loss function as in Definition 3.1.
The output matrix of the multi-layer transformer needs to pass an additional linear layer to map the hidden dimension $d$ to the vocabulary size $d_{\mathrm{voc}}$. Assuming $d_{\mathrm{voc}}$ is a constant, the weight matrix dimensions for this additional MLP layer are $d\times d_{\mathrm{voc}}$.
The probability tensor $Y_{\mathrm{pred}}\in\mathbb{R}^{n\times d_{\mathrm{voc}}}$ is the final output.
We denote the ground truth as $Y_{\mathrm{gt}}\in\mathbb{R}^{n\times d_{\mathrm{voc}}}$ corresponding to $Y_{\mathrm{pred}}$.
According to the cross-entropy loss definition, the formula is expressed as

$$ $\displaystyle L_{\mathrm{cross-entropy}}(X)=-\sum_{j=1}^{n}\sum_{k=1}^{d_{ \mathrm{voc}}}(Y_{\mathrm{gt}})_{j,k}\log((Y_{\mathrm{pred}})_{j,k})$ $$

where the summation iterates over all elements, and the ground truth $(Y_{\mathrm{gt}})_{j,k}=1$ for the correct class and 0 0 otherwise.

###### Definition 3.1 (Loss function L ⁢ ( X ) 𝐿 𝑋 L(X) italic_L ( italic_X ) ) .

###### Remark 3.2 .

### 3.2 Closed forms of gradient components

In training large language models (LLMs), updating the model necessitates computing the gradient of weights for every layer.
Consequently, it becomes essential to derive the closed-form expressions for all corresponding gradient components with respect to the weights of the query, key, and value matrices in the transformer model.
We first define some intermediate variables before detailing these gradient components in each self-attention transformer layer.

###### Definition 3.3 (Intermediate variables $T_{i}$ ) .

Let $m$ denote the number of transformer layers in the model.
Let $m$-layer self-attention transformer be defined as Definition 1.3.
Let $d$ denote the hidden dimension.
Let $n$ denote the sequence length.
Let $X\in\mathbb{R}^{n\times d}$ be the input sentence.
Let $g_{i}$ denote components other than self-attention in the $i$-th transformer layer.
Let $\mathsf{Attn}_{i}$ denote the self-attention module in the $i$-th transformer layer (see also Definition 1.2).

For $i\in\{0,1,2,\cdots,m\}$, we define $T_{i}(X)\in\mathbb{R}^{n\times d}$ be the intermediate variable (hidden states) output by $i$-th layer self-attention transformer. Namely, we have

$$ $\displaystyle T_{i}(X)=\begin{cases}g_{0}(X),&~{}i=0;\\ (g_{i}\circ\mathsf{Attn}_{i})(T_{i-1}(X)),&~{}i\in[m].\end{cases}$ $$

Here, we use $\circ$ to denote function composition.

Then, we are ready to introduce the closed forms of the three gradient components in a single self-attention transformer layer.
Notably, according to the chain rule, the gradient of the $k$-th transformer layer in LLMs depends on the gradient components from the $(k+1)$-th transformer layer.
The gradient can be calculated for every transformer layer by combining the upstream and local gradients.
The closed forms of the gradients for each layer in multi-layer transformers are formalized in the following lemma (Lemma 3.4).

###### Lemma 3.4 (Closed form of gradient components, informal version of Lemma C.4 ) .

Let $L(X)$ be defined as in Definition 3.1, and the $m$-layer transformer defined as in Definition 1.3.
Let $W_{Q_{i}},W_{K_{i}},W_{V_{i}}\in\mathbb{R}^{d\times d}$ denote the attention weight in the $i$-th attention.
Let $T_{i}(X)$ denote the intermediate variable output by $i$-th self-attention transformer layer (see Definition 3.3).
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
For $j\in[n],k\in[d]$, let $G_{i}(j,k)$ denote the $(j,k)$-th entry of $G_{i}$,
let $\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))_{j,k}}{\mathrm{d}T_{i-1}(X)}\in
\mathbb{R}^{n\times d}$ denote the gradient of $(j,k)$-th entry of $\mathsf{Attn}_{i}(T_{i-1}(X))$.
Then, we can show that

- •
Part 1.
$\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=\sum_{j=1}^{n}\sum_{k
=1}^{d}G_{i}(j,k)\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))_{j,k}}{
\mathrm{d}T_{i-1}(X)}.$
- •
Part 2. Let $W_{*_{i}}$ be $W_{Q_{i}},W_{K_{i}}$ or $W_{V_{i}}$, then
$\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{*_{i}}}=\sum_{j=1}^{n}\sum_{k=
1}^{d}G_{i}(j,k)\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))_{j,k}}{
\mathrm{d}W_{*_{i}}}.$

Our main results are based on the above closed forms of four gradient components.

###### Definition 3.3 (Intermediate variables T i subscript 𝑇 𝑖 T_{i} italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) .

###### Lemma 3.4 (Closed form of gradient components, informal version of Lemma C.4 ) .

## 4 Main Results

Now, we present our main findings.
We will work through this section in the following order:
In Section 4.1, we delineate the computational efficiency of our gradient calculation methods in each single layer.
In Section 4.2, we introduce our main theorem (Theorem 4.2) for multi-layer transformer by integrating the preceding results and provide our main algorithm (Algorithm 1).
Section 4.3 discusses how we transcend the previous works.

### 4.1 Fast computing for single layer

In the case of single-layer attention, we provide our theorem that state the three gradient components can be calculated in almost linear time with negligible error.

###### Theorem 4.1 (Single-layer gradient approximation) .

We assume $d=O(\log n)$ and each number in matrices can be written using $O(\log n)$ bits. Assume the number of layers $m=n^{o(1)}$.
Let $L(X)$ be defined as Definition 3.1.
Suppose we have a single-layer self-attention transformer model ($m=1$ in Definition 1.3).
We can approximate one-layer self-attention for three gradient components, i.e. $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$, $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{Q}W_{K}^{\top}}$ and $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V}}$, in $n^{1+o(1)}$ time with $1/\operatorname{poly}(n)$ error.

###### Proof.

We finish the proof by Lemma 5.1, 5.2 and 5.3.
∎

Figure: Algorithm 1 Almost Linear Time (ALT) Multi-layer Transformer Gradient Approximation

###### Theorem 4.1 (Single-layer gradient approximation) .

###### Proof.

### 4.2 Fast computing for multi-layer transformers

Based on the results demonstrated in previous sections, we are ready to introduce our main result: the gradients of the whole transformer model can be approximated in almost linear time.

###### Theorem 4.2 (Main result, formal version of Theorem 1.4 ) .

Let $m$ denote the number of transformer layers.
We assume $d=O(\log n)$ and each number in matrices can be written using $O(\log n)$ bits. Assume the number of layers $m=n^{o(1)}$.
We can show that, for any $i\in[m]$, all the gradient components (see also Lemma 3.4) of the $i$-th layer can be computed by Algorithm 1 in almost linear time $n^{1+o(1)}$, and the approximation error of the entire $m$ layer transformer model can be bounded by $1/\operatorname{poly}(n)$.

###### Proof.

We prove the theorem by directly combining Theorem 4.1 and Lemma 5.5.
∎

Theorem 4.2 demonstrates that, during the training of a multi-layer transformer model, at each training iteration, the gradient computation for the weight matrices of each layer can be performed in almost linear time $n^{1+o(1)}$. This result supports the feasibility of fast training for any transformer-based large language models (LLMs).
In Algorithm 1, we illustrate the process of back-propagating gradients from the $m$-th transformer layer back to the first layer.
This algorithm highlights the significance of the gradient with respect to the intermediate variables $T_{i}(X)$.
Due to the application of the chain rule in gradient computation, the gradient of $T_{i}(X)$ is indispensable for determining the gradients of the weight matrices $W_{Q_{i}},W_{K_{i}}$ and $W_{V_{i}}$ at the $i$-th layer.
Consequently, by iteratively computing the gradient for $T_{i}(X)$, we systematically propagate the gradient through to the initial transformer layer.
Additionally, our algorithm is capable of computing the gradient with respect to the input data $X$. Therefore, our algorithm also supports fast prompt tuning. For a more in-depth discussion on this topic, please refer to Section 6.

###### Theorem 4.2 (Main result, formal version of Theorem 1.4 ) .

###### Proof.

### 4.3 Beyond the previous work

Our algorithm exhibits significant advancements over two seminal prior studies, namely AS [] and AS24a []. In AS [], the authors proposed an almost linear time algorithm for computing the forward process of the attention mechanism. In contrast, AS24a [] introduced an almost linear time algorithm for the backward of attention mechanism.
However, AS24a [] has the following limitations:
(1) only computing gradients for a single layer of the attention mechanism, which cannot extend to multiple layers;
(2) calculating gradients with respect to a specific loss, namely the $\ell_{2}$ loss;
(3) computing gradients only for the weight matrix $W_{Q_{i}},W_{K_{i}}$ (as defined in Definition 1.2), but ignore other crucial components such as the MLP layer following attention computation and the activation function. These limitations are inherent in their technique and prevents the applicability of the method in multiple layer transformers.

In our work, we have the following improvements beyond previous work:
(1) we enable almost linear time gradient computation across an entire transformer layer, incorporating both the MLP layer and the activation function;
(2) our algorithm supports gradient calculation for general loss function $L(X)$ (see Definition 3.1);
(3) we extend the gradient calculation to include not only $W_{Q_{i}},W_{K_{i}}$ but also $T_{i}(X)$ and $W_{V_{i}}$.
These advancements collectively demonstrate a substantial leap forward from the methodologies in AS [] and AS24a [].

## 5 Technical Overview

### 5.1 Low-rank approximation for attention matrix

In this section, we delve into the crucial techniques behind our work: the low-rank approximation of the attention matrix, which is achieved through the polynomial method [, ].
Drawing inspiration from AS [], the intuition of this approximation lies in the fact that the attention matrix $f(X)\in\mathbb{R}^{n\times n}$ (as defined in Definition 1.2), also referred to as the similarity matrix in attention mechanism, can be effectively approximated by low-rank matrices $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$, where $k_{1}=n^{o(1)}$.
The naive method for calculating the attention matrix $f(X)$ has a time complexity of $O(n^{2})$, whereas the input data $X\in\mathbb{R}^{n\times d}$ contains only $d\cdot n=n^{1+o(1)}$ entries. This discrepancy suggests the potential of using low-rank representations of $f(X)$ to design a fast algorithm.

An example of how to use the low-rank representations is the attention forward.
First note that approximating $f(X)$ alone does not lead to a fast algorithm, since $U_{1}V_{1}^{\top}$ still requires $n\times n$ entries. But by using the structure of the attention $\mathsf{Attn}(X):=f(X)V$ where $V=XW_{V}$, we can do it faster.
By expressing $f(X)$ as $U_{1}V_{1}^{\top}$, the attention forward becomes
$\underbrace{U_{1}}_{n\times k_{1}}\underbrace{V_{1}^{\top}}_{k_{1}\times n}
\underbrace{V}_{n\times d}.$
It is well known that different multiplication sequences can lead to dramatically different numbers of operations required, so the order of matrix multiplications matters, which is indeed the case here.
We first perform $V_{1}^{\top}V\in\mathbb{R}^{k_{1}\times d}$ and this cost $O(k_{1}nd)=n^{1+o(1)}$ time.
Then we can compute $U_{1}V_{1}^{\top}V$ within $O(nk_{1}d)=n^{1+o(1)}$ time.

This method significantly reduces the computation time of the attention forward from $O(n^{2})$ to almost linear time, $n^{1+o(1)}$.
Driven by this technique and analyzing the close forms of the gradients, we extend the acceleration to the gradient of the entire model.

### 5.2 Accelerating gradient computation of T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X )

Based on the low-rank approximation method mentioned in Section 5.1, we compute the gradient of $L(X)$ with respect to the intermediate variable $T_{i}(X)$, which denotes the output of the $i$-th transformer layer.
This computation is critical as it enables us to calculate gradients for other gradient components because of the chain rule.

#### Extending to general loss functions.

According to the findings in DSXY [], the gradient $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$ can be decomposed into five components, namely $C_{2}(X),C_{4}(X),C_{6}(X),C_{7}(X),C_{8}(X)$, as detailed in Lemma D.1.
However, the gradient result presented in DSXY [] is tailored to a specific loss function, the $\ell_{2}$ loss, limiting its applicability to a narrow range of scenarios.
The primary challenge in extending the scope to encompass general loss functions is the absence of a unified analytical framework. Previous analyses [, ] are limited to individual, specific loss functions. In this work, we introduce a comprehensive analysis framework (Definition 3.1) and we have demonstrated its applicability to the cross-entropy loss (see Remark 3.2). Consequently, by utilizing this generalized analysis framework, we extend the notation $L(X)$ to include a wide range of general loss functions.

#### Accelerating the gradient computation.

An important step in accelerating the gradient computation for the entire multi-layer transformer model is to accelerate the computation of the gradient on the intermediate variables $T_{i}(X)$.
The key challenge is that, to compute the gradient on $T_{i}(X)$, we need to compute the gradient on other components of one transformer layer, such as residual connection, multi-head attention, and causal attention mask (see more details in Section 6). We conduct comprehensive analysis on those components in the transformer layer, and prove that, by using low-rank approximation technique, the computation of gradient $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$ can be computed in almost linear time $n^{1+o(1)}$ (Lemma 5.1).

A crucial aspect of speeding up gradient computation for the entire multi-layer transformer model involves accelerating the calculation of gradients with respect to the intermediate variables $T_{i}(X)$. The main challenge lies in the fact that computing the gradient of $T_{i}(X)$ requires calculating the gradients for other components within a transformer layer, including the residual connection, multi-head attention, and causal attention mask (see Section 6). We have conducted an extensive analysis of these components within the transformer layer (see Section I, J, and K) and demonstrated that, through the application of low-rank approximation techniques, the gradient $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$ can be computed in almost linear time $n^{1+o(1)}$ (Lemma 5.1).
In particular, we apply the low-rank approximation technique on the five terms $C_{2}(X),C_{4}(X),C_{6}(X),C_{7}(X),C_{8}(X)$ respectively, demonstrating
that each term can be computed in almost linear time, $n^{1+o(1)}$, as shown in Section E.
Then we aggregate those terms, as described in Section E.6.
Since all five terms are $n\times d$ matrices, the summation of these terms remains almost linear in complexity.
We then conclude that for any single-layer transformer, the gradient computation with respect to the input can be performed in almost linear time $n^{1+o(1)}$, as stated in Lemma 5.1.

The statement made for a single transformer layer can be readily generalized to any layer within an $m$-layer transformer model.
For instance, consider the intermediate variables $T_{i}(X)$ and $T_{i-1}(X)$ (as defined in Definition 3.3), where $T_{i}(X)=(g_{i}\circ\mathsf{Attn}_{i})(T_{i-1}(X))$.
Given the gradient $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$, as established in the previous paragraph, we compute the gradient with respect to $T_{i-1}(X)$, namely $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$, in almost linear time $n^{1+o(1)}$.
For a multi-layer transformer model, the above process can be conducted recursively. Thus, we can compute the gradient of the loss function $L(X)$ on *any* $T_{i}(X)$ in almost linear time $n^{1+o(1)}$.

###### Lemma 5.1 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$ , informal version of Lemma E.11 ) .

Let $L(X)$ be defined as Definition 3.1.
Let $m$ denote the number of self-attention transformer layers (see Definition 1.3).
Let $T_{i}(X)$ denote the intermediate variable output by $i$-th self-attention transformer layer (see Definition 3.3).
We show that $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.

###### Lemma 5.1 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ T i ⁢ ( X ) d 𝐿 𝑋 d subscript 𝑇 𝑖 𝑋 \frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X ) end_ARG , informal version of Lemma E.11 ) .

### 5.3 Accelerating gradient computation of W i subscript 𝑊 𝑖 W_{i} italic_W start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT and W V i subscript 𝑊 subscript 𝑉 𝑖 W_{V_{i}} italic_W start_POSTSUBSCRIPT italic_V start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUBSCRIPT

In Section 5.2, we detailed the fast computation of gradients for intermediate variables $T_{i}(X)$.
Let $W_{i}:=W_{Q_{i}}W_{K_{i}}^{\top}$, with $W_{Q_{i}}$ and $W_{K_{i}}$ representing the query and key weight matrices, respectively, the gradients of $W_{i}$ and $W_{V_{i}}$ represent *all* trainable weight matrices in a transformer layer.
Consequently, by determining the gradients for $W_{i}$ and $W_{V_{i}}$ across each layer, we achieve almost linear time gradient back-propagation throughout multi-layer transformer models.

#### Fast gradient computation.

The prior study in AS24a [] demonstrated that the gradient of $W_{i}$ can be computed in almost linear time.
We extend their findings by adapting their approach to accommodate general loss function $L(X)$ (as defined in Definition 3.1) and further generalize their results to include the gradient computation for both $W_{i}$ and $W_{V_{i}}$ in each transformer layer (Lemma 5.2 and 5.3).

###### Lemma 5.2 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ , informal version of Lemma F.5 ) .

Let $L(X)$ be defined as Definition 3.1, and $m$ be the number of self-attention transformer layers (Definition 1.3).
For any $i\in[m]$, let $W_{i}=W_{Q_{i}}W_{K_{i}}^{\top},W_{V_{i}}\in\mathbb{R}^{d\times d}$ denote the attention weight in the $i$-th transformer layer.
We show that $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.

###### Lemma 5.3 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}$ , informal version of Lemma G.4 ) .

Let $L(X)$ be defined as Definition 3.1, and $m$ be the number of self-attention transformer layers (Definition 1.3).
For any $i\in[m]$, let $W_{i}=W_{Q_{i}}W_{K_{i}}^{\top},W_{V_{i}}\in\mathbb{R}^{d\times d}$ denote the attention weight in the $i$-th transformer layer.
We show that $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.

###### Lemma 5.2 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ W i d 𝐿 𝑋 d subscript 𝑊 𝑖 \frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d italic_W start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG , informal version of Lemma F.5 ) .

###### Lemma 5.3 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ W V i d 𝐿 𝑋 d subscript 𝑊 subscript 𝑉 𝑖 \frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d italic_W start_POSTSUBSCRIPT italic_V start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUBSCRIPT end_ARG , informal version of Lemma G.4 ) .

### 5.4 Accelerating gradient computation for multi-Layer transformers

In this section, our focus turns to extending the single-layer transformer result from the previous section to a multi-layer transformer.

#### Running time analysis.

We derive the closed-form gradient for the non-attention components within a transformer layer, namely the $g_{i}$ function defined in Definition 1.3.
With the closed-form gradient of $g_{i}$ established in Lemma H.1, we then demonstrate in Lemma 5.4 that the gradient computation for $g_{i}$ can also be achieved in almost linear time.
Given that the number of layers $m=n^{o(1)}$ is much smaller than $n$
and the computation time for gradients on each layer is $n^{1+o(1)}$, we only need to iteratively repeat this procedure for $m$ times.
Therefore, the overall running time for computing gradients across the entire model is $m\cdot n^{1+o(1)}=n^{1+o(1)}$.

###### Lemma 5.4 (Computation time for $G_{i}$ , informal version of Lemma H.2 ) .

Let $T_{i}(X)$ be defined as Definition 3.3, i.e. $T_{i}(X)=(g_{i}\circ\mathsf{Attn}_{i})(T_{i-1}(X))$.
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
Assume we already have $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$.
Assuming for any $Z\in\mathbb{R}^{n\times d}$, we have $g_{i}(Z)\in\mathbb{R}^{n\times d}$, and $g_{i}(Z)=\phi(Z\cdot W_{g})$, where $W_{g}\in\mathbb{R}^{d\times d}$ and $\phi:\mathbb{R}\rightarrow\mathbb{R}$ denotes any element-wise activation function. Let $\phi^{\prime}$ denote the derivative of $\phi$.
Then, we show that $G_{i}$ can be computed in $n^{1+o(1)}$ time.

###### Lemma 5.4 (Computation time for G i subscript 𝐺 𝑖 G_{i} italic_G start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT , informal version of Lemma H.2 ) .

#### Error propagation analysis.

Here, we consider the approximation error. In our setting, the approximation error originates from the low-rank approximation of the attention matrix, as detailed in Lemma C.13.
As discussed in previous sections, the approximation error in each layer can be bounded by $1/\operatorname{poly}(n)$. Then, we only need to focus on how error propagates in different layers.

We first prove that our $1/\operatorname{poly}(n)$ approximation error statement holds for a single-layer transformer, as evidenced in Lemma H.3.
Subsequently, through mathematical induction and leveraging the results of error propagation over the gradient of $g_{i}$, we show that the approximation error can be bounded by $1/\operatorname{poly}(n)$ for any $m$-layer transformer (Lemma 5.5), where the number of layers $m$ is considered small.

###### Lemma 5.5 (Multi-layer transformer gradient approximation, informal version of Theorem H.4 ) .

Let $L(X)$ be defined as Definition 3.1.
Let $X$ be defined as Definition 1.2.
Suppose we have a $m$-layer transformer (see Definition 1.3). Then, for any $i\in[m]$, we can show that:
($i$) Running time:
Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$, $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$, and $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}$ in $n^{1+o(1)}$ time;
($ii$) Error bound:
The approximation of the entire transformer model can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}$ satisfies
$\|\widetilde{g}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}\leq 1/
\operatorname{poly}(n)$.

###### Lemma 5.5 (Multi-layer transformer gradient approximation, informal version of Theorem H.4 ) .

## 6 Extensions

#### Multi-head attention and residual connections.

Multi-head attention and residual connections are important components in attention mechanisms. While these components were not involved in our initial analysis for simplicity, incorporating them into our algorithm is straightforward, as detailed in Sections B.1 and B.2.
Our algorithm maintains the capability to compute gradients for multi-layer transformers with multi-head attention and residual connection in almost linear time, suggesting that it can be readily adapted to more practical transformer models.
The detailed analysis of incorporating residual connection with our framework can be found in Section J and Lemma J.3.
For the synergy with multi-head attention, we provide comprehensive analysis in Section K and Lemma K.2.

#### Causal attention mask.

The causal attention mask is critical to prevent transformers from “cheating” during training by ensuring future information is not used. The full-rank characteristic of the causal attention mask poses challenges for low-rank approximations. Nevertheless, we have identified a method to accelerate the computation of causal masked attention by exploiting its inherent properties,
showing almost linear time complexity. A comprehensive explanation is provided in Section B.3.
More detailed analysis
can be found in Section I and Lemma I.7 and I.8.

#### Prompt tuning.

Prompt tuning (or prefix learning) is a prevalent approach in parameter-efficient fine-tuning (PEFT), which requires the calculation of gradients on input data $X$.
Given our algorithm’s ability to compute gradients for intermediate variables $T_{i}$ in approximately linear time, we can similarly accelerate the gradient computation for input data $X$, thus enhancing the efficiency of the prompt tuning process. Additional details are provided in Section B.5.

#### Synergy with system-level attention acceleration.

Many contemporary works focus on system-level acceleration of attention mechanisms, often by leveraging caching and mitigating I/O bottlenecks.
Our algorithm has the potential to integrate with such advancements. By combining our theoretical improvements in computation time (from $O(n^{2})$ to $n^{1+o(1)}$) with system-level optimizations, the overall efficiency of attention mechanism computation may improve further.
We leave the implementation of our method on GPU as future work since there are several coding challenges. More details can be found in Section B.4.

## 7 Conclusion

The attention mechanism in transformer models has quadratic time complexity with respect to the input token length.
In this work, we proposed a novel Algorithm 1, which can approximately train a multi-layer transformer model in almost linear time, introducing only a small error.
Importantly, our algorithm is designed to be compatible with general loss functions, practical sub-modules (residual connection, casual mask, multi-head attention), and general gradient-based algorithms. It may be seamlessly integrated with other system-level acceleration techniques.
While we lack enterprise-scale computational resources for training large language models to provide empirical support, our theoretical findings suggest that we can accelerate the training of LLMs in practice.

## Acknowledgement

Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.

## Appendix A More Related Work

#### Attention mechanism.

Attention mechanisms, including self-attention and cross-attention, are pivotal techniques employed in state-of-the-art neural networks. Since it was introduced in VSP^+ [], it has gained widespread adoption across various domains. In particular, it is integral to decoder-only LLMs [] and the Vision Transformer (ViT) architecture []. The former has been instrumental in the remarkable success of LLMs, while the latter has significantly advanced the field of computer vision, encompassing applications such as image generation [, , ], detection [], segmentation [], and layout generation [, , ].
Moreover, attention mechanism can be integrated into multi-modal models [, , , ], math reasoning [], diffusion models [, , , , , ], differential privacy [, , , , , , , , ] and many other techniques [, , , , , , , ].

#### Attention theory.

BCB [] introduced attention mechanisms in NLP, enhancing encoder-decoder architecture with variable-length vectors to improve machine translation. Building on this, LPM [] developed local and global attention variants, further refining NLP tasks.
Recent Large Language Model research has focused extensively on attention computation [, , ]. Studies by ZHDK [], CLP^+ [], KKL [] use Locality Sensitive Hashing for attention approximation, with ZHDK [] offering efficient dot-product attention. BSZ [] and AS [] explore static and dynamic attention calculations, while LSZ [] investigates hyperbolic regression regularization. DMS [] proposes algorithms for reducing attention matrix dimensionality in LLMs.
Attention has also been examined from optimization and convergence perspectives [, , , ], investigating word co-occurrence learning [], regression problems with exponential activation functions [], attention mechanism evolution during training [], and the impact of heavy-tailed noise on stochastic gradient descent []. Theoretical explorations of attention variants include quantum attention [], tensor attention [, ], and differentially private attention [, , ].

#### More methods for model acceleration.

Various techniques have been developed for model acceleration.
One approach involves modifying model architectures to enable faster inference, such as Mamba [], Linearizing Transformers [], PolySketchFormer [], and the Hopfield Model [, , , , , , , ] and so on.
Another line of work is to prune the weights in a neural network to reduce running time and memory consumption [, , , , , , ].
In addition, specific techniques have been developed to accelerate LLM generation [, , , ].

## Appendix B Discussion and Extension Details

In Section B.1, we argue that our framework can easily adapt to the multi-head attention mechanism.
In Section B.2, we introduce how to integrate residual connection to our framework.
In Section B.3, we detail the integration of the causal attention mask into our algorithm.
In Section B.4, we discuss the possibility of the synergy between our theoretical side attention acceleration and the existing system-level attention acceleration mechanism.
In Section B.5, we show how to expedite prompt tuning using our results.

### B.1 Multi-head attention

The multi-head attention mechanism was first introduced by VSP^+ [].
This innovation allows a token to simultaneously attend to multiple positions within the same layer, thereby enriching the model’s capacity for capturing various dependencies.
However, this enhanced capability comes with an increase in the size of the attention matrix $f(X)$ from $1\times n\times n$ to $h\times n\times n$, where $h$ is the number of attention heads.
To mitigate the computational burden, each head’s vector is derived by splitting the original vector, reducing the dimensionality of each head to $d_{h}:=d/h$.
To summarize, the key distinctions between multi-head and single-head attention are (1) an enlarged attention matrix $f(X)$ and (2) a reduced dimensionality $d_{h}$ within each attention head.

#### Enlarged attention matrix.

As previously discussed, the attention matrix’s dimensionality increases with the number of heads, $h$.
Despite this expansion, the application of the low-rank approximation technique, as outlined in Section 5.1, ensures that the computation time for the attention matrix remains almost linear.
Specifically, for a constant number of heads $h$ in the multi-head mechanism, the time complexity for computing $f(X)\in\mathbb{R}^{h\times n\times n}$ is $h\cdot n^{1+o(1)}=n^{1+o(1)}$.

#### Reduced dimensionality.

Another differentiating factor of multi-head attention is the lower dimensionality processed by each head, i.e. $d_{h}:=d/h$, compared the full $d$ in single-head attention.
This reduction ensures that the gradient computation time does not increase with the introduction of multiple attention heads.

We provide comprehensive analysis of the synergy of our algorithm with multi-head attention in Section K. We first prove in Lemma K.2, with the addition of multi-head attention, the gradient over the attention mechanism can be computed in almost linear time.
Then, we further prove that for any multi-layer transformer, with multi-head attention, the gradient can be computed in almost linear time as well.

### B.2 Residual connection

Residual connection is a pivotal technique in deep neural network architectures, effectively addressing issues such as vanishing and exploding gradients during training process, and facilitating faster convergence of the model.
Residual connection is also integrated into the standard attention mechanism.
Formally, given the intermediate variable $T_{i}(X)$ output by the $i$-th transformer layer as defined in Definition 3.3, we provide the formal definition of residual connection in Definition J.1 and J.2.
Since the residual connection only brings an additional add operation to each component and
with $T_{i}(X)$ belonging to the space $\mathbb{R}^{n\times d}$, the residual connection introduces only a marginal computational overhead of $O(n\cdot d)$ per layer. Consequently, the total computational cost for each layer is $O(n\cdot d)+n^{1+o(1)}=n^{1+o(1)}$.
Hence, by intuition, the inclusion of residual connections does not compromise the overall complexity of our method.

The detailed analysis is provided in Section J, where we first prove in Lemma J.3, that if the gradient over one structure can be computed in almost linear time, then with the addition of the residual connection, the gradient can also be computed in almost linear time.
Then we use math induction to extend our result to the entire multi-layer transformer model.

### B.3 Causal attention mask

In transformer training, attention mask is a crucial component, designed to prevent a given token from attending to future tokens in the sequence.
Causal attention mask is a widely used attention mask, which is configured as a lower triangular matrix, where elements on or below the main diagonal are ones, with all other entries being zeros.

Now we describe how to incorporate this into our algorithm.
Let $M\in\{0,1\}^{n\times n}$ represent the causal attention mask (see Definition I.2).
Let $\widehat{f}(X):=D^{-1}(M\odot A)$ where $A=\exp(XWX^{\top}/d)$ and $D:=\operatorname{diag}((M\odot A)\cdot{\bf 1}_{n})$.
Lemma I.1 reveals that $A$ has a low-rank representation given by $U_{0}V_{0}^{\top}$.
Using Lemma I.3, we know $(M\odot(U_{0}V_{0}^{\top}))\cdot v$ for any vector $v\in\mathbb{R}^{n}$ can be computed in almost linear time.

To integrate the causal mask into the gradient computation within each transformer layer, we first find all instances that have the structure of $f(X)\cdot H$ or $(f(X)\odot(UV^{\top}))\cdot H$, where $H,U,V$ are low rank matrices.
Then, we replace $f(X)$ with $\widehat{f}(X)$ in these instances.
More detailed analysis of causal attention can be found in Section I.
To be more specific, we group the gradient components for $T_{i},W_{i},W_{V_{i}}$ into two categories, one for dot product (Lemma I.7), another for Hadamard product (Lemma I.8).
After showing each component can be calculated in almost linear time, the overall gradient computation remains $n^{1+o(1)}$ time.
Thus, our framework can seamlessly accommodate causal attention masks.

### B.4 System-level attention acceleration

The attention computing acceleration involves a two-pronged strategy that leverages both system-level improvements (e.g. Flash Attention [, , ]) and the theoretical time complexity improvements (e.g. our work and HJK^+ []).

Numerous efforts have been made in the literature to accelerate attention calculations at the system level. For instance, Flash Attention [, , ] targets the I/O bottleneck inherent in attention mechanisms. Studies such as block-wise parallel decoding [] focus on implementing parallel decoding within transformer models to enhance inference speed. Additionally, recent advancements in the field of speculative decoding, such as Medusa [], leverage a smaller, more efficient model to generate predictions, with the larger model only responsible for validating, the smaller model’s outputs [].

Despite these innovations, the aforementioned methods do not address the fundamental quadratic time complexity $O(n^{2})$ of the attention mechanisms. This presents an opportunity to complement our low-rank approximation technique, with these system-level optimizations, thereby achieving an even greater acceleration in attention computation.
For instance, we could design an I/O-aware algorithm for Algorithm 1, similar to the approach taken by Flash Attention, to effectively leverage GPU acceleration.

To implement our algorithm practically on GPU, we have some coding challenges to fix: (1) we need to define some new tensor operations in PyTorch, e.g. Eq. (E.2), Eq. (E.4); (2) we need to systematically re-implement some back-propagation function of the current PyTorch function; (3) we need to implement some CUDA function to run our algorithm in parallel for the casual mask, see discussion in Section B.3.
We may leave this as our future work.

### B.5 Prompt tuning

Prompt tuning, as introduced by various studies [, , , , , ], has emerged as a parameter-efficient fine-tuning strategy for large language models (LLMs).
Specifically, prompt tuning involves adjusting “soft prompts” conditioned on frozen LLMs.
This method requires relatively small number of tuneable parameters compared with fine-tuning the entire LLMs, making it a popular choice for conserving training resources, including data and computational power.

The analysis reveals that the essence of prompt tuning involves computing gradients with respect to the soft prompts $X_{p}$ across the entire model. In both prompt tuning and full fine-tuning, the quadratic $O(n^{2})$ computational complexity of gradient calculation remains the same due to the self-attention mechanism inherent in LLMs.

In this work, leveraging the low-rank approximation technique discussed in Section 5.1, our algorithm (Algorithm 1) efficiently computes gradients on soft prompts $X_{p}$ over the entire model in almost linear time. This suggests that our method is universal and can also be applied within traditional prompt tuning frameworks.

## Appendix C Preliminary on Gradient Calculation

In Section C.1, we list several useful math facts used in the following sections of this paper.
In Section C.2, we provide the close forms of the gradient components.
In Section C.3, we introduce some mathematical definitions to facilitate understanding of gradient calculations.
In Section C.4, we list some low rank approximation technique introduced in AS [] and AS24a [].
In Section C.5, we demonstrate that the entries of matrices defined in Section C.3 are bounded.

#### Notations.

For two vectors $x\in\mathbb{R}^{n}$ and $y\in\mathbb{R}^{n}$, we use $\langle x,y\rangle$ to denote the inner product between $x,y$. Namely, $\langle x,y\rangle=\sum_{i=1}^{n}x_{i}y_{i}$.
We use $e_{i}$ to denote a vector where only $i$-th coordinate is $1$, and other entries are 0 0.
For each $a,b\in\mathbb{R}^{n}$, we use $a\odot b\in\mathbb{R}^{n}$ to denote the Hardamard product, i.e. the $i$-th entry of $(a\odot b)$ is $a_{i}b_{i}$ for all $i\in[n]$.
We use ${\bf 1}_{n}$ to denote a length-$n$ vector where all the entries are ones.
We use $\|A\|_{\infty}$ to denote the $\ell_{\infty}$ norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e. $\|A\|_{\infty}:=\max_{i\in[n],j\in[d]}|A_{i,j}|$.
We use $\operatorname{poly}(n)$ to denote polynomial time complexity with respective to $n$.

### C.1 Basic math facts

In this section, we provide some useful basic math facts,

###### Fact C.1 .

Let $x,y,z\in\mathbb{R}^{n}$. Then we have

- •
$\langle x\odot y,z\rangle=x^{\top}\operatorname{diag}(y)z$.
- •
$\langle x,(y\odot z)\rangle=\langle y,(x\odot z)\rangle=\langle z,(y\odot x)\rangle$
- •
$\langle x,y\rangle=\langle x\odot y,{\bf 1}_{n}\rangle$.

Then, we introduce a classical folklore used for the Hadamard product of two matrices.

###### Fact C.2 (Folklore, [ 6 ] ) .

Let $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$. Let $U_{2},V_{2}\in\mathbb{R}^{n\times k_{2}}$. Then we have

$$ $\displaystyle(\underbrace{U_{1}}_{n\times k_{1}}\underbrace{V_{1}^{\top}}_{k_{ 1}\times n})\odot(\underbrace{U_{2}}_{n\times k_{2}}\underbrace{V_{2}^{\top}}_ {k_{2}\times n})=\underbrace{(U_{1}\oslash U_{2})}_{n\times k_{1}k_{2}} \underbrace{(V_{1}\oslash V_{2})^{\top}}_{k_{1}k_{2}\times n}$ $$

Here, given $U_{1}\in\mathbb{R}^{n\times k_{1}}$ and $U_{2}\in\mathbb{R}^{n\times k_{2}}$, the $U_{1}\oslash U_{2}\in\mathbb{R}^{n\times k_{1}k_{2}}$ is the row-wise Kronecker product, i.e., $(U_{1}\oslash U_{2})_{i,l_{1}+(l_{2}-1)k_{1}}:=(U_{1})_{i,l_{1}}U_{i,l_{2}}$ for all $i\in[n]$, $l_{1}\in[k_{1}]$ and $l_{2}\in[k_{2}]$.

###### Fact C.1 .

###### Fact C.2 (Folklore, [ 6 ] ) .

### C.2 Close form of three gradient components

We first restate the definition of self-attention, where we denote $W:=W_{Q}W_{K}^{\top}\in\mathbb{R}^{d\times d}$ for simplicity.

###### Definition C.3 (Self-attention module) .

Let $X\in\mathbb{R}^{n\times d}$ denote the input sequence, where $n$ is the number of input tokens and $d$ is the hidden dimension size.
Let $W_{V}\in\mathbb{R}^{d\times d}$ be the value weight matrix, and let $W:=W_{Q}W_{K}^{\top}\in\mathbb{R}^{d\times d}$ be the key-query weight matrix.
The self-attention function $\mathsf{Attn}(X)$ with weights $W,W_{V}$ is:

$$ $\displaystyle\mathsf{Attn}(X)=\mathsf{Softmax}(XWX^{\top}/d)\cdot X\cdot W_{V}.$ $$

where $\mathsf{Softmax}$ is applied to each row of its input matrix.
The attention can be re-written as:

$$ $\displaystyle\mathsf{Attn}(X)=f(X)\cdot X\cdot W_{V},$ $$

where (1) $A:=\exp(XWX^{\top}/d)\in\mathbb{R}^{n\times n}$ and $\exp$ is applied element-wise, (2) $D:=\operatorname{diag}(A{\bf 1}_{n})\in\mathbb{R}^{n\times n}$, and (3) $f(X):=D^{-1}A\in\mathbb{R}^{n\times n}$ is the attention matrix.

Note that the gradient of $W_{Q}$ and $W_{K}$ can easily be calculated from the gradient of $W$, i.e.,

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{Q}}$ $\displaystyle~{}=\frac{\mathrm{d}L(X)}{\mathrm{d}W}\cdot\frac{\mathrm{d}W}{ \mathrm{d}W_{Q}}$ $\displaystyle~{}=\frac{\mathrm{d}L(X)}{\mathrm{d}W}\cdot W_{K}$ $$

where the first step follows from the chain rule, and the second step follows from basic calculus.

Then, we show how to derive the close form for the gradient components within each layer of a multi-layer transformer.

###### Lemma C.4 (Close form of gradient components, formal version of Lemma 3.4 ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $W_{i}:=W_{Q_{i}}W_{K_{i}}^{\top}\in\mathbb{R}^{d\times d}$ be the key-query weight matrix, $W_{V_{i}}\in\mathbb{R}^{d\times d}$ be the value weight matrix for the $i$-th transformer layer.
- •
Let $T_{i}(X)$ denote the intermediate variable output by $i$-th self-attention transformer layer (see Definition 3.3).
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$, let $\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))_{i_{2},j_{2}}}{\mathrm{d}T_{i-1}
(X)}\in\mathbb{R}^{n\times d}$ denote the gradient of $(i_{2},j_{2})$-th entry of $\mathsf{Attn}_{i}(T_{i-1}(X))$.

Then, we can show that

- •
Part 1.
$\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=\sum_{i_{2}=1}^{n}
\sum_{j_{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{
i-1}(X))_{i_{2},j_{2}}}{\mathrm{d}T_{i-1}(X)}.$
- •
Part 2.
$\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}=\sum_{i_{2}=1}^{n}\sum_{j_
{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))
_{i_{2},j_{2}}}{\mathrm{d}W_{i}}.$
- •
Part 3.
$\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}=\sum_{i_{2}=1}^{n}\sum
_{j_{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}
(X))_{i_{2},j_{2}}}{\mathrm{d}W_{V_{i}}}.$

###### Proof.

We have

- •
$L(X)\in\mathbb{R}$.
- •
$\mathsf{Attn}_{i}(T_{i-1}(X))\in\mathbb{R}^{n\times d},T_{i-1}(X)\in\mathbb{R}
^{n\times d}$.
- •
$W_{i}\in\mathbb{R}^{d\times d},W_{V_{i}}\in\mathbb{R}^{d\times d}$.

Therefore, we have

- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}\in\mathbb{R}^{n\times d},~{}~{}
\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}{\mathrm{d}T_{i-1}(X)}\in\mathbb
{R}^{(n\times d)\times(n\times d)}$.
- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}\in\mathbb{R}^{d\times d},~{}~{}\frac{
\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}{\mathrm{d}W_{i}}\in\mathbb{R}^{(n
\times d)\times(d\times d)}$.
- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}\in\mathbb{R}^{d\times d},~{}~{}
\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}{\mathrm{d}W_{V_{i}}}\in\mathbb{
R}^{(n\times d)\times(d\times d)}$.

Then, simply applying chain rule, we can get the final results.
∎

###### Definition C.3 (Self-attention module) .

###### Lemma C.4 (Close form of gradient components, formal version of Lemma 3.4 ) .

###### Proof.

### C.3 Basic notations for computing gradients

Before we move on to compute gradients, we need to define some useful notations.

We begin with introducing the index for a matrix.

###### Definition C.5 (Simplified notations) .

For any matrix $Z\in\mathbb{R}^{n\times d}$, for $i\in[n],j\in[d]$, we have following definitions:

- •
Let $\underbrace{Z_{i,j}}_{\mathrm{scalar}}$ and $Z(i,j)$ denote the $(i,j)$-th entry of $Z$.
- •
Let $\underbrace{Z_{i,*}}_{d\times 1}$ and $Z(i,*)$ denote the $i$-th row of $Z$.
- •
Let $\underbrace{Z_{*,j}}_{n\times 1}$ and $Z(*,j)$ denote the $j$-th column of $Z$.

Then, we define the exponential matrix in the attention mechanism.

###### Definition C.6 (Exponential function $u$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d}$
- •
Let $W:=W_{Q}W_{K}^{\top}\in\mathbb{R}^{d\times d}$

We define $u(X)\in\mathbb{R}^{n\times n}$ as follows

$$ $\displaystyle u(X):=\exp(XWX^{\top})$ $$

Then, we introduce the summation vector of the aforementioned exponential matrix.

###### Definition C.7 (Sum function of softmax $\alpha$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d}$
- •
Let $u(X)$ be defined as Definition C.6

We define $\alpha(X)\in\mathbb{R}^{n}$ as follows

$$ $\displaystyle\alpha(X):=u(X)\cdot{\bf 1}_{n}$ $$

Then, with the help of the summation vector, we are ready to normalize the exponential matrix and get the softmax probability matrix.

###### Definition C.8 (Softmax probability function $f$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d}$
- •
Let $u(X)\in\mathbb{R}^{n\times n}$ be defined as Definition C.6
- •
Let $\alpha(X)\in\mathbb{R}^{n}$ be defined as Definition C.7

We define $f(X)\in\mathbb{R}^{n\times n}$ as follows

$$ $\displaystyle f(X):=\operatorname{diag}(\alpha(X))^{-1}u(X)$ $$

where we define $f(X)_{j_{0}}^{\top}\in\mathbb{R}^{n}$ is the $j_{0}$-th row of $f(X)$.

Besides the probability matrix introduced above, we introduce the value matrix in the following definition.

###### Definition C.9 (Value function $h$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d}$
- •
Let $W_{V}\in\mathbb{R}^{d\times d}$

We define $h(X)\in\mathbb{R}^{n\times d}$ as follows

$$ $\displaystyle h(X)=XW_{V}$ $$

Then, we introduce $s(X)$ to represent the output of the attention mechanism.

###### Definition C.10 (Self-attention output $s$ ) .

If we have the below conditions,

- •
Let $f(X)$ be defined as Definition C.8
- •
Let $h(X)$ be defined as Definition C.9

We define $s(X)\in\mathbb{R}^{n\times d}$ as follows

$$ $\displaystyle s(X)=f(X)h(X)$ $$

Then, we introduce $q(X)$ and $p(X)$ to facilitate the calculation of the gradient on $W$.

###### Definition C.11 (Definition of $q(X)$ ) .

If we have the below conditions,

- •
Let $h(X)\in\mathbb{R}^{n\times d}$ be defined as in Definition C.9.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

We define $q(X)\in\mathbb{R}^{n\times n}$ as

$$ $\displaystyle q(X)=\underbrace{G_{i}}_{n\times d}\underbrace{h(X)^{\top}}_{d \times n}.$ $$

where we define $q(X)_{j_{0}}^{\top}\in\mathbb{R}^{n}$ is the $j_{0}$-th row of $q(X)$.

###### Definition C.12 (Definition of $p(X)$ , Definition C.5 in AS24a [ 6 ] ) .

For every index $j_{0}\in[n]$, we define $p(X)_{j_{0}}\in\mathbb{R}^{n}$ as

$$ $\displaystyle p(X)_{j_{0}}:=(\operatorname{diag}(f(X)_{j_{0}})-f(X)_{j_{0}}f(X )_{j_{0}}^{\top})q(X)_{j_{0}}$ $$

where we have $p(X)\in\mathbb{R}^{n\times n}$ and we define $p(X)_{j_{0}}^{\top}\in\mathbb{R}^{n}$ is the $j_{0}$-th row of $p(X)$.

Furthermore, we define $p_{1}(X)=f(X)\odot q(X)$ and $p_{2}(X)=\operatorname{diag}(p_{1}(X)\cdot{\bf 1}_{n})f(X)$.
Additionally, we can calculate $p(X)$ as

$$ $\displaystyle p(X)=p_{1}(X)-p_{2}(X)$ $$

###### Definition C.5 (Simplified notations) .

###### Definition C.6 (Exponential function u 𝑢 u italic_u ) .

###### Definition C.7 (Sum function of softmax α 𝛼 \alpha italic_α ) .

###### Definition C.8 (Softmax probability function f 𝑓 f italic_f ) .

###### Definition C.9 (Value function h ℎ h italic_h ) .

###### Definition C.10 (Self-attention output s 𝑠 s italic_s ) .

###### Definition C.11 (Definition of q ⁢ ( X ) 𝑞 𝑋 q(X) italic_q ( italic_X ) ) .

###### Definition C.12 (Definition of p ⁢ ( X ) 𝑝 𝑋 p(X) italic_p ( italic_X ) , Definition C.5 in AS24a [ 6 ] ) .

### C.4 Low rank representations

Using AS []’s polynomial method techniques, we can obtain the following low-rank representation result.

###### Lemma C.13 (Low rank representation to $f$ , Section 3 of AS [ 5 ] , Lemma D.1 of AS24a [ 6 ] ) .

For any $A=o(\sqrt{\log n})$, there exists a $k_{1}=n^{o(1)}$ such that:
Let $X\in\mathbb{R}^{n\times d}$ and $W\in\mathbb{R}^{d\times d}$ be a square matrix.
It holds that $\|XW\|_{\infty}\leq R,\|X\|_{\infty}\leq R$,
then there are two matrices $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$ such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.
Here $f(X)=D^{-1}\exp(XWX^{\top})$ (see also Definition C.8) and we define $D=\operatorname{diag}(\exp(XWX^{\top}){\bf 1}_{n})$ (see also Definition C.7).
Moreover, these matrices $U_{1},V_{1}$ can be explicitly constructed in $n^{1+o(1)}$ time.

A similar technique can be applied to $s(X)$.

###### Lemma C.14 (Low rank representation to $s$ ) .

Let $d=O(\log n)$. Assume that each number in the $n\times d$ matrices $h(X)\in\mathbb{R}^{n\times d}$ can be written using $O(\log n)$ bits. Let $n\times d$ matrix $s(X)\in\mathbb{R}^{n\times d}$ be defined as Definition C.10. Then, there are two matrices $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$ we have $\|U_{1}V_{1}^{\top}h(X)-s(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.

###### Proof.

We can show that

$$ $\displaystyle\|U_{1}V_{1}^{\top}h(X)-s(X)\|_{\infty}=$ $\displaystyle~{}\|U_{1}V_{1}^{\top}h(X)-f(X)h(X)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(\underbrace{U_{1}V_{1}^{\top}}_{n\times n}-\underbrace{f(X) }_{n\times n})\underbrace{h(X)}_{n\times d}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|\underbrace{U_{1}V_{1}^{\top}}_{n\times n}-\underbrace{f(X) }_{n\times n}\|_{\infty}\|\underbrace{h(X)}_{n\times d}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|\underbrace{U_{1}V_{1}^{\top}}_{n\times n}-\underbrace{f(X) }_{n\times n}\|_{\infty}\cdot\operatorname{poly}(n)$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from the choice of $s(X)$, the 2nd step comes from $AC-BC=(A-B)C$ holds for any matrices $A$, $B$, and $C$, the 3rd step is because of basic linear algebra, the 4th step is due to each number in $h(X)$ can be written using $O(\log(n))$ bits, the fifth step follows from $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.

∎

We can also get a low-rank representation of $p_{1}(x)$ and $p_{2}(x)$.

###### Lemma C.15 (Low rank representation to $p_{1}(X)$ , Lemma D.4 of AS24a [ 6 ] ) .

Let $k_{1}=n^{o(1)}$. Let $k_{2}=n^{o(1)}$.
Assume that $p_{1}(X):=f(X)\odot q(X)$. Assume $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$ approximates the $f(X)$ such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. Assume $U_{2},V_{2}\in\mathbb{R}^{n\times k_{2}}$ approximates the $q(X)\in\mathbb{R}^{n\times n}$ such that $\|U_{2}V_{2}^{\top}-q(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. Then there are matrices $U_{3},V_{3}\in\mathbb{R}^{n\times k_{3}}$ such that $\|U_{3}V_{3}^{\top}-p_{1}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{3},V_{3}$ can be explicitly constructed in $n^{1+o(1)}$ time.

###### Lemma C.16 (Low rank representation $p_{2}(X)$ , Lemma D.5 of AS24a [ 6 ] ) .

Let $k_{1}=n^{o(1)}$. Let $k_{2}=n^{o(1)}$. Let $k_{4}=n^{o(1)}$.
Assume that $p_{2}(X)$ is an $n\times n$ where $j_{0}$-th row $p_{2}(X)_{j_{0}}=f(X)_{j_{0}}f(X)^{\top}_{j_{0}}q(X)_{j_{0}}$ for each $j_{0}\in[n]$. Assume $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$ approximates the $f(X)$ such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. Assume $U_{2},V_{2}\in\mathbb{R}^{n\times k_{2}}$ approximates the $q(X)\in\mathbb{R}^{n\times n}$ such that $\|U_{2}V_{2}^{\top}-q(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. Then there are matrices $U_{4},V_{4}\in\mathbb{R}^{n\times k_{4}}$ such that $\|U_{4}V_{4}^{\top}-p_{2}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{4},V_{4}$ can be explicitly constructed in $n^{1+o(1)}$ time.

###### Lemma C.13 (Low rank representation to f 𝑓 f italic_f , Section 3 of AS [ 5 ] , Lemma D.1 of AS24a [ 6 ] ) .

###### Lemma C.14 (Low rank representation to s 𝑠 s italic_s ) .

###### Proof.

###### Lemma C.15 (Low rank representation to p 1 ⁢ ( X ) subscript 𝑝 1 𝑋 p_{1}(X) italic_p start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ( italic_X ) , Lemma D.4 of AS24a [ 6 ] ) .

###### Lemma C.16 (Low rank representation p 2 ⁢ ( X ) subscript 𝑝 2 𝑋 p_{2}(X) italic_p start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) , Lemma D.5 of AS24a [ 6 ] ) .

### C.5 Bounded entries of matrices

In this section, we provide proof that entries of matrices are bounded.

We begin with the exponential matrix $f(X)$.

###### Lemma C.17 (Bounded entries of $f(X)$ ) .

If we have the below conditions,

- •
Let $f(X)\in\mathbb{R}^{n\times n}$ be defined in Definition C.8.

Then, we can show that

$$ $\displaystyle\|f(X)\|_{\infty}\leq 1$ $$

###### Proof.

By Definition C.8, we have

$$ $\displaystyle f(X)=\operatorname{diag}(\alpha(X))^{-1}u(X)$ $$

By Definition C.7, we have

$$ $\displaystyle\alpha(X)=u(X){\bf 1}_{n}$ $$

Combining above two equations, we have

$$ $\displaystyle\|f(X)\|_{\infty}\leq 1$ $$

∎

A similar analysis can be applied to $h(X)$ and $s(X)$ as well.

###### Lemma C.18 (Bounded entries of $h(X)$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V}$ can be re represented using $O(\log(n))$ bits.
- •
Let $h(X)\in\mathbb{R}^{n\times d}$ be defined in Definition C.9.

Then, we can show that

$$ $\displaystyle\|h(X)\|_{\infty}\leq\operatorname{poly}(n)$ $$

###### Proof.

By Definition C.9, we have

$$ $\displaystyle h(X):=XW_{V}$ $$

Then, we have

$$ $\displaystyle\|h(X)\|_{\infty}=$ $\displaystyle~{}\|XW_{V}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|X\|_{\infty}\|W_{V}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\operatorname{poly}(n)$ $$

where the 1st step is from the definition of $h(X)$, the 2nd step comes from basic linear algebra, the 3rd step is because of each entry in $X$ and $W_{V}$ can be represented by $O(\log(n))$ bits.
∎

###### Lemma C.19 (Bounded entries of $s(X)$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V}$ can be re represented using $O(\log(n))$ bits.
- •
Let $s(X)\in\mathbb{R}^{n\times d}$ be defined in Definition C.10.

Then, we can show that

$$ $\displaystyle\|s(X)\|_{\infty}\leq\operatorname{poly}(n)$ $$

###### Proof.

By Definition C.10, we have

$$ $\displaystyle\underbrace{s(X)}_{n\times d}=\underbrace{f(X)}_{n\times n} \underbrace{h(X)}_{n\times d}$ $$

Then, we have

$$ $\displaystyle\|s(X)\|_{\infty}=$ $\displaystyle~{}\|f(X)h(X)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|f(X)\|_{\infty}\|h(X)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\operatorname{poly}(n)$ $$

where the 1st step is from the definition of $c(X)$, the 2nd step comes from basic linear algebra, the 3rd step is because of Lemma C.17, C.18.
∎

###### Lemma C.17 (Bounded entries of f ⁢ ( X ) 𝑓 𝑋 f(X) italic_f ( italic_X ) ) .

###### Proof.

###### Lemma C.18 (Bounded entries of h ⁢ ( X ) ℎ 𝑋 h(X) italic_h ( italic_X ) ) .

###### Proof.

###### Lemma C.19 (Bounded entries of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X ) ) .

###### Proof.

## Appendix D Matrix View

In this section, we dive into analyzing the gradient of $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$.

In Section D.1, we give the gradient of $s(X)$ with respective to $X$.
In Section D.2, we show the close form of the gradient on $T_{i}(X)$ via the chain rule.
In Section D.3, we integrate each $C_{i}(X)$ to its corresponding matrix term $B_{i}(X)$.
In Section D.4, applying the similar technique used in the previous section, we integrate the gradient on $T_{i}(X)$ into its corresponding matrix view.
In Section D.5, we further apply matrix integration on each matrix term in the gradient on $T_{i}(X)$ calculated in the previous section.
In Section D.6, we give the matrix view of all gradient components.

### D.1 Gradient of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X )

In this section, we give the gradient of $s(X)$ with respective to $X$.

The results from DSXY [] give the gradient of $c(X)$.
By chain rule, the gradient of $s(X)$ is equivalent to the gradient of $c(X)$ from DSXY [], since $c(X)=s(X)-B$ where $B$ is a constant matrix.

###### Lemma D.1 (Gradient of $s(X)_{i_{0},j_{0}}$ , Lemma B.16 in DSXY [ 23 ] ) .

If we have the below conditions,

- •
Let $s(X)\in\mathbb{R}^{n\times d}$ be defined as Definition C.10

Then, we have

- •
Part 1. For all $i_{0}=i_{1}\in[n]$, $j_{0},j_{1}\in[d]$,
$\displaystyle\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}=C_
{1}(X)+C_{2}(X)+C_{3}(X)+C_{4}(X)+C_{5}(X)$
where we have definitions:
–
$C_{1}(X):=-s(X)_{i_{0},j_{0}}\cdot f(X)_{i_{0},i_{0}}\cdot\langle W_{j_{1},*},
X_{i_{0},*}\rangle$
–
$C_{2}(X):=-s(X)_{i_{0},j_{0}}\cdot\langle f(X)_{i_{0},*},XW_{*,j_{1}}\rangle$
–
$C_{3}(X):=f(X)_{i_{0},i_{0}}\cdot h(X)_{i_{0},j_{0}}\cdot\langle W_{j_{1},*},X
_{i_{0},*}\rangle$
–
$C_{4}(X):=\langle f(X)_{i_{0},*}\odot(XW_{*,j_{1}}),h(X)_{*,j_{0}}\rangle$
–
$C_{5}(X):=f(X)_{i_{0},i_{0}}\cdot(W_{V})_{j_{1},j_{0}}$
- •
Part 2. For all $i_{0}\neq i_{1}\in[n]$, $j_{0},j_{1}\in[d]$,
$\displaystyle\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}=C_
{6}(X)+C_{7}(X)+C_{8}(X)$
where we have definitions:
–
$C_{6}(X):=-s(X)_{i_{0},j_{0}}\cdot f(X)_{i_{1},i_{0}}\cdot\langle W_{j_{1},*},
X_{i_{0},*}\rangle$
*
This is corresponding to $C_{1}(X)$
–
$C_{7}(X):=f(X)_{i_{1},i_{0}}\cdot h(X)_{i_{1},j_{0}}\cdot\langle W_{j_{1},*},X
_{i_{0},*}\rangle$
*
This is corresponding to $C_{3}(X)$
–
$C_{8}(X):=f(X)_{i_{1},i_{0}}\cdot(W_{V})_{j_{1},j_{0}}$
*
This is corresponding to $C_{5}(X)$

###### Lemma D.1 (Gradient of s ⁢ ( X ) i 0 , j 0 𝑠 subscript 𝑋 subscript 𝑖 0 subscript 𝑗 0 s(X)_{i_{0},j_{0}} italic_s ( italic_X ) start_POSTSUBSCRIPT italic_i start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT , italic_j start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT end_POSTSUBSCRIPT , Lemma B.16 in DSXY [ 23 ] ) .

### D.2 Gradient on T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X )

In the Lemma D.2, we use the chain rule to calculate the close form of the gradient on $T_{i}(X)$.

###### Lemma D.2 (Gradient for $T_{i}(X)$ ) .

If we have the below conditions,

- •
Let $\mathsf{Attn}_{i}$ be defined as Definition C.3.
- •
Let $T_{i}(X)\in\mathbb{R}^{n\times d}$ be defined as Definition 3.3.
- •
Let $s(X)$ be defined as Definition C.10.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, for $i_{1}\in[n]$, $j_{1}\in[d]$, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)_{i_{1},j_{1}}}=\sum_{i _{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0 },j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}$ $$

###### Proof.

By Lemma C.4, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=\sum_{i_{2}=1}^{n} \sum_{j_{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{ i-1}(X))_{i_{2},j_{2}}}{\mathrm{d}T_{i-1}(X)}.$ $$

By Definition C.3 and Definition C.10, we have

$$ $\displaystyle\mathsf{Attn}_{i}(T_{i-1}(X))=s(T_{i-1}(X))$ $$

Therefore, by combining above two equations and substituting variable $T_{i-1}(X)=X$, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)_{i_{1},j_{1}}}=\sum_{i _{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0 },j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}$ $$

∎

###### Lemma D.2 (Gradient for T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

### D.3 Matrix view of C ⁢ ( X ) 𝐶 𝑋 C(X) italic_C ( italic_X )

In this section, we will provide the matrix view of $C_{i}(X)\in\mathbb{R}$, for $i\in\{6,7,8,2,4\}$.
We will consider each $C_{i}(X)$ one by one.
We begin with $C_{6}(X)$.

###### Lemma D.3 (Matrix view of $C_{6}(X)$ ) .

If we have the below conditions,

- •
Let $C_{6}(X,i_{1},j_{1}):=-s(X)_{i_{0},j_{0}}\cdot f(X)_{i_{1},i_{0}}\cdot\langle W
_{j_{1},*},X_{i_{0},*}\rangle$ be defined as in Lemma D.1.
- •
We define a matrix $B_{6}(X)\in\mathbb{R}^{n\times d}$.
For all $i_{1}\in[n],j_{1}\in[d]$, let $B_{6}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $B_{6}(X)$. We define $B_{6}(i_{1},j_{1})=C_{6}(X,i_{1},j_{1})$.

Then, we can show that

$$ $\displaystyle\underbrace{B_{6}(X)}_{n\times d}=\underbrace{-s(X)_{i_{0},j_{0}} }_{1\times 1}\underbrace{f(X)_{*,i_{0}}}_{n\times 1}\underbrace{(W\cdot X_{i_{ 0},*})^{\top}}_{1\times d}$ $$

###### Proof.

We have

$$ $\displaystyle C_{6}(X,i_{1},j_{1})=$ $\displaystyle~{}-s(X)_{i_{0},j_{0}}\cdot f(X)_{i_{1},i_{0}}\cdot\langle W_{j_{ 1},*},X_{i_{0},*}\rangle$ $\displaystyle=$ $\displaystyle~{}-s(X)_{i_{0},j_{0}}\cdot f(X)_{i_{1},i_{0}}\cdot X_{i_{0},*}^{ \top}W_{j_{1},*}$ $$

where the 1st step is from the choice of $C_{6}(X)$, the 2nd step comes from $\langle a,b\rangle=a^{\top}b$ holds for any $a,b\in\mathbb{R}^{d}$.

We have

$$ $\displaystyle\underbrace{B_{6}(X)(i_{1},*)}_{d\times 1}=$ $\displaystyle~{}-\underbrace{s(X)_{i_{0},j_{0}}}_{1\times 1}\underbrace{f(X)_{ i_{1},i_{0}}}_{1\times 1}\underbrace{W}_{d\times d}\underbrace{X_{i_{0},*}}_{d \times 1}$ $$

Then, we have

$$ $\displaystyle\underbrace{B_{6}(X)}_{n\times d}=\underbrace{-s(X)_{i_{0},j_{0}} }_{1\times 1}\underbrace{f(X)_{*,i_{0}}}_{n\times 1}\underbrace{(W\cdot X_{i_{ 0},*})^{\top}}_{1\times d}$ $$

∎

A similar analysis procedure can also be applied on $C_{7}(X)$.

###### Lemma D.4 (Matrix view of $C_{7}(X)$ ) .

If we have the below conditions,

- •
Let $C_{7}(X,i_{1},j_{1}):=f(X)_{i_{1},i_{0}}\cdot h(X)_{j_{0},i_{1}}\cdot\langle W
_{j_{1},*},X_{i_{0},*}\rangle$ be defined as in Lemma D.1.
- •
We define a matrix $B_{7}(X)\in\mathbb{R}^{n\times d}$.
For all $i_{1}\in[n],j_{1}\in[d]$, let $B_{7}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $B_{7}(X)$. We define $B_{7}(i_{1},j_{1})=C_{7}(X,i_{1},j_{1})$.

Then, we can show that

$$ $\displaystyle\underbrace{B_{7}(X)}_{n\times d}=$ $\displaystyle~{}\underbrace{(f(X)_{*,i_{0}}\odot h(X)_{*,j_{0}})}_{n\times 1} \cdot\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $$

###### Proof.

We have

$$ $\displaystyle C_{7}(X,i_{1},j_{1})=$ $\displaystyle~{}f(X)_{i_{1},i_{0}}\cdot h(X)_{i_{1},j_{0}}\cdot\langle W_{j_{1 },*},X_{i_{0},*}\rangle$ $\displaystyle=$ $\displaystyle~{}f(X)_{i_{1},i_{0}}\cdot h(X)_{i_{1},j_{0}}\cdot W_{j_{1},*}^{ \top}X_{i_{0},*}$ $$

where the 1st step is from the choice of $C_{7}(X)$,
the 2nd step comes from $\langle a,b\rangle=a^{\top}b$ holds for any $a,b\in\mathbb{R}^{d}$.

We have

$$ $\displaystyle B_{7}(X)(i_{1},*)=f(X)_{i_{1},i_{0}}\cdot h(X)_{i_{1},j_{0}} \cdot W\cdot X_{i_{0},*}$ $$

Then, we have

$$ $\displaystyle\underbrace{B_{7}(X)}_{n\times d}=$ $\displaystyle~{}\underbrace{(f(X)_{*,i_{0}}\odot h(X)_{*,j_{0}})}_{n\times 1} \cdot\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $$

∎

Then, we provide an analysis of $C_{8}(X)$.

###### Lemma D.5 (Matrix view of $C_{8}(X)$ ) .

If we have the below conditions,

- •
Let $C_{8}(X,i_{1},j_{1}):=f(X)_{i_{1},i_{0}}\cdot(W_{V})_{j_{1},j_{0}}$ be defined as in Lemma D.1.
- •
We define a matrix $B_{8}(X)\in\mathbb{R}^{n\times d}$.
For all $i_{1}\in[n],j_{1}\in[d]$, let $B_{8}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $B_{8}(X)$. We define $B_{8}(i_{1},j_{1})=C_{8}(X,i_{1},j_{1})$.

Then, we can show that

$$ $\displaystyle\underbrace{B_{8}(X)}_{n\times d}=\underbrace{f(X)_{*,i_{0}}}_{n \times 1}\underbrace{(W_{V})_{*,j_{0}}^{\top}}_{1\times d}$ $$

###### Proof.

We have

$$ $\displaystyle C_{8}(X,i_{1},j_{1})=$ $\displaystyle~{}f(X)_{i_{1},i_{0}}\cdot(W_{V})_{j_{1},j_{0}}$ $$

where the 1st step is from the choice of $C_{7}(X)$.

We have

$$ $\displaystyle B_{8}(X)(i_{1},*)=f(X)_{i_{1},i_{0}}\cdot(W_{V})_{*,j_{0}}$ $$

Then, we have

$$ $\displaystyle\underbrace{B_{8}(X)}_{n\times d}=\underbrace{f(X)_{*,i_{0}}}_{n \times 1}\underbrace{(W_{V})_{*,j_{0}}^{\top}}_{1\times d}$ $$

∎

Now, we consider $C_{2}(X)$.

###### Lemma D.6 (Matrix view of $C_{2}(X)$ ) .

If we have the below conditions,

- •
Let $C_{2}(X,j_{1}):=-s(X)_{i_{0},j_{0}}\cdot\langle f(X)_{i_{0},*},XW_{*,j_{1}}\rangle$ be defined as in Lemma D.1.
- •
We define a matrix $B_{2}(X)\in\mathbb{R}^{d}$. For all $j_{1}\in[d]$, the $j_{1}$-th entry of $B_{2}(X)$ is defined as $C_{2}(X,j_{1})$.

Then, we can show that

$$ $\displaystyle\underbrace{B_{2}(X)}_{d\times 1}=\underbrace{-s(X)_{i_{0},j_{0}} }_{1\times 1}\underbrace{W^{\top}}_{d\times d}\underbrace{X^{\top}}_{d\times n }\underbrace{f(X)_{i_{0},*}}_{n\times 1}$ $$

###### Proof.

We have

$$ $\displaystyle C_{2}(X,j_{1})=$ $\displaystyle~{}-s(X)_{i_{0},j_{0}}\cdot\langle f(X)_{i_{0},*},XW_{*,j_{1}}\rangle$ $\displaystyle=$ $\displaystyle~{}-s(X)_{i_{0},j_{0}}\cdot(XW_{*,j_{1}})^{\top}f(X)_{i_{0},*}$ $\displaystyle=$ $\displaystyle~{}\underbrace{-s(X)_{i_{0},j_{0}}}_{1\times 1}\underbrace{W_{*,j _{1}}^{\top}}_{1\times d}\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_ {0},*}}_{n\times 1}$ $$

where the 1st step is from the choice of $C_{2}(X)$, the second step follows
from $\langle a,b\rangle=a^{\top}b$, for any $a,b\in\mathbb{R}^{n}$.

Then, we have

$$ $\displaystyle\underbrace{B_{2}(X)}_{d\times 1}=\underbrace{-s(X)_{i_{0},j_{0}} }_{1\times 1}\underbrace{W^{\top}}_{d\times d}\underbrace{X^{\top}}_{d\times n }\underbrace{f(X)_{i_{0},*}}_{n\times 1}$ $$

∎

Finally, we analyze $C_{4}(X)$, which is the last term we need to compute.

###### Lemma D.7 (Matrix view of $C_{4}(X)$ ) .

If we have the below conditions,

- •
Let $C_{4}(X,j_{1}):=\langle f(X)_{i_{0},*}\odot(XW_{*,j_{1}}),h(X)_{*,j_{0}}\rangle$ be defined as in Lemma D.1.
- •
We define a matrix $B_{4}(X)\in\mathbb{R}^{d}$. For all $j_{1}\in[d]$, the $j_{1}$-th entry of $B_{4}(X)$ is defined as $C_{4}(X,j_{1})$.

Then, we can show that

$$ $\displaystyle\underbrace{B_{4}(X)}_{d\times 1}=\underbrace{W^{\top}}_{d\times d }\underbrace{X^{\top}}_{d\times n}\underbrace{(f(X)_{i_{0},*}\odot h(X)_{*,j_{ 0}})}_{n\times 1}$ $$

###### Proof.

We have

$$ $\displaystyle C_{4}(X,j_{1})=$ $\displaystyle~{}\langle f(X)_{i_{0},*}\odot(XW_{*,j_{1}}),h(X)_{*,j_{0}}\rangle$ $\displaystyle=$ $\displaystyle~{}\langle f(X)_{i_{0},*}\odot h(X)_{*,j_{0}},(XW_{*,j_{1}})\rangle$ $\displaystyle=$ $\displaystyle~{}(XW_{*,j_{1}})^{\top}(f(X)_{i_{0},*}\odot h(X)_{*,j_{0}})$ $$

where the 1st step is from the choice of $C_{4}(X)$, the 2nd step comes from Fact C.1, and the last step follows from basic linear algebra.
∎

###### Lemma D.3 (Matrix view of C 6 ⁢ ( X ) subscript 𝐶 6 𝑋 C_{6}(X) italic_C start_POSTSUBSCRIPT 6 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma D.4 (Matrix view of C 7 ⁢ ( X ) subscript 𝐶 7 𝑋 C_{7}(X) italic_C start_POSTSUBSCRIPT 7 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma D.5 (Matrix view of C 8 ⁢ ( X ) subscript 𝐶 8 𝑋 C_{8}(X) italic_C start_POSTSUBSCRIPT 8 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma D.6 (Matrix view of C 2 ⁢ ( X ) subscript 𝐶 2 𝑋 C_{2}(X) italic_C start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma D.7 (Matrix view of C 4 ⁢ ( X ) subscript 𝐶 4 𝑋 C_{4}(X) italic_C start_POSTSUBSCRIPT 4 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

### D.4 Matrix view of gradient on T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X )

Since we have got the matrix view of each $C_{i}(X)$ term in the previous section, we can get the matrix view of the gradient on $T_{i}(X)$ in Lemma D.8.

###### Lemma D.8 (Matrix view of single entry of gradient) .

If we have the below conditions,

- •
Let $s(X)$ be defined as Definition C.10.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
Let $B_{6}(X),B_{7}(X),B_{8}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.3, Lemma D.4, and Lemma D.5
- •
Let $B_{2}(X),B_{4}(X)\in\mathbb{R}^{d}$ be defined in Lemma D.6 and Lemma D.7.

For any $i_{0}\in[n],j_{0}\in[d]$, we have

$$ $\displaystyle G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{ \mathrm{d}X}=\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot(\underbrace{B_{6 }(X)+B_{7}(X)+B_{8}(X)}_{n\times d}+\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{(B_{2}(X)+B_{4}(X))^{\top}}_{1\times d})$ $$

###### Proof.

By Lemma D.1, we have

- •
Part 1. For all $i_{0}=i_{1}\in[n]$, $j_{0},j_{1}\in[d]$,
$\displaystyle\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}=C_
{1}(X)+C_{2}(X)+C_{3}(X)+C_{4}(X)+C_{5}(X)$
(1)
- •
Part 2. For all $i_{0}\neq i_{1}\in[n]$, $j_{0},j_{1}\in[d]$,
$\displaystyle\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X_{i_{1},j_{1}}}=C_
{6}(X)+C_{7}(X)+C_{8}(X)$
(2)

Since for any $i_{1}\in[n],j_{1}\in[d]$, let $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X_{i_{1},
j_{1}}}$ denote the $(i_{1},j_{1})$-th entry of $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X}$, we consider the following two cases:

- •
Case 1. The $i_{0}$-th row of $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X}$.
- •
Case 2. The other $n-1$ rows of $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X}$ where $i_{1}\neq i_{0}$.

We first consider Case 1.

Recall that the matrix view of $C_{2}(X),C_{4}(X)\in\mathbb{R}$ are $B_{2}(X),B_{4}(X)\in\mathbb{R}^{d}$, and the matrix view of $C_{6}(X),C_{7}(X),C_{8}(X)\in\mathbb{R}$ are $B_{6}(X),B_{7}(X),B_{8}(X)\in\mathbb{R}^{n\times d}$, respectively.

For $k\in\{6,7,8\}$, we use $B_{k}(X)(s,*)\in\mathbb{R}^{d}$ to denote the $s$-th row of $B_{k}(X)$.

We use $(G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X})(i_{0
},*)\in\mathbb{R}^{d}$ to denote the $i_{0}$-th row of $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X}$.

Since $C_{6}(X),C_{7}(X),C_{8}(X)$ are the corresponding parts of $C_{1}(X),C_{3}(X),C_{5}(X)$, and by Eq. (1), then we can have the following

$$ $\displaystyle~{}(G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{ \mathrm{d}X})(i_{0},*)$ $\displaystyle=$ $\displaystyle~{}\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot\underbrace{(B _{6}(X)(i_{0},*)+B_{7}(X)(i_{0},*)+B_{8}(X)(i_{0},*)+B_{2}(X)+B_{4}(X))}_{d \times 1}$ $$

We then consider Case 2.

For $k\in\{6,7,8\}$, we use $B_{k}(X)(\neq s,*)\in\mathbb{R}^{(n-1)\times d}$ to denote the matrix $B_{k}(X)$ with the $s$-th row removed.

Similarly, we use $(G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X})(\neq
i
_{0},*)\in\mathbb{R}^{(n-1)\times d}$ to denote the matrix $G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{\mathrm{d}X}$ with the $i_{0}$-th row removed.

By Eq. (2), we have

$$ $\displaystyle(G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{ \mathrm{d}X})(\neq i_{0},*)=\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot \underbrace{(B_{6}(X)(\neq i_{0},*)+B_{7}(X)(\neq i_{0},*)+B_{8}(X)(\neq i_{0} ,*))}_{d\times(n-1)}$ $$

Combining Case 1 and Case 2 together, we have

$$ $\displaystyle G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{ \mathrm{d}X}=\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot(\underbrace{B_{6 }(X)+B_{7}(X)+B_{8}(X)}_{n\times d}+\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{(B_{2}(X)+B_{4}(X))^{\top}}_{1\times d})$ $$

∎

Then, we have the matrix view of $T_{i}(X)$ gradient.

###### Lemma D.9 (Matrix view of $T_{i}(X)$ gradient) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $T(X)$ be defined as Definition 3.3.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
Let $B_{6}(X),B_{7}(X),B_{8}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.3, Lemma D.4, and Lemma D.5
- •
Let $B_{2}(X),B_{4}(X)\in\mathbb{R}^{d}$ be defined in Lemma D.6 and Lemma D.7.

Then, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=\sum_{i_{0}=1}^{n} \sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot(\underbrace {B_{6}(X)+B_{7}(X)+B_{8}(X)}_{n\times d}+\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{(B_{2}(X)+B_{4}(X))^{\top}}_{1\times d})$ $$

###### Proof.

By Lemma D.8, we have

$$ $\displaystyle G_{i}(i_{0},j_{0})\cdot\frac{\mathrm{d}s(X)_{i_{0},j_{0}}}{ \mathrm{d}X}=\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\cdot(\underbrace{B_{6 }(X)+B_{7}(X)+B_{8}(X)}_{n\times d}+\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{(B_{2}(X)+B_{4}(X))^{\top}}_{1\times d})$ $$

Then, by Lemma C.4
we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=\sum_{i_{2}=1}^{n} \sum_{j_{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{ i-1}(X))_{i_{2},j_{2}}}{\mathrm{d}T_{i-1}(X)}.$ $$

After combining the above two equations, we are done.
∎

###### Lemma D.8 (Matrix view of single entry of gradient) .

###### Proof.

###### Lemma D.9 (Matrix view of T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X ) gradient) .

###### Proof.

### D.5 Matrix view of each term in gradient on T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X )

In this subsection, we reduce the double summation to a matrix product for easy and clear analysis.

We first work on the $B_{6}$ term.

###### Lemma D.10 (Matrix view of $B_{6}(X)$ term) .

If we have the below conditions,

- •
Let $\underbrace{B_{6}(X)}_{n\times d}=\underbrace{-s(X)_{i_{0},j_{0}}}_{1\times 1}
\underbrace{f(X)_{*,i_{0}}}_{n\times 1}\underbrace{(W\cdot X_{i_{0},*})^{\top}
}_{1\times d}$ be defined in Lemma D.3.
- •
We define $z_{6}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{6}(X)_{*,i_{0}}}_{n\times 1}=(\underbrace{G_{i}(i_
{0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{
f(X)_{*,i_{0}}}_{n\times 1}$
- •
Let $f(X)\in\mathbb{R}^{n\times n}$ be defined in Definition C.8.
- •
Let $W\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{6}(X)}_{n\times d}=-\underbrace{z_{6}(X)}_{n \times n}\underbrace{X}_{n\times d}\underbrace{W^{\top}}_{d\times d}$ $$

###### Proof.

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})B_{6}(X)=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j _{0})}_{1\times 1}\underbrace{s(X)_{i_{0},j_{0}}}_{1\times 1}\underbrace{f(X)_ {*,i_{0}}}_{n\times 1}\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0}, j_{0})}_{1\times 1}\underbrace{s(X)_{i_{0},j_{0}}}_{1\times 1})\underbrace{f(X )_{*,i_{0}}}_{n\times 1}\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\underbrace{G_{i}(i_{0},*)^{\top}}_{1 \times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{f(X)_{*,i_{0}}}_{ n\times 1}\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\underbrace{G_{i}(i_{0},*)^{\top}}_{1 \times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{f(X)_{*,i_{0}}}_{ n\times 1}\underbrace{X_{i_{0},*}^{\top}}_{1\times d}\underbrace{W^{\top}}_{d \times d}$ $$

where the 1st step is from the choice of $B_{6}(X)$, the 2nd step comes from basic algebra, the 3rd step is because of $a^{\top}b=\sum_{i=1}^{d}a_{i}\cdot b_{i}$ holds for any $a,b\in\mathbb{R}^{d}$, the 4th step is due to $(AB)^{\top}=B^{\top}A^{\top}$ for any matrices $A$ and $B$.

Recall that we have $\underbrace{z_{6}(X)_{*,i_{0}}}_{n\times 1}=(\underbrace{G_{i}(i_{0},*)^{\top}
}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{f(X)_{*,i_{0}
}}_{n\times 1}$.

Then, we have

$$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\underbrace{G_{i}(i_{0},*)^{\top}}_{1 \times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{f(X)_{*,i_{0}}}_{ n\times 1}\underbrace{X_{i_{0},*}^{\top}}_{1\times d}\underbrace{W^{\top}}_{d \times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\underbrace{z_{6}(X)_{*,i_{0}}}_{n\times 1} \underbrace{X_{i_{0},*}^{\top}}_{1\times d}\underbrace{W^{\top}}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\underbrace{z_{6}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W^{\top}}_{d\times d}$ $$

where the 1st step is from the choice of $z_{6}(X)$, the 2nd step comes from basic linear algebra.
∎

Then, we can get the matrix view of $B_{7}(X)$ term.

###### Lemma D.11 (Matrix view of $B_{7}(X)$ term) .

If we have the below conditions,

- •
Let $\underbrace{B_{7}(X)}_{n\times d}=\underbrace{(f(X)_{*,i_{0}}\odot h(X)_{*,j_{
0}})}_{n\times 1}\cdot\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ be defined in Lemma D.4.
- •
We define $z_{7}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{7}(X)_{*,i_{0}}}_{n\times 1}=\underbrace{f(X)_{*,i
_{0}}}_{n\times 1}\odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}(i_{0},*
)}_{d\times 1}).$
- •
Let $X\in\mathbb{R}^{n\times d},W\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{7}(X)}_{n\times d}=\underbrace{z_{7}(X)}_{n\times n }\underbrace{X}_{n\times d}\underbrace{W^{\top}}_{d\times d}$ $$

###### Proof.

We have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{7}(X)}_{n\times d}=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_ {0})}_{1\times 1}\underbrace{(f(X)_{*,i_{0}}\odot h(X)_{*,j_{0}})}_{n\times 1} \cdot\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}(\underbrace{f(X)_{*,i_{0}}}_{n\times 1} \odot(\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\underbrace {h(X)_{*,j_{0}}}_{n\times 1}))\cdot\underbrace{(W\cdot X_{i_{0},*})^{\top}}_{1 \times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}(\underbrace{f(X)_{*,i_{0}}}_{n\times 1} \odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}(i_{0},*)}_{d\times 1})) \cdot\underbrace{(X_{i_{0},*}^{\top}W^{\top})}_{1\times d}$ $$

where the 1st step is from the choice of $B_{7}(X)$, the 2nd step comes from basic algebra, the 3rd step is because of basic linear algebra.

Recall that we have $\underbrace{z_{7}(X)_{*,i_{0}}}_{n\times 1}=\underbrace{f(X)_{*,i_{0}}}_{n
\times 1}\odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}(i_{0},*)}_{d
\times 1})$.

Then we have

$$ $\displaystyle~{}\sum_{i_{0}=1}^{n}(\underbrace{f(X)_{*,i_{0}}}_{n\times 1} \odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}(i_{0},*)}_{d\times 1})) \cdot\underbrace{(X_{i_{0},*}^{\top}W^{\top})}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{z_{7}(X)_{*,i_{0}}}_{n\times 1} \underbrace{X_{i_{0},*}^{\top}}_{1\times d}\underbrace{W^{\top}}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\underbrace{z_{7}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W^{\top}}_{d\times d}$ $$

where the 1st step is from the choice of $z_{7}(X)$, the 2nd step comes from basic linear algebra.
∎

Then, we consider $B_{8}(X)$.

###### Lemma D.12 (Matrix view of $B_{8}(X)$ term) .

If we have the below conditions,

- •
Let $\underbrace{B_{8}(X)}_{n\times d}=\underbrace{f(X)_{*,i_{0}}}_{n\times 1}
\underbrace{(W_{V})_{*,j_{0}}^{\top}}_{1\times d}$ be defined in Lemma D.5.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{8}(X)}_{n\times d}=\underbrace{f(X)}_{n\times n} \underbrace{G_{i}}_{n\times d}\underbrace{W_{V}^{\top}}_{d\times d}$ $$

###### Proof.

We have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{8}(X)}_{n\times d}=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_ {0})}_{1\times 1}\underbrace{f(X)_{*,i_{0}}}_{n\times 1}\underbrace{(W_{V})_{* ,j_{0}}^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{f(X)_{*,i_{0}}}_{n\times 1}(\sum _{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\underbrace{(W_{V})_{ *,j_{0}}^{\top}}_{1\times d})$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{f(X)_{*,i_{0}}}_{n\times 1} \underbrace{G_{i}(i_{0},*)^{\top}}_{1\times d}\underbrace{W_{V}^{\top}}_{d \times d}$ $\displaystyle=$ $\displaystyle~{}\underbrace{f(X)}_{n\times n}\underbrace{G_{i}}_{n\times d} \underbrace{W_{V}^{\top}}_{d\times d}$ $$

where the 1st step is from the choice of $B_{8}(X)$, the 2nd step comes from basic algebra, the 3rd step is because of basic linear algebra, the 4th step is due to basic linear algebra.

∎

Now, we can do the matrix view of $B_{2}(X)$ term.

###### Lemma D.13 (Matrix view of $B_{2}(X)$ term) .

If we have the below conditions,

- •
Let $\underbrace{B_{2}(X)}_{d\times 1}=\underbrace{-s(X)_{i_{0},j_{0}}}_{1\times 1}
\underbrace{W^{\top}}_{d\times d}\underbrace{X^{\top}}_{d\times n}\underbrace{
f(X)_{i_{0},*}}_{n\times 1}$ be defined in Lemma D.6
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
We define $z_{2}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{2}(X)_{i_{0},*}}_{n\times 1}=(\underbrace{G_{i}(i_
{0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{
f(X)_{i_{0},*}}_{n\times 1}$
- •
Let $X\in\mathbb{R}^{n\times d},W\in\mathbb{R}^{d\times d}$ be defined in Definition C.3

Then we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{2}(X)^{\top}}_{ 1\times d}=-\underbrace{z_{2}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

###### Proof.

We have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{2}(X)^{\top}}_{ 1\times d}=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j _{0})}_{1\times 1}\underbrace{s(X)_{i_{0},j_{0}}}_{1\times 1}\underbrace{e_{i_ {0}}}_{n\times 1}\underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_ {n\times d}\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0}, j_{0})}_{1\times 1}\underbrace{s(X)_{i_{0},j_{0}}}_{1\times 1})\underbrace{e_{ i_{0}}}_{n\times 1}\underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X }_{n\times d}\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}(\underbrace{G_{i}(i_{0},*)^{\top}}_{1 \times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{e_{i_{0}}}_{n \times 1}\underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_{n\times d }\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1}( \underbrace{G_{i}(i_{0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d \times 1})\underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_{n \times d}\underbrace{W}_{d\times d}$ $$

where the 1st step is from the choice of $B_{2}(X)$, the 2nd step comes from basic algebra, the 3rd step is because of $a^{\top}b=\sum_{i=1}^{d}a_{i}\cdot b_{i}$ holds for any $a,b\in\mathbb{R}^{d}$, the 4th step is due to $(AB)^{\top}=B^{\top}A^{\top}$ holds for any matrix $A,B$.

Recall that we have
$\underbrace{z_{2}(X)_{i_{0},*}}_{n\times 1}=(\underbrace{G_{i}(i_{0},*)^{\top}
}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{f(X)_{i_{0},*
}}_{n\times 1}$.

Then, we have

$$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1}( \underbrace{G_{i}(i_{0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d \times 1})\underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_{n \times d}\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{z_{2}(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}-\underbrace{z_{2}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

where the 1st step is from the choice of $z_{2}(X)$, the 2nd step comes from basic linear algebra.
∎

Finally, we do a similar analysis for the term $B_{4}(X)$. Then, we get all the matrix views we need.

###### Lemma D.14 (Matrix view of $B_{4}(X)$ term) .

If we have the below conditions,

- •
Let $\underbrace{B_{4}(X)}_{d\times 1}=\underbrace{W^{\top}}_{d\times d}\underbrace
{X^{\top}}_{d\times n}\underbrace{(f(X)_{i_{0},*}\odot h(X)_{*,j_{0}})}_{n
\times 1}$ be defined in Lemma D.7.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
We define $z_{4}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{4}(X)_{i_{0},*}}_{n\times 1}=\underbrace{f(X)_{i_{
0},*}}_{n\times 1}\odot\underbrace{(h(X)G_{i}(i_{0},*))}_{n\times 1}$

Then we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{4}(X)^{\top}}_{ 1\times d}=\underbrace{z_{4}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

###### Proof.

We have

$$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_ {0})}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{4}(X)^{\top} }_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_ {0})}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{(f(X)_{i_{0},*} ^{\top}\odot h(X)_{*,j_{0}}^{\top})}_{1\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1}( \underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\odot(\sum_{j_{0}=1}^{d} \underbrace{G_{i}(i_{0},j_{0})}_{1\times 1}\underbrace{h(X)_{*,j_{0}}^{\top}}_ {1\times n}))\underbrace{X}_{n\times d}\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1}( \underbrace{f(X)_{i_{0},*}^{\top}}_{1\times n}\odot\underbrace{(h(X)G_{i}(i_{0 },*))^{\top}}_{1\times n})\underbrace{X}_{n\times d}\underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\underbrace{e_{i_{0}}}_{n\times 1} \underbrace{z_{4}(X)_{i_{0},*}^{\top}}_{1\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\underbrace{z_{4}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

where the 1st step is from the choice of $B_{4}(X)$, the 2nd step comes from basic algebra, the 3rd step is because of basic linear algebra, the 4th step is due to the choice of $z_{4}(X)$, the 5th step follows from basic linear algebra.
∎

###### Lemma D.10 (Matrix view of B 6 ⁢ ( X ) subscript 𝐵 6 𝑋 B_{6}(X) italic_B start_POSTSUBSCRIPT 6 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

###### Lemma D.11 (Matrix view of B 7 ⁢ ( X ) subscript 𝐵 7 𝑋 B_{7}(X) italic_B start_POSTSUBSCRIPT 7 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

###### Lemma D.12 (Matrix view of B 8 ⁢ ( X ) subscript 𝐵 8 𝑋 B_{8}(X) italic_B start_POSTSUBSCRIPT 8 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

###### Lemma D.13 (Matrix view of B 2 ⁢ ( X ) subscript 𝐵 2 𝑋 B_{2}(X) italic_B start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

###### Lemma D.14 (Matrix view of B 4 ⁢ ( X ) subscript 𝐵 4 𝑋 B_{4}(X) italic_B start_POSTSUBSCRIPT 4 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### D.6 Components of gradient on T i ⁢ ( X ) subscript 𝑇 𝑖 𝑋 T_{i}(X) italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( italic_X )

###### Definition D.15 (Definition of $D_{k}$ ) .

If we have the below conditions,

- •
For $k_{1}\in\{6,7,8\}$, let $B_{k_{1}}(X)\in\mathbb{R}^{n\times d}$ be defined as Lemma D.3, D.4, and D.5, respectively.
- •
For $k_{2}\in\{2,4\}$, let $B_{k_{2}}(X)\in\mathbb{R}^{d\times 1}$ be defined as Lemma D.6 and D.7, respectively.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.

We define $D_{k}\in\mathbb{R}^{n\times d}$ as follows:

- •
For $k_{1}\in\{6,7,8\}$, we define
$\displaystyle D_{k_{1}}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}
(i_{0},j_{0})}_{1\times 1}\underbrace{B_{k_{1}}(X)}_{n\times d}$
- •
For $k_{2}\in\{2,4\}$, we define
$\displaystyle D_{k_{2}}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}
(i_{0},j_{0})}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{k_{
2}}(X)^{\top}}_{1\times d}$

###### Definition D.16 (Definition of $K$ ) .

If we have the below conditions,

- •
Let $s(X)\in\mathbb{R}^{n\times d}$ be defined as Definition C.10.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.

We define $K\in\mathbb{R}^{n}$, where for each $i_{0}\in[n]$, we define

$$ $\displaystyle\underbrace{K_{i_{0}}}_{1\times 1}=\underbrace{G_{i}(i_{0},*)^{ \top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1}$ $$

Furthermore, we have

$$ $\displaystyle\underbrace{K}_{n\times 1}=\underbrace{(G_{i}\odot s(X))}_{n \times d}\underbrace{{\bf 1}_{d}}_{d\times 1}$ $$

###### Lemma D.17 (Close form of $D_{k}$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d},W\in\mathbb{R}^{d\times d}$ be defined as Definition C.3.
- •
For $k\in\{6,7,8,2,4\}$, let $D_{k}\in\mathbb{R}^{n\times d}$ be defined as Definition D.15.
- •
For $k_{3}\in\{6,7,2,4\}$, let $z_{k_{3}}(X)\in\mathbb{R}^{n\times n}$ be defined as Lemma D.10, D.11, D.13, and D.14, respectively.
- •
Let $K\in\mathbb{R}^{n}$ be defined as Definition D.16.
- •
We define $z_{6}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{6}(X)}_{n\times n}=\underbrace{f(X)}_{n\times n}
\underbrace{\operatorname{diag}(K)}_{n\times n}.$
- •
We define $z_{7}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{7}(X)}_{n\times n}=\underbrace{f(X)}_{n\times n}
\odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}^{\top}}_{d\times n})$
- •
We define $z_{2}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{2}(X)}_{n\times n}=\underbrace{\operatorname{diag}
(K)}_{n\times n}\underbrace{f(X)}_{n\times n}$
- •
We define $z_{4}(X)\in\mathbb{R}^{n\times n}$, which satisfies
$\displaystyle\underbrace{z_{4}(X)}_{n\times n}=\underbrace{f(X)}_{n\times n}
\odot(\underbrace{G_{i}}_{n\times d}\underbrace{h(X)^{\top}}_{d\times n})$

Then, we can show that the close forms of $D_{k}$ can be written as follows:

- •
$D_{6}=-\underbrace{z_{6}(X)}_{n\times n}\underbrace{X}_{n\times d}\underbrace{
W^{\top}}_{d\times d}$.
- •
$D_{7}=\underbrace{z_{7}(X)}_{n\times n}\underbrace{X}_{n\times d}\underbrace{W
^{\top}}_{d\times d}$.
- •
$D_{8}=\underbrace{f(X)}_{n\times n}\underbrace{G_{i}}_{n\times d}\underbrace{W
_{V}^{\top}}_{d\times d}$.
- •
$D_{2}=-\underbrace{z_{2}(X)}_{n\times n}\underbrace{X}_{n\times d}\underbrace{
W}_{d\times d}$.
- •
$D_{4}=\underbrace{z_{4}(X)}_{n\times n}\underbrace{X}_{n\times d}\underbrace{W
}_{d\times d}$.

###### Proof.

We finish the proof by parts.

- •
By Lemma D.10, we have the close form of $D_{6}$.
- •
By Lemma D.11, we have the close form of $D_{7}$.
- •
By Lemma D.12, we have the close form of $D_{8}$.
- •
By Lemma D.13, we have the close form of $D_{2}$.
- •
By Lemma D.14, we have the close form of $D_{4}$.

∎

###### Definition D.15 (Definition of D k subscript 𝐷 𝑘 D_{k} italic_D start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) .

###### Definition D.16 (Definition of K 𝐾 K italic_K ) .

###### Lemma D.17 (Close form of D k subscript 𝐷 𝑘 D_{k} italic_D start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) .

###### Proof.

## Appendix E Fast Computation for Gradient on T ⁢ ( X ) 𝑇 𝑋 T(X) italic_T ( italic_X )

In this section, we give an almost linear time $n^{1+o(1)}$ algorithm for each $B_{i}(X)$ term.
Namely, we consider $B_{6}(X),B_{7}(X),B_{8}(X),B_{2}(X),B_{4}(X)$ in Section E.1, E.2, E.3, E.4, and E.5, respectively.

### E.1 Fast computation for B 6 ⁢ ( X ) subscript 𝐵 6 𝑋 B_{6}(X) italic_B start_POSTSUBSCRIPT 6 end_POSTSUBSCRIPT ( italic_X ) term

Before we introduce the almost linear time algorithm for $B_{6}(X)$ term, we need to introduce the accelerated algorithm for the key component term, $z_{6}(X)$, in Lemma E.2.

We first compute $K$, which is defined in Definition D.16

###### Lemma E.1 (Computation time for $K$ ) .

If we have the below conditions,

- •
Let $K\in\mathbb{R}^{n}$ be defined as Definition D.16.

Then, we can show that $K$ can be computed in $O(n\cdot d)$ time.

###### Proof.

Since for each $i_{0}\in[n]$, we have

$$ $\displaystyle\underbrace{K_{i_{0}}}_{1\times 1}=\underbrace{G_{i}(i_{0},*)^{ \top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1}$ $$

Then, we have that it takes $O(d)$ time for calculating each entry.

Since there are total $n$ entries in $K$, the overall computation time for $K$ is $O(n\cdot d)$.
∎

We now compute $z_{6}(X)$.

###### Lemma E.2 (Fast computation for $z_{6}(X)$ ) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $z_{6}(X)\in\mathbb{R}^{n\times n}$ be defined in Lemma D.10.

Then, for some $k_{6}=n^{o(1)}$, there are matrices $U_{6},V_{6}\in\mathbb{R}^{n\times k_{6}}$ such that $\|U_{6}V_{6}^{\top}-z_{6}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{6},V_{6}$ can be constructed in $n^{1+o(1)}$ time.

###### Proof.

Recall in Lemma D.10, we have define $z_{6}(X)$ satisfying the following equation

$$ $\displaystyle\underbrace{z_{6}(X)_{*,i_{0}}}_{n\times 1}=(\underbrace{G_{i}(i_ {0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{ f(X)_{*,i_{0}}}_{n\times 1}$ (3) $$

Recall that $K\in\mathbb{R}^{n}$ has been defined in Definition D.16.
By Lemma E.1, we have $K$ can be computed in $O(n\cdot d)$ time.

We also have

$$ $\displaystyle\underbrace{z_{6}(X)}_{n\times n}=\underbrace{f(X)}_{n\times n} \underbrace{\operatorname{diag}(K)}_{n\times n}$ $$

By Lemma C.13, we have $U_{1},V_{1}\in\mathbb{R}^{n\times k_{1}}$ such that

$$ $\displaystyle\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

Let $U_{6}=U_{1}$, $V_{6}=\operatorname{diag}(K)V_{1}$.

We have $V_{6}=\underbrace{\operatorname{diag}(K)}_{n\times n}\underbrace{V_{1}}_{n
\times k_{1}}$ can be computed in $nk_{1}$ time.

The overall running time for constructing $U_{6}$ and $V_{6}$ is $n^{1+o(1)}$.

Then, we consider the error bound.

We have

$$ $\displaystyle\|U_{6}V_{6}^{\top}-z_{6}(X)\|_{\infty}=$ $\displaystyle~{}\|U_{1}V_{1}^{\top}\operatorname{diag}(K)-f(X)\operatorname{ diag}(K)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\|\operatorname{diag}(K)\| _{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n(\epsilon/\operatorname{poly}(n))\|\operatorname{diag}(K)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from the choice of $U_{6}$, $V_{6}$, the 2nd step comes from basic linear algebra, the 3rd step is because of Lemma C.13, the 4th step is due to $\|\operatorname{diag}(K)\|_{\infty}\leq\operatorname{poly}(n)$.

∎

Then, we are ready to introduce the almost linear time algorithm for $B_{6}(X)$ term.

###### Lemma E.3 (Fast computation for $B_{6}(X)$ term) .

If we have the below conditions,

- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $B_{6}(X)\in\mathbb{R}^{n\times n}$ be defined in Lemma D.3.
- •
We define $D_{6}\in\mathbb{R}^{n\times d}$, where $D_{6}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})B_{6}(X)$.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, there is an algorithm to approximate $D_{6}$ in $n^{1+o(1)}$ time, and it can achieve $\epsilon/\operatorname{poly}(n)$ accuracy.

Namely, the algorithm output $\widetilde{D}_{6}$ satisfying

$$ $\displaystyle\|D_{6}-\widetilde{D}_{6}\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

###### Proof.

Recall that in Lemma D.10, we have defined $z_{6}(X)\in\mathbb{R}^{n\times n}$, which satisfies

$$ $\displaystyle\underbrace{z_{6}(X)_{*,i_{0}}}_{n\times 1}=(\underbrace{G_{i}(i_ {0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{ f(X)_{*,i_{0}}}_{n\times 1}$ $$

And, in that Lemma, we also have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{6}(X)}_{n\times d}=-\underbrace{z_{6}(X)}_{n \times n}\underbrace{X}_{n\times d}\underbrace{W^{\top}}_{d\times d}$ $$

Let $U_{6},V_{6}\in\mathbb{R}^{n\times k_{6}}$ be defined as Lemma E.2.

Let $\widetilde{z}_{6}(X)=U_{6}V_{6}^{\top}$.

By Lemma E.2, we have

$$ $\displaystyle\|\widetilde{z}_{6}(X)-z_{6}(X)\|_{\infty}\leq\epsilon/ \operatorname{poly}(n)$ (4) $$

Proof of running time.

We compute in the following way:

- •
Compute $\underbrace{V_{6}^{\top}}_{k_{6}\times n}\underbrace{X}_{n\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{V_{6}^{\top}X}_{k_{6}\times d}\underbrace{W^{\top}}_{d\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{U_{6}}_{n\times k_{6}}\underbrace{V_{6}^{\top}XW^{\top}}_{k_{6}
\times d}$, which takes $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|\widetilde{z}_{6}(X)XW^{\top}-z_{6}(X)XW^{\top}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n\|\widetilde{z}_{6}(X)-z_{6}(X)\|_{\infty}\|X\|_{ \infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from Eq.(4), the 3rd step is because of $\|W\|_{\infty}\leq\operatorname{poly}(n)$ and $\|X\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.1 (Computation time for K 𝐾 K italic_K ) .

###### Proof.

###### Lemma E.2 (Fast computation for z 6 ⁢ ( X ) subscript 𝑧 6 𝑋 z_{6}(X) italic_z start_POSTSUBSCRIPT 6 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma E.3 (Fast computation for B 6 ⁢ ( X ) subscript 𝐵 6 𝑋 B_{6}(X) italic_B start_POSTSUBSCRIPT 6 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### E.2 Fast computation for B 7 ⁢ ( X ) subscript 𝐵 7 𝑋 B_{7}(X) italic_B start_POSTSUBSCRIPT 7 end_POSTSUBSCRIPT ( italic_X ) term

Similar to the analysis process of $B_{6}(X)$ term, we first provide the almost linear time algorithm for $z_{7}(X)$, then provide that algorithm for $B_{7}(X)$.

###### Lemma E.4 (Fast computation for $z_{7}(X)$ ) .

If we have the below conditions,

- •
Let $z_{7}(X)\in\mathbb{R}^{n\times n}$ be defined in Lemma D.11.
- •
By Lemma C.13, let $U_{1},V_{1}$ be the low rank approximation of $f(X)$, such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, for some $k_{7}=n^{o(1)}$, there are matrices $U_{7},V_{7}\in\mathbb{R}^{n\times k_{7}}$ such that $\|U_{7}V_{7}^{\top}-z_{7}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{7},V_{7}$ can be constructed in $n^{1+o(1)}$ time.

###### Proof.

Recall that in Lemma D.11, we have defined $z_{7}(X)\in\mathbb{R}^{n\times n}$, where the $i_{0}$-th column of $z_{7}(X)$ satisfies

$$ $\displaystyle\underbrace{z_{7}(X)_{*,i_{0}}}_{n\times 1}=\underbrace{f(X)_{*,i _{0}}}_{n\times 1}\odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}(i_{0},* )}_{d\times 1})$ $$

which is equivalent to

$$ $\displaystyle\underbrace{z_{7}(X)}_{n\times n}=\underbrace{f(X)}_{n\times n} \odot(\underbrace{h(X)}_{n\times d}\underbrace{G_{i}^{\top}}_{d\times n})$ $$

By Lemma C.13, we know $\widetilde{f}(X):=U_{1}V_{1}^{\top}$ is a good approximation for $f(X)$.

We choose $U_{7}=U_{1}\oslash h(X)$ and $V_{7}=V_{1}\oslash G_{i}$,
where $U_{7},V_{7}\in\mathbb{R}^{n\times k_{1}d}$.

Proof of running time.

For $U_{7}=U_{1}\oslash h(X)$, since $U_{1}\in\mathbb{R}^{n\times k_{1}},h(X)\in\mathbb{R}^{n\times d}$, constructing $U_{7}$ takes $O(ndk_{1})=O(n^{1+o(1)})$ time.

Similarly, constructing $V_{7}$ takes $O(n^{1+o(1)})$ time.

Proof of error bound.

Using Fact C.2, we have

$$ $\displaystyle\|U_{7}V_{7}^{\top}-z_{7}(X)\|_{\infty}=$ $\displaystyle~{}\|U_{7}V_{7}^{\top}-f(X)\odot(h(X)G_{i}^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}\oslash h(X))(V_{1}\oslash G_{i})^{\top}-f(X)\odot(h(X )G_{i}^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}V_{1}^{\top})\odot(h(X)G_{i}^{\top})-f(X)\odot(h(X)G_{ i}^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|\widetilde{f}(X)\odot(h(X)G_{i}^{\top})-f(X)\odot(h(X)G_{i}^ {\top})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\|h(X)\|_{\infty}\|G_{i}\|_{\infty}\cdot\epsilon/ \operatorname{poly}(n)$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ (5) $$

where the 1st step is from the definition of $z_{7}(X)$, the 2nd step comes from the choice of $U_{7}$ and $V_{7}$, the 3rd step is because of Fact C.2, the 4th step is due to the definition of $\widetilde{f}(X)$, the 5th step follows from $\|\widetilde{f}(X)-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$, the sixth step follows from Lemma C.18 and $\|G_{i}\|_{\infty}\leq\operatorname{poly}(n)$.

∎

Then, we can do similarly fast computation for $B_{7}$ term.

###### Lemma E.5 (Fast computation for $B_{7}(X)$ term) .

If we have the below conditions,

- •
Let $B_{7}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.4.
- •
We define $D_{7}\in\mathbb{R}^{n\times d}$, where $D_{7}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})B_{7}(X)$.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d},B\in\mathbb{R}^{n
\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, there is an algorithm to approximate $D_{7}$ in $n^{1+o(1)}$ time, and it can achieve $\epsilon/\operatorname{poly}(n)$ accuracy.

Namely, the algorithm output $\widetilde{D}_{7}$ satisfies

$$ $\displaystyle\|D_{7}-\widetilde{D}_{7}\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

###### Proof.

In Lemma D.11, we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{7}(X)}_{n\times d}=\underbrace{z_{7}(X)}_{n\times n }\underbrace{X}_{n\times d}\underbrace{W^{\top}}_{d\times d}$ $$

Let $U_{7},V_{7}\in\mathbb{R}^{n\times k_{7}}$ be defined in Lemma E.4.

Let $\widetilde{z}_{7}(X):=U_{7}V_{7}^{\top}$.

By Lemma E.4, we have

$$ $\displaystyle\|\widetilde{z}_{7}(X)-z_{7}(X)\|_{\infty}\leq\epsilon/ \operatorname{poly}(n)$ (6) $$

Proof of running time.

We compute in the following way:

- •
Compute $\underbrace{V_{7}^{\top}}_{k_{7}\times n}\underbrace{X}_{n\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{V_{7}^{\top}X}_{k_{7}\times d}\underbrace{W^{\top}}_{d\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{U_{7}}_{n\times k_{7}}\underbrace{V_{7}^{\top}XW^{\top}}_{k_{7}
\times d}$, which takes $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|\widetilde{z}_{7}(X)XW^{\top}-z_{7}(X)XW^{\top}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n\|\widetilde{z}_{7}(X)-z_{7}(X)\|_{\infty}\|X\|_{ \infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from Eq. (6), the 3rd step is because of $\|W\|_{\infty}\leq\operatorname{poly}(n)$ and $\|X\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.4 (Fast computation for z 7 ⁢ ( X ) subscript 𝑧 7 𝑋 z_{7}(X) italic_z start_POSTSUBSCRIPT 7 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma E.5 (Fast computation for B 7 ⁢ ( X ) subscript 𝐵 7 𝑋 B_{7}(X) italic_B start_POSTSUBSCRIPT 7 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### E.3 Fast computation for B 8 ⁢ ( X ) subscript 𝐵 8 𝑋 B_{8}(X) italic_B start_POSTSUBSCRIPT 8 end_POSTSUBSCRIPT ( italic_X ) term

Then, we can do fast computations on $B_{8}(X)$ term.

###### Lemma E.6 (Fast computation for $B_{8}(X)$ term) .

If we have the below conditions,

- •
Let $B_{8}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.5.
- •
We define $D_{8}\in\mathbb{R}^{n\times d}$, where $D_{8}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}G_{i}(i_{0},j_{0})B_{8}(X)$.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, there is an algorithm to approximate $D_{8}$ in $n^{1+o(1)}$ time, and it can achieve $\epsilon/\operatorname{poly}(n)$ accuracy.

Namely, the algorithm output $\widetilde{D}_{8}$ satisfies

$$ $\displaystyle\|D_{8}-\widetilde{D}_{8}\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

###### Proof.

Recall that in Lemma D.12, we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{B_{8}(X)}_{n\times d}=\underbrace{f(X)}_{n\times n} \underbrace{G_{i}}_{n\times d}\underbrace{W_{V}^{\top}}_{d\times d}$ $$

Let $\widetilde{f}(X):=U_{1}V_{1}^{\top}$ denote the approximation of $f(X)$.

By Lemma C.13, we have

$$ $\displaystyle\|f(X)-\widetilde{f}(X)\|_{\infty}\leq\epsilon/\operatorname{poly }(n)$ (7) $$

Proof of running time.

We compute in the following way:

- •
Compute $\underbrace{V_{1}^{\top}}_{k_{1}\times n}\underbrace{G_{i}}_{n\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{V_{1}^{\top}G_{i}}_{k_{1}\times d}\underbrace{W_{V}^{\top}}_{d
\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{U_{1}}_{n\times k_{1}}\underbrace{V_{1}^{\top}G_{i}W_{V}^{\top}}_{
k_{1}\times d}$, which takes $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|\widetilde{f}(X)G_{i}W_{V}^{\top}-f(X)G_{i}W_{V}^{\top}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n\|\widetilde{f}(X)-f(X)\|_{\infty}\|G_{i}\|_{\infty}\| W_{V}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n(\epsilon/\operatorname{poly}(n))\|G_{i}\|_{\infty}\|W _{V}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from Eq.(7), the 3rd step is because of $\|G_{i}\|_{\infty}\leq\operatorname{poly}(n)$ and $\|W_{V}\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.6 (Fast computation for B 8 ⁢ ( X ) subscript 𝐵 8 𝑋 B_{8}(X) italic_B start_POSTSUBSCRIPT 8 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### E.4 Fast computation for B 2 ⁢ ( X ) subscript 𝐵 2 𝑋 B_{2}(X) italic_B start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) term

Then, we provide the proof of how to do fast computation on $B_{2}(X)$.

###### Lemma E.7 (Fast computation for $z_{2}(X)$ ) .

If we have the below conditions,

- •
Let $z_{2}(X)\in\mathbb{R}^{n\times n}$ be defined as in Lemma D.13.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, for some $k_{9}=n^{o(1)}$, there are matrices $U_{9},V_{9}\in\mathbb{R}^{n\times k_{9}}$ such that $\|U_{9}V_{9}^{\top}-z_{2}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{9},V_{9}$ can be constructed in $n^{1+o(1)}$ time.

###### Proof.

Recall that in Lemma D.13, we have defined $z_{2}(X)\in\mathbb{R}^{n\times n}$, where the $i_{0}$-th row of $z_{2}(X)$ satisfies

$$ $\displaystyle\underbrace{z_{2}(X)_{i_{0},*}}_{n\times 1}=(\underbrace{G_{i}(i_ {0},*)^{\top}}_{1\times d}\underbrace{s(X)_{i_{0},*}}_{d\times 1})\underbrace{ f(X)_{i_{0},*}}_{n\times 1}$ $$

Recall that $K\in\mathbb{R}^{n}$ has been defined in Definition D.16.

By Lemma E.1, we have $K$ can be computed in $O(n\cdot d)$ time.

We also have

$$ $\displaystyle\underbrace{z_{2}(X)}_{n\times n}=\underbrace{\operatorname{diag} (K)}_{n\times n}\underbrace{f(X)}_{n\times n}$ $$

By Lemma C.13, let $U_{1},V_{1}$ be the low rank approximation of $f(X)$, such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.

Let $U_{9}=\operatorname{diag}(K)U_{1}$, $V_{6}=V_{1}$.

We have $U_{9}=\underbrace{\operatorname{diag}(K)}_{n\times n}\underbrace{U_{1}}_{n
\times k_{1}}$ can be computed in $nk_{1}$ time.

The overall running time for constructing $U_{9}$ and $V_{9}$ is $n^{1+o(1)}$.

Then, we consider the error bound.

We have

$$ $\displaystyle\|U_{9}V_{9}^{\top}-z_{2}(X)\|_{\infty}=$ $\displaystyle~{}\|\operatorname{diag}(K)U_{1}V_{1}^{\top}-\operatorname{diag}( K)f(X)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\|\operatorname{diag}(K)\| _{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n(\epsilon/\operatorname{poly}(n))\|\operatorname{diag}(K)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ (8) $$

where the 1st step is from the choice of $U_{6}$, $V_{6}$, the 2nd step comes from basic linear algebra, the 3rd step is because of Lemma C.13, the 4th step is due to $\|\operatorname{diag}(K)\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.8 (Fast computation for $B_{2}(X)$ term) .

If we have the below conditions,

- •
Let $B_{2}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.6.
- •
We define $D_{2}\in\mathbb{R}^{n\times d}$, where $D_{2}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0})}_{1
\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{2}(X)^{\top}}_{1
\times d}$.
- •
Let $X\in\mathbb{R}^{d\times n},W,W_{V}\in\mathbb{R}^{d\times d},B\in\mathbb{R}^{n
\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},B,G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, there is an algorithm to approximate $D_{2}$ in $n^{1+o(1)}$ time, and it can achieve $\epsilon/\operatorname{poly}(n)$ accuracy.

Namely, the algorithm output $\widetilde{D}_{2}$ satisfies

$$ $\displaystyle\|D_{2}-\widetilde{D}_{2}\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

###### Proof.

In Lemma D.13, we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{2}(X)^{\top}}_{ 1\times d}=-\underbrace{z_{2}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

Let $U_{9},V_{9}\in\mathbb{R}^{n\times k_{9}}$ be defined in Lemma E.7.

Let $\widetilde{z}_{2}(X):=U_{9}V_{9}^{\top}$.

By Lemma E.7, we have

$$ $\displaystyle\|\widetilde{z}_{2}(X)-z_{2}(X)\|_{\infty}\leq\epsilon/ \operatorname{poly}(n)$ (9) $$

Proof of running time.

We compute in the following way:

- •
Compute $\underbrace{V_{9}^{\top}}_{k_{9}\times n}\underbrace{X}_{n\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{V_{9}^{\top}X}_{k_{9}\times d}\underbrace{W}_{d\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{U_{9}}_{n\times k_{9}}\underbrace{V_{9}^{\top}XW}_{k_{9}\times d}$, which takes $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|\widetilde{z}_{2}(X)XW-z_{2}(X)XW\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n\|\widetilde{z}_{2}(X)-z_{2}(X)\|_{\infty}\|X\|_{ \infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from Eq.(9), the 3rd step is because of $\|W\|_{\infty}\leq\operatorname{poly}(n)$ and $\|X\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.7 (Fast computation for z 2 ⁢ ( X ) subscript 𝑧 2 𝑋 z_{2}(X) italic_z start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma E.8 (Fast computation for B 2 ⁢ ( X ) subscript 𝐵 2 𝑋 B_{2}(X) italic_B start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### E.5 Fast computation for B 4 ⁢ ( X ) subscript 𝐵 4 𝑋 B_{4}(X) italic_B start_POSTSUBSCRIPT 4 end_POSTSUBSCRIPT ( italic_X ) term

Finally, our analysis shows that we can do fast computations for $B_{4}(X)$ term. After that, we showed that all terms can be computed quickly.

###### Lemma E.9 (Fast computation for $z_{4}(X)$ ) .

If we have the below conditions,

- •
Let $z_{4}(X)\in\mathbb{R}^{n\times n}$ be defined in Lemma D.14.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, for some $k_{10}=n^{o(1)}$, there are matrices $U_{10},V_{10}\in\mathbb{R}^{n\times k_{10}}$, let $\widetilde{z}_{4}(X):=U_{10}V_{10}^{\top}$, such that $\|\widetilde{z}_{4}(X)-z_{4}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$. The matrices $U_{10},V_{10}$ can be constructed in $n^{1+o(1)}$ time.

###### Proof.

In Lemma D.14, we have defined $z_{4}(X)\in\mathbb{R}^{n\times n}$, where the $i_{0}$-th column of $z_{4}(X)$ satisfies

$$ $\displaystyle\underbrace{z_{4}(X)_{i_{0},*}}_{n\times 1}=(\underbrace{f(X)_{i_ {0},*}}_{n\times 1}\odot\underbrace{(h(X)G_{i}(i_{0},*))}_{n\times 1})$ $$

which is equivalent to

$$ $\displaystyle\underbrace{z_{4}(X)}_{n\times n}=(\underbrace{f(X)}_{n\times n} \odot\underbrace{G_{i}}_{n\times d}\underbrace{h(X)^{\top}}_{d\times n})$ $$

By Lemma C.13, let $U_{1},V_{1}$ be the low rank approximation of $f(X)$, such that $\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.

We choose $U_{10}=U_{1}\oslash G_{i}$ and $V_{10}=V_{1}\oslash h(X)$,
where $U_{10},V_{10}\in\mathbb{R}^{n\times k_{1}d}$.

Proof of running time.

For $U_{10}=U_{1}\oslash G_{i}$, since $U_{1}\in\mathbb{R}^{n\times k_{1}},G_{i}\in\mathbb{R}^{n\times d}$, constructing $U_{10}$ takes $O(ndk_{1})=O(n^{1+o(1)})$ time.

Similarly, constructing $V_{10}$ takes $O(n^{1+o(1)})$ time.

Proof of error bound.

Let $\widetilde{f}(X):=U_{1}V_{1}^{\top}$.

Using Fact C.2, we have

$$ $\displaystyle~{}\|\widetilde{z}_{4}(X)-z_{4}(X)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|U_{10}V_{10}^{\top}-f(X)\odot(G_{i}\cdot h(X)^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}\oslash G_{i})(V_{1}\oslash h(X))^{\top}-f(X)\odot(G_{ i}\cdot h(X)^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}V_{1}^{\top})\odot(G_{i}\cdot h(X)^{\top})-f(X)\odot(G _{i}\cdot h(X)^{\top})\|_{\infty}$ $$

where the 1st step is from the definition of $\widetilde{z}_{4}(X),z_{4}(X)$,
the 2nd step comes from the choice of $U_{10}$ and $V_{10}$, the 3rd step is because of Fact C.2.

$$ $\displaystyle~{}\|(U_{1}V_{1}^{\top})\odot(G_{i}\cdot h(X)^{\top})-f(X)\odot(G _{i}\cdot h(X)^{\top})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|U_{1}V_{1}^{\top}-f(X)\|_{\infty}\|G_{i}\cdot h(X)^{\top}\|_ {\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot(\epsilon/\operatorname{poly}(n))\|h(X)\|_{\infty}\|G_{i }\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from $\|U_{1}V_{1}-f(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$, the 3rd step is because of Lemma C.18 and $\|G_{i}\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.10 (Fast computation for $B_{4}(X)$ term) .

If we have the below conditions,

- •
Let $B_{4}(X)\in\mathbb{R}^{n\times d}$ be defined in Lemma D.7.
- •
We define $D_{4}\in\mathbb{R}^{n\times d}$, where $D_{4}:=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0})}_{1
\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{4}(X)^{\top}}_{1
\times d}$.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.

Then, we can show that, there is an algorithm to approximate $D_{4}$ in $n^{1+o(1)}$ time, and it can achieve $\epsilon/\operatorname{poly}(n)$ accuracy.

Namely, the algorithm output $\widetilde{D}_{4}$ satisfies

$$ $\displaystyle\|D_{4}-\widetilde{D}_{4}\|_{\infty}\leq\epsilon/\operatorname{ poly}(n)$ $$

###### Proof.

In Lemma D.14, we have

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_{0} )}_{1\times 1}\underbrace{e_{i_{0}}}_{n\times 1}\underbrace{B_{4}(X)^{\top}}_{ 1\times d}=\underbrace{z_{4}(X)}_{n\times n}\underbrace{X}_{n\times d} \underbrace{W}_{d\times d}$ $$

Let $\widetilde{z}_{4}(X):=U_{10}V_{10}^{\top}$.

By Lemma E.9, we have

$$ $\displaystyle\|\widetilde{z}_{4}(X)-z_{4}(X)\|_{\infty}\leq\epsilon/ \operatorname{poly}(n)$ (10) $$

Proof of running time.

We compute in the following way:

- •
Compute $\underbrace{V_{10}^{\top}}_{k_{10}\times n}\underbrace{X}_{n\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{V_{10}^{\top}X}_{k_{10}\times d}\underbrace{W}_{d\times d}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{U_{10}}_{n\times k_{10}}\underbrace{V_{10}^{\top}XW}_{k_{10}\times
d}$, which takes $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|\widetilde{z}_{4}(X)XW-z_{4}(X)XW\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n\|\widetilde{z}_{4}(X)-z_{4}(X)\|_{\infty}\|X\|_{ \infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot n(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}\|W\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic linear algebra, the 2nd step comes from Eq.(10), the 3rd step is because of $\|W\|_{\infty}\leq\operatorname{poly}(n)$ and $\|X\|_{\infty}\leq\operatorname{poly}(n)$.

∎

###### Lemma E.9 (Fast computation for z 4 ⁢ ( X ) subscript 𝑧 4 𝑋 z_{4}(X) italic_z start_POSTSUBSCRIPT 4 end_POSTSUBSCRIPT ( italic_X ) ) .

###### Proof.

###### Lemma E.10 (Fast computation for B 4 ⁢ ( X ) subscript 𝐵 4 𝑋 B_{4}(X) italic_B start_POSTSUBSCRIPT 4 end_POSTSUBSCRIPT ( italic_X ) term) .

###### Proof.

### E.6 Putting everything together

After we have analyzed each $B_{i}(X)$ term in the previous section, we put them together in this section, to analyze the overall running time and error bound of the gradient of $L(X)$ on $T_{i}(X)$ in Lemma E.11.

###### Lemma E.11 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$ , formal version of Lemma 5.1 ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $m$ denote the number of self-attention transformer model (see Definition 1.3).
- •
For any $i\in[m]$, let $T_{i}(X)$ be defined as Definition 3.3.
- •
Let $X\in\mathbb{R}^{n\times d},W,W_{V}\in\mathbb{R}^{d\times d}$ be defined in Definition C.3.
- •
Assuming each entry of $X,W,W_{V},G_{i}$ can be re represented using $O(\log(n))$ bits.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
Assume $G_{i}$ can be computed in $n^{1+o(1)}$ time.

We can show that $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.
Namely, our algorithm can output $\widetilde{g}_{t}$ in $n^{1+o(1)}$ time, which satisfies

$$ $\displaystyle\|\widetilde{g}_{t}-\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}\| _{\infty}\leq 1/\operatorname{poly}(n)$ $$

###### Proof.

By Lemma D.9, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{d}\underbrace{G_{i}(i_{0},j_ {0})}_{1\times 1}\cdot(\underbrace{B_{6}(X)+B_{7}(X)+B_{8}(X)}_{n\times d}+ \underbrace{e_{i_{0}}}_{n\times 1}\underbrace{(B_{2}(X)+B_{4}(X))^{\top}}_{1 \times d})$ $\displaystyle=$ $\displaystyle~{}\sum_{i\in\{2,4,6,7,8\}}D_{i}$ $$

where the 1st step is from Lemma D.9, the 2nd step comes from the definition of $D_{6},D_{7},D_{8},D_{2},D_{4}$.

Then, by Lemma E.3, E.5, E.6, E.8, E.10, we have $D_{6},D_{7},D_{8},D_{2},D_{4}\in\mathbb{R}^{n\times d}$ can be approximated in $n^{1+o(1)}$ time, with up to $\epsilon/\operatorname{poly}(n)$ error.

Namely, for $i\in\{2,4,6,7,8\}$, let $\widetilde{D}_{i}\in\mathbb{R}^{n\times d}$ denote the approximated version of $D$, we have

$$ $\displaystyle\|\widetilde{D}_{i}-D\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$ $$

Let $\widetilde{g}_{t}=\sum_{i\in\{2,4,6,7,8\}}\widetilde{D}_{i}$.

Proof of running time.

The running time for $\widetilde{g}_{t}=\sum_{i\in\{2,4,6,7,8\}}\widetilde{D}_{i}$ is $5nd$.

Therefore, the overall running time for computing $\widetilde{g}_{t}$ is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle\|\widetilde{g}_{t}-\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}\| _{\infty}=$ $\displaystyle~{}\|\sum_{i\in\{2,4,6,7,8\}}(\widetilde{D}_{i}-D_{i})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\sum_{i\in\{2,4,6,7,8\}}\|(\widetilde{D}_{i}-D_{i})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from the definition of $\widetilde{g}_{t}$ and $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}$, the 2nd step comes from basic algebra, the 3rd step is because of $\|\widetilde{D}_{i}-D\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$.

Then, choose $\epsilon=1/\operatorname{poly}(n)$, we have

$$ $\displaystyle\|\widetilde{g}_{t}-\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)}\| _{\infty}\leq 1/\operatorname{poly}(n)$ $$

∎

###### Lemma E.11 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ T i − 1 ⁢ ( X ) d 𝐿 𝑋 d subscript 𝑇 𝑖 1 𝑋 \frac{\mathrm{d}L(X)}{\mathrm{d}T_{i-1}(X)} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d italic_T start_POSTSUBSCRIPT italic_i - 1 end_POSTSUBSCRIPT ( italic_X ) end_ARG , formal version of Lemma 5.1 ) .

###### Proof.

## Appendix F Fast Computation for Gradient on W 𝑊 W italic_W

In Section F.1, we introduce some essential notations used in this section.
In Section F.2, we offer the gradient of $s(X)$ on $W$, which is equivalent to the gradient of the output of the attention mechanism on $W$.
In Section F.3, we illustrate the gradient of $L(X)$ on $W$.
In Section F.4, we introduce the almost linear time algorithm for calculating the gradient of $L(X)$ on $W$, along with the error bound analysis.

### F.1 Key concepts

###### Definition F.1 (Definition of $\operatorname{\mathsf{A}}$ , [ 6 ] ) .

Let $A_{1},A_{2}\in\mathbb{R}^{n\times d}$ be two matrices. Suppose that $\operatorname{\mathsf{A}}=A_{1}\otimes A_{2}\in\mathbb{R}^{n^{2}\times d^{2}}$. We define $\operatorname{\mathsf{A}}_{j_{0}}\in\mathbb{R}^{n\times d^{2}}$ be a $n\times d^{2}$ size sub-block from $\operatorname{\mathsf{A}}$. Note that there are $n$ such sub-blocks.

###### Remark F.2 .

Note that the $A_{1},A_{2}$ matrices in Definition F.1 is $X$ in our setting.
Since in AS24a [], they consider a more general setting, where $A_{1},A_{2}$ can be difference matrices, while in our problem, we consider self-attention. Therefore, in our paper, we have $A_{1}=A_{2}=X$.

###### Definition F.1 (Definition of 𝖠 𝖠 \operatorname{\mathsf{A}} sansserif_A , [ 6 ] ) .

###### Remark F.2 .

### F.2 Gradient of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X ) on W 𝑊 W italic_W

We begin with introducing the close form of the gradient of $s(X)$.

AS24a [] proved the close form of the gradient of $c(X)=s(X)-B$ with respect to $W$ for a constant matrix $B$.
By chain rule, this is equivalent to the gradient of $s(X)$ with respect to $W$.

###### Lemma F.3 (Gradient of $s(X)$ on $W$ , Lemma B.1 in AS24a [ 6 ] ) .

If we have the below conditions,

- •
Let $\operatorname{\mathsf{A}}$ be defined as Definition F.1. For every $i\in[d^{2}]$, define $\operatorname{\mathsf{A}}_{j_{0},i}\in\mathbb{R}^{n}$ to be the $i$-th column for $\operatorname{\mathsf{A}}_{j_{0}}\in\mathbb{R}^{n\times d^{2}}$.
- •
Let $f(X),h(X),s(X)$ be defined as Definition C.8, C.9, C.10.
- •
Let $W\in\mathbb{R}^{d\times d}$ be defined as Definition C.3. Let $w\in\mathbb{R}^{d^{2}}$ denote the vector representation of $W$.

Then, for each $i\in[d^{2}]$, we have
For each $j_{0}\in[n]$, for every $i_{0}\in[d]$

$$ $\displaystyle\frac{\mathrm{d}s(X)_{j_{0},i_{0}}}{\mathrm{d}w_{i}}=\langle \operatorname{\mathsf{A}}_{j_{0},i}\odot f(X)_{j_{0}},h(X)_{i_{0}}\rangle- \langle f(X)_{j_{0}},h(X)_{i_{0}}\rangle\cdot\langle\operatorname{\mathsf{A}}_ {j_{0},i},f(X)_{j_{0}}\rangle$ $$

###### Lemma F.3 (Gradient of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X ) on W 𝑊 W italic_W , Lemma B.1 in AS24a [ 6 ] ) .

### F.3 Gradient of L ⁢ ( X ) 𝐿 𝑋 L(X) italic_L ( italic_X ) on W 𝑊 W italic_W

Differing from the $\ell_{2}$ loss function used in AS24a [], our framework supports arbitrary loss functions.
Therefore, we use Lemma F.4 to illustrate the gradient of $L(X)$ on $W$.

###### Lemma F.4 (Gradient of $L(X)$ on $W$ ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $W\in\mathbb{R}^{d\times d},X\in\mathbb{R}^{n\times d}$ be Defined as Definition C.3.
- •
Let $p(X)$ be defined as Definition C.12.

Then, we can show that

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}=X^{\top}\cdot p(X)\cdot X$ $$

###### Proof.

By Lemma F.3, we have, for each $i\in[d^{2}]$, we have
For each $j_{0}\in[n]$, for every $i_{0}\in[d]$

$$ $\displaystyle\frac{\mathrm{d}s(X)_{j_{0},i_{0}}}{\mathrm{d}w_{i}}=\langle \underbrace{\operatorname{\mathsf{A}}_{j_{0},i}}_{n\times 1}\odot\underbrace{f (X)_{j_{0}}}_{n\times 1},\underbrace{h(X)_{i_{0}}}_{n\times 1}\rangle-\langle \underbrace{f(X)_{j_{0}}}_{n\times 1},\underbrace{h(X)_{i_{0}}}_{n\times 1} \rangle\cdot\langle\underbrace{\operatorname{\mathsf{A}}_{j_{0},i}}_{n\times 1 },\underbrace{f(X)_{j_{0}}}_{n\times 1}\rangle$ (11) $$

By Fact C.1, we have

$$ $\displaystyle\langle\operatorname{\mathsf{A}}_{j_{0},i}\odot f(X)_{j_{0}},h(X) _{i_{0}}\rangle=\operatorname{\mathsf{A}}_{j_{0},i}^{\top}\operatorname{diag}( f(X)_{j_{0}})h(X)_{i_{0}}$ $$

and

$$ $\displaystyle\langle f(X)_{j_{0}},h(X)_{i_{0}}\rangle\cdot\langle f(X)_{j_{0}} ,\operatorname{\mathsf{A}}_{j_{0},i}\rangle=\operatorname{\mathsf{A}}_{j_{0},i }^{\top}f(X)_{j_{0}}f(X)_{j_{0}}^{\top}h(X)_{i_{0}}$ $$

By Eq. (11), for each $i\in[d^{2}]$, we have
For each $j_{0}\in[n]$, for every $i_{0}\in[d]$, we have

$$ $\displaystyle\frac{\mathrm{d}s(X)_{j_{0},i_{0}}}{\mathrm{d}w_{i}}= \operatorname{\mathsf{A}}_{j_{0},i}^{\top}(\operatorname{diag}(f(X)_{j_{0}})-f (X)_{j_{0}}f(X)_{j_{0}}^{\top})h(X)_{i_{0}}$ $$

which implies,

$$ $\displaystyle\frac{\mathrm{d}s(X)_{j_{0},i_{0}}}{\mathrm{d}W}=\underbrace{ \operatorname{\mathsf{A}}_{j_{0}}^{\top}}_{d^{2}\times n}\underbrace{( \operatorname{diag}(f(X)_{j_{0}})-f(X)_{j_{0}}f(X)_{j_{0}}^{\top})}_{n\times n }\underbrace{h(X)_{i_{0}}}_{n\times 1}$ (12) $$

By Lemma C.4, for $i\in[m]$, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}=\sum_{i_{2}=1}^{n}\sum_{j_ {2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X)) _{i_{2},j_{2}}}{\mathrm{d}W_{i}}.$ (13) $$

By the definition of $s(X)$ (Definition C.10), we have

$$ $\displaystyle s(X)=\mathsf{Attn}_{i}(T_{i-1}(X))$ $$

Combining Eq. (12) and Eq. (13), for each $i\in[m]$, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}=\sum_{j_{0}=1}^{n}\sum_{i_ {0}=1}^{d}\underbrace{G_{i}(j_{0},i_{0})}_{1\times 1}\cdot\underbrace{ \operatorname{\mathsf{A}}_{j_{0}}^{\top}}_{d^{2}\times n}\underbrace{( \operatorname{diag}(f(X)_{j_{0}})-f(X)_{j_{0}}f(X)_{j_{0}}^{\top})}_{n\times n }\underbrace{h(X)_{i_{0}}}_{n\times 1}$ (14) $$

Recall that we have defined $q(X)$ in Definition C.11,

$$ $\displaystyle q(X)_{j_{0}}:=\sum_{i_{0}=1}^{d}G_{i}(j_{0},i_{0})\cdot h(X)_{i_ {0}}$ (15) $$

Recall that $p(x)_{j_{0}}\in\mathbb{R}^{n}$ is define as Definition C.12,

$$ $\displaystyle p(x)_{j_{0}}:=(\operatorname{diag}(f(x)_{j_{0}})-f(x)_{j_{0}}f(x )_{j_{0}}^{\top})q(x)_{j_{0}}.$ (16) $$

Then, we have

$$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\sum_{i_{0}=1}^{d}\underbrace{G_{i}(j_{0},i_ {0})}_{1\times 1}\cdot\underbrace{\operatorname{\mathsf{A}}_{j_{0}}^{\top}}_{d ^{2}\times n}\underbrace{(\operatorname{diag}(f(X)_{j_{0}})-f(X)_{j_{0}}f(X)_{ j_{0}}^{\top})}_{n\times n}\underbrace{h(X)_{i_{0}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\underbrace{\operatorname{\mathsf{A}}_{j_{0} }^{\top}}_{d^{2}\times n}\underbrace{(\operatorname{diag}(f(X)_{j_{0}})-f(X)_{ j_{0}}f(X)_{j_{0}}^{\top})}_{n\times n}\underbrace{q(X)_{j_{0}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\operatorname{\mathsf{A}}_{j_{0}}^{\top}p_{j _{0}}(X)$ $\displaystyle=$ $\displaystyle~{}\underbrace{X^{\top}}_{d\times n}\underbrace{p(X)}_{n\times n} \underbrace{X}_{n\times d}$ $$

where the 1st step is from Eq. (14), the 2nd step comes from Eq. (15), the 3rd step is because of Eq. (16), the 4th step is due to the tensor tricks.

∎

###### Lemma F.4 (Gradient of L ⁢ ( X ) 𝐿 𝑋 L(X) italic_L ( italic_X ) on W 𝑊 W italic_W ) .

###### Proof.

### F.4 Fast computation

Finally, we introduce the almost linear time algorithm and its error analysis of the gradient of $L(X)$ on $W$ in Lemma F.5.

###### Lemma F.5 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $m$ denote the number of self-attention transformer layers (see Definition 1.3).
- •
For any $i\in[m]$, let $W_{i}=W_{Q_{i}}W_{K_{i}}^{\top}$ denote the attention weight in the $i$-th transformer layer.

We can show that $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.
Namely, our algorithm can output $\widetilde{g}_{w}$ in $n^{1+o(1)}$ time, which satisfies

$$ $\displaystyle\|\widetilde{g}_{w}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}\|_{ \infty}\leq 1/\operatorname{poly}(n)$ $$

###### Proof.

Recall by Lemma C.15, C.16, we have defined $p_{1}(X),p_{2}(X)\in\mathbb{R}^{n\times n}$.

In those Lemmas, we have $p_{1}(X),p_{2}(X)$ have low rank approximation $U_{3}V_{3}^{\top}$ and $U_{4}V_{4}^{\top}$, respectively.

By the definition of $p(X)$ (Definition C.12), we have

$$ $\displaystyle p(X)=p_{1}(X)-p_{2}(X)$ (17) $$

Then, by Lemma F.4, we have

$$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}$ $\displaystyle=$ $\displaystyle~{}X^{\top}p(X)X$ $\displaystyle=$ $\displaystyle~{}X^{\top}(p_{1}(X)-p_{2}(X))X$ $$

where the 1st step is from Lemma F.4, the 2nd step comes from Eq. (17).

Let $\widetilde{p}_{1}(X),\widetilde{p}_{2}(X)$ denote the low rank approximations for $p_{1}(X),p_{2}(X)$, respectively.

Proof of running time.
We first compute $X^{\top}\widetilde{p}_{1}(X)X$ in following order

- •
Compute $\underbrace{X^{\top}}_{d\times n}\underbrace{U_{3}}_{n\times k_{3}}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{X^{\top}U_{3}}_{d\times k_{3}}\underbrace{V_{3}^{\top}}_{k_{3}
\times n}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{X^{\top}U_{3}V_{3}^{\top}}_{d\times n}\underbrace{X}_{n\times d}$, which takes $n^{1+o(1)}$ time.

The overall running time for $X^{\top}\widetilde{p}_{1}(X)X$ is $n^{1+o(1)}$.

Similarly, the overall running time for $X^{\top}\widetilde{p}_{2}(X)X$ is $n^{1+o(1)}$.

Since $X^{\top}\widetilde{p}_{1}(X)X,X^{\top}\widetilde{p}_{2}(X)X\in\mathbb{R}^{d
\times d}$, the computation time for $X^{\top}(\widetilde{p}_{1}(X)-\widetilde{p}_{2}(X))X$ is $O(d^{2})$.

Therefore, the overall running time for $X^{\top}(\widetilde{p}_{1}(X)-\widetilde{p}_{2}(X))X$ is $n^{1+o(1)}$.

Proof of error bound.

We consider the error for $X^{\top}\widetilde{p}_{1}(X)X$ first.

$$ $\displaystyle~{}\|X^{\top}\widetilde{p}_{1}(X)X-X^{\top}p_{1}(X)X\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|X^{\top}(\widetilde{p}_{1}(X)-p_{1}(X))X\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n^{2}\|X\|_{\infty}^{2}\|\widetilde{p}_{1}(X)-p_{1}(X)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n^{2}(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}^{2}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ (18) $$

where the 1st step is from basic algebra, the 2nd step comes from basic linear algebra, the 3rd step is because of $\|\widetilde{p}_{1}(X)-p_{1}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$, the 4th step is due to $\|X\|_{\infty}\leq\operatorname{poly}(n)$.

Similarly, we can have

$$ $\displaystyle\|X^{\top}\widetilde{p}_{2}(X)X-X^{\top}p_{2}(X)X\|_{\infty}\leq \epsilon/\operatorname{poly}(n)$ (19) $$

Therefore, we have

$$ $\displaystyle~{}\|X^{\top}\widetilde{p}(X)X-X^{\top}p(X)X\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|X^{\top}\widetilde{p}_{1}(X)X-X^{\top}p_{1}(X)X+X^{\top} \widetilde{p}_{2}(X)X-X^{\top}p_{2}(X)X\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|X^{\top}\widetilde{p}_{1}(X)X-X^{\top}p_{1}(X)X\|_{\infty}+ \|X^{\top}\widetilde{p}_{2}(X)X-X^{\top}p_{2}(X)X\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}(\epsilon/\operatorname{poly}(n))+(\epsilon/\operatorname{poly }(n))$ $\displaystyle=$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic algebra, the 2nd step comes from triangle inequality, the 3rd step is because of Eq. (F.4) and Eq. (19), the 4th step is due to basic algebra.

Then, we choose $\epsilon=1/\operatorname{poly}(n)$, we have

$$ $\displaystyle\|\widetilde{g}_{w}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}}\|_{ \infty}\leq 1/\operatorname{poly}(n)$ $$

∎

###### Lemma F.5 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ W i d 𝐿 𝑋 d subscript 𝑊 𝑖 \frac{\mathrm{d}L(X)}{\mathrm{d}W_{i}} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d italic_W start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG ) .

###### Proof.

## Appendix G Fast Computation for Gradient on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT

In Section G.1, we introduce the close form of the gradient of $s(X)$ on $W_{V}$.
In Section G.2, we provide the close form of the gradient of $L(X)$ on $W_{V}$.
In Section G.3, based on the close form calculated in the previous section, we introduce the almost linear time algorithm for computing the gradient of $L(X)$ on $W_{V}$.

### G.1 Gradient of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X ) on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT

Since $s(X)=f(X)h(X)$, we begin with considering the gradient of $h(X)$ on $W_{V}$ in Lemma G.1.

###### Lemma G.1 (Gradient of $h(X)$ on $W_{V}$ ) .

If we have the below conditions,

- •
Let $h(X)$ be defined as Definition C.9.
- •
Let $W_{V}$ be defined as Definition C.3.

Then, for any $i_{0}\in[n],j_{0}\in[d]$ and any $i_{1},j_{1}\in[d]$, we have

$$ $\displaystyle\frac{\mathrm{d}h(X)_{i_{0},j_{0}}}{\mathrm{d}(W_{V})_{i_{1},j_{1 }}}=\begin{cases}X_{i_{0},i_{1}}&~{}j_{0}=j_{1}\\ 0&~{}j_{0}\neq j_{1}\end{cases}$ $$

###### Proof.

Since $h_{i_{0},j_{0}}$ satisfies

$$ $\displaystyle h_{i_{0},j_{0}}=X_{i_{0},*}^{\top}(W_{V})_{*,j_{0}},$ $$

we have $h_{i_{0},j_{0}}$ only depends on $(W_{V})_{*,j_{0}}$.

Hence, we have, for $j_{0}\neq j_{1}$,

$$ $\displaystyle\frac{\mathrm{d}h(X)_{i_{0},j_{0}}}{\mathrm{d}(W_{V})_{i_{1},j_{1 }}}=0$ $$

For $j_{0}=j_{1}$ case, we have

$$ $\displaystyle\frac{\mathrm{d}h(X)_{i_{0},j_{0}}}{\mathrm{d}(W_{V})_{i_{1},j_{0 }}}=X_{i_{0},i_{1}}$ $$

∎

Combining the result in the previous Lemma and the chain rule, we can have the gradient of $s(X)$ on $W_{V}$ in Lemma G.2.

###### Lemma G.2 (Gradient of $s(X)$ on $W_{V}$ ) .

If we have the below conditions,

- •
Let $s(X)$ be defined as Definition C.10.
- •
Let $W_{V}$ be defined as Definition C.3.

Then, for any $i_{2}\in[n],j_{2}\in[d]$ and any $i_{1},j_{1}\in[d]$, we have

- •
Part 1.
$\displaystyle\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}(W_{V})_{i_{1},j_{1
}}}=\begin{cases}f(X)_{i_{2},*}^{\top}X_{*,i_{1}}&~{}j_{2}=j_{1}\\
0&~{}j_{2}\neq j_{1}\end{cases}$
- •
Part 2.
$\displaystyle\underbrace{\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}W_{V}}}
_{d\times d}=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_{2},*}}_{n
\times 1}\underbrace{e_{j_{2}}^{\top}}_{1\times d}$

###### Proof.

Proof of Part 1.

By Definition C.10, we have

$$ $\displaystyle s(X)_{i_{2},j_{2}}:=f(X)_{i_{2},*}^{\top}h(X)_{*,j_{2}}$ (20) $$

Therefore, $s(X)_{i_{2},j_{2}}$ is only depends on $h(X)_{*,j_{2}}$, which further means $s(X)_{i_{2},j_{2}}$ is only depends on $(W_{V})_{*,j_{2}}$.

Hence, for $j_{1}\neq j_{2}$, we have

$$ $\displaystyle\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}(W_{V})_{i_{1},j_{2 }}}=0$ $$

We consider $j_{1}=j_{2}$ case.

By, Eq. (20), we can derive that

$$ $\displaystyle\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}h(X)_{i_{3},j_{2}}} =f(X)_{i_{2},i_{3}}$ (21) $$

By chain rule, we have

$$ $\displaystyle~{}\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}(W_{V})_{i_{1},j _{2}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{3}=1}^{d}\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{ d}h(X)_{i_{3},j_{2}}}\frac{\mathrm{d}h(X)_{i_{3},j_{2}}}{\mathrm{d}(W_{V})_{i_ {1},j_{2}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{3}=1}^{d}f(X)_{i_{2},i_{3}}\frac{\mathrm{d}h(X)_{i_{3 },j_{2}}}{\mathrm{d}(W_{V})_{i_{1},j_{2}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{3}=1}^{d}f(X)_{i_{2},i_{3}}X_{i_{3},i_{1}}$ $\displaystyle=$ $\displaystyle~{}f(X)_{i_{2},*}^{\top}X_{*,i_{1}}$ (22) $$

where the 1st step is from chain rule, the 2nd step comes from Eq. (21), the 3rd step is because of Lemma G.1, the 4th step is due to basic linear algebra.

Proof of Part 2.

By Eq (G.1), we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}(W_{V}) _{*,j_{2}}}}_{d\times 1}=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_ {2},*}}_{n\times 1}$ $$

which implies

$$ $\displaystyle\underbrace{\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}W_{V}}} _{d\times d}=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_{2},*}}_{n \times 1}\underbrace{e_{j_{2}}^{\top}}_{1\times d}$ $$

∎

###### Lemma G.1 (Gradient of h ⁢ ( X ) ℎ 𝑋 h(X) italic_h ( italic_X ) on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT ) .

###### Proof.

###### Lemma G.2 (Gradient of s ⁢ ( X ) 𝑠 𝑋 s(X) italic_s ( italic_X ) on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT ) .

###### Proof.

### G.2 Gradient of L ⁢ ( X ) 𝐿 𝑋 L(X) italic_L ( italic_X ) on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT

Since we have already got the close form of the gradient of $s(X)$ on $W_{V}$, we can easily extend it and get the close form of the gradient of $L(X)$ on $W_{V}$ in Lemma G.3.

###### Lemma G.3 (Gradient of $L(X)$ on $W_{V}$ ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $W_{V}$ be defined as Definition C.3.

Then, we can show that

$$ $\displaystyle\underbrace{\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}}_{d\times d }=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)}_{n\times n}\underbrace{G_ {i}}_{n\times d}$ $$

###### Proof.

We slightly abuse the notation, using $W_{V}$ to represent $V_{i}$ in Lemma G.1, G.2.

By Lemma G.2, we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}s(X)_{i_{2},j_{2}}}{\mathrm{d}W_{V}}} _{d\times d}=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_{2},*}}_{n \times 1}\underbrace{e_{j_{2}}^{\top}}_{1\times d}$ (23) $$

By Lemma C.4, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}=\sum_{i_{2}=1}^{n}\sum _{j_{2}=1}^{d}G_{i}(i_{2},j_{2})\cdot\frac{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1} (X))_{i_{2},j_{2}}}{\mathrm{d}W_{V_{i}}}.$ (24) $$

By Definition C.10 and Definition C.3, we have

$$ $\displaystyle s(X)=\mathsf{Attn}_{i}(T_{i-1}(X))$ $$

Therefore, combining Eq. (23) and Eq. (24), we have

$$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{2}=1}^{n}\sum_{j_{2}=1}^{d}\underbrace{G_{i}(i_{2},j_ {2})}_{1\times 1}\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)_{i_{2},*}}_ {n\times 1}\underbrace{e_{j_{2}}^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{2}=1}^{n}\underbrace{X^{\top}}_{d\times n}\underbrace {f(X)_{i_{2},*}}_{n\times 1}\sum_{j_{2}=1}^{d}\underbrace{G_{i}(i_{2},j_{2})}_ {1\times 1}\underbrace{e_{j_{2}}^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{2}=1}^{n}\underbrace{X^{\top}}_{d\times n}\underbrace {f(X)_{i_{2},*}}_{n\times 1}\underbrace{G_{i}(i_{2},*)^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)}_{n\times n} \underbrace{G_{i}}_{n\times d}$ $$

where the 1st step is from Eq. (23) and Eq. (24), the 2nd step comes from basic algebra, the 3rd step is because of basic linear algebra, the 4th step is due to basic linear algebra.

∎

###### Lemma G.3 (Gradient of L ⁢ ( X ) 𝐿 𝑋 L(X) italic_L ( italic_X ) on W V subscript 𝑊 𝑉 W_{V} italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT ) .

###### Proof.

### G.3 Fast computation

Finally, we can introduce our almost linear time algorithm for computing the $L(X)$ gradient on $W_{V}$.

###### Lemma G.4 (Fast computation for $\frac{\mathrm{d}L(X)}{\mathrm{d}(W_{V})_{i}}$ ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $m$ denote the number of self-attention transformer layers (see Definition 1.3).
- •
For any $i\in[m]$, let $W_{V_{i}}\in\mathbb{R}^{d\times d}$ denote the attention weight in the $i$-th transformer layer.

We can show that $\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}$ can be approximated in $n^{1+o(1)}$ time, with $1/\operatorname{poly}(n)$ approximation error.
Namely, our algorithm can output $\widetilde{g}_{v}$ in $n^{1+o(1)}$ time, which satisfies

$$ $\displaystyle\|\widetilde{g}_{v}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}\|_ {\infty}\leq 1/\operatorname{poly}(n)$ $$

###### Proof.

Recall in Lemma C.13, $U_{1}V_{1}^{\top}$ is the low rank approximation of $f(X)$.

Let $\widetilde{f}(X):=U_{1}V_{1}^{\top}$ denote the low rank approximation of $f(X)$.

Recall in Lemma G.3, we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}}_{d\times d }=\underbrace{X^{\top}}_{d\times n}\underbrace{f(X)}_{n\times n}\underbrace{G_ {i}}_{n\times d}$ $$

Proof of running time.

We compute $X^{\top}\widetilde{f}(X)G_{i}$ in following order

- •
Compute $\underbrace{X^{\top}}_{d\times n}\cdot\underbrace{U_{1}}_{n\times k_{1}}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{X^{\top}\cdot U_{1}}_{d\times k_{1}}\cdot\underbrace{V_{1}^{\top}}
_{k_{1}\times n}$, which takes $n^{1+o(1)}$ time.
- •
Compute $\underbrace{X^{\top}\cdot U_{1}\cdot V_{1}^{\top}}_{d\times n}\cdot\underbrace
{G_{i}}_{n\times d}$, which takes $d^{2}\cdot n$ time.

The overall running time is $n^{1+o(1)}$.

Proof of error bound.

We have

$$ $\displaystyle~{}\|X^{\top}\cdot f(X)\cdot G_{i}-X^{\top}\cdot\widetilde{f}(X) \cdot G_{i}\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|X^{\top}\cdot(f(X)-\widetilde{f}(X))\cdot G_{i}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n^{2}\|X\|_{\infty}\|f(X)-\widetilde{f}(X)\|_{\infty}\|G_{i}\| _{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n^{2}(\epsilon/\operatorname{poly}(n))\|X\|_{\infty}\|G_{i}\|_ {\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is from basic algebra, the 2nd step comes from basic linear algebra, the 3rd step is because of $\|f(X)-\widetilde{f}(X)\|_{\infty}\leq\epsilon/\operatorname{poly}(n)$, the 4th step is due to $\|X\|_{\infty}\leq\operatorname{poly}(n)$ and $\|G_{i}\|_{\infty}\leq\operatorname{poly}(n)$.

Let $\widetilde{g}_{v}=X^{\top}\cdot\widetilde{f}(X)\cdot G_{i}$.

We choose $\epsilon=1/\operatorname{poly}(n)$. Then, we have

$$ $\displaystyle\|\widetilde{g}_{v}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{i}}}\|_ {\infty}\leq 1/\operatorname{poly}(n)$ $$

∎

###### Lemma G.4 (Fast computation for d ⁢ L ⁢ ( X ) d ⁢ ( W V ) i d 𝐿 𝑋 d subscript subscript 𝑊 𝑉 𝑖 \frac{\mathrm{d}L(X)}{\mathrm{d}(W_{V})_{i}} divide start_ARG roman_d italic_L ( italic_X ) end_ARG start_ARG roman_d ( italic_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG ) .

###### Proof.

## Appendix H Gradient Approximation for Entire Model

In Section H.1, we introduce the close form of $G_{i}$ and argue that $G_{i}$ can be computed in almost linear time $n^{1+o(1)}$.
In Section H.2, we provide the almost linear time algorithm for gradient computing on a single-layer transformer.
In Section H.3, with the help of math induction, we introduce the almost linear time algorithm for computing the gradient of the multi-layer transformer, along with its approximation error.

### H.1 Computation time for G i subscript 𝐺 𝑖 G_{i} italic_G start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT

Here we consider $g_{i}$ in Definition 1.3 as a linear layer with an arbitrary non-linear activation $\phi$.
Since $g_{i}$ can be viewed as a composition of an MLP and an activation function, we begin with analyzing the $T_{i}$ gradient on $\mathsf{Attn_{i}}$.

###### Lemma H.1 (Gradient of $T_{i}$ on $\mathsf{Attn}_{i}$ ) .

If we have the below conditions,

- •
Let $T_{i}(X)$ be defined as Definition 3.3.
- •
Assuming for any $Z\in\mathbb{R}^{n\times d}$, we have $g_{i}(Z)\in\mathbb{R}^{n\times d}$, and $g_{i}(Z)=\phi(ZW_{g})$, where $W_{g}\in\mathbb{R}^{d\times d}$ and $\phi:\mathbb{R}\rightarrow\mathbb{R}$ denotes any element-wise activation function. Let $\phi^{\prime}$ denote the derivative of $\phi$.
- •
We simplify the notation, using $T_{i}$ and $\mathsf{Attn}_{i}$ to represent $T_{i}(X)$ and $\mathsf{Attn}_{i}(T_{i-1}(X))$, respectively.
- •
For any matrix $Z\in\mathbb{R}^{n\times d}$, we use $Z(i,j)$ to denote the $(i,j)$-th entry of $Z$.

Then, we can show that, for any $i_{4},i_{5}\in[n],j_{4},j_{5}\in[d]$,

- •
Part 1.
$\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i
_{5},j_{5})}=\begin{cases}\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)
^{\top}W_{g}(*,j_{4}))}_{1\times 1}\underbrace{W_{g}(j_{5},j_{4})}_{1\times 1}
&~{}i_{4}=i_{5}\\
0&~{}i_{4}\neq i_{5}\end{cases}$
- •
Part 2.
$\displaystyle\underbrace{\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf
{Attn}_{i}}}_{n\times d}=\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^
{\top}W_{g}(*,j_{4}))}_{1\times 1}\underbrace{e_{i_{4}}}_{n\times 1}
\underbrace{W_{g}(*,j_{4})^{\top}}_{1\times d}$

###### Proof.

Proof of Part 1.

By the definition of $T_{i}$ (Definition 3.3), for $i_{4}\in[d],j_{4}\in[n]$, we have

$$ $\displaystyle T_{i}(i_{4},j_{4})=\phi(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}(* ,j_{4}))$ $$

Therefore, for any $i_{5}\neq i_{4}$, we have

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i _{5},j_{5})}=0$ $$

Then, we consider $i_{4}=i_{5}$ case.

By basic calculus, we have

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i _{4},j_{5})}=\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}( *,j_{4}))}_{1\times 1}\underbrace{W_{g}(j_{5},j_{4})}_{1\times 1}$ $$

Combining two equations mentioned above, we have the result for Part 1.

Proof of Part 2.

By result of Part 1, for $i_{5}=i_{4}$, we have

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i _{4},j_{5})}=\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}( *,j_{4}))}_{1\times 1}\underbrace{W_{g}(j_{5},j_{4})}_{1\times 1}$ $$

which implies

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i _{4},*)}=\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}(*,j_ {4}))}_{1\times 1}\underbrace{W_{g}(*,j_{4})}_{d\times 1}$ $$

By result of Part 1, for $i_{5}\neq i_{4}$, we have

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}(i _{5},*)}=0$ $$

By basic linear algebra, combining the two equations mentioned above, we have

$$ $\displaystyle\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}}= \underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}(*,j_{4}))}_{1 \times 1}\underbrace{e_{i_{4}}}_{n\times 1}\underbrace{W_{g}(*,j_{4})^{\top}}_ {1\times d}$ $$

∎

Then, we can argue that the computation for $G_{i}$ can be done in almost linear time $n^{1+o(1)}$.

###### Lemma H.2 (Computation time for $G_{i}$ , formal version of Lemma 5.4 ) .

If we have the below conditions,

- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
Assuming we already have $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}(X)}$.
- •
Assuming for any $Z\in\mathbb{R}^{n\times d}$, we have $g_{i}(Z)\in\mathbb{R}^{n\times d}$, and $g_{i}(Z)=\phi(ZW_{g})$, where $W_{g}\in\mathbb{R}^{d\times d}$ and $\phi:\mathbb{R}\rightarrow\mathbb{R}$ denotes any element-wise activation function. Let $\phi^{\prime}$ denote the derivative of $\phi$.
- •
We simplify the notation, using $T_{i}$ and $\mathsf{Attn}_{i}$ to represent $T_{i}(X)$ and $\mathsf{Attn}_{i}(T_{i-1}(X))$, respectively.
- •
For any matrix $Z\in\mathbb{R}^{n\times d}$, we use $Z(i,j)$ to denote the $(i,j)$-th entry of $Z$.

Then, we can show that $G_{i}$ can be computed in $n^{1+o(1)}$ time.

###### Proof.

Let $g_{T_{i}}:=\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}}$, and for any $i_{4}\in[n],j_{4}\in[d]$, let $g_{T_{i}}(i_{4},j_{4})$ denote the $(i_{4},j_{4})$-th entry of $g_{T_{i}}$.

Similarly, for any $i_{5}\in[n],j_{5}\in[d]$, let $T_{i}(i_{5},j_{5})$ denote the $(i_{5},j_{5})$-th entry of $T_{i}$.

We can have

$$ $\displaystyle G_{i}=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}T_{i}}\cdot\frac{\mathrm{d}T_{ i}}{\mathrm{d}\mathsf{Attn}_{i}}$ $\displaystyle=$ $\displaystyle~{}g_{T_{i}}\cdot\frac{\mathrm{d}T_{i}}{\mathrm{d}\mathsf{Attn}_{ i}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{4}=1}^{n}\sum_{j_{4}=1}^{d}g_{T_{i}}(i_{4},j_{4}) \cdot\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}}$ $$

where the 1st step is from the definition of $G_{i}$, the 2nd step comes from chain rule, the 3rd step is because of the definition of $g_{T_{i}}$, the 4th step is due to chain rule.

$$ $\displaystyle~{}\sum_{i_{4}=1}^{n}\sum_{j_{4}=1}^{d}g_{T_{i}}(i_{4},j_{4}) \cdot\frac{\mathrm{d}T_{i}(i_{4},j_{4})}{\mathrm{d}\mathsf{Attn}_{i}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{4}=1}^{n}\sum_{j_{4}=1}^{d}\underbrace{g_{T_{i}}(i_{4 },j_{4})}_{1\times 1}\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{ \top}W_{g}(*,j_{4}))}_{1\times 1}\underbrace{e_{i_{4}}}_{n\times 1}\underbrace {W_{g}(*,j_{4})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{4}=1}^{n}\underbrace{e_{i_{4}}}_{n\times 1}\sum_{j_{4 }=1}^{d}\underbrace{g_{T_{i}}(i_{4},j_{4})}_{1\times 1}\underbrace{\phi^{ \prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g}(*,j_{4}))}_{1\times 1} \underbrace{W_{g}(*,j_{4})^{\top}}_{1\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{4}=1}^{n}\underbrace{e_{i_{4}}}_{n\times 1}( \underbrace{W_{g}}_{d\times d}(\underbrace{g_{T_{i}}(i_{4},*)}_{d\times 1} \odot\underbrace{\phi^{\prime}(\mathsf{Attn}_{i}(i_{4},*)^{\top}W_{g})}_{d \times 1}))^{\top}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(g_{T_{i}}\odot\phi^{\prime}(\mathsf{Attn}_{i}W_{g }))}_{n\times d}\underbrace{W_{g}^{\top}}_{d\times d}$ (25) $$

where the 1st step is from Lemma H.1, the 2nd step comes from basic algebra, the 3rd step is because of basic linear algebra, the 4th step is due to basic linear algebra.

By Eq. (H.1), we have the close form of $G_{i}$.

We can compute $G_{i}$ in the following order

- •
Compute $\underbrace{(g_{T_{i}}\odot\phi^{\prime}(\mathsf{Attn}_{i}W_{g}))}_{n\times d}$, which takes $n\cdot d$ time.
- •
Compute $\underbrace{(g_{T_{i}}\odot\phi^{\prime}(\mathsf{Attn}_{i}W_{g}))}_{n\times d}
\underbrace{W_{g}^{\top}}_{d\times d}$, which takes $d^{2}\cdot n$ time.

Therefore, the overall running time for $G_{i}$ is $n^{1+o(1)}$.

∎

###### Lemma H.1 (Gradient of T i subscript 𝑇 𝑖 T_{i} italic_T start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT on 𝖠𝗍𝗍𝗇 i subscript 𝖠𝗍𝗍𝗇 𝑖 \mathsf{Attn}_{i} sansserif_Attn start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) .

###### Proof.

###### Lemma H.2 (Computation time for G i subscript 𝐺 𝑖 G_{i} italic_G start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT , formal version of Lemma 5.4 ) .

###### Proof.

### H.2 Fast computation for single-layer transformer

In this section, we dive into the computation time and approximation error of the gradient of a single-layer transformer.
We demonstrate in the following Lemma that the gradient of a single-layer transformer can be computed in almost linear time $n^{1+o(1)}$, and its error can be bounded by $1/\operatorname{poly}(n)$.

###### Lemma H.3 (Single-layer transformer gradient approximation) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $X$ be defined as Definition C.3.
- •
Let the gradient matrix $G_{i}\in\mathbb{R}^{n\times d}$ be defined as $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
Assuming for any $Z\in\mathbb{R}^{n\times d}$, we have $g_{i}(Z)\in\mathbb{R}^{n\times d}$, and $g_{i}(Z)=\phi(Z\cdot W_{g})$, where $W_{g}\in\mathbb{R}^{d\times d}$ and $\phi:\mathbb{R}\rightarrow\mathbb{R}$ denotes any element-wise activation function. Let $\phi^{\prime}$ denote the derivative of $\phi$.
- •
Suppose we have a single-layer transformer (see Definition 1.3).

Then, we can show that,

- •
Part 1: running time. Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $n^{1+o(1)}$ time.
- •
Part 2: error bound. The approximation error of the single-layer transformer can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}_{1}$ satisfies
$\displaystyle\|\widetilde{g}_{1}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}
\leq 1/\operatorname{poly}(n)$

###### Proof.

By Definition 1.3, a single-layer transformer has following structure:

$$ $\displaystyle g_{1}\circ\mathsf{Attn}_{1}\circ g_{0}(X)$ $$

By the definition of $G_{i}$, we have

$$ $\displaystyle G_{1}=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{1}(T_{0}(X))}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}T_{1}(X)}\cdot\frac{\mathrm{d} T_{1}(X)}{\mathrm{d}\mathsf{Attn}_{1}(T_{0}(X))}$ (26) $$

By Lemma H.2, we have $G_{1}$ can be computed in $n^{1+o(1)}$ time.

Proof of Part 1: running time.

For less confusion, in this part of the proof, we ignore the approximation error temporarily.

Since we have got $G_{1}$, we use methods mentioned in Lemma E.11, F.5, G.4 to compute $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}(X)},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{1
}},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{1}}}$, respectively, which takes $n^{1+o(1)}$ time for each.

Then, since we have $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}(X)}$, again by Lemma H.2, we have $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ can be computed in $n^{1+o(1)}$ time.

Therefore, the overall running time is $n^{1+o(1)}$.

Proof of Part 2: error bound.

Then, we move on to the error bound.

By Lemma H.2 and Eq. (H.2), there is no approximation error when computing $G_{1}$.

By Lemma E.11, F.5, G.4, we have there is $1/\operatorname{poly}(n)$ approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}(X)},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{1
}},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{1}}}$, respectively.

Let $\widetilde{g}_{t_{0}},\widetilde{g}_{w_{1}},\widetilde{g}_{v_{1}}$ denote the approximation results of $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}(X)},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{1
}},\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{1}}}$, respectively.

We have

$$ $\displaystyle\|\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}(X)} \|_{\infty}\leq 1/\operatorname{poly}(n)$ (27) $$

and

$$ $\displaystyle\|\widetilde{g}_{w_{1}}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{1}}\|_ {\infty}\leq 1/\operatorname{poly}(n)$ $$

and

$$ $\displaystyle\|\widetilde{g}_{v_{1}}-\frac{\mathrm{d}L(X)}{\mathrm{d}W_{V_{1}} }\|_{\infty}\leq 1/\operatorname{poly}(n)$ $$

Let $\widetilde{G}_{0}=\widetilde{g}_{t_{0}}\cdot\frac{\mathrm{d}T_{0}(X)}{\mathrm{
d}X}$ denote the approximated version of $G_{0}$.

We have

$$ $\displaystyle~{}\|\widetilde{G}_{0}-G_{0}\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0} (X)})\cdot\frac{\mathrm{d}T_{0}(X)}{\mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\cdot d\|\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{ d}T_{0}(X)}\|_{\infty}\|\frac{\mathrm{d}T_{0}(X)}{\mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\cdot d(1/\operatorname{poly}(n))\|\frac{\mathrm{d}T_{0}(X)}{ \mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}1/\operatorname{poly}(n)$ $$

where the 1st step is from the definition of $\widetilde{G}_{0}$, the 2nd step comes from basic linear algebra, the 3rd step is because of Eq. (27), the 4th step is due to each entry can be written by $O(\log n)$ bits.

Let $\widetilde{g}_{1}=\widetilde{G}_{0}$.

Therefore, we have

$$ $\displaystyle\|\widetilde{g}_{1}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty} \leq 1/\operatorname{poly}(n)$ $$

∎

###### Lemma H.3 (Single-layer transformer gradient approximation) .

###### Proof.

### H.3 Fast computation for multi-layer transformer

Since we have already demonstrated that almost linear time gradient computation can be applied to a single-layer transformer, with the help of math induction, we can easily generalize that result to the multi-layer transformer.
In the following Lemma, we display that the gradient of the multi-layer transformer can be computed in almost linear time, and its approximation error can be bounded by $1/\operatorname{poly}(n)$.

###### Lemma H.4 (Multi-layer transformer gradient approximation, formal version of Lemma 5.5 ) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $X$ be defined as Definition C.3.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
For $i_{2}\in[n],j_{2}\in[d]$, let $G_{i}(i_{2},j_{2})$ denote the $(i_{2},j_{2})$-th entry of $G_{i}$.
- •
Let gradient components for each layer be computed according to Lemma E.11, F.5, G.4.
- •
Assuming for any $Z\in\mathbb{R}^{n\times d}$, we have $g_{i}(Z)\in\mathbb{R}^{n\times d}$, and $g_{i}(Z)=\phi(Z\cdot W_{g})$, where $W_{g}\in\mathbb{R}^{d\times d}$ and $\phi:\mathbb{R}\rightarrow\mathbb{R}$ denotes any element-wise activation function. Let $\phi^{\prime}$ denote the derivative of $\phi$.
- •
Suppose we have a $m$-layer transformer (see Definition 1.3).

Then, we can show that,

- •
Part 1: running time. Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $n^{1+o(1)}$ time.
- •
Part 2: error bound. The approximation error of the multi-layer transformer can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}$ satisfies
$\displaystyle\|\widetilde{g}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}\leq
1
/\operatorname{poly}(n)$

###### Proof.

We use math induction to prove this Lemma.

Step 1: Proof of a single-layer transformer.

Firstly, by Lemma H.3, we have that for one-layer transformer, our conclusion is established.

Step 2: Assumption for $k$-layer transformer.

Secondly, we assume for any $k$, for $k$-layer transformer model, we have

- •
Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $O(n^{1+o(1)})$ time.
- •
The approximation error of the $k$-layer transformer can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}$ satisfies
$\displaystyle\|\widetilde{g}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}\leq
1
/\operatorname{poly}(n)$

Step 3: Proof of $(k+1)$-layer transformer.

Thirdly, we consider the $(k+1)$-layer transformer model.

Without loss of generality, we assume that the additional transformer layer is added at the beginning of the model.

Namely, let $\mathsf{F}_{k}$ denote a $k$-layer transformer model.
We have

$$ $\displaystyle\mathsf{F}_{k}(X)=g_{k}\circ\mathsf{Attn}_{k}\circ\cdots\circ g_{ 1}\circ\mathsf{Attn}_{1}\circ g_{0}(X)$ $$

Let the $(k+1)$-layer transformer model have the following structure:

$$ $\displaystyle\mathsf{F}_{k+1}(X)=\mathsf{F}_{k}\circ\mathsf{Attn}\circ g(X)$ (28) $$

Let $T_{0}:=g(X)$.

By assumption, we have

- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}(T_{0})}$ can be approximated in $n^{1+o(1)}$ time.
- •
Let $\widetilde{g}_{k}$ denote the approximated version of $\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}(T_{0})}$. We have
$\displaystyle\|\widetilde{g}_{k}-\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}
(T_{0})}\|_{\infty}\leq 1/\operatorname{poly}(n)$
(29)

Step 3.1: Proof of the running time for $(k+1)$-layer transformer

For less confusion, in this part of the proof, we ignore the approximation error temporarily.

By the assumption, we have $\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}(T_{0})}$ can be approximated in $n^{1+o(1)}$ time.

We compute $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in following order:

- •
Since we already have $\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}(T_{0})}$, by Lemma E.11, the computation time for $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}}$ is $n^{1+o(1)}$.
- •
Since we have $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}}$, by Lemma H.2, the computation time for $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ is $n^{1+o(1)}$.

Therefore, for $(k+1)$-layer transformer, the overall running time for $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ is $n^{1+o(1)}$.

Step 3.2: Proof of the error bound for $(k+1)$-layer transformer

By Lemma E.11, during the process of solving the approximated version of $\frac{\mathrm{d}L(X)}{\mathrm{d}g(X)}$, the approximation error will not be magnified by more than $\operatorname{poly}(n)$.

Let $\widetilde{g}_{t_{0}}$ denote the approximated version of $\frac{\mathrm{d}L(X)}{\mathrm{d}g(X)}$, we have

$$ $\displaystyle~{}\|\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{d}g(X)} \|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\operatorname{poly}(n)\|\widetilde{g}_{k}-\frac{\mathrm{d}L(X) }{\mathrm{d}T(X)}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}1/\operatorname{poly}(n)$ (30) $$

where the 1st step is from the above statement, the 2nd step comes from Eq. (29), the 3rd step is because of basic algebra.

Then, we consider

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}X}=\frac{\mathrm{d}L(X)}{\mathrm{ d}g(X)}\cdot\frac{\mathrm{d}g(X)}{\mathrm{d}X}$ (31) $$

Recall that we have $\widetilde{g}=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$.
Then, we have

$$ $\displaystyle~{}\|\widetilde{g}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{d}g(X)} )\cdot\frac{\mathrm{d}g(X)}{\mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\cdot d\|\widetilde{g}_{t_{0}}-\frac{\mathrm{d}L(X)}{\mathrm{ d}g(X)}\|_{\infty}\|\frac{\mathrm{d}g(X)}{\mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}n\cdot d(1/\operatorname{poly}(n))\|\frac{\mathrm{d}g(X)}{ \mathrm{d}X}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}1/\operatorname{poly}(n)$ $$

where the 1st step is from Eq. (31), the 2nd step comes from basic linear algebra, the 3rd step is because of Eq. (H.3), the 4th step is due to each entry can be written by $O(\log n)$ bits.

Step 4: Use math induction.

So far, with the assumption that our statement holds under $k$-layer transformer, we have proved that our statement still holds under $(k+1)$-layer transformer.

Therefore, by math induction, our statement holds for any $m$-layer transformer.

∎

###### Lemma H.4 (Multi-layer transformer gradient approximation, formal version of Lemma 5.5 ) .

###### Proof.

## Appendix I Causal Attention Mask

This section will discuss how to combine the causal attention mask with our framework.
We argue that even with the causal attention mask, we can also achieve almost linear time gradient computing for the multi-layer transformer.

In Section I.1, we introduce essential tools from literature to deal with the causal mask added on the attention matrix.
In Section I.2, we show that with the addition of causal mask, our framework can still achieve almost linear time gradient computation.

### I.1 Tools from previous work

Firstly, we restate a classical low-rank approximation method in the literature.

###### Lemma I.1 (Low-rank approximation, [ 5 ] ) .

Suppose $Q,K\in\mathbb{R}^{n\times d}$, with $\|Q\|_{\infty}\leq R$, and $\|K\|_{\infty}\leq R$. Let $A:=\exp(QK^{\top}/d)\in\mathbb{R}^{n\times n}$. For accuracy parameter $\epsilon\in(0,1)$, there is a positive integer $g$ bounded above by

$$ $\displaystyle g=O\Big{(}\max\Big{\{}\frac{\log(1/\epsilon)}{\log(\log(1/ \epsilon)/R)},R^{2}\Big{\}}\Big{)},$ $$

and a positive integer $r$ bounded above by

$$ $\displaystyle r\leq{2(g+d)\choose 2g}$ $$

such that:
There is a matrix $\widetilde{A}\in\mathbb{R}^{n\times n}$ that is an $(\epsilon,r)$-approximation of $A\in\mathbb{R}^{n\times n}$.
Furthermore, the matrices $U_{0}$ and $V_{0}$ defining $\widetilde{A}$ can be computed in $O(n\cdot r)$ time.

Then, we provide the formal definition for the causal attention mask.

###### Definition I.2 (Causal attention mask, [ 57 ] ) .

We define the causal attention mask as $M\in\{0,1\}^{n\times n}$, where $M_{i,j}=1$ if $i\geq j$ and $M_{i,j}=0$ otherwise.

Figure: Algorithm 2 Causal attention mask algorithm, Algorithm 4 in LLS^+24d []

In previous work [], they point out there exists an algorithm (Algorithm 2) that can calculate low-rank matrices (with the causal attention mask) multiplication with any vector $v$ in almost linear time.
We restate their results in Lemma I.3.

###### Lemma I.3 (Fast computation for causal attention mask on tensor, [ 57 ] ) .

Let $M\in\{0,1\}^{n\times n}$ be a causal attention mask defined in Definition I.2.
Let $U_{0},V_{0}\in\mathbb{R}^{n\times k}$.
Let $v\in\mathbb{R}^{n}$.
Then, there exists an algorithm (see Algorithm 2) whose output satisfies that

$$ $\displaystyle Y=(M\odot(U_{0}V_{0}^{\top}))v,$ $$

which takes $O(nk)$ time.

We extend their results to the multiplication of matrix with $n^{o(1)}$ columns.

###### Lemma I.4 (Fast computation for causal attention mask on matrix) .

If we have the below conditions,

- •
Let $M\in\{0,1\}^{n\times n}$ be a causal attention mask defined in Definition I.2.
- •
Let $U_{0},V_{0}\in\mathbb{R}^{n\times k}$ where $k=n^{o(1)}$.
- •
Let $H\in\mathbb{R}^{n\times k_{H}}$ where $k_{H}=n^{o(1)}$.

Then, there exists an algorithm, whose output satisfies that

$$ $\displaystyle Z=(M\odot(U_{0}V_{0}^{\top}))H,$ $$

which takes $n^{1+o(1)}$ time.

###### Proof.

For $j\in[k_{H}]$, let $H_{*,j}\in\mathbb{R}^{n}$ denote the $j$-th column of $H$.

By Lemma I.3, we can compute $(M\odot(U_{0}V_{0}^{\top}))H_{*,j}$ in $O(nk)$ time.

There are $k_{H}$ columns in total.
Therefore, the overall running time is $O(nkk_{H})=O(n\cdot n^{o(1)}\cdot n^{o(1)})=n^{1+o(1)}$.
∎

###### Lemma I.1 (Low-rank approximation, [ 5 ] ) .

###### Definition I.2 (Causal attention mask, [ 57 ] ) .

###### Lemma I.3 (Fast computation for causal attention mask on tensor, [ 57 ] ) .

###### Lemma I.4 (Fast computation for causal attention mask on matrix) .

###### Proof.

### I.2 Fast computation with causal mask

We can easily change all low-rank matrices multiplication to the algorithm mentioned in Lemma I.4.
Then, our framework can support the causal attention mask and still achieves almost linear time gradient computing for the multi-layer transformer.

The causal mask directly affects the attention matrix, so it’s necessary to define the attention matrix with the causal mask applied.

###### Definition I.5 .

Let $M\in\{0,1\}^{n\times n}$ be a causal attention mask defined in Definition I.2.
We define attention matrix with causal mask as:

$$ $\displaystyle\widehat{f}(X):=D^{-1}(M\odot A)$ $$

where $A:=\exp(XWX^{\top}/d)$ and $D:=\operatorname{diag}((M\odot A)\cdot{\bf 1}_{n})$.

After analyzing the components of gradients on $T_{i}(X),W_{i},W_{V_{i}}$ in Section E, F and G, we categorize them into two groups: one involving the dot product and the other involving the Hadamard product of the attention matrix.
Then, we can show $\widehat{f}(X)H$ and $(\widehat{f}(X)\odot(UV^{\top}))H$ for low rank matrices $U,V,H$ can be approximated in almost linear time.

###### Lemma I.6 .

If we have the below conditions,

- •
Let $\widehat{f}(X)$ be defined in Definition I.5.
- •
Let $U,V\in\mathbb{R}^{n\times k}$ where $k=n^{o(1)}$.
- •
Let $H\in\mathbb{R}^{n\times k_{H}}$ where $k_{H}=n^{o(1)}$.

Then, approximating the following takes $n^{1+o(1)}$ time:

- •
Part 1. $\widehat{f}(X)H$
- •
Part 2. $(\widehat{f}(X)\odot(UV^{\top}))H$

###### Proof.

From Definition I.5, we know

$$ $\displaystyle\widehat{f}(X):=D^{-1}(M\odot A)$ $$

where $D:=\operatorname{diag}((M\odot A)\cdot{\bf 1}_{n})$.

By Lemma I.1, $U_{0}V_{0}^{\top}$ is a good approximation for $A$.
Then, we can approximate $\widehat{f}(X)$ by:

$$ $\displaystyle D^{-1}(M\odot(U_{0}V_{0}^{\top}))$ $$

where $D:=\operatorname{diag}((M\odot(U_{0}V_{0}^{\top}))\cdot{\bf 1}_{n})$.

Using Lemma I.3, we know $(M\odot(U_{0}V_{0}^{\top}))\cdot v$ for any vector $v\in\mathbb{R}^{n}$ can be computed in almost linear time.

We begin by examining the normalization matrix $D^{-1}$.
Calling Lemma I.3, we compute $(M\odot(U_{0}V_{0}^{\top}))\cdot{\bf 1}_{n}$ in almost linear time.
Then, it takes $O(n)$ time to make $(M\odot(U_{0}V_{0}^{\top}))\cdot{\bf 1}_{n}$ diagonal.
Given that $D$ is diagonal, its inverse $D^{-1}$ can be determined in $O(n)$ time.
Thus, we can compute $D^{-1}$ in almost linear time.

Proof of Part 1.
$H$ can be viewed as a combination of $k_{H}$ vectors, each of size $n$.
Calling Lemma I.4, we can compute $(M\odot(U_{0}V_{0}^{\top}))H$ in $n^{1+o(1)}$ time.

Finally, we compute $\underbrace{D^{-1}}_{n\times n}\underbrace{(M\odot(U_{0}V_{0}^{\top}))H}_{n
\times k_{H}}$, which takes $n^{1+o(1)}$ time since $D^{-1}$ is diagonal.
The overall gradient computation remains $n^{1+o(1)}$ time.

Proof of Part 2.
The proof for this part involves Fact C.2.
We can show

$$ $\displaystyle~{}((D^{-1}(M\odot(U_{0}V_{0}^{\top})))\odot(UV^{\top}))H$ $\displaystyle=$ $\displaystyle~{}((M\odot(D^{-1}U_{0}V_{0}^{\top}))\odot(UV^{\top}))H$ $\displaystyle=$ $\displaystyle~{}(M\odot((D^{-1}U_{0}V_{0}^{\top})\odot(UV^{\top})))H$ $\displaystyle=$ $\displaystyle~{}(M\odot((D^{-1}U_{0})\oslash U)(V_{0}\oslash V)^{\top})H$ $$

where the 1st step is from $D(A\odot B)=(DA)\odot B=A\odot(DB)$ for diagonal matrix $D\in\mathbb{R}^{m\times m}$ and $A,B\in\mathbb{R}^{m\times n}$, the 2nd step comes from $(A\odot B)\odot C=A\odot(B\odot C)$ for $A,B,C\in\mathbb{R}^{m\times n}$, and the last step follows from Fact C.2.

Let $U_{M}:=(D^{-1}U_{0})\oslash U$ and $V_{M}:=V_{0}\oslash V$.

For $U_{M}$, we compute $\underbrace{D^{-1}}_{n\times n}\underbrace{U_{0}}_{n\times k}$ which takes $nk$ time.
We then compute $\underbrace{(D^{-1}U_{0})}_{n\times k}\oslash\underbrace{U}_{n\times k}$ which takes $O(nk^{2})$ time.

For $V_{M}$, we compute $\underbrace{V_{0}}_{n\times k}\oslash\underbrace{V}_{n\times k}$ which takes $O(nk^{2})$ time.

We now have $(M\odot(U_{M}V_{M}^{\top})H$. Calling Lemma I.4, we finish the proof.
∎

We now prove for gradient components that have dot product.

###### Lemma I.7 (Components for dot product) .

If we have the below conditions,

- •
Let $\widehat{f}(X)$ be defined in Definition I.5.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
Let $D_{6}=-f(X)\operatorname{diag}(K)XW^{\top}$ be defined in Lemma D.17.
- •
Let $D_{2}=-\operatorname{diag}(K)f(X)XW$ be defined in Lemma D.17.
- •
Let $D_{8}=f(X)G_{i}W_{V}^{\top}$ be defined in Lemma D.17.
- •
Let $g_{v}:=X^{\top}f(X)G_{i}$ be the gradient on $W_{V_{i}}$ and defined in Lemma G.3.

Then, we can show the following can be approximated in almost linear time:

- •
Part 1. $\widehat{D}_{6}=-\widehat{f}(X)\operatorname{diag}(K)XW^{\top}$
- •
Part 2. $\widehat{D}_{2}=-\operatorname{diag}(K)\widehat{f}(X)XW$
- •
Part 3. $\widehat{D}_{8}=\widehat{f}(X)G_{i}W_{V}^{\top}$
- •
Part 4. $\widehat{g}_{v}:=X^{\top}\widehat{f}(X)G_{i}$

###### Proof.

Proof of Part 1.
For $\widehat{D}_{6}$, we compute $\underbrace{\operatorname{diag}(K)}_{n\times n}\underbrace{X}_{n\times d}$ first, which takes $nd$ time.

Then, we compute $\underbrace{\widehat{f}(X)}_{n\times n}\underbrace{\operatorname{diag}(K)X}_{n
\times d}$ using Part 1. of Lemma I.6, which takes $n^{1+o(1)}$ time.

Finally, we compute $\underbrace{\widehat{f}(X)\operatorname{diag}(K)X}_{n\times d}\underbrace{W^{
\top}}_{d\times d}$, which takes $n^{1+o(1)}$ time.

Proof of Part 2.
For $\widehat{D}_{2}$,
we compute $\underbrace{\widehat{f}(X)}_{n\times n}\underbrace{X}_{n\times d}$ using Part 1. of Lemma I.6, which takes $n^{1+o(1)}$ time.

Then, we compute $\underbrace{\operatorname{diag}(K)}_{n\times n}\underbrace{\widehat{f}(X)X}_{n
\times d}$, which takes $nd$ time.

After that, we compute $\underbrace{\operatorname{diag}(K)\widehat{f}(X)X}_{n\times d}\underbrace{W}_{
d\times d}$, which takes $n^{1+o(1)}$ time.

Proof of Part 3.
For $\widehat{D}_{8}$, we compute in the following steps:

We compute $\underbrace{\widehat{f}(X)}_{n\times n}\underbrace{G_{i}}_{n\times d}$ using Part 1. of Lemma I.6, which takes $n^{1+o(1)}$ time.

Then, we compute $\underbrace{\widehat{f}(X)G_{i}}_{n\times d}\underbrace{W_{V}^{\top}}_{d\times
d}$, which takes $n\cdot d^{2}$ time.

Proof of Part 4.
For $\widehat{g}_{v}$, we compute in the following steps:

We compute $\underbrace{\widehat{f}(X)}_{n\times n}\underbrace{G_{i}}_{n\times d}$ using Part 1. of Lemma I.6, which takes $n^{1+o(1)}$ time.

Then, we compute $\underbrace{X^{\top}}_{d\times n}\underbrace{\widehat{f}(X)G_{i}}_{n\times d}$, which takes $n\cdot d^{2}$ time.
∎

We then prove for gradient components that have Hadamard product.

###### Lemma I.8 (Components for Hadamard product) .

If we have the below conditions,

- •
Let $\widehat{f}(X)$ be defined in Definition I.5.
- •
Let $G_{i}\in\mathbb{R}^{n\times d}$ denote the gradient matrix resulting from the application of the chain rule up to the function $g_{i}$, i.e., $G_{i}=\frac{\mathrm{d}L(X)}{\mathrm{d}\mathsf{Attn}_{i}(T_{i-1}(X))}$.
- •
Let $D_{7}=(f(X)\odot(h(X)G_{i}^{\top}))XW^{\top}$ be defined in Lemma D.17.
- •
Let $D_{4}=(f(X)\odot(G_{i}h(X)^{\top}))XW$ be defined in Lemma D.17.
- •
Let $g_{w}:=X^{\top}p(X)X=X^{\top}(p_{1}(X)-p_{2}(X))X$ be the gradient on $W_{i}$ and defined in Definition C.12 and Lemma F.5 where $p_{1}(X)=f(X)\odot q(X)$ and $p_{2}(X)=\operatorname{diag}(p_{1}(X)\cdot{\bf 1}_{n})f(X)$.

Then, we can show the following can be approximated in almost linear time:

- •
Part 1. $\widehat{D}_{7}=(\widehat{f}(X)\odot(h(X)G_{i}^{\top}))XW^{\top}$
- •
Part 2. $\widehat{D}_{4}=(\widehat{f}(X)\odot(G_{i}h(X)^{\top}))XW$
- •
Part 3. $\widehat{g}_{w}:=X^{\top}(\widehat{p}_{1}(X)-\widehat{p}_{2}(X))X$ where $\widehat{p}_{1}(X)=\widehat{f}(X)\odot q(X)$ and $p_{2}(X)=\operatorname{diag}(\widehat{p}_{1}(X)\cdot{\bf 1}_{n})\widehat{f}(X)$.

###### Proof.

Proof of Part 1.
For $\widehat{D}_{7}$, we can compute $\underbrace{(\widehat{f}(X)\odot(h(X)G_{i}^{\top}))}_{n\times n}\underbrace{X}
_{n\times d}$ using Part 2. of Lemma I.6, which takes $n^{1+o(1)}$ time.

We then compute $\underbrace{(\widehat{f}(X)\odot(h(X)G_{i}^{\top}))X}_{n\times d}\underbrace{W
^{\top}}_{d\times d}$, which takes $nd^{2}$ time.

Proof of Part 2.
For $\widehat{D}_{7}$, we can compute $\underbrace{(\widehat{f}(X)\odot(G_{i}h(X)^{\top}))}_{n\times n}\underbrace{X}
_{n\times d}$ using Part 2. of Lemma I.6, which takes $n^{1+o(1)}$ time.

We then compute $\underbrace{(\widehat{f}(X)\odot(G_{i}h(X)^{\top}))X}_{n\times d}\underbrace{W
}_{d\times d}$, which takes $nd^{2}$ time.

Proof of Part 3.
For $\widehat{g}_{w}$, we consider $X^{\top}\widehat{p}_{1}(X)X$ first.
Based on Definition C.11, we have $\widehat{p}_{1}(X)=\widehat{f}(X)\odot q(X)=\widehat{f}(X)\odot(G_{i}h(X)^{
\top})$.
We then compute $(\widehat{f}(X)\odot(G_{i}h(X)^{\top}))X$ using Part 2. of Lemma I.6, which takes $n^{1+o(1)}$ time.
After that, we compute $\underbrace{X^{\top}}_{d\times n}\underbrace{(\widehat{f}(X)\odot(G_{i}h(X)^{
\top}))X}_{n\times d}$, which takes $nd^{2}$ time.

Now we consider $X^{\top}\widehat{p}_{2}(X)X$.
By definition, $\widehat{p}_{2}(X)=\operatorname{diag}(\widehat{p}_{1}(X)\cdot{\bf 1}_{n})
\widehat{f}(X)$.
We first compute $\widehat{p}_{1}(X)\cdot{\bf 1}_{n}=(\widehat{f}(X)\odot(G_{i}h(X)^{\top}))
\cdot{\bf 1}_{n}$ using Part 2. of Lemma I.6, which takes $n^{1+o(1)}$ time.
Meanwhile, we compute $\widehat{f}(X)X$ using Part 1. of Lemma I.6, which takes $n^{1+o(1)}$ time.
We then have $\underbrace{\operatorname{diag}(\widehat{p}_{1}(X)\cdot{\bf 1}_{n})}_{n\times n
}\underbrace{\widehat{f}(X)X}_{n\times d}$, which takes $nd$ time.
Finally, we compute $\underbrace{X^{\top}}_{d\times n}\underbrace{\operatorname{diag}(\widehat{p}_{
1}(X)\cdot{\bf 1}_{n})\widehat{f}(X)X}_{n\times d}$, which takes $nd^{2}$ time.

Together, $\underbrace{X^{\top}\widehat{p}_{1}(X)X}_{d\times d}-\underbrace{X^{\top}
\widehat{p}_{2}(X)X}_{d\times d}$ takes $d^{2}$ time.
∎

Thus, we show that our framework can support causal attention masks.

###### Definition I.5 .

###### Lemma I.6 .

###### Proof.

###### Lemma I.7 (Components for dot product) .

###### Proof.

###### Lemma I.8 (Components for Hadamard product) .

###### Proof.

## Appendix J Residual Connection

In this section, we discuss how to adapt our framework to the attention mechanism with the residual connection.

In Section J.1, we provide a formalized definition of the two residual connections used in the attention mechanism.
In Section J.2, we argue that with the addition of the residual connection, the gradient over the attention mechanism can be computed in almost linear time $n^{1+o(1)}$ and the approximation error can be bound by $1/\operatorname{poly}(n)$.
In Section J.3, we use math induction to show that the gradient over the entire transformer with the residual connection can also be computed in almost linear time $n^{1+o(1)}$.

### J.1 Key concepts

Recall that in Definition 3.3, we have defined $T_{i}(X)\in\mathbb{R}^{n\times d}$ as the intermediate variable output by the $i$-th transformer layer. For simplicity, we use $T_{i}$ to represent $T_{i}(X)$ in the rest part of this section.
Namely, we have

$$ $\displaystyle T_{i}=(g_{i}\circ\mathsf{Attn}_{i})(T_{i-1})$ $$

Then, we consider adding the residual connection to our framework.
Note that there are two residual connection operations in one transformer layer.
We first define the residual connection over the $\mathsf{Attn}_{i}$ in Definition J.1.

###### Definition J.1 (Residual connection over $\mathsf{Attn_{i}}$ ) .

If we have the below conditions,

- •
Let $T_{i}$ be defined as Definition 3.3.
- •
Let $\mathsf{Attn}_{i}$ be defined as Definition C.3.

We define $Z_{i}\in\mathbb{R}^{n\times d}$ as the output with the residual connection of $\mathsf{Attn}_{i}$. Namely, we have

$$ $\displaystyle Z_{i}=T_{i-1}+\mathsf{Attn}_{i}(T_{i-1})$ $$

Then, we consider the second residual connection over the MLP layer $g_{i}$, where we have the formal definition for this in Definition J.2.

###### Definition J.2 (Residual connection over $g_{i}$ ) .

If we have the below conditions,

- •
Let the multi-layer transformer be defined as Definition 1.3.
- •
Let the intermediate variable $T_{i}$ be defined as Definition 3.3.
- •
Let $g_{i}$ denote the components other than self-attention in the $i$-th transformer layer.
- •
Let $Z_{i}\in\mathbb{R}^{n\times d}$ be defined as Definition J.1.

Then $T_{i}$, the output of $i$-th layer transformer with the residual connection, should have the following form:

$$ $\displaystyle T_{i}=Z_{i}+g_{i}(Z_{i})$ $$

###### Definition J.1 (Residual connection over 𝖠𝗍𝗍𝗇 𝗂 subscript 𝖠𝗍𝗍𝗇 𝗂 \mathsf{Attn_{i}} sansserif_Attn start_POSTSUBSCRIPT sansserif_i end_POSTSUBSCRIPT ) .

###### Definition J.2 (Residual connection over g i subscript 𝑔 𝑖 g_{i} italic_g start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) .

### J.2 Analysis of the residual connection

In the previous section, we have defined the two residual connection operations.

In this section, we argue that if the gradient computation can be done in almost linear time without the residual connection, then with the addition of the residual connection, the gradient computation can also be completed in almost linear time.

###### Lemma J.3 (Analysis of the residual connection) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $Y_{R}\in\mathbb{R}^{n\times d}$ and $X_{R}\in\mathbb{R}^{n\times d}$ denote the output and input of the residual connection, respectively.
- •
Let $\mathsf{H}:\mathbb{R}^{n\times d}\rightarrow\mathbb{R}^{n\times d}$ denote some layer in the transformer, such as MLP, $\mathsf{Attn}$, etc.
- •
Suppose the residual connection can be written as
$\displaystyle Y_{R}=X_{R}+\mathsf{H}(X_{R}).$
- •
Assuming we have $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}\in\mathbb{R}^{n\times d}$, then we can calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}\frac{\mathrm{d}\mathsf{H}(X_{R})}{
\mathrm{d}X_{R}}$ in almost linear time $n^{1+o(1)}$.

Then, we can show that,

- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}X_{R}}$ can be calculated in almost linear time $n^{1+o(1)}$.
- •
If $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}$ has $1/\operatorname{poly}(n)$ approximation error, then the approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{R}}$ is still $1/\operatorname{poly}(n)$.

###### Proof.

By the chain rule, we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}X_{R}}=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}\frac{\mathrm{d}Y_{R}}{ \mathrm{d}X_{R}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}(I+\frac{\mathrm{d} \mathsf{H}(X_{R})}{\mathrm{d}X_{R}})$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}+\frac{\mathrm{d}L(X)}{ \mathrm{d}Y_{R}}\frac{\mathrm{d}\mathsf{H}(X_{R})}{\mathrm{d}X_{R}}$ (32) $$

where the 1st step is from the chain rule, the 2nd step comes from basic calculus, the 3rd step is because of basic algebra.

By the assumption, we already have $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}$, and $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}\frac{\mathrm{d}\mathsf{H}(X_{R})}{
\mathrm{d}X_{R}}$ can be computed in almost linear time $n^{1+o(1)}$.

The addition operation between $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}$ and $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{R}}\frac{\mathrm{d}\mathsf{H}(X_{R})}{
\mathrm{d}X_{R}}$ takes $n\cdot d$ time.

Therefore, the overall running time for $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{R}}$ is $n^{1+o(1)}$.

Then, we consider the approximation error.

By Eq. (J.2) and basic linear algebra, the approximation error will not be magnified by more than $(n\cdot d\operatorname{poly}(n)+1)$.
Since $(n\cdot d\operatorname{poly}(n)+1)(1/\operatorname{poly}(n))=\operatorname{
poly}(n)$, the approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{R}}$ can be bounded by $1/\operatorname{poly}(n)$.

∎

###### Lemma J.3 (Analysis of the residual connection) .

###### Proof.

### J.3 Analysis for the entire model with the residual connection

In the previous section, we have shown that, with the addition of the residual connection on a single component, the gradient computation time can still be done in almost linear time.
We will apply this finding to the entire model.

We begin by single layer proof.

###### Lemma J.4 (Fast gradient computation for single-layer transformer with residual connection) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $X\in\mathbb{R}^{n\times d}$ be defined as Definition C.3.
- •
Suppose we have a single-layer transformer (see Definition 1.3).
- •
Let the residual connection be defined as Definition J.1 and J.2.

Then, we can show that,

- •
Part 1: running time. Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $n^{1+o(1)}$ time.
- •
Part 2: error bound. The approximation error of the single-layer transformer with the residual connection can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}_{r_{1}}$ satisfies
$\displaystyle\|\widetilde{g}_{r_{1}}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{
\infty}\leq 1/\operatorname{poly}(n)$

###### Proof.

We use $T_{i}$ to represent $T_{i}(X)$ for simplicity.
By the definition of $T_{i}$ (see also Definition 3.3), we have the following equations

$$ $\displaystyle T_{0}=g_{0}(X)$ $$

Follow Definition J.1 and J.2, we have

$$ $\displaystyle Z_{1}=T_{0}+\mathsf{Attn}_{1}(T_{0})$ $$

and

$$ $\displaystyle T_{1}=Z_{1}+g_{1}(Z_{1})$ $$

Then we calculate the gradient by the following steps:

- •
Step 1: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{1}}$. By the definition of $L(X)$ (see also Definition 3.1), we have $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{1}}$ can be computed in $n\cdot d$ time.
- •
Step 2: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}Z_{1}}$. By Lemma H.2, the assumption in Lemma J.3 is satisfied. Therefore, we have $\frac{\mathrm{d}L(X)}{\mathrm{d}Z_{1}}$ can be computed in almost linear time $n^{1+o(1)}$.
- •
Step 3: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}}$. By Lemma E.11, the assumption in Lemma J.3 is satisfied. Hence, $\frac{\mathrm{d}L(X)}{\mathrm{d}T_{0}}$ can be computed in almost linear time. By Lemma E.11, the approximation error is $1/\operatorname{poly}(n)$.
- •
Step 4: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$. By Lemma H.2, $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ can be computed in $n^{1+o(1)}$. The approximation error is $(n\cdot d)(1/\operatorname{poly}(n))=(1/\operatorname{poly}(n))$.

To sum up, we can show that the overall running time for $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ is $n^{1+o(1)}$ and the approximation error is $1/\operatorname{poly}(n)$.

Let $\widetilde{g}_{r_{1}}$ be the output of Step 4. Then we are done.

∎

We now prove for multi-layer.

###### Lemma J.5 (Fast gradient computation for multi-layer transformer with residual connection) .

If we have the below conditions,

- •
Let $L(X)$ be defined as Definition 3.1.
- •
Let $X\in\mathbb{R}^{n\times d}$ be defined as Definition C.3.
- •
Let the residual connection be defined as Definition J.1 and J.2.
- •
Suppose we have a $m$-layer transformer (see Definition 1.3).

Then, we can show that,

- •
Part 1: running time. Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $n^{1+o(1)}$ time.
- •
Part 2: error bound. The approximation error of the $m$-layer transformer with the residual connection can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}_{r}$ satisfies
$\displaystyle\|\widetilde{g}_{r}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}
\leq 1/\operatorname{poly}(n)$

###### Proof.

We use math induction in this proof.

Step 1: Proof of a single-layer transformer.

Firstly, by Lemma J.4, we have the statement holds for a single-layer transformer.

Step 2: Assumption for $k$-layer transformer.

Secondly, we assume for any $k$, for $k$-layer transformer model, we have

- •
Part 1: running time. Our algorithm can approximate $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in $O(n^{1+o(1)})$ time.
- •
Part 2: error bound. The approximation error of the $k$-layer transformer can be bounded by $1/\operatorname{poly}(n)$.
Namely, our algorithm output $\widetilde{g}$ satisfies
$\displaystyle\|\widetilde{g}-\frac{\mathrm{d}L(X)}{\mathrm{d}X}\|_{\infty}\leq
1
/\operatorname{poly}(n)$

Step 3: Proof of $(k+1)$-layer transformer.

Thirdly, we consider the $(k+1)$-layer transformer model.

Let $\mathsf{F}_{k}$ denote a $k$-layer transformer with the residual connection.

Then, the entire model can be written as

$$ $\displaystyle(\mathsf{F}_{k}\circ g_{0})(X)$ $$

By the definition of $T_{i}$, we have

$$ $\displaystyle T_{0}=g_{0}(X)$ $$

Then, by definition of $Z_{i}$ (see also Definition J.1), we have

$$ $\displaystyle Z_{1}=T_{0}+\mathsf{Attn}_{1}(T_{0})$ $$

By Definition J.2, we have

$$ $\displaystyle T_{1}=Z_{1}+g_{1}(Z_{1})$ $$

Without loss of generality, we assume that the additional transformer layer is added at the beginning of the model.
Then, the $(k+1)$-layer transformer model has the following structure:

$$ $\displaystyle\mathsf{F}_{k+1}(X)=\mathsf{F}_{k}(T_{1})$ $$

By the assumption for $k$-layer transformer, we have
$\frac{\mathrm{d}L(X)}{\mathrm{d}T_{1}}$ can be computed in almost linear time $n^{1+o(1)}$ and the approximation error can be bounded by $1/\operatorname{poly}(n)$.

We apply similar proof of Lemma J.4, then we can show that, we can compute $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ in almost linear time $n^{1+o(1)}$ and the approximation error can be bounded by $1/\operatorname{poly}(n)$.

∎

###### Lemma J.4 (Fast gradient computation for single-layer transformer with residual connection) .

###### Proof.

###### Lemma J.5 (Fast gradient computation for multi-layer transformer with residual connection) .

###### Proof.

## Appendix K Multi-head Attention

Following the notation used in Section B.1, we use $h$ to denote the number of heads, and $d_{h}=d/h$ to denote the dimension of each head.

###### Definition K.1 (Multi-head attention) .

If we have the below conditions,

- •
Let $h$ denote the number of heads.
- •
Let $d$ denote the hidden dimension. Let $d_{h}=d/h$ denote the dimension of each attention head.
- •
Let $Q,K,V\in\mathbb{R}^{n\times d}$ be defined as Definition C.3.
- •
Let $f(X)$ be defined as Definition C.8.
- •
Let $s(X)$ be defined as Definition C.10.

The multi-head attention can be formalized as follows:

- •
Step 1. Split the hidden dimension $d$ of $Q,K,V\in\mathbb{R}^{n\times d}$ into $h$ parts. Then, for each $l\in[h]$, we have $Q_{l},K_{l},V_{l}\in\mathbb{R}^{n\times d_{h}}$.
- •
Step 2. For each $l\in[h]$, calculate the attention matrix $f_{l}:=\mathsf{Softmax}(Q_{l}K_{l}^{\top}/d_{h})\in\mathbb{R}^{n\times n}$, and calculate the corresponding attention result $s_{l}:=f_{l}V_{l}\in\mathbb{R}^{n\times d_{h}}$.
- •
Step 3. Concatenate $s_{l}\in\mathbb{R}^{n\times d_{h}}$ together, then we have the final multi-head attention output $s\in\mathbb{R}^{n\times d}$.

Then, we dive into the analysis of the gradient computation process over the attention mechanism with multi-head attention.

###### Lemma K.2 (Analysis of the multi-head attention) .

If we have the below conditions,

- •
Let $\mathsf{Attn}(X)$ be defined as Definition C.3.
- •
Let multi-head attention mechanism be defined as Definition K.1.
- •
Let $Y_{m},X_{m}\in\mathbb{R}^{n\times d}$ denote the output and input of the multi-head attention, respectively.

Then, we can show that,

- •
$\frac{\mathrm{d}L(X)}{\mathrm{d}X_{m}}$ can be calculated in almost linear time $n^{1+o(1)}$.
- •
If $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}$ has $1/\operatorname{poly}(n)$ approximation error, then the approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{m}}$ is still $1/\operatorname{poly}(n)$.

###### Proof.

Following the notations used in Definition K.1, for $l\in[h]$, we use $s_{l}\in\mathbb{R}^{n\times d_{h}}$ to denote the output by each attention head.
And we use $s\in\mathbb{R}^{n\times d}$ to denote the concatenated version of the output of the multi-head attention.

By the chain rule and the definition of $L(X)$ (see also Definition 3.1), we have

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}X_{m}}=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{ m}}{\mathrm{d}s}\frac{\mathrm{d}s}{\mathrm{d}X_{m}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{ m}}{\mathrm{d}s}\sum_{l=1}^{h}\frac{\mathrm{d}s_{l}}{\mathrm{d}X_{m}}$ $$

where the 1st step is from the chain rule, the 2nd step comes from $s\in\mathbb{R}^{n\times d}$ is the concatenated version of $s_{l}\in\mathbb{R}^{n\times d_{h}}$.

We calculate the gradient in the following steps:

- •
Step 1: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}$. By the definition of $L(X)$ (Definition 3.1), we have that $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}$ can be calculated in $n\cdot d$ time.
- •
Step 2: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{m}}{\mathrm{d}s}$. Since we already have $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}$, by Lemma H.2, we have $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{m}}{\mathrm{d}s}$ can be computed in almost linear time $n^{1+o(1)}$.
- •
Step 3: Calculate $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{m}}{\mathrm{d}s}
\sum_{l=1}^{h}\frac{\mathrm{d}s_{l}}{\mathrm{d}X_{m}}$. For each $l\in[h]$, by Lemma E.11, $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{m}}{\mathrm{d}s}
\cdot\frac{\mathrm{d}s_{l}}{\mathrm{d}X_{m}}$ can be computed in $n^{1+o(1)}$. Since the number of heads $h$ can be viewed as a constant here, it takes $n^{1+o(1)}$ time to compute the gradients on $h$ heads.

Therefore, the overall running time for $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{m}}$ is $n^{1+o(1)}$.

Then, we consider the error bound.

By assumption, there is $1/\operatorname{poly}(n)$ approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}$.
For each $l\in[h]$, the approximation error will not be magnified by more than $n^{2}\cdot d\cdot d_{h}\cdot\operatorname{poly}(n)$ on $\frac{\mathrm{d}L(X)}{\mathrm{d}Y_{m}}\cdot\frac{\mathrm{d}Y_{m}}{\mathrm{d}s}
\cdot\frac{\mathrm{d}s_{l}}{\mathrm{d}X_{m}}$.

Then, since there is total $h$ heads, the approximation error on $\frac{\mathrm{d}L(X)}{\mathrm{d}X_{m}}$ can be bound by

$$ $\displaystyle h\cdot n^{2}\cdot d\cdot d_{h}\cdot\operatorname{poly}(n)\cdot(1 /\operatorname{poly}(n))=1/\operatorname{poly}(n)$ $$

∎

Similar to the proof of Lemma H.3 and H.4, we apply Lemma K.2 to deal with the multi-head attention in each transformer layer.
Then, we can show that $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ can be computed in almost linear time $n^{1+o(1)}$ and the approximation error can be bounded by $1/\operatorname{poly}(n)$.

###### Definition K.1 (Multi-head attention) .

###### Lemma K.2 (Analysis of the multi-head attention) .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: See Weaknesses section.
soundness: 3
strengths: The work focuses on a highly relevant topic, given the growing context length in the training of LLMs, and highlights the need for algorithms with solid theoretical backing to improve the quadratic complexity of gradient computations in the attention mechanism.

The results seem broadly applicable, as they hold across different loss functions (unlike previous studies) and can incorporate important elements of multi-layer transformers, such as MLPs, residual connections, and causal masking.
summary: The paper addresses the important problem of reducing the quadratic computational complexity of gradient calculation in transformers, proposing an approximation algorithm that computes gradients in almost linear time with a bounded polynomial error. The work builds on recent work in reducing the quadratic complexity of the forward pass and gradient computation for single-layer transformers to linear complexity.
weaknesses: The work is very interesting, but I have a few questions and points to raise.

In general, I feel that there are missing intuitions about the approach in the main body. For example, the technical overview section mentions that the attention matrix can be approximated by low-rank matrices. The related work section (line 157) states that the work uses polynomial approximation techniques (Aggarwal and Alman, 2022), but there is no further discussion about these techniques in the main body. I strongly believe that a small discussion section should be included, as this addresses the heart of the challenge.

Similarly, for Algorithm 1, all the steps reference lemmas in the appendix, but there is no discussion of these steps in Section 5. While Section 5 discusses accelerating the gradient for individual variables, there should be a connection and reference to Algorithm 1 in this discussion. Furthermore, I strongly suggest including a small proof sketch for the key lemma, Lemma 5.1, which shows the linear time complexity of the gradient computation with respect to $T_i$.

Lastly, paragraphs 401-407 and 409-413 essentially say the same thing with minor rewording—please address this repetition.

Some questions:

—What assumptions do the results of Theorem 4.2 (linear time complexity and $1/poly(n)$ error) require to hold?

—The results operate in a very practical regime of $d≪n$, but $d=O(log⁡ (n))$ seems too small, considering practical values like 1024–4096. I understand that such assumptions are often necessary for theoretical results, but how sensitive do you think the results are to this choice? The same question applies to the sub-polynomial choice for the number of layers $m$ and the logarithmic bit precision. Does this imply a practical limit on the number of layers for the algorithm to remain efficient?

—It would be useful to see the algorithm in practice and compare it to quadratic attention. However, I feel that the benefits will only be apparent when n is very large (e.g., at least 100k, which requires significant computational resources and may be difficult to run. Have the authors had a chance to experiment with this?

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Could the authors provide examples or scenarios where this method outperforms existing quadratic complexity methods in practical applications?
2. For cases where the multi-layer approximation results in slight errors, how do these impact tasks with high sensitivity to precision (e.g., in medical or financial text processing)?
3. Can the authors clarify the compatibility of this method with common pre-training setups in large language models? Specifically, how would this method perform in tandem with models like LLaMA or GPT-4?
4. What happens if the attention matrix is full-rank? Does the proposed method lose its advantage in such cases?
5. A more detailed, quantitative comparison with recent works on transformer efficiency, such as FlashAttention[1], hyperattention[2], or tensor-based approximation methods, would enhance the positioning of the proposed method.

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
[2] HyperAttention: Long-context Attention in Near-Linear Time
soundness: 3
strengths: Strengths:

1.Innovation in Complexity Reduction: The primary contribution of an almost linear time approximation method for gradient computation addresses a significant bottleneck in transformer models, advancing efficiency in training large language models.
2.Comprehensive Theoretical Support: The paper offers a well-developed theoretical foundation, demonstrating that the proposed method maintains accuracy within a polynomially small approximation error.
3.General Applicability: The approach is compatible with various sub-modules (e.g., residual connections, multi-head attention) and loss functions, enhancing its utility across diverse transformer-based applications.
4.Potential Impact on Long-Context Models: Given the paper’s focus on scaling efficiency for long input sequences, the findings are likely to benefit future developments in large language models that need to handle extensive contexts effectively.
summary: This paper proposes a method to compute gradients in multi-layer transformers in almost linear time, addressing the quadratic complexity bottleneck of traditional self-attention mechanisms. By leveraging a low-rank approximation technique, the approach aims to make large language models more efficient, particularly for long-context tasks, while maintaining a small, polynomial approximation error. The algorithm is versatile, supporting general loss functions and widely used transformer components such as residual connections, multi-head attention, and causal masks, making it theoretically applicable to a broad range of transformer-based architectures.
weaknesses: 1.Lack of Empirical Validation: Although the paper establishes a strong theoretical foundation, it lacks empirical evaluation to validate the proposed method's performance in real-world applications. While there is discussion of related methods, a deeper quantitative or empirical comparison (e.g., speed, memory usage) with current acceleration techniques could further support the claims of the paper. Without experiments or benchmarks, it is difficult to assess the practical effectiveness of the method, especially in comparison with other acceleration techniques.
2.Unclear Approximation Error Impact: While the approximation error is stated to be polynomially small, the paper does not explore how this error might affect model performance in downstream tasks. In applications requiring high precision, even small errors could be significant; this is particularly relevant in domains where precision is critical, such as healthcare or finance.
3.Limitations of Low-Rank Approximation: The proposed approach relies on low-rank approximation of the attention matrix. However, if the attention matrix is effectively full-rank, this assumption may not hold, potentially negating the efficiency gains of the method. Further clarification on the handling of full-rank attention maps would improve the paper.
4.Potential Implementation Challenges: The paper notes that implementing the approach, particularly on GPUs, may involve re-defining tensor operations and re-implementing specific backpropagation steps. This could pose a barrier to immediate adoption by practitioners, as it requires significant engineering effort. This is a good paper and I would like to raise my score if the author can provide more empirical supports.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In Theorem 1.6 of [1], authors show that there is a $n^{1+o(1)}$ time algorithm that computes the gradient up to 1/poly(n) accuracy. Can the authors please comment on the difference between this result and Theorem 4.1? Both in terms of guarantees, as well as in terms of proof technique. Does the present paper's proof use Theorem 1.6 in any way?
2. Can the authors elaborate on the proof of Theorem H.4 for the multi-layer analysis? Specifically, I refer to line 3646-3650: how did $n/poly(n)$ end up becoming $1/poly(n)$? Since the authors are using an inductive proof, I don't think one can rigorously prove the inductive step with explicitly writing down the degree in $poly(n)$. Related to this problem, I don't see where the authors use the assumed bound on $m$, surely, the inductive proof has to use this fact?
3. Can the authors please comment on what are the key challenges and theoretical contributions in proving their result on causal masks? In particular, please highlight what additional work was done on top of the results by [2].

**minor suggestions that did not affect score**
1. I suggest to remove mention of $m =n^{o(1)}$ in Theorem 4.1 if you are going to assume $m=1$ in the same theorem.
2. line 401-410 seems to be a duplication of 409-416. 

[1] [The Fine-Grained Complexity of Gradient Computation for Training Large Language Models](https://arxiv.org/pdf/2402.04497), Alman and Song 2024

[2] [Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers.](https://arxiv.org/pdf/2405.05219) Liang et al 2024a
soundness: 3
strengths: Overall, I think being able to compute the attention (both forward and backward) is a question of great theoretical and practical interest.

I think the extensions to a $m$-layer Transformer (Theorem 4.2) is interesting, well-motivated, and proof is non-trivial. (However, I do have trouble verifying some parts of the proof, see below).
summary: In this paper, the authors provide a number of algorithms that enable one to approximately compute the gradient (and forward pass) of a multi-layer Transformer in time that is nearly-linear in $n$, the input sequence length. The key challenge is in dealing with the self-similarity matrix, which has $O(n^2)$ entries. The dimension $d$ is assumed to be $O(log(n))$.

In Theorem 4.1, authors establish nearly-linear time guarantees for computing the gradient of a single-layer Transformer.

In Theorem 4.2, the authors generalize this result to a Transformer with up to $m=n^{o(1)}$ layers.

In both Theorem 4.1 and 4.2, the guarantee is roughly $1/poly(n)$ error in $n^{1+o(1)}$ time. 

In section 6, the authors discuss further extensions of their results to various settings, including multi-headed attention, causal mask, and prompt tuning.
weaknesses: My main criticism of the paper is that, while the results are interesting and important, I am uncertain about the contributions of the present paper, given the existing papers on this topic. 

Theorem 4.1 shows fast gradient computation for a single layer. This result seems identical (at least in big-O notation) to the result in Theorem 1.6 of [1]. On line 446, authors claim that they generalize the results in [1] in two ways: (a) accomodate general loss function L(X), and (b) generalize their result to both $W_i$ and $W_{V_i}$. However, I do not see why (a) is meaningful -- for the purpose of computing gradients, one can wlog consider a linear $L(X)$, which is exactly handled by [1]. I also do not see why (b) is meaningful -- [1] already handles derivative wrt $W_{i}$, whereas the derivative of $W_{V_i}$ is just a minor extension of existing results (see Remark 1.3 of [1]).

I feel that the generalization of 4.1 to the multi-layer setting is potentially quite nontrivial -- the error from a single-layer's grradient computation can blow up over layers, and thus handling this is non-trivial. However, I had a bit of trouble verifying the part of the proof that bounds this error propagation. The main result seems to be Lemma 5.5 (Theorem H.4). In this proof, the key seems to be bounding error accumulation in the gradient computation over $m$ layers, whose proof I have some trouble verifying. (see item 2 under questions below)

Finally, the authors discuss a number of extensions of their work to handling things like multi-headed attention, and causal masks. I think the handling of causal mask is also interesting, but the result seem to be almost immediate following [2].

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: * Could the authors provide empirical results or benchmarks to demonstrate the practical efficiency and effectiveness of the proposed gradient approximation method on real-world large language models? How does it compare to state-of-the-art methods in terms of both speed and accuracy?

* The paper claims that the approximation error is polynomially small, but how does this error affect downstream tasks such as text generation or classification? Are there scenarios where this approximation might lead to a noticeable degradation in performance, especially in sensitive applications?

* The authors mention potential coding challenges for implementing this method, particularly on GPUs. Have the authors considered these challenges in detail, and can they provide more concrete guidance or solutions for implementing this algorithm efficiently in real-world systems, especially for large-scale models like GPT-4 or LLaMA?
soundness: 2
strengths: *  The paper introduces a method to approximate gradient computation in multi-layer transformers in almost linear time, breaking the quadratic time complexity barrier. This efficiency boost is crucial for training large language models, especially for long-context tasks, and could reduce both computational and energy costs.

* The proposed algorithm is versatile, supporting general loss functions and working with commonly used transformer components such as residual connections, multi-head attention, and causal masks. This ensures that the method can be applied to a wide range of practical transformer architectures and tasks without major adjustments.
summary: This paper presents a novel method to approximate the gradient computation in multi-layer transformer models with almost linear time complexity. Traditionally, the self-attention mechanism in transformers has a quadratic time complexity with respect to the input sequence length, which becomes a bottleneck during training, especially for models with long input contexts. The authors propose a fast approximation algorithm that calculates gradients in time proportional to $𝑛^{1+o(1)}$ , where n is the input length, and achieves a polynomially small approximation error. This method is applicable to general loss functions and transformers with practical components such as residual connections, multi-head attention, and causal masks. The results hold significant potential for improving the training efficiency of large language models without sacrificing much accuracy.
weaknesses: * While the paper provides a strong theoretical foundation for the proposed method, there is no empirical evidence or experimental results presented to validate its performance in real-world transformer models. This leaves a gap in understanding how well the method performs in practice compared to theoretical expectations.

* Although the approximation error is stated to be polynomially small, the paper does not deeply explore how this approximation might affect model performance in different tasks. There could be edge cases where the accuracy trade-offs are more significant, especially in critical applications requiring precise outcomes.

* The paper acknowledges that practical implementation, especially on GPUs, may face coding and optimization challenges, such as needing to redefine tensor operations and re-implement certain backpropagation functions. This could limit the immediate adoption of the method by practitioners looking for a ready-to-use solution.

Meta-review: The reviewers and AC found this work to be an interesting contribution. In terms of rating, this paper was borderline, slightly below the acceptance threshold. However, the final decision is rejection for the following reasons:

- As initially raised by Reviewer o3CW, this work appears somewhat incremental over [Alman & Song, 2024a]. The authors state: 
"Fast gradient computation. The prior study in Alman & Song (2024a) demonstrated that the gradient of $W_i$ can be computed in almost linear time. We extend their findings by adapting their approach to accommodate general loss function $L(X)$ (as defined in Definition 3.1) and further generalize their results to include the gradient computation for both $W_i$ and $W_{V_i}$ in each transformer layer (Lemma 5.2 and 5.3)". First, I am not sure how much complexity $W_{V_i}$ brings. Second, I don't see either how loss function itself make a major difference in gradient compute complexity. At the minimum, the authors should have done a much better job distinguishing novelty and proof techniques over [Alman & Song, 2024a].

- As pointed out by Reviewer FvhV, the constraint on $d=log(n)$ is not very realistic and not fully in line with practical LLMs. There should have been comprehensive discussions of practical considerations: What the typical scaling of $d$ and the number of heads are, and how the results change if $d$ grows polynomially rather than logarithmic.

- Additional drawbacks of this work include the lack of thorough empirical validation, which is partly addressed during the author response, and connections to practical LLM training.

Overall, this submission could be strengthened by addressing the above points and will represent a strong contribution to an upcoming venue. However, current weaknesses prevent its acceptance as is.

Decision: Reject