Tensor Attention Training: Provably Efficient Learning of Higher-order Transformers

## Abstract

Tensor Attention, a multi-view attention that is able to capture high-order correlations among multiple modalities, can overcome the representational limitations of classical matrix attention. However, the O ⁢ ( n 3 ) 𝑂 superscript 𝑛 3 O(n^{3}) italic_O ( italic_n start_POSTSUPERSCRIPT 3 end_POSTSUPERSCRIPT ) time complexity of tensor attention poses a significant obstacle to its utilization in transformers, where n 𝑛 n italic_n is the input sequence length. In this work, we prove that the backward gradient of tensor attention training can be computed 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 , the same complexity as its forward computation under the bounded entries assumption. We provide a closed-form solution for the gradient and propose a fast computation method utilizing polynomial approximation methods and tensor algebraic techniques. Furthermore, we prove the necessity and tightness of our assumption through hardness analysis, showing that slightly weakening it renders the gradient problem unsolvable in truly subcubic time. Our theoretical results establish the feasibility of efficient higher-order transformer training and may facilitate practical applications of tensor attention architectures.

## 1 Introduction

The generative large language models (LLMs), such as Mistral [], Llama [], Llama2 [], Llama3 [], Gemma [], GPT-3 [], GPT-4 [], Claude3 [],
Grok-1 [] and many more, have been widely involved in people’s living and work in these two years, such as bio-informatics [], coding [], education [], finance [], law [], and even writing NeurIPS conference reviews [].
The success of LLMs is based on the transformer architecture introduced by [], which also has been introduced into other modality [], such as vision-language models, e.g., CLIP [], Flamingo [], LLaMA-Adapter [, ], LLava [, ], BLIP [, ], MiniGPT-4 [], Qwen [, ], Gemini [], MM1 [].

The above open-sourced large models use two-view matrix attention, i.e., each attention score/entry is related to two tokens (one query token and one key token) to capture the data correlation.
More specifically, let $Z$ be hidden representations and $Q=ZW_{Q},K=ZW_{K},V=ZW_{V}$ be the corresponding query, key, and value matrices after projections using weights $W_{Q},W_{K},W_{V}$ respectively.
Then, the classical/matrix attention head can be written as $\mathsf{Att}(Z)=\mathsf{Softmax}(QK^{\top})V$.

On the other hand, many studies find that multi-view is crucial for high-order correlation in various kinds of data, e.g., math [], graph [, ], and multi-modality [].
For example, recently, OpenAI released GPT-4o [], and Google released Project Astra [], two flagship multi-modality models that aim to reason across three views, i.e., audio, vision, and text in real-time.

However, [] theoretically and empirically shows that classical attention can capture pairwise correlation but not triple-wise correlation due to the representational limitations of matrix attention.
In other words, one classical matrix attention head “cannot” capture the information relevant to multi-views simultaneously unless using multiple layers with careful architecture design.
This poses a fundamental technical obstacle for multi-modality models to efficiently fuse multiple representations/views to capture the high-order correlation among them.

**Table 1: Comparison to previous works.**
| Reference | For/Backward | Matrix/Tensor |
| --- | --- | --- |
| [] | Forward | Matrix |
| [] | Forward | Matrix |
| [] | Forward | Matrix |
| [] | Backward | Matrix |
| [] | Forward | Tensor |
| Ours (Theorem 5.2) | Backward | Tensor |

To fundamentally solve the above obstacle, [] and [] introduce Tensor Attention, which is a higher-order generalization of matrix attention that can capture high-order/multi-view information intrinsically. More specifically, it is defined as $\mathsf{Softmax}(Q(K_{1}\oslash K_{2})^{\top})(V_{1}\oslash V_{2})$ (Definition 3.5, and illustrated in Figure 1), where $\oslash$ is one kind of tensor operation, and $Q$, $K_{1}/V_{1}$, $K_{2}/V_{2}$ can be from different views/modalities.
However, to implement Tensor Attention practically, we must overcome the complexity bottleneck.
Let the input token length be $n$, then the forward and backward time complexity of tensor attention will be $O(n^{3})$ as $Q(K_{1}\oslash K_{2})^{\top}\in\mathbb{R}^{n\times n^{2}}$ [], while the time complexity of matrix attention is $O(n^{2})$ only as $QK^{\top}\in\mathbb{R}^{n\times n}$ [].
For example, the input length of Llama2 [] is 4096. So, intuitively, if we put tensor attention in Llama2, the input length will reduce to 256 to keep the same complexity in running speed and memory consumption.

Figure: Figure 1: The visualization of tensor attention with $\mathsf{Softmax}$ activation function (Definition 3.5). We give an example of input token length $n=8$, feature dimension $d=4$.
Refer to caption: x1.png

There are several recent works to overcome the time complexity bottleneck above, e.g., $O(n^{2})$ for matrix attention and $O(n^{3})$ for tensor attention.
[] accelerate matrix attention forward via kernel density estimation and get truly sub-quadratic time running time.
[] uses the polynomial approximation method to map the matrix attention into low-rank matrices during forward computation, leading to an almost linear time complexity $n^{1+o(1)}$ when entries are bounded.
Similarly, under sparsity assumptions, [] achieves nearly linear time computation for matrix attention forward by identifying the larger entries in the attention matrix.
On the one hand, with fine-grained analysis, [] proposes a new backward algorithm to compute the gradient of matrix attention in almost linear time complexity $n^{1+o(1)}$ as well, under the same bounded entry assumption.
On the other hand, [] surprisingly finds that the *forward* computation of tensor attention can also be achieved in almost linear time $n^{1+o(1)}$ rather than almost quadratic time $n^{2+o(1)}$, under similar assumptions as []. See a summary in Table 1.
Thus, it is natural to ask,

Can we achieve almost linear time for gradient computation in Tensor Attention Training?

We provide a positive answer in this work. Under the same bounded entries assumption as [], we propose Algorithm 1 to fast compute the backward gradient of Tensor Attention Training in almost linear time $n^{1+o(1)}$ as its forward computation. Thus, our results may make the tensor attention practical, as we can get around the $O(n^{3})$ complexity barrier both in its forward and backward computation.
Our contributions are summarized as follows:

- •
Under fine-grained analysis, we give the closed-form solution of the gradient computation of tensor attention (Lemma 4.1) and its time complexity without acceleration (Theorem 4.3).
- •
Based on the closed-form solution, by utilizing polynomial approximation methods and tensor computation techniques, we propose Algorithm 1 to fast compute the backward gradient of tensor attention training in almost linear time as its forward computation (Theorem 5.2).
- •
Furthermore, we prove that our assumption is necessary and “tight” by hardness analysis, i.e., if we slightly weaken the assumption, there is no algorithm that can solve the tensor attention gradient computation in truly sub-cubic complexity (Theorem 6.3).

## 2 Related Work

#### Fast attention computation.

In recent years, significant advances have been made in the development of efficient attention computation. One research direction involves employing low-rank approximations, polynomial kernel, or random features for the attention matrix [, , , , , , , , , , , , ], which scales the computational complexity sub-quadratically with sequence length. Another method explores patterns of sparse attention that lessen the computational load [, , , , , ]. Additionally, using linear attention as an alternative to softmax attention has emerged as a substantial area of study [, , , , , , , , , ]. These innovations have enhanced the capability of transformer-based models to handle longer sequences, thereby broadening their potential applications across various fields [, , , , , , , , , , , , , ].

#### Tensor computation for high-order representation.

Tensors excel over matrices in capturing higher-order relationships within data. Calculating low-rank factorizations or approximations of tensors is essential in a wide range of computer science applications, such as natural language processing [, , , ], computer vision [, , , ], computer graphics [, , ], security [, , ], and data mining [, , , ]. Moreover, tensors are crucial in numerous machine learning applications [, , , , , , ] and other diverse fields [, , , , ].

Roadmap.
In Section 3, we introduce the notations, several useful definitions, and our loss function.
In Section 4, we give the closed form of the gradient of our loss function, and also its computational time complexity.
In Section 5, we prove that we can compute the gradient in almost linear time.
In Section 6, we provide the hardness analysis.
In Section 7, we give the conclusion of our paper.

## 3 Preliminary

In this section, we first provide the notations we use. In Section 3.1, we provide general definitions related to tensor operation. In Section 3.2, we provide key definitions that we utilize in this paper.

Basic notations.
We use $[n]$ to denote $\{1,2,\dots,n\}$.
We use $e_{i}$ to denote a column vector where only $i$-th location is $1$ and zeros everywhere else.
We denote an all $1$ vector using ${\bf 1}_{n}\in\mathbb{R}^{n}$ .
We use $\langle a,b\rangle$ to denote the inner product of $a,b\in\mathbb{R}^{d}$ i.e. $\langle a,b\rangle:=\sum_{i=1}^{d}a_{i}b_{i}$.
We use $\|x\|_{p}$ to denote the $\ell_{p}$ norm of a vector $x\in\mathbb{R}^{n}$, i.e. $\|x\|_{2}:=(\sum_{i=1}^{n}x_{i}^{2})^{1/2}$, and $\|x\|_{\infty}:=\max_{i\in[n]}|x_{i}|$. We use $\circ$ to denote the Hadamard product i.e. the $(i,j)$-entry of $A\circ B$ is $A_{i,j}B_{i,j}$.

We use $\operatorname{tr}[A]:=\sum_{i=1}^{n}A_{i,i}$ to denote the trace of a matrix $A\in\mathbb{R}^{n\times n}$.
We use $\exp(A)$ to denote a matrix where $\exp(A)_{i,j}:=\exp(A_{i,j})$ for a matrix $A\in\mathbb{R}^{n\times d}$.
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 $\|A\|_{F}$ to denote the Frobenius norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e. $\|A\|_{F}:=\sqrt{\sum_{i\in[n]}\sum_{j\in[d]}|A_{i,j}|^{2}}$.
We use $\operatorname{poly}(n)$ to denote polynomial time complexity w.r.t $n$.

#### Tensor related notations.

Let $A\in\mathbb{R}^{n\times d}$. We use $a:=\operatorname{vec}(A)$ to denote a length $nd$ vector. We stack rows of $A$ into a column vector, i.e. $\operatorname{vec}(A):=[a_{1}^{\top},a_{2}^{\top},\dots,a_{n}^{\top}]^{\top}$ where $a_{i}^{\top}$ is the $i$-th row of $A$, or simply $\operatorname{vec}(A)_{j+(i-1)d}:=A_{i,j}$ for any $i\in[n],j\in[d]$, visualized in Fig. 2.
Let $I_{d}\in\mathbb{R}^{d\times d}$ denote an identity matrix. Let $\mathsf{I}_{d}\in\mathbb{R}^{d\times d\times d}$ denote an identity tensor, i.e., the diagonal entries are $1$ and zeros everywhere else.
Let $X\in\mathbb{R}^{d\times d^{2}}$. Let $x\in\mathbb{R}^{d^{3}}$ denote the vectorization of $X\in\mathbb{R}^{d\times d^{2}}$. Let $\mathsf{X}\in\mathbb{R}^{d\times d\times d}$ be the tesnorization of $X\in\mathbb{R}^{d\times d^{2}}$,
where $\mathsf{X}_{a,b,c}=X_{a,(b-1)d+c}$ for any $a,b,c\in[d]$.
And we have $X=\mathrm{mat}(\mathsf{X})$.

Figure: Figure 2: The visualization of vectorization operator $\operatorname{vec}(\cdot)$, which stacks rows of a matrix $A\in\mathbb{R}^{n\times d}$ into a column vector $a\in\mathbb{R}^{nd}$. In this figure, we give an example of $n=3,d=4$. The components of $A$ and $a$ are also given for easier understanding.
Refer to caption: x2.png

### 3.1 Definition of Tensor Operations

We define some operations like the Kronecker product, which is a matrix operation applied to two matrices of any size, producing a block matrix.
It is different from regular matrix multiplication and will be useful for our introduction and analysis of tensor attention.

###### Definition 3.1 ( $\otimes$ Kronecker product) .

Given $K_{1}\in\mathbb{R}^{n_{1}\times d_{1}}$ and $K_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$, for any $i_{1}\in[n_{1}],i_{2}\in[n_{2}],j_{1}\in[d_{1}],j_{2}\in[d_{2}]$, we define $K:=K_{1}\otimes K_{2}\in\mathbb{R}^{n_{1}n_{2}\times d_{1}d_{2}}$ as follows

$$ $\displaystyle K_{i_{1}+(i_{2}-1)n_{1},j_{1}+(j_{2}-1)d_{1}}=(K_{1})_{i_{1},j_{ 1}}\cdot(K_{2})_{i_{2},j_{2}}.$ $$

In this work, we will primarily use the following column-wise and row-wise versions of the Kronecker product, which are special kinds of Kronecker product.

###### Definition 3.2 ( $\oslash$ column-wise Kronecker product) .

Given matrices $K_{1}\in\mathbb{R}^{n_{1}\times d},K_{2}\in\mathbb{R}^{n_{2}\times d}$, we define matrix $K:=K_{1}\oslash K_{2}\in\mathbb{R}^{n_{1}n_{2}\times d}$ as follows

$$ $\displaystyle K_{i_{1}+(i_{2}-1)n_{1},j}:=(K_{1})_{i_{1},j}\cdot(K_{2})_{i_{2} ,j},~{}~{}~{}\forall i_{1}\in[n_{1}],i_{2}\in[n_{2}],j\in[d].$ $$

###### Definition 3.3 ( $\ominus$ row-wise Kronecker product) .

Given matrices $K_{1}\in\mathbb{R}^{n\times d_{1}},K_{2}\in\mathbb{R}^{n\times d_{2}}$, we define matrix $K:=K_{1}\ominus K_{2}\in\mathbb{R}^{n\times d_{1}d_{2}}$
as follows

$$ $\displaystyle K_{i,j_{1}+(j_{2}-1)d_{1}}:=(K_{1})_{i,j_{1}}\cdot(K_{2})_{i,j_{ 2}},~{}~{}~{}\forall i\in[n],j_{1}\in[d_{1}],j_{2}\in[d_{2}].$ $$

###### Definition 3.1 ( ⊗ tensor-product \otimes ⊗ Kronecker product) .

###### Definition 3.2 ( ⊘ ⊘ \oslash ⊘ column-wise Kronecker product) .

###### Definition 3.3 ( ⊖ symmetric-difference \ominus ⊖ row-wise Kronecker product) .

### 3.2 Key Definitions of Tensor Attention

Now, we are ready to introduce the tensor attention. First, we introduce the parameters and input.

###### Definition 3.4 (Input and weight matrix) .

We define the input sequence as $Z\in\mathbb{R}^{n\times d}$ and the key, query, and value weight matrix as $W_{K_{1}},W_{K_{2}},W_{Q},W_{V_{1}},W_{V_{2}}\in\mathbb{R}^{d\times d}$. Then, we define the key, query, and value matrix as $K_{1}:=ZW_{K_{1}}\in\mathbb{R}^{n\times d}$, $K_{2}:=ZW_{K_{2}}\in\mathbb{R}^{n\times d}$, $Q:=ZW_{Q}\in\mathbb{R}^{n\times d}$, $V_{1}:=ZW_{V_{1}}\in\mathbb{R}^{n\times d},V_{2}:=ZW_{V_{2}}\in\mathbb{R}^{n
\times d}$.

Then, based on Kronecker product, we define tensor attention in the following way.

###### Definition 3.5 (Tensor attention, Definition 7 in [ 96 ] , Definition 1.1 in [ 15 ] ) .

Given input matrices $Q,K_{1},K_{2},$ $V_{1},$ $V_{2}\in\mathbb{R}^{n\times d}$, compute the following matrix

$$ $\displaystyle\underbrace{D^{-1}}_{n\times n}\underbrace{A}_{n\times n^{2}} \underbrace{V}_{n^{2}\times d}\in\mathbb{R}^{n\times d},$ $$

where (1) $A:=\exp(QK^{\top}/d)\in\mathbb{R}^{n\times n^{2}}$ and $K:=K_{1}\oslash K_{2}\in\mathbb{R}^{n^{2}\times d}$, (2) $D:=\operatorname{diag}(A{\bf 1}_{n^{2}})\in\mathbb{R}^{n\times n}$, and (3) $V:=V_{1}\oslash V_{2}\in\mathbb{R}^{n^{2}\times d}$.

###### Remark 3.6 .

In Definition 3.5, on the one hand, we separate the $\mathsf{Softmax}$ operation into an element-wise $\exp$ operation and a diagonal normalization matrix $D$ for a more transparent formulation.
On the other hand, we change $K,V\in\mathbb{R}^{n\times d}$ in classical attention to $K_{1}\oslash K_{2},V_{1}\oslash V_{2}\in\mathbb{R}^{n^{2}\times d}$ in tensor attention, where $\oslash$ is a tensor operation defined in Definition 3.2.

Our Definition 3.5 covers the self-attention setting, when the query/key/values $Q,K_{1},K_{2},V_{1},V_{2}$ follow Definition 3.4 where they share the same input. It is then a tensor self-attention, which can capture high-order information of the input $Z$. When the query/key/values have different inputs, it is then a tensor cross-attention that can capture high-order relationships among multiple inputs.

Also, note that we have $A\in\mathbb{R}^{n\times n^{2}}$ in Definition 3.5. Although $QK^{\top}$ is a low-rank matrix with rank at most $d$, $\exp(QK^{\top})$ may be a full-rank matrix in general.
Thus, it is clear to see the exact forward computation of tensor attention takes $O(n^{3})$ time. Here, we introduce a forward tensor attention approximation task, which will help us formulate the tensor attention gradient approximation task later. Furthermore, [] show that they can solve this approximation task in almost linear time $n^{1+o(1)}$ (Lemma 5.1).

###### Definition 3.7 (Approximate Tensor Attention Computation ( $\mathsf{ATAttC}(n,d,B,\epsilon)$ ), Definition 1.2 in [ 15 ] ) .

Given input matrices $Q,K_{1},K_{2},V_{1},V_{2}\in\mathbb{R}^{n\times d}$ and parameters $\epsilon,B>0$, where $\max\{\|Q\|_{\infty},$ $\|K_{1}\|_{\infty},$ $\|K_{2}\|_{\infty},\|V_{1}\|_{\infty},$ $\|V_{2}\|_{\infty}\}\leq B$. Let $A,D,V$ be defined in Definition 3.5.
Then, our target is to output a matrix $T\in\mathbb{R}^{n\times d}$ satisfying

$$ $\displaystyle\|\underbrace{T}_{n\times d}-\underbrace{D^{-1}}_{n\times n} \underbrace{A}_{n\times n^{2}}\underbrace{V}_{n^{2}\times d}\|_{\infty}\leq\epsilon.$ $$

Figure: Figure 3: The visualization of loss function defined in Definition 3.8. Let $A_{1},A_{2},A_{3},A_{4},A_{5}$ and $E$ be $n\times d$ input matrices. Let $Y$ be a given matrix with size $d^{2}\times d$. The Kronecker product operator $\otimes$ is defined in Definition 3.1. We minimize matrix $X\in\mathbb{R}^{d\times d^{2}}$ in our loss function. We first compute $\exp(A_{1}X(A_{2}\otimes A_{3})^{\top})$. Then, we compute $D(X):=\operatorname{diag}(\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}){\bf 1}_{n^{2 }})$. Afterwards, we compute $D(X)^{-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top})(A_{4}\otimes A_{5})Y-E$. Finally, we optimize $X$ to compute the minimum of its Frobenius norm with a scaling factor $0.5$.
Refer to caption: x3.png

For our focus, tensor attention training, we would like to find weights to fit the tensor attention to a desired output $E$. We first simplify the attention expression of Definition 3.5, whose inputs are from Definition 3.4 with weight matrices $W_{Q},W_{K_{1}},W_{K_{2}},W_{V_{1}},W_{V_{2}}\in\mathbb{R}^{d\times d}$. Let $X:=W_{Q}\cdot(W_{K_{1}}\oslash W_{K_{2}})^{\top}\in\mathbb{R}^{d\times d^{2}}$ and $Y:=W_{V_{1}}\oslash W_{V_{2}}\in\mathbb{R}^{d^{2}\times d}$. It can be verified that the tensor attention equals $D^{-1}\exp(ZX(Z\otimes Z)^{\top}/d)(Z\otimes Z)Y$. Then we model the tensor attention training as the following tensor attention optimization problem (where $A_{1},A_{2},A_{3},A_{4},A_{5}$ are introduced to replace $Z$ to capture more general settings such as cross-attention). See Figure 3 for an illustration.

###### Definition 3.8 (Tensor attention optimization) .

Suppose that $A_{1},A_{2},A_{3},A_{4},A_{5},E\in\mathbb{R}^{n\times d}$ and $Y_{1},Y_{2}\in\mathbb{R}^{d\times d}$ are given. We formulate the attention optimization problem as

$$ $\displaystyle\min_{X\in\mathbb{R}^{d\times d^{2}}}\mathsf{Loss}(X):=0.5\|D(X)^ {-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d)(A_{4}\otimes A_{5})Y-E\|_{F}^{2},$ $$

where (1) $A_{2}\otimes A_{3}\in\mathbb{R}^{n^{2}\times d^{2}}$ is the tensor product between $A_{2}$ and $A_{3}$, (2) $D(X)=\operatorname{diag}(\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d){\bf 1}_{n^{
2}})\in\mathbb{R}^{n\times n}$, and (3) $Y=Y_{1}\oslash Y_{2}\in\mathbb{R}^{d^{2}\times d}$.

The naive gradient computation for the above tensor attention training takes $\Omega(n^{3})$ time. The gradient for $X$ is the bottleneck while that for $Y$ is not, since $A_{1}X(A_{2}\otimes A_{3})^{\top}\in\mathbb{R}^{n\times n^{2}}$ lies in the non-linear function $\mathsf{Softmax}$. Also, note that with gradients of $X$ and $Y$, it is easy to get the gradients of the weight matrices $W_{Q},W_{K_{1}},W_{K_{2}},W_{V_{1}},W_{V_{2}}$. Therefore, our main focus is the following Approximate Tensor Attention Loss Gradient Computation task.

###### Definition 3.9 (Approximate Tensor Attention Loss Gradient Computation ( $\mathsf{ATAttLGC}(n,d,B,\epsilon)$ )) .

Let $\epsilon,B>0$. Let $A_{1},A_{2},A_{3},A_{4},A_{5},E\in\mathbb{R}^{n\times d}$ and let $X_{1},X_{2},X_{3},Y_{1},Y_{2}\in\mathbb{R}^{d\times d}$ (see Definition 3.8).
Let $X=X_{1}\cdot(X_{2}\oslash X_{3})^{\top}\in\mathbb{R}^{d\times d^{2}}$.
Assume that $\max\{\|A_{1}X_{1}\|_{\infty}$, $\|A_{2}X_{2}\|_{\infty}$, $\|A_{3}X_{3}\|_{\infty}$, $\|A_{4}Y_{1}\|_{\infty}$, $\|A_{5}Y_{2}\|_{\infty}\}\leq B$. Let us assume that any numbers in the previous matrices are in the $\log(n)$ bits model.
We define $\mathsf{Loss}(X)$ the same as Definition 3.8.
Let the gradient of loss function $\mathsf{Loss}(X)$ be $\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}\in\mathbb{R}^{d\times d^{2}}$.
Then, our target is to output a matrix $\widetilde{g}\in\mathbb{R}^{d\times d^{2}}$ satisfying
$\|\widetilde{g}-\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}\|_{\infty}\leq\epsilon.$

###### Definition 3.4 (Input and weight matrix) .

###### Definition 3.5 (Tensor attention, Definition 7 in [ 96 ] , Definition 1.1 in [ 15 ] ) .

###### Remark 3.6 .

###### Definition 3.7 (Approximate Tensor Attention Computation ( 𝖠𝖳𝖠𝗍𝗍𝖢 ⁢ ( n , d , B , ϵ ) 𝖠𝖳𝖠𝗍𝗍𝖢 𝑛 𝑑 𝐵 italic-ϵ \mathsf{ATAttC}(n,d,B,\epsilon) sansserif_ATAttC ( italic_n , italic_d , italic_B , italic_ϵ ) ), Definition 1.2 in [ 15 ] ) .

###### Definition 3.8 (Tensor attention optimization) .

###### Definition 3.9 (Approximate Tensor Attention Loss Gradient Computation ( 𝖠𝖳𝖠𝗍𝗍𝖫𝖦𝖢 ⁢ ( n , d , B , ϵ ) 𝖠𝖳𝖠𝗍𝗍𝖫𝖦𝖢 𝑛 𝑑 𝐵 italic-ϵ \mathsf{ATAttLGC}(n,d,B,\epsilon) sansserif_ATAttLGC ( italic_n , italic_d , italic_B , italic_ϵ ) )) .

## 4 Exact Tensor Attention Gradient Computation and Complexity

In this section, we provide the closed form of the tensor attention gradient of the loss function (Definition 3.8) and also its computational time.
First, we calculate the closed form of the gradient in the following lemma, whose proof is in Appendix E.5.

###### Lemma 4.1 (Closed form of gradient, informal version of Lemma E.6 ) .

Define the function $\mathsf{F}(x)\in\mathbb{R}^{n\times n^{2}}$ as in Definition D.6. Suppose that $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ are three given matrices. Suppose that $\mathsf{Loss}(x)$ is defined as Definition 3.8, where $x=\operatorname{vec}(X)$.
Then, we have

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x)}{\mathrm{d}x}=\operatorname{vec} (A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3}))\in\mathbb{R}^{d^{3}}.$ $$

Note that, $\mathsf{F}(x)$ is a size $n\times n^{2}$ matrix which is the bottleneck obstacle in time complexity.

###### Definition 4.2 .

Let ${\cal T}_{\mathrm{mat}}(a,b,c)$ denote the time of multiplying $a\times b$ matrix and $b\times c$ matrix.

Then, with straightforward analysis, we get the following theorem about the time complexity of naive computation. The complete proof is in Appendix E.6.

###### Theorem 4.3 (Tensor attention gradient computation, informal version of Theorem E.7 ) .

Suppose that $A_{1},A_{2},A_{3},$ $A_{4},A_{5},E\in\mathbb{R}^{n\times d}$ are input fixed matrices. We denote matrix variables as $X\in\mathbb{R}^{d\times d^{2}}$ and $Y\in\mathbb{R}^{d^{2}\times d}$ (gradient computation is w.r.t. $X$ ). Let $g=\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}\in\mathbb{R}^{d\times d^{2}}$ (for definition of $\mathsf{Loss}(X)$, see Definition 3.8). Then, we show that computing the gradient $g\in\mathbb{R}^{d\times d^{2}}$ requires ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})$ time.

Note that ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})\geq\Omega(n^{3})$. Thus, the naive tensor attention gradient computation is a complexity obstacle in practice as discussed in Section 1.
Based on the closed formulation in Lemma 4.1, we derive our acceleration method, which will be introduced in the following section.

###### Lemma 4.1 (Closed form of gradient, informal version of Lemma E.6 ) .

###### Definition 4.2 .

###### Theorem 4.3 (Tensor attention gradient computation, informal version of Theorem E.7 ) .

## 5 Fast Tensor Attention Gradient Computation

In this section, we show how to compute the tensor attention matrix gradient in almost linear time. In Section 5.1, we demonstrate our main results. In Section 5.2, we introduce some key tensor techniques used in our proof.

Figure: Algorithm 1 Almost Linear Time Tensor Attention Gradient Computation

### 5.1 Main Results for Fast Gradient Computation

Polynomial approximation methods involve representing complex functions through simpler polynomial forms to facilitate easier analysis and computation.
They are crucial in numerical analysis, aiding in the efficient solution of differential equations and optimization problems, and are widely used in simulations and machine learning [, ].

Based on the polynomial approximation methods, [] get the following result about tensor attention acceleration, which will be used to prove our main result.

###### Lemma 5.1 (Theorem 1.4 in [ 15 ] ) .

There is an algorithm that solves $\mathsf{ATAttC}(n,d=O(\log n),B=o(\sqrt[3]{\log n}),\epsilon=1/\operatorname{
poly}(n))$ (see Definition 3.7) in time $n^{1+o(1)}$.

Using similar polynomial approximation methods, and combined with a series of tensor analysis techniques (Section 5.2), we get our main acceleration results.

###### Theorem 5.2 (Main result for fast gradient computation) .

Assuming the entries of $A_{1},A_{2},$ $A_{3},$ $A_{4},A_{5},E\in\mathbb{R}^{n\times d}$ and $X_{1},X_{2},X_{3},Y_{1},Y_{2}\in\mathbb{R}^{d\times d}$ are represented using $O(\log n)$ bits.
Then, there exist an algorithm (Algorithm 1) that runs in $n^{1+o(1)}$ time to solve $\mathsf{ATAttLGC}(n,d=O(\log n),B=o(\sqrt[3]{\log n}),\epsilon=1/\operatorname
{poly}(n))$ (see Definition 3.9), i.e., our algorithm computes a gradient matrix $\widetilde{g}\in\mathbb{R}^{d\times d^{2}}$ satisfying $\|\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}-\widetilde{g}\|_{\infty}\leq 1
/\operatorname{poly}(n)$.

###### Proof sketch of Theorem 5.2 .

The complete proof can be found in Appendix F.6.

We use the polynomial approximation method to obtain low-rank approximation results for $D^{-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d)$ in Lemma F.1. However, this cannot be directly used for the closed form of the tensor attention gradient solution in Theorem 4.3. Utilizing a series of tensor techniques (Section 5.2 and Appendix C), we smartly convey these low rank properties throughout the gradient formulation and computation, where two key steps are fixed in Lemma F.5 and Lemma F.7.
∎

###### Remark 5.3 .

The assumption in Theorem 5.2 is practical. In practice, especially in recent long context tasks, the $n$ is large, e.g, $n=2\times 10^{6}$ for Google’s Gemini 1.5 Pro
[], while the model training uses a half-precision floating-point format, e.g., the bit number is $16$. Furthermore, our assumption is “tight”, where if we slightly weaken the assumption, there is no algorithm that can solve the tensor attention gradient computation in truly sub-cubic complexity (Theorem 6.3).

Our Theorem 5.2 accurately approximates ($\epsilon=1/\operatorname{poly}(n)$) the tensor attention gradient computation in almost linear time $n^{1+o(1)}$ under practical assumptions (see the above Remark 5.3).
Thus, our methods solve the last puzzle of tensor attention acceleration. Combined with previous work on tensor attention inference, this may make tensor attention practical, as we overcome the theoretical cubic time complexity barrier both in inference and training.

We provide Algorithm 1 for our almost linear time tensor attention training method. In the detailed algorithm, first, we construct $U_{1},V_{1},W_{1}$ in Lemma F.1. Then, we construct $U_{2},V_{2},W_{2}$ in Lemma F.3 and $U_{3},V_{3},W_{3}$ in Lemma F.5. We show how to construct $U_{4},V_{4},W_{4}$ in Lemma F.7. Finally, we construct $U_{5},V_{5},W_{5}$ and compute the gradient $g$ in almost linear time in Theorem F.8.

###### Lemma 5.1 (Theorem 1.4 in [ 15 ] ) .

###### Theorem 5.2 (Main result for fast gradient computation) .

###### Proof sketch of Theorem 5.2 .

###### Remark 5.3 .

### 5.2 Tensor Operation Analysis Techniques

Here, we introduce some key techniques for proving Theorem 5.2. These techniques make it possible to convey the low-rank property even during the tensor operations, solving the novel technical challenges in tensor attention gradient computation.

We first introduce a swap rule and a distributed rule, where both proofs are in Appendix C.2.

###### Fact 5.4 (Swap rule for tensor and matrix product) .

Let $W_{1},W_{2}\in\mathbb{R}^{d\times d},A_{1},A_{2}\in\mathbb{R}^{n\times d}$. We have

$$ $\displaystyle\underbrace{(A_{1}\otimes A_{2})}_{{n^{2}\times d^{2}}}\cdot \underbrace{(W_{1}\oslash W_{2})}_{{d^{2}\times d}}=\underbrace{(A_{1}\cdot W_ {1})}_{{n\times d}}\oslash\underbrace{(A_{2}\cdot W_{2})}_{{n\times d}}.$ $$

Fact 5.4 tells us that we can swap the order of tensor operation and matrix multiplication, allowing us to always compute the low dimension first to reduce the complexity.

###### Fact 5.5 .

Let $U_{1}\in\mathbb{R}^{n_{1}\times d}$ and $U_{2}\in\mathbb{R}^{n_{1}\times k}$.
Let $V_{1}\in\mathbb{R}^{n_{2}\times d}$ and $V_{2}\in\mathbb{R}^{n_{2}\times k}$. Let $W_{1}\in\mathbb{R}^{n_{3}\times d}$ and $W_{2}\in\mathbb{R}^{n_{3}\times k}$.
We have

$$ $\displaystyle\underbrace{(U_{1}\ominus U_{2})}_{n_{1}\times dk}\cdot( \underbrace{(V_{1}\ominus V_{2})}_{n_{2}\times dk}\oslash\underbrace{(W_{1} \ominus W_{2})}_{n_{3}\times dk})^{\top}=(\underbrace{U_{1}}_{n_{1}\times d}( \underbrace{V_{1}}_{n_{2}\times d}\oslash\underbrace{W_{1}}_{n_{3}\times d})^{ \top})\circ(\underbrace{U_{2}}_{n_{1}\times k}(\underbrace{V_{2}}_{n_{2}\times k }\oslash\underbrace{W_{2}}_{n_{3}\times k})^{\top})$ $$

Fact 5.5 tells us that the multiple tensor operation can be distributed to a different format. If we have some low-rank matrix/tensor, we can distribute them into each component, so that each component can be accelerated via the low-rank property. Intuitively, this allows us to borrow some low-rank benefits from other terms to fix the bottleneck terms.

We provide an important tool whose proof is in Appendix C.2.

###### Lemma 5.6 ( Informal version of Lemma C.13 ) .

Given $A_{1}\in\mathbb{R}^{n_{1}\times d_{1}}$, $A_{2}\in\mathbb{R}^{n_{2}\times d_{1}}$, let $A:=(A_{1}\oslash A_{2})\in\mathbb{R}^{n_{1}n_{2}\times d_{1}}$. Given $B_{1}\in\mathbb{R}^{n_{1}\times d_{2}}$, $B_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$, let $B:=(B_{1}\oslash B_{2})\in\mathbb{R}^{n_{1}n_{2}\times d_{2}}$.
We define $C\in\mathbb{R}^{d_{1}\times d_{2}}$ as $C:=A^{\top}B$ and $C_{1}:=A_{1}^{\top}B_{1}\in\mathbb{R}^{d_{1}\times d_{2}},C_{2}:=A_{2}^{\top}B
_{2}\in\mathbb{R}^{d_{1}\times d_{2}}$.
Then, we have $C_{1}\circ C_{2}=C$ and given $A_{1},A_{2},B_{1},B_{2}$, we can get $C$ in ${\cal T}_{\mathrm{mat}}(d_{1},\max\{n_{1},n_{2}\},d_{2})$ time.

Lemma 5.6 is a highly non-trivial method to handle tensor operation, $\circ$ and matrix multiplication together. By using the method, we save the computation time from ${\cal T}_{\mathrm{mat}}(d,n^{2},d)$ to ${\cal T}_{\mathrm{mat}}(d,n,d)$, which gets rid of the bottleneck quadratic term $n^{2}$.

Lastly, we introduce a tensor trick, which can reduce a tensor operation to a matrix multiplication operation. The proof is in Appendix C.3.

###### Fact 5.7 (Tensor-trick) .

Given matrices $A_{1}\in\mathbb{R}^{n_{1}\times d_{1}},A_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$ and $X\in\mathbb{R}^{d_{1}\times d_{2}}$, we have $\operatorname{vec}(A_{1}XA_{2}^{\top})=(A_{1}\otimes A_{2})\operatorname{vec}(
X)\in\mathbb{R}^{n_{1}n_{2}}$.

Technical novelty over previous works.
We generalize beyond the results of AS24b [], which only provide methods for tensor attention forward. Our paper presents a detailed analysis for tensor attention backward, providing both upper bound and lower bound.
Though we build on some results from AS24b [] and AS24a [], generalizing to tensor attention backward posed many technical challenges. These challenges are unique to our setting and not presented in previous settings like matrix attention [, ] or tensor attention forward []. To be more specific, we prove many key properties for tensor operation needed for backward though not needed for forward, including Facts 5.4 (swap rule for tensor and matrix product), 5.5 (distribution rule for tensor and matrix product), C.11 (tensor computation reduction to matrix product), C.12 (distribution rule for tensor computation), Claim C.20 (tensor product to matrix product). Fact 5.4, used as a key part to prove Lemmas F.1 and F.3, gives the swap rule for tensor operations. Lemma 5.6 supports the proof of Fact 5.5 and helps bypass the $O(n^{3}d^{2})$ time complexity bottleneck in the fast computation of $U_{2}$. Fact 5.5, crucial in proving Lemma F.5, shows the distributive nature of tensor operations. Using Facts C.11, C.12, and Claim C.20, we leverage the structure of low-rank matrices $U_{5},V_{5},W_{5}$ to prove Theorem 5.2.

###### Fact 5.4 (Swap rule for tensor and matrix product) .

###### Fact 5.5 .

###### Lemma 5.6 (
Informal version of Lemma C.13 ) .

###### Fact 5.7 (Tensor-trick) .

## 6 Tensor Attention Gradient Complexity Lower Bound

In this section, we show that our assumption is necessary.
First, we introduce some hardness analysis background in Section 6.1. Then, we introduce our main hardness result in Section 6.2.

### 6.1 𝖲𝖤𝖳𝖧 𝖲𝖤𝖳𝖧 \mathsf{SETH} sansserif_SETH and Tensor Attention Forward Hardness

We provide the findings that our results are based on. We first introduce a well-known hypothesis in hardness analysis.
The Strong Exponential Time Hypothesis ($\mathsf{SETH}$), a well-established conjecture, has been instrumental in establishing fine-grained lower bounds for numerous algorithmic problems, as highlighted in the survey by []. More than two decades ago, [] introduced $\mathsf{SETH}$ as an enhanced version of the well-known $\mathsf{P}\neq\mathsf{NP}$ conjecture, positing that current algorithms solving the $\mathsf{SAT}$ problem are nearly optimal in terms of efficiency.

###### Hypothesis 6.1 (Strong Exponential Time Hypothesis ( $\mathsf{SETH}$ ), [ 46 ] ) .

Given $\epsilon>0$, there exists $k\geq 3\in\operatorname*{\mathbb{Z}}$ such that it is impossible to solve $k$-$\mathsf{SAT}$ problem with $n$ variables in $O(2^{(1-\epsilon)n})$ time, including using any randomized algorithms.

We will critically utilize the hardness result of the forward tensor attention computation.

###### Lemma 6.2 (Theorem 1.3 in [ 15 ] ) .

Assuming $\mathsf{SETH}$, for any constant $\delta>0$, no algorithm can solve $\mathsf{ATAttC}(n,d=\Theta(\log n),B=\Theta(\sqrt[3]{(1+\gamma)\log n}),
\epsilon=n^{\gamma-O(1)})$ (Definition 3.7) in $O(n^{3-\delta})$ time, even if the inputs meet the following conditions for any $\gamma\geq 0$: (1) $V\in\{0,1\}^{n^{2}\times d}$, (2) There exists $B_{a}\leq O((1+\gamma)\log^{2}n)=O(d(\sqrt[3]{(1+\gamma)\log n})^{3})$ where all entries of $Q(K_{1}\oslash K_{2})^{\top}$ are within the range $[1,B_{a}]$ and more than half entries in each row of $Q(K_{1}\oslash K_{2})^{\top}$ are equal to $B_{a}$.

This result shows that assuming $\mathsf{SETH}$, if we just slightly weaken the assumption from $B=O(\sqrt[3]{\log n})$ to $B=\Theta(\sqrt[3]{(1+\gamma)\log n})$ with $\gamma=\omega(1)$, then the tenor attention forward computation is hard, i.e., no algorithm can solve it in truly sub-cubic time.

###### Hypothesis 6.1 (Strong Exponential Time Hypothesis ( 𝖲𝖤𝖳𝖧 𝖲𝖤𝖳𝖧 \mathsf{SETH} sansserif_SETH ), [ 46 ] ) .

###### Lemma 6.2 (Theorem 1.3 in [ 15 ] ) .

### 6.2 Main Result for Hardness

Based on the above observation (Lemma 6.2), we prove our main result for tensor attention gradient computation hardness.

###### Theorem 6.3 (Main result for hardness) .

Let $\gamma:\mathbb{N}\to\mathbb{N}$ be any function with $\gamma(n)=o(\log n)$ and $\gamma(n)=\omega(1)$.
Assuming $\mathsf{SETH}$, for any constant $\delta>0$, it is impossible to solve $\mathsf{ATAttLGC}(n,d=\Theta(\log n),B=\Theta(\sqrt[3]{\gamma(n)\cdot\log n}),
\epsilon=O(1/(\log n)^{4}))$ (Definition 3.9) in time $O(n^{{{}{}{}3}-\delta})$
when $E=0$, $\mathsf{Y}=\mathsf{I}_{d}$, $\mathsf{X}=\lambda\mathsf{I}_{d}$ for some scalar $\lambda\in[0,1]$.

See the formal proof in Appendix G.2.
The intuition is that if we can solve $\mathsf{ATAttLGC}$ in $O(t)$ time, then we can solve $\mathsf{ATAttC}$ in $O(t\cdot\log^{11}(n))$ time by interpolation and “integral”.
We see a similar sharp complexity
transition as forward computation (Lemma 6.2): assuming $\mathsf{SETH}$, if we slightly weaken the assumption from $B=O(\sqrt[3]{\log n})$ to $B=\Theta(\sqrt[3]{(1+\gamma)\log n})$ with $\gamma=\omega(1)$, then the tensor attention gradient computation will be unsolvable in truly sub-cubic time as well.

###### Theorem 6.3 (Main result for hardness) .

## 7 Discussion and Conclusion

In this work, we proved that the backward gradient of tensor attention training can be computed in almost linear $n^{1+o(1)}$ time, the same complexity as its forward computation, under a bounded entries assumption. We provided a closed-form solution for the gradient and proposed a fast computation method utilizing polynomial approximation and tensor algebraic techniques. Furthermore, we proved the necessity and tightness of our assumption through hardness analysis, showing that slightly weakening it renders the tensor attention gradient problem unsolvable in truly subcubic time. Our theoretical results establish the feasibility of efficient higher-order transformer training and may facilitate practical applications of tensor attention architectures. Due to space limits, we provide our further discussion and extension in Appendix A. Future work can perform empirical evaluations of the method in practical large language models, and explore how these findings can be implemented in real-world scenarios to enable the development of powerful higher-order models.

## 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 Further Discussion and Extension

#### Connection to real applications.

There are some empirical studies attempting to implement similar tensor attention (three order) in language modeling [] and 3D medical image segmentation []. However, due to cubic time complexity, their models remain relatively small, e.g, 12M parameters in []. Although small scale, [, ] demonstrates the significant potential of tensor attention. Our work proves that an almost linear time algorithm for tensor attention mechanisms exists (Algorithm 1). This advancement could enable the scaling up of tensor attention and facilitate novel model designs in multi-modality, 3D imaging, and beyond.
On the other hand, we abstract the most challenging part (the highest time complexity operation) in high-order attention into a clear mathematical problem and provide a solution. Our work introduces a new concept to the community, suggesting that cubic time complexity may not be the bottleneck in implementing three-order attention during training.
Practical implementation poses additional significant challenges, considering numerous other techniques and operations, such as dropout, layer normalization, residual connections, position encoding, and many others.
We hope our work inspires further algorithmic design.

#### Feasibility when the large value exists in the matrices.

If there exist many large entries in $Q,K_{1},K_{2},V_{1},V_{2}$, our hardness results (Theorem 6.3) indicate that no algorithm can accelerate the attention computation. However, several exciting works [, ] have shown that large entries are very sparse in the attention matrix. This suggests that our Algorithm 1 could inspire many potential practical implementations. One straightforward approach is to handle large entries separately, as in [], and apply our algorithm to the remaining parts. There is undoubtedly a broad algorithm design space, and we hope our work provides valuable insights.

#### Extend our technique to compute the module-wise gradient.

Let $n$ be the input toke length, and $d$ be the hidden dimension. At the $i$-th layer of transformer model, let $G_{i}\in\mathbb{R}^{n\times d}$ denote the output of upstream gradient, $X_{i}\in\mathbb{R}^{n\times d}$ be defined in Definition 3.8, and $\mathsf{Attn}_{i}:=D^{-1}AV$ be the tensor attention model where $D,A,V$ are defined in Definition 3.5. Let $\mathsf{Loss}$ be some loss function. Then, by the chain rule, we have the module-wise gradient $\frac{\mathrm{d}\mathsf{Loss}}{\mathrm{d}X_{i}}=\mathrm{vec}(G_{i})\frac{
\mathrm{d}\mathsf{Attn}_{i}}{\mathrm{d}X_{i}}$.

#### Extend our technique to the multi-head attention.

The gradient computation for each attention head in the same layer is independent of the others; each head only depends on its upstream gradient and its current module-wise gradient according to the chain rule. Therefore, our analysis can be directly applied to multi-head attention.

#### Generalize to scenarios involving multiple modalities

In our three-order attention, one attention module can handle three modalities simultaneously, i.e., $Q,K_{1},K_{2}$. For more modality, e.g., $m>3$ modality, there are two potential solutions in our minds.
First, we could use $m$-order attention, i.e., $Q,K_{1},K_{2},\dots,K_{m-1}$. The inference and training time complexity for this approach are still unknown, and we leave it as our future work.
Second, we could use multiple modules of three-order attention. Note that one layer of standard attention may introduce one more modality $K_{1}$ each time, while one layer of three-order attention may introduce two more modalities $K_{1},K_{2}$ each time. Thus, if we have $m+1$ modality and $Q$ is from one modality, say text, then the standard attention may need $m$ layers to merge all modalities together, whereas three-order attention may only need $\log(m)$ layers to merge them all together.

#### Societal impacts.

We delve into and offer a deeper understanding of the attention mechanism, introducing a novel approach to integrate multi-modality into attention through the tensor attention algorithm. We also demonstrate that the computation of both forward and backward tensor attention can be achieved with almost linear time complexity.
Regarding the negative societal impact, since our work is completely theoretical in nature, we do not foresee any potential negative societal impacts which worth pointing out.

## Appendix B More Related Work

#### Large language models and transformer.

The foundation of the success of generative large language models (LLMs) lies in the decoder-only transformer architecture, as introduced by []. This architecture has become critical for many leading models in natural language processing (NLP) []. These models have already demonstrated their capabilities in various real-world applications, including language translation [], sentiment analysis [], and language modeling [], due to their emergent ability, e.g., compositional ability [, , ], in-context learning [, , ]. The transformer leverages a self-attention mechanism, which enables the model to identify long-range dependencies within the input sequence. Self-attention calculates a weighted sum of input tokens, with weights based on the similarity between token pairs. This allows the model to focus on pertinent information from various parts of the sequence during output generation.

## Appendix C Tensor Operation Background

In Section C.1, we define the notation of computational time and the tensor operation.
In Section C.2, we provide some helpful facts of tensor operation.
In Section C.3, we provide some helpful facts of vectorization operation.
In Section C.4, we provide some helpful facts about the tensor product.

### C.1 General definitions and tensor operation

###### Fact C.1 ( [ 19 , 20 ] ) .

We can show that
${\cal T}_{\mathrm{mat}}(a,b,c)=O({\cal T}_{\mathrm{mat}}(a,c,b))=O({\cal T}_{
\mathrm{mat}}(b,a,c))=O({\cal T}_{\mathrm{mat}}(b,c,a))=O({\cal T}_{\mathrm{
mat}}(c,a,b))=O({\cal T}_{\mathrm{mat}}(c,b,a))$.

We define the third mode tensor product, which is the core operator of tensor operations.

###### Definition C.2 (Third mode tensor product $(\cdot,\cdot,\cdot)$ ) .

Let $\mathsf{X}\in\mathbb{R}^{d\times d\times d}$. Given matrices $A_{1}\in\mathbb{R}^{n\times d}$, $A_{2}\in\mathbb{R}^{n\times d}$ and $A_{3}\in\mathbb{R}^{n\times d}$.
Let operator $\mathsf{X}(A_{1},A_{2},A_{3})\in\mathbb{R}^{n\times n\times n}$ satisfying

$$ $\displaystyle\mathsf{X}(A_{1},A_{2},A_{3})_{i,j,l}:=\sum_{a=1}^{d}\sum_{b=1}^{ d}\sum_{c=1}^{d}\mathsf{X}_{a,b,c}(A_{1})_{i,a}(A_{2})_{j,b}(A_{3})_{l,c},~{}~ {}~{}~{}\forall i\in[n],j\in[n],l\in[n].$ $$

###### Definition C.3 ( $\odot$ tensor computation) .

Given matrices $A\in\mathbb{R}^{n\times d}$, $B\in\mathbb{R}^{n\times d}$, $C\in\mathbb{R}^{n\times d}$, we use $T=A\odot B\odot C\in\mathbb{R}^{n\times n\times n}$ to denote an tensor whose entries are given by

$$ $\displaystyle T_{i,j,l}:=\sum_{a=1}^{d}A_{i,a}B_{j,a}C_{l,a},~{}~{}~{}\forall i \in[n],j\in[n],l\in[n].$ $$

We note that a tensor $T$ can be written in the form $A\odot B\odot C$ like this if and only if its tensor rank is at most $d$.

###### Fact C.1 ( [ 19 , 20 ] ) .

###### Definition C.2 (Third mode tensor product ( ⋅ , ⋅ , ⋅ ) ⋅ ⋅ ⋅ (\cdot,\cdot,\cdot) ( ⋅ , ⋅ , ⋅ ) ) .

###### Definition C.3 ( ⊙ direct-product \odot ⊙ tensor computation) .

### C.2 Facts for tensor operation

###### Fact C.4 (Transpose rule) .

We show the results below

- •
Suppose that $\underbrace{K}_{n_{1}n_{2}\times d}=\underbrace{K_{1}}_{n_{1}\times d}\oslash
\underbrace{K_{2}}_{n_{2}\times d}$.
We have $\underbrace{K^{\top}}_{d\times n_{1}n_{2}}=\underbrace{K_{1}^{\top}}_{d\times n
_{1}}\ominus\underbrace{K_{2}^{\top}}_{d\times n_{2}}$.
- •
Suppose that $\underbrace{Q}_{n\times d_{1}d_{2}}=\underbrace{Q_{1}}_{n\times d_{1}}\ominus
\underbrace{Q_{2}}_{n\times d_{2}}$.
We have $\underbrace{Q^{\top}}_{d_{1}d_{2}\times n}=\underbrace{Q_{1}^{\top}}_{d_{1}
\times n}\oslash\underbrace{Q_{2}^{\top}}_{d_{2}\times n}$.
- •
Suppose that $\underbrace{V}_{n_{1}n_{2}\times d_{1}d_{2}}=\underbrace{V_{1}}_{n_{1}\times d
_{1}}\otimes\underbrace{V_{2}}_{n_{2}\times d_{2}}$.
We have $\underbrace{V^{\top}}_{d_{1}d_{2}\times n_{1}n_{2}}=\underbrace{V_{1}^{\top}}_
{d_{1}\times n_{1}}\otimes\underbrace{V_{2}^{\top}}_{d_{2}\times n_{2}}$.

###### Proof.

The proof is very straightforward.
∎

###### Fact C.5 (Swap rule) .

Let $V_{1}\in\mathbb{R}^{n\times d}$. Let $V_{2}\in\mathbb{R}^{n\times k}$.
Let $W_{1}\in\mathbb{R}^{m\times d}$. Let $W_{2}\in\mathbb{R}^{m\times k}$.
We can show swap rule for $\oslash$ and $\ominus$,

$$ $\displaystyle\underbrace{(V_{1}\ominus V_{2})}_{n\times dk}\oslash\underbrace{ (W_{1}\ominus W_{2})}_{m\times dk}=\underbrace{(V_{1}\oslash W_{1})}_{mn\times d }\ominus\underbrace{(V_{2}\oslash W_{2})}_{mn\times k}$ $$

And we can show swap rule for $\otimes$ and $\ominus$,

$$ $\displaystyle\underbrace{(V_{1}\ominus V_{2})}_{n\times dk}\otimes\underbrace{ (W_{1}\ominus W_{2})}_{m\times dk}=\underbrace{(V_{1}\otimes W_{1})}_{mn\times dk }\ominus\underbrace{(V_{2}\otimes W_{2})}_{mn\times dk}$ $$

###### Proof.

The proof is trivially following from definition of $\oslash$ and $\ominus$.

Note that for any $i_{1}\in[n],i_{2}\in[m],j_{1}\in[d],j_{2}\in[k]$

$$ $\displaystyle~{}((V_{1}\ominus V_{2})\oslash(W_{1}\ominus W_{2}))_{i_{1}+(i_{2 }-1)n,j_{1}+(j_{2}-1)d}$ $\displaystyle=$ $\displaystyle~{}(V_{1})_{i_{1},j_{1}}(V_{2})_{i_{1},j_{2}}(W_{1})_{i_{2},j_{1} }(W_{2})_{i_{2},j_{2}}$ $\displaystyle=$ $\displaystyle~{}((V_{1}\oslash W_{1})\ominus(V_{2}\oslash W_{2}))_{i_{1}+(i_{2 }-1)n,j_{1}+(j_{2}-1)d}$ $$

Thus, we complete the proof.
∎

###### Remark C.6 .

In Fact C.5, due to definition $V_{1}$ and $V_{2}$ need to have the same number of rows. $W_{1}$ and $W_{2}$ also need to have the same number of rows. $V_{1}$ and $W_{1}$ need to have same number of columns, and $V_{2}$ and $V_{2}$ need to have same number of columns.

###### Fact C.7 (Swap rule for tensor product and matrix product, Restatement of Fact 5.4 ) .

Let $W_{1},W_{2}\in\mathbb{R}^{d\times d}$ and $A_{1},A_{2}\in\mathbb{R}^{n\times d}$. We have

$$ $\displaystyle\underbrace{(A_{1}\otimes A_{2})}_{{n^{2}\times d^{2}}}\cdot \underbrace{(W_{1}\oslash W_{2})}_{{d^{2}\times d}}=\underbrace{(A_{1}\cdot W_ {1})}_{{n\times d}}\oslash\underbrace{(A_{2}\cdot W_{2})}_{{n\times d}}.$ $$

###### Proof of Fact 5.4 .

For any $i_{1},i_{2}\in[n],j\in[d]$, we have

$$ $\displaystyle~{}((A_{1}\otimes A_{2})\cdot(W_{1}\oslash W_{2}))_{i_{1}+(i_{2}- 1)n,j}$ $\displaystyle=$ $\displaystyle~{}\sum_{k_{1}\in[d],k_{2}\in[d]}(A_{1}\otimes A_{2})_{i_{1}+(i_{ 2}-1)n,k_{1}+(k_{2}-1)d}(W_{1}\oslash W_{2})_{k_{1}+(k_{2}-1)d,j}$ $\displaystyle=$ $\displaystyle~{}\sum_{k_{1}\in[d],k_{2}\in[d]}(A_{1}\otimes A_{2})_{i_{1}+(i_{ 2}-1)n,k_{1}+(k_{2}-1)d}\cdot(W_{1})_{k_{1},j}\cdot(W_{2})_{k_{2},j}$ $\displaystyle=$ $\displaystyle~{}\sum_{k_{1}\in[d],k_{2}\in[d]}(A_{1})_{i_{1},k_{1}}\cdot(A_{2} )_{i_{2},k_{2}}\cdot(W_{1})_{k_{1},j}\cdot(W_{2})_{k_{2},j}$ $\displaystyle=$ $\displaystyle~{}(\sum_{k_{1}\in[d]}(A_{1})_{i_{1},k_{1}}\cdot(W_{1})_{k_{1},j} )\cdot(\sum_{k_{2}\in[d]}(A_{2})_{i_{2},k_{2}}\cdot(W_{2})_{k_{2},j})$ $\displaystyle=$ $\displaystyle~{}(A_{1}\cdot W_{1})_{i_{1},j}\cdot(A_{2}\cdot W_{2})_{i_{2},j}$ $\displaystyle=$ $\displaystyle~{}((A_{1}\cdot W_{1})\oslash(A_{2}\cdot W_{2}))_{i_{1}+(i_{2}-1) n,j},$ $$

where the first step follows matrix multiplication, the second step follows Definition 3.2, the third step follows Definition 3.1, the fourth step follows simple algebra, the fifth step follows matrix multiplication and the last step follows Definition 3.2.
∎

###### Fact C.8 (Restatement of Fact 5.5 ) .

Let $U_{1}\in\mathbb{R}^{n_{1}\times d}$ and $U_{2}\in\mathbb{R}^{n_{1}\times k}$.
Let $V_{1}\in\mathbb{R}^{n_{2}\times d}$ and $V_{2}\in\mathbb{R}^{n_{2}\times k}$. Let $W_{1}\in\mathbb{R}^{n_{3}\times d}$ and $W_{2}\in\mathbb{R}^{n_{3}\times k}$.
We have

$$ $\displaystyle\underbrace{(U_{1}\ominus U_{2})}_{n_{1}\times dk}\cdot( \underbrace{(V_{1}\ominus V_{2})}_{n_{2}\times dk}\oslash\underbrace{(W_{1} \ominus W_{2})}_{n_{3}\times dk})^{\top}=(\underbrace{U_{1}}_{n_{1}\times d}( \underbrace{V_{1}}_{n_{2}\times d}\oslash\underbrace{W_{1}}_{n_{3}\times d})^{ \top})\circ(\underbrace{U_{2}}_{n_{1}\times k}(\underbrace{V_{2}}_{n_{2}\times k }\oslash\underbrace{W_{2}}_{n_{3}\times k})^{\top})$ $$

###### Proof of Fact 5.5 .

We can show that

$$ $\displaystyle(U_{1}\ominus U_{2})((V_{1}\ominus V_{2})\oslash(W_{1}\ominus W_{ 2}))^{\top}=$ $\displaystyle~{}(U_{1}\ominus U_{2})((V_{1}\oslash W_{1})\ominus(V_{2}\oslash W _{2}))^{\top}$ $\displaystyle=$ $\displaystyle~{}(U_{1}\ominus U_{2})((V_{1}\oslash W_{1})^{\top}\oslash(V_{2} \oslash W_{2})^{\top})$ $\displaystyle=$ $\displaystyle~{}(U_{1}^{\top}\oslash U_{2}^{\top})^{\top}((V_{1}\oslash W_{1}) ^{\top}\oslash(V_{2}\oslash W_{2})^{\top})$ $\displaystyle=$ $\displaystyle~{}(U_{1}(V_{1}\oslash W_{1})^{\top})\circ(U_{2}(V_{2}\oslash W_{ 2})^{\top})$ $$

where first step is due to swapping rule for $\oslash$ and $\ominus$ (see Fact C.5), the second step follows from Fact C.4, the third step follows from Fact C.4, and
the last step follows from Lemma C.13.
∎

###### Fact C.9 .

Let $U_{1}\in\mathbb{R}^{n_{1}\times d^{2}}$ and $U_{2}\in\mathbb{R}^{n_{1}\times k^{2}}$.
Let $V_{1}\in\mathbb{R}^{n_{2}\times d}$ and $V_{2}\in\mathbb{R}^{n_{2}\times k}$. Let $W_{1}\in\mathbb{R}^{n_{3}\times d}$ and $W_{2}\in\mathbb{R}^{n_{3}\times k}$.
We have

$$ $\displaystyle\underbrace{(U_{1}\ominus U_{2})}_{n_{1}\times d^{2}k^{2}}\cdot( \underbrace{(V_{1}\ominus V_{2})}_{n_{2}\times dk}\otimes\underbrace{(W_{1} \ominus W_{2})}_{n_{3}\times dk})^{\top}=(\underbrace{U_{1}}_{n_{1}\times d^{2 }}(\underbrace{V_{1}}_{n_{2}\times d}\otimes\underbrace{W_{1}}_{n_{3}\times d} )^{\top})\circ(\underbrace{U_{2}}_{n_{1}\times k^{2}}(\underbrace{V_{2}}_{n_{2 }\times k}\otimes\underbrace{W_{2}}_{n_{3}\times k})^{\top})$ $$

###### Proof.

We can show that,

$$ $\displaystyle~{}\underbrace{(U_{1}\ominus U_{2})}_{n_{1}\times d^{2}k^{2}} \cdot(\underbrace{(V_{1}\ominus V_{2})}_{n_{2}\times dk}\otimes\underbrace{(W_ {1}\ominus W_{2})}_{n_{3}\times dk})^{\top}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(U_{1}\ominus U_{2})}_{n_{1}\times d^{2}k^{2}} \cdot((V_{1}\otimes W_{1})\ominus(V_{2}\otimes W_{2}))^{\top}$ $\displaystyle=$ $\displaystyle~{}(U_{1}\ominus U_{2})\cdot((V_{1}\otimes W_{1})^{\top}\oslash(V _{2}\otimes W_{2})^{\top})$ $\displaystyle=$ $\displaystyle~{}(U_{1}^{\top}\oslash U_{2}^{\top})^{\top}\cdot((V_{1}\otimes W _{1})^{\top}\oslash(V_{2}\otimes W_{2})^{\top})$ $\displaystyle=$ $\displaystyle~{}(\underbrace{U_{1}}_{n_{1}\times d^{2}}(\underbrace{V_{1}}_{n_ {2}\times d}\otimes\underbrace{W_{1}}_{n_{3}\times d})^{\top})\circ( \underbrace{U_{2}}_{n_{1}\times k^{2}}(\underbrace{V_{2}}_{n_{2}\times k} \otimes\underbrace{W_{2}}_{n_{3}\times k})^{\top})$ $$

where the first step is because of the swap rule for $\otimes$ and $\ominus$ (see Fact C.5), the second step follows from Fact C.4, the third step follows from Fact C.4, and the last step follows from Lemma C.13.
∎

###### Claim C.10 .

Let $A,B,C\in\mathbb{R}^{n\times d}$.

Part 1.
Let $I_{d}\in\mathbb{R}^{d\times d}$ denote an identity matrix. Then, we have

$$ $\displaystyle AI_{d}B^{\top}=AB^{\top}.$ $$

Part 2.
Let $\mathsf{I}_{d}\in\mathbb{R}^{d\times d\times d}$ denote an identity tensor. Then we can show that

$$ $\displaystyle\mathsf{I}_{d}(A,B,C)=A\odot B\odot C$ $$

###### Proof.

Now we prove for each part.

Proof of Part1.
Using the property of identity matrix, it’s easy to see this holds.

Proof of Part2.

$$ $\displaystyle\mathsf{I}_{d}(A,B,C)=$ $\displaystyle~{}\sum_{a=1}^{d}\sum_{b=1}^{d}\sum_{c=1}^{d}(\mathsf{I}_{d})_{a, b,c}(A)_{i,a}(B)_{j,b}(C)_{l,c}$ $\displaystyle=$ $\displaystyle~{}\sum_{a=1}^{d}(A)_{i,a}(B)_{j,a}(C)_{l,a}$ $\displaystyle=$ $\displaystyle~{}A\odot B\odot C$ $$

where the first step follows from Definition C.2, the second step follows from the property of identity tensor $(\mathsf{I}_{d})_{i,j,k}$, which equals $1$ only when $i=j=k$ and 0 0 elsewhere, and the last step follows from Definition C.3.
∎

###### Fact C.11 .

Let $U,V,W\in\mathbb{R}^{n\times d}$, we have

$$ $\displaystyle\underbrace{U(V\oslash W)^{\top}}_{{n\times n^{2}}}=\underbrace{ \mathsf{mat}(U\odot V\odot W)}_{{n\times n^{2}}}.$ $$

###### Proof.

For any $i,j,k\in[n]$, we have

$$ $\displaystyle\mathsf{mat}(U\odot V\odot W)_{i,(j-1)n+k}=$ $\displaystyle~{}(U\odot V\odot W)_{i,j,k}$ $\displaystyle=$ $\displaystyle~{}\sum_{a\in[d]}U_{i,a}V_{j,a}W_{k,a}$ $\displaystyle=$ $\displaystyle~{}\sum_{a\in[d]}U_{i,a}(V\oslash W)_{(j-1)n+k,a}$ $\displaystyle=$ $\displaystyle~{}\sum_{a\in[d]}U_{i,a}((V\oslash W)^{\top})_{a,(j-1)n+k}$ $\displaystyle=$ $\displaystyle~{}(U(V\oslash W)^{\top})_{i,(j-1)n+k},$ $$

where the first step by definition of $\mathsf{mat}$, the second step follows Definition C.3, the third step follows Definition 3.2, the fourth step follows from transpose, and the last step follows from matrix multiplication.
∎

###### Fact C.12 .

Given $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ and $W_{1},W_{2},W_{3}\in\mathbb{R}^{n\times k}$,
we have

$$ $\displaystyle\underbrace{[W_{1}\odot W_{2}\odot W_{3}](A_{1}^{\top},A_{2}^{ \top},A_{3}^{\top})}_{{d\times d\times d}}=\underbrace{((A_{1}^{\top}W_{1}) \odot(A_{2}^{\top}W_{2})\odot(A_{3}^{\top}W_{3}))}_{{d\times d\times d}}.$ $$

###### Proof.

The proof is trivial by Definition C.3 and Definition C.2.
∎

We prove an important tool, which will be used in analyzing the running time of our algorithm.

###### Lemma C.13 ( Formal version of Lemma 5.6 ) .

If the following condition holds

- •
Let $\oslash$ be defined as Definition 3.2.
- •
Given $A_{1}\in\mathbb{R}^{n_{1}\times d_{1}}$, $A_{2}\in\mathbb{R}^{n_{2}\times d_{1}}$, let $A:=(A_{1}\oslash A_{2})\in\mathbb{R}^{n_{1}n_{2}\times d_{1}}$.
- •
Given $B_{1}\in\mathbb{R}^{n_{1}\times d_{2}}$, $B_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$, let $B:=(B_{1}\oslash B_{2})\in\mathbb{R}^{n_{1}n_{2}\times d_{2}}$.
- •
We define $C\in\mathbb{R}^{d_{1}\times d_{2}}$ as $C:=A^{\top}B$
- •
We define $\underbrace{C_{1}}_{d_{1}\times d_{2}}:=A_{1}^{\top}B_{1},\underbrace{C_{2}}_{
d_{1}\times d_{2}}:=A_{2}^{\top}B_{2}$

Then, we have

- •
Part 1. $C_{1}\circ C_{2}=C$
- •
Part 2. Given as input $A_{1},A_{2},B_{1},B_{2}$, we can get $C$ in ${\cal T}_{\mathrm{mat}}(d_{1},\max\{n_{1},n_{2}\},d_{2})$ time.

###### Proof.

For each $i\in[n_{1}]$, let $a_{1,i}^{\top}$ denote the $i$-th row of $A_{1}\in\mathbb{R}^{n_{1}\times d_{1}}$.

For each $i\in[n_{2}]$, let $a_{2,i}^{\top}$ denote the $i$-th row of $A_{2}\in\mathbb{R}^{n_{2}\times d_{1}}$.

For each $i\in[n_{1}]$, let $b_{1,i}^{\top}$ denote the $i$-th row of $B_{1}\in\mathbb{R}^{n_{1}\times d_{2}}$.

For each $i\in[n_{2}]$, let $b_{2,i}^{\top}$ denote the $i$-th row of $B_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$.

Recall that $C_{1}\in\mathbb{R}^{d_{1}\times d_{2}}$ and $C_{2}\in\mathbb{R}^{d_{1}\times d_{2}}$,

$$ $\displaystyle C_{1}:=A_{1}^{\top}B_{1},~{}~{}~{}C_{2}:=A_{2}^{\top}B_{2}$ $$

Thus, we see that for all $\forall k_{1}\in[d_{1}],k_{2}\in[d_{2}]$

$$ $\displaystyle(C_{1})_{k_{1},k_{2}}=$ $\displaystyle~{}\sum_{i=1}^{n_{1}}a_{1,i,k_{1}}b_{1,i,k_{2}}$ $\displaystyle(C_{2})_{k_{1},k_{2}}=$ $\displaystyle~{}\sum_{j=1}^{n_{2}}a_{2,j,k_{1}}b_{2,j,k_{2}}$ $$

Then, we can write $C\in\mathbb{R}^{d_{1}\times d_{2}}$ as

$$ $\displaystyle\underbrace{C}_{d_{1}\times d_{2}}=$ $\displaystyle~{}\underbrace{A^{\top}}_{d_{1}\times n_{1}n_{2}}\underbrace{B}_{ n_{1}n_{2}\times d_{2}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{n_{1}n_{2}}\underbrace{A_{i,*}}_{d_{1}\times 1} \underbrace{(B_{i,*})^{\top}}_{1\times d_{2}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{n_{1}}\sum_{j=1}^{n_{2}}\underbrace{A_{i+(j-1)n_{1 },*}}_{d_{1}\times 1}\cdot\underbrace{(B_{i+(j-1)n_{1},*})^{\top}}_{1\times d_ {2}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{n_{1}}\sum_{j=1}^{n_{2}}\underbrace{(a_{1,i}\circ a _{2,j})}_{d_{1}\times 1}\cdot\underbrace{(b_{1,i}\circ b_{2,j})^{\top}}_{1 \times d_{2}}$ (1) $$

where the first step follows from definition of $C\in\mathbb{R}^{d\times d}$, the second step follows from the matrix can written as the summation of $n_{1}n_{2}$ rank-$1$ matrices, the third step follows from changing the index, the forth step follows from $\underbrace{A_{i+(j-1)n_{1},*}}_{d_{1}\times 1}=\underbrace{a_{1,i}}_{d_{1}
\times 1}\circ\underbrace{a_{2,j}}_{d_{1}\times 1}$ by Definition 3.2.

From the above, we can calculate that the entry of $C$ in location $k_{1}\in[d_{1}],k_{2}\in[d_{2}]$ is

$$ $\displaystyle C_{k_{1},k_{2}}=$ $\displaystyle~{}\sum_{i=1}^{n_{1}}\sum_{j=1}^{n_{2}}(a_{1,i}\circ a_{2,j})_{k_ {1}}\cdot(b_{1,i}\circ b_{2,j})_{k_{2}}^{\top}$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{n_{1}}\sum_{j=1}^{n_{2}}a_{1,i,k_{1}}a_{2,j,k_{1}} b_{1,i,k_{2}}b_{2,j,k_{2}}$ $\displaystyle=$ $\displaystyle~{}(\sum_{i=1}^{n_{1}}a_{1,i,k_{1}}b_{1,i,k_{2}})\cdot(\sum_{j=1} ^{n_{2}}a_{2,j,k_{1}}b_{2,j,k_{2}})$ $\displaystyle=$ $\displaystyle~{}(C_{1})_{k_{1},k_{2}}\cdot(C_{2})_{k_{1},k_{2}}$ $$

where the first step follows from Eq. (C.2), the second step follows from simple algebra, the third step follows from separating the summation over $i$ and the summation over $j$, and the last step follows from definition of matrices $C_{1}$ and $C_{2}$.

Thus, we can conclude

$$ $\displaystyle C=C_{1}\circ C_{2}.$ $$

The algorithm will first compute $C_{1}$ and $C_{2}$, which takes ${\cal T}_{\mathrm{mat}}(d_{1},\max\{n_{1},n_{2}\},d_{2})$ time. Then it calculates $C_{1}\circ C_{2}$, which takes $O(d_{1}d_{2})$ time.
∎

###### Fact C.4 (Transpose rule) .

###### Proof.

###### Fact C.5 (Swap rule) .

###### Proof.

###### Remark C.6 .

###### Fact C.7 (Swap rule for tensor product and matrix product, Restatement of Fact 5.4 ) .

###### Proof of Fact 5.4 .

###### Fact C.8 (Restatement of Fact 5.5 ) .

###### Proof of Fact 5.5 .

###### Fact C.9 .

###### Proof.

###### Claim C.10 .

###### Proof.

###### Fact C.11 .

###### Proof.

###### Fact C.12 .

###### Proof.

###### Lemma C.13 (
Formal version of Lemma 5.6 ) .

###### Proof.

### C.3 Facts for vectorization operation

###### Fact C.14 .

Let $A,B\in\mathbb{R}^{n\times d}$. Then,

$$ $\displaystyle\operatorname{tr}[A^{\top}B]=\operatorname{vec}(A)^{\top} \operatorname{vec}(B)$ $$

###### Proof.

We can show

$$ $\displaystyle\operatorname{tr}[A^{\top}B]=$ $\displaystyle~{}\sum_{i=1}^{n}\sum_{j=1}^{d}A_{i,j}B_{i,j}$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(A)^{\top}\operatorname{vec}(B)$ $$

where the first step is due to the definition of trace, and the second step is because of the definition of $\operatorname{vec}$ operator.
∎

###### Fact C.15 .

Let $a\in\mathbb{R}^{n},b\in\mathbb{R}^{d}$. Then,

$$ $\displaystyle\operatorname{vec}(ab^{\top})=a\otimes b$ $$

###### Proof.

We can show

$$ $\displaystyle\operatorname{vec}(ab^{\top})=$ $\displaystyle~{}\operatorname{vec}(\begin{bmatrix}a_{1}b^{\top}\\ a_{2}b^{\top}\\ \dots\\ a_{n}b^{\top}\end{bmatrix})$ $\displaystyle=$ $\displaystyle~{}[a_{1}b^{\top},a_{2}b^{\top},\dots,a_{n}b^{\top}]^{\top}$ $\displaystyle=$ $\displaystyle~{}a\otimes b$ $$

where the first step follows from the definition of the outer product, the second step follows from the definition of vectorization operator $\operatorname{vec}(\cdot)$ which stacks rows of a matrix into a column vector, and the last step follows from Definition 3.1.
∎

###### Fact C.16 (Tensor-trick, Restatement of Fact 5.7 ) .

Given matrices $A_{1}\in\mathbb{R}^{n_{1}\times d_{1}},A_{2}\in\mathbb{R}^{n_{2}\times d_{2}}$ and $X\in\mathbb{R}^{d_{1}\times d_{2}}$, we have $\operatorname{vec}(A_{1}XA_{2}^{\top})=(A_{1}\otimes A_{2})\operatorname{vec}(
X)\in\mathbb{R}^{n_{1}n_{2}}$.

###### Proof of Fact 5.7 .

We can show

$$ $\displaystyle\operatorname{vec}(A_{1}XA_{2}^{\top})=$ $\displaystyle~{}\sum_{i=1}^{d_{1}}\sum_{j=1}^{d_{2}}X_{i,j}\operatorname{vec}( A_{1,*,i}(A_{2,*,j})^{\top})$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{d_{1}}\sum_{j=1}^{d_{2}}X_{i,j}(\underbrace{A_{1,* ,i}}_{n_{1}\times 1}\otimes\underbrace{A_{2,*,j}}_{n_{2}\times 1})$ $\displaystyle=$ $\displaystyle~{}\sum_{i=1}^{d_{1}}(\underbrace{A_{1,*,i}}_{n_{1}\times 1} \otimes\underbrace{A_{2}}_{n_{2}\times d_{2}})\underbrace{X_{i,*}}_{d_{2} \times 1}$ $\displaystyle=$ $\displaystyle~{}(A_{1}\otimes A_{2})\operatorname{vec}(X)$ $$

where the first step is due to the matrix being able to be written as a summation of vectors, the second step follows from Fact C.15, the third step follows from that matrix can be written as a summation of vectors, and the last step follows from the definition of vectorization operator $\operatorname{vec}(\cdot)$.
∎

###### Fact C.17 .

Let $A\in\mathbb{R}^{n_{1}\times n_{2}}$, $B\in\mathbb{R}^{n_{2}\times n_{3}}$, $C\in\mathbb{R}^{n_{3}\times n_{4}}$, $D\in\mathbb{R}^{n_{4}\times n_{5}}$.

We have

$$ $\displaystyle\operatorname{tr}[ABCD]=\operatorname{vec}(A^{\top})^{\top}(B \otimes D^{\top})\operatorname{vec}(C)$ $$

###### Proof.

We can show

$$ $\displaystyle\operatorname{tr}[ABCD]=$ $\displaystyle~{}\operatorname{vec}(A^{\top})^{\top}\operatorname{vec}(BCD)$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(A^{\top})^{\top}(B\otimes D^{\top}) \operatorname{vec}(C)$ $$

where the first step follows from Fact C.14, and the second step follows from Fact C.16.
∎

###### Fact C.18 .

Let $A,B\in\mathbb{R}^{n\times n}$ be two $n\times n$ symmetric matrices. Let $X$ and $Y$ denote two $n\times n$ matrices. Then we have

$$ $\displaystyle\operatorname{vec}(A)^{\top}(X\otimes Y)\operatorname{vec}(B)= \operatorname{vec}(A)^{\top}(Y\otimes X)\operatorname{vec}(B)$ $$

###### Proof.

We can show that

$$ $\displaystyle\operatorname{vec}(A)^{\top}(X\otimes Y)\operatorname{vec}(B)=$ $\displaystyle~{}\operatorname{tr}[A^{\top}XBY^{\top}]$ $\displaystyle=$ $\displaystyle~{}\operatorname{tr}[BY^{\top}A^{\top}X]$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(B^{\top})^{\top}(Y^{\top}\otimes X^{\top}) \operatorname{vec}(A^{\top})$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(B)^{\top}(Y^{\top}\otimes X^{\top}) \operatorname{vec}(A)$ $\displaystyle=$ $\displaystyle~{}((Y^{\top}\otimes X^{\top})\operatorname{vec}(A))^{\top} \operatorname{vec}(B)$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(A)^{\top}(Y\otimes X)\operatorname{vec}(B)$ $$

where the first step follows from Fact C.17, the second step follows from the cyclic property of trace, the third step follows from Fact C.17, the fourth step follows from $A,B$ is symmetric, the fifth step is due to the definition of inner product, and the last step is due to Fact C.4.
∎

###### Fact C.14 .

###### Proof.

###### Fact C.15 .

###### Proof.

###### Fact C.16 (Tensor-trick, Restatement of Fact 5.7 ) .

###### Proof of Fact 5.7 .

###### Fact C.17 .

###### Proof.

###### Fact C.18 .

###### Proof.

### C.4 Facts for tensor product

###### Fact C.19 .

Let $X=\underbrace{\mathrm{mat}(\mathsf{I}_{d})}_{d\times d^{2}}$, where $\mathsf{I}_{d}\in\mathbb{R}^{d\times d\times d}$ and $A_{1},A_{2}\in\mathbb{R}^{n\times d}$. We have

$$ $\displaystyle\underbrace{(A_{1}\otimes A_{2})}_{{n^{2}\times d^{2}}} \underbrace{X^{\top}}_{{d^{2}\times d}}=\underbrace{A_{1}\oslash A_{2}}_{{n^{2 }\times d}}.$ $$

###### Proof.

For any $i_{1},i_{2}\in[n],j\in[d]$, we have

$$ $\displaystyle((A_{1}\otimes A_{2})X^{\top})_{i_{1}+(i_{2}-1)n,j}=$ $\displaystyle~{}\sum_{k_{1}\in[d],k_{2}\in[d]}(A_{1}\otimes A_{2})_{i_{1}+(i_{ 2}-1)n,k_{1}+(k_{2}-1)d}X_{j,k_{1}+(k_{2}-1)d}$ $\displaystyle=$ $\displaystyle~{}\sum_{k_{1}\in[d],k_{2}\in[d]}(A_{1})_{i_{1},k_{1}}\cdot(A_{2} )_{i_{2},k_{2}}X_{j,k_{1}+(k_{2}-1)d}$ $\displaystyle=$ $\displaystyle~{}(A_{1})_{i_{1},j}\cdot(A_{2})_{i_{2},j}$ $\displaystyle=$ $\displaystyle~{}(A_{1}\oslash A_{2})_{i_{1}+(i_{2}-1)n,j},$ $$

where the first step is due to matrix multiplication, the second step follows Definition 3.1, the third step follows $X_{j,k_{1}+(k_{2}-1)d}=1$ when $j=k_{1}=k_{2}$, and $X_{j,k_{1}+(k_{2}-1)d}=0$ otherwise, and the last step is because of Definition 3.2.
∎

###### Claim C.20 .

Given $X\in\mathbb{R}^{d\times d^{2}}$. Note $\mathsf{X}\in\mathbb{R}^{d\times d\times d}$ denotes its tensor version.
Given matrices $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$. Following Definition C.2, we can show

$$ $\displaystyle(\underbrace{A_{1}}_{n\times d}\underbrace{X}_{d\times d^{2}} \underbrace{(A_{2}\otimes A_{3})^{\top}}_{d^{2}\times n^{2}})_{i,(j-1)n+l}=( \underbrace{\mathsf{X}(A_{1},A_{2},A_{3})}_{n\times n\times n})_{i,j,l},~{}~{} ~{}~{}\forall i\in[n],j\in[n],l\in[n]$ $$

and

$$ $\displaystyle\operatorname{vec}(\underbrace{A_{1}}_{n\times d}\underbrace{X}_{ d\times d^{2}}\underbrace{(A_{2}\otimes A_{3})^{\top}}_{d^{2}\times n^{2}})= \operatorname{vec}(\underbrace{\mathsf{X}(A_{1},A_{2},A_{3})}_{n\times n\times n }).$ $$

###### Proof.

We can show that

$$ $\displaystyle(A_{1}X(A_{2}\otimes A_{3})^{\top})_{i,(j-1)n+l}=$ $\displaystyle~{}\sum_{a=1}^{d}\sum_{b=1}^{d}\sum_{c=1}^{d}(A_{1})_{i,a}X_{a,(b -1)d+c}(A_{2})_{j,b}(A_{3})_{l,c}$ $\displaystyle=$ $\displaystyle~{}\sum_{a=1}^{d}\sum_{b=1}^{d}\sum_{c=1}^{d}\mathsf{X}_{a,b,c}(A _{1})_{i,a}(A_{2})_{j,b}(A_{3})_{l,c}$ $\displaystyle=$ $\displaystyle~{}\mathsf{X}(A_{1},A_{2},A_{3})_{i,j,l},$ $$

where the first step follows the Kronecker product Definition 3.1, the second step follows $\mathsf{X}_{a,b,c}=X_{a,(b-1)d+c}$, and the last step is due to Definition C.2.
∎

Now, we introduce a key claim that can reduce the tensor product to matrix multiplication and Kronecker product to make calculation easy.

###### Claim C.21 .

Let $\mathsf{I}_{d}\in\mathbb{R}^{d\times d\times d}$ and $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$. We have
$\mathrm{mat}(\mathsf{I}_{d}(A_{1},A_{2},A_{3}))=A_{1}\mathrm{mat}(\mathsf{I}_{
d})(A_{2}\otimes A_{3})^{\top}=A_{1}(A_{2}\oslash A_{3})^{\top}\in\mathbb{R}^{
n\times n^{2}}$ .

###### Proof.

The proof follows from Claim C.20 and Fact C.19.
∎

###### Fact C.19 .

###### Proof.

###### Claim C.20 .

###### Proof.

###### Claim C.21 .

###### Proof.

## Appendix D Gradient Formulation and Analysis

In Section D.1, we define some useful function that will help further calculation. In Section D.2, we define the expression for the loss function. In Section D.3, we give detailed gradient computation.

### D.1 Definitions for useful functions

We will introduce the definition of $\mathsf{K}$, $\alpha$, $\mathsf{S}$, and $\mathsf{L}$ used in loss formulation.

###### Definition D.1 .

We define $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ to be three matrices in size $n\times d$. Suppose that $\mathsf{A}=A_{1}\otimes A_{2}\otimes A_{3}\in\mathbb{R}^{n^{3}\times d^{3}}$.
Let $\mathsf{A}_{j_{0}}\in\mathbb{R}^{n^{2}\times d^{3}}$ represent an $n^{2}\times d^{3}$ sub-block from $\mathsf{A}$. There are $n$ such sub-blocks, i.e. the $(i+(j_{0}-1)\cdot n^{2})$-th row, $j$-th column of $\mathsf{A}$ is the $i$-th row, $j$-th column of $\mathsf{A}_{j_{0}}$, for $i\in[n^{2}],j\in[d^{3}],j_{0}\in[n]$.

For all $j_{0}\in[n]$, we denote function $\mathsf{K}(x)_{j_{0}}:\mathbb{R}^{d^{3}}\rightarrow\mathbb{R}^{n^{2}}$ as below:

$$ $\displaystyle\mathsf{K}(x)_{j_{0}}:=\underbrace{\exp(\mathsf{A}_{j_{0}}x)}_{n^ {2}\times 1}.$ $$

###### Definition D.2 .

Let three matrices $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ in size $n\times d$. We define $\mathsf{A}_{j_{0}}\in\mathbb{R}^{n^{2}\times d^{3}}$ be a $n^{2}\times d^{3}$ size sub-block from $\mathsf{A}$ (see as Definition D.1 ). (Recall that $\mathsf{A}=A_{1}\otimes A_{2}\otimes A_{3}\in\mathbb{R}^{n^{3}\times d^{3}}$.)

For any index $j_{0}\in[n]$, we denote function $\alpha(x)_{j_{0}}:\mathbb{R}^{d^{3}}\rightarrow\mathbb{R}$ as follows:

$$ $\displaystyle\alpha(x)_{j_{0}}:=\langle\underbrace{\exp(\mathsf{A}_{j_{0}}x)}_ {n^{2}\times 1},\underbrace{{\bf 1}_{n^{2}}}_{n^{2}\times 1}\rangle.$ $$

###### Definition D.3 .

Suppose that $\alpha(x)_{j_{0}}\in\mathbb{R}$ (see Definition D.2).

Recall $\mathsf{K}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ (see Definition D.1).

For a fixed $j_{0}\in[n]$, we define function $\mathsf{S}(x)_{j_{0}}:\mathbb{R}^{d^{3}}\rightarrow\mathbb{R}^{n^{2}}$ as follows:

$$ $\displaystyle\mathsf{S}(x)_{j_{0}}:=\underbrace{\alpha(x)_{j_{0}}^{-1}}_{ \mathrm{scalar}}\underbrace{\mathsf{K}(x)_{j_{0}}}_{n^{2}\times 1}.$ $$

We use $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ to denote the matrix where $j_{0}$-th row is $(\mathsf{S}(x)_{j_{0}})^{\top}$. (Note that we can rewrite $\mathsf{S}(x)=D^{-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d)\in\mathbb{R}^{n
\times n^{2}}$ and where $D=\operatorname{diag}(\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d){\bf 1}_{n^{2}})$.)

###### Definition D.4 .

Let $\mathsf{A}_{3}=A_{4}\otimes A_{5}\in\mathbb{R}^{n^{2}\times d^{2}}$, where $A_{4},A_{5},\in\mathbb{R}^{n\times d}$. Let $Y_{1},Y_{2}\in\mathbb{R}^{d\times d}$.
Let $Y=Y_{1}\oslash Y_{2}\in\mathbb{R}^{d^{2}\times d}$ denote the matrix representation of $y\in\mathbb{R}^{d^{3}}$.
For all $i_{0}\in[d]$, we define $\mathsf{L}()_{i_{0}}:\mathbb{R}^{d^{3}}\rightarrow\mathbb{R}^{n^{2}}$ as follows:

$$ $\displaystyle\mathsf{L}(y)_{i_{0}}:=\underbrace{\mathsf{A}_{3}}_{n^{2}\times d ^{2}}\underbrace{Y_{*,i_{0}}}_{d^{2}\times 1}.$ $$

Let $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ matrix where $i_{0}$ column is $\mathsf{L}(y)_{i_{0}}$. (Note that we can rewrite $\mathsf{L}(y)=(A_{4}\otimes A_{5})Y$.)

We will define $\mathsf{W}$ and $\mathsf{F}$ used in gradient analysis.

###### Definition D.5 .

Let $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$ (see Definition D.7). Let $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ (see Definition D.4).

We define $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ to be

$$ $\displaystyle\mathsf{W}(x):=\underbrace{\mathsf{V}(x)}_{n\times d}\underbrace{ \mathsf{L}(y)^{\top}}_{d\times n^{2}}$ $$

We denote $\mathsf{W}(x)_{j_{0}}^{\top}$ as the $j_{0}$-th row of $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$.

###### Definition D.6 .

For all index $j_{0}\in[n]$, let us define $\mathsf{F}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ to be

$$ $\displaystyle\underbrace{\mathsf{F}(x)_{j_{0}}}_{{n^{2}\times 1}}:=\underbrace {(\operatorname{diag}(\mathsf{S}(x)_{j_{0}})-\mathsf{S}(x)_{j_{0}}\mathsf{S}(x )_{j_{0}}^{\top})}_{{n^{2}\times n^{2}}}\underbrace{\mathsf{W}(x)_{j_{0}}}_{{n ^{2}\times 1}}.$ $$

We define $\mathsf{F}(x)\in\mathbb{R}^{n\times n^{2}}$ in the sense that $\mathsf{F}(x)_{j_{0}}^{\top}$ is the $j_{0}$-th row of $\mathsf{F}(x)$.

###### Definition D.1 .

###### Definition D.2 .

###### Definition D.3 .

###### Definition D.4 .

###### Definition D.5 .

###### Definition D.6 .

### D.2 Definitions for loss function

We now present some useful definitions pertaining to $x\in\mathbb{R}^{d^{3}}$.

###### Definition D.7 .

For all $j_{0}\in[n]$, we denote $\mathsf{S}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$
as the normalized vector (see Definition D.3). For all $i_{0}\in[d]$, we denote $\mathsf{L}(y)_{i_{0}}$ to be the same in Definition D.4.

Consider every $j_{0}\in[n]$, every $i_{0}\in[d]$. Let us consider $\mathsf{V}(x)_{j_{0},i_{0}}:\mathbb{R}^{d^{3}}\rightarrow\mathbb{R}$
as follows:

$$ $\displaystyle\mathsf{V}(x)_{j_{0},i_{0}}:=\langle\mathsf{S}(x)_{j_{0}},\mathsf {L}(y)_{i_{0}}\rangle-E_{j_{0},i_{0}},$ $$

where $E_{j_{0},i_{0}}$ is the $(j_{0},i_{0})$-th coordinate of $E\in\mathbb{R}^{n\times d}$ for $j_{0}\in[n],i_{0}\in[d]$.
This is the same as $\underbrace{\mathsf{V}(x)}_{n\times d}=\underbrace{\mathsf{S}(x)}_{n\times n^{
2}}\underbrace{\mathsf{L}(y)}_{n^{2}\times d}-\underbrace{E}_{n\times d}$.

###### Definition D.8 .

For all $j_{0}\in[n]$, for all $i_{0}\in[d]$. We define
$\mathsf{Loss}(x)_{j_{0},i_{0}}$ to be $:=0.5\mathsf{V}(x)_{j_{0},i_{0}}^{2}$.

###### Definition D.7 .

###### Definition D.8 .

### D.3 Further information on gradient computation

In this section, we offer detailed analysis to help the computations of gradient and derivative. It is noted that, for the sake of convenience in deriving a closed-form expression for our gradient, we omit the $1/d$ normalization factor in $\mathsf{S}$. As this factor merely scales the result, it does not impact the overall computation of these matrices.

###### Remark D.9 .

Recall that in Definition 3.8, we consider $X\in\mathbb{R}^{d\times d\times d}$ for gradient computation, which has $d^{3}$ number of parameters.
On the other hand, in Definition 3.9, we have $X=X_{1}\cdot(X_{2}^{\top}\ominus X_{3}^{\top})\in\mathbb{R}^{d\times d^{2}}$ which has $3d^{2}$ number of parameters, which indeed guarantee computation acceleration.

###### Lemma D.10 (The gradient computation for various functions w.r.t. $x_{i}$ ) .

Let $x\in\mathbb{R}^{d^{3}}$. Let $j_{0}\in[n],i_{0}\in[d]$.
For all $i\in[d^{3}]$, we define $\mathsf{A}_{j_{0},i}\in\mathbb{R}^{n^{2}}$ to be the $i$-th column for $\mathsf{A}_{j_{0}}\in\mathbb{R}^{n^{2}\times d^{3}}$. Recall that $\mathsf{K}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ is defined in Definitions D.1. The scalar function $\alpha(x)_{j_{0}}\in\mathbb{R}$ is defined in Definitions D.2 . Column function $\mathsf{S}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ is defined in Definitions D.3. Scalar function $\mathsf{V}(x)_{j_{0},i_{0}}\in\mathbb{R}$ is defined in Definitions D.7. Scalar function $\mathsf{Loss}(x)_{j_{0},i_{0}}\in\mathbb{R}$ is defined in Definitions D.8.

Then, for each $i\in[d^{3}]$, we have

- •
Part 1.
$\displaystyle\frac{\mathrm{d}x}{\mathrm{d}x_{i}}=e_{i}$
- •
Part 2. For any $j_{0}\in[n]$,
$\displaystyle\frac{\mathrm{d}\mathsf{A}_{j_{0}}x}{\mathrm{d}x_{i}}=\mathsf{A}_
{j_{0},i}$
- •
Part 3. For any $j_{0}\in[n]$
$\displaystyle\frac{\mathrm{d}\mathsf{K}(x)_{j_{0}}}{\mathrm{d}x_{i}}=\mathsf{A
}_{j_{0},i}\circ\mathsf{K}(x)_{j_{0}}$
- •
Part 4. For any $j_{0}\in[n]$,
$\displaystyle\frac{\mathrm{d}\alpha(x)_{j_{0}}}{\mathrm{d}x_{i}}=\langle
\mathsf{A}_{j_{0},i},\mathsf{K}(x)_{j_{0}}\rangle$
- •
Part 5. For any $j_{0}\in[n]$,
$\displaystyle\frac{\mathrm{d}\mathsf{S}(x)_{j_{0}}}{\mathrm{d}x_{i}}=\mathsf{A
}_{j_{0},i}\circ\mathsf{S}(x)_{j_{0}}-\langle\mathsf{A}_{j_{0},i},\mathsf{S}(x
)_{j_{0}}\rangle\cdot\mathsf{S}(x)_{j_{0}}$
- •
Part 6. For any $j_{0}\in[n]$, for any $i_{0}\in[d]$,
$\displaystyle\frac{\mathrm{d}\langle\mathsf{S}(x)_{j_{0}},\mathsf{L}(y)_{i_{0}
}\rangle}{\mathrm{d}x_{i}}=\langle\mathsf{L}(y)_{i_{0}},\mathsf{A}_{j_{0},i}
\circ\mathsf{S}(x)_{j_{0}}\rangle-\langle\mathsf{L}(y)_{i_{0}},\mathsf{S}(x)_{
j_{0}}\rangle\cdot\langle\mathsf{A}_{j_{0},i},\mathsf{S}(x)_{j_{0}}\rangle$
- •
Part 7. For any $j_{0}\in[n]$, for each $i_{0}\in[d]$
$\displaystyle\frac{\mathrm{d}\mathsf{V}(x)_{j_{0},i_{0}}}{\mathrm{d}x_{i}}=
\langle\mathsf{A}_{j_{0},i}\circ\mathsf{S}(x)_{j_{0}},\mathsf{L}(y)_{i_{0}}
\rangle-\langle\mathsf{S}(x)_{j_{0}},\mathsf{L}(y)_{i_{0}}\rangle\cdot\langle
\mathsf{A}_{j_{0},i},\mathsf{S}(x)_{j_{0}}\rangle$
- •
Part 8. For any $j_{0}\in[n]$, for each $i_{0}\in[d]$
$\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x)_{j_{0},i_{0}}}{\mathrm{d}x_{i}}=
(\langle\mathsf{L}(y)_{i_{0}},\mathsf{A}_{j_{0},i}\circ\mathsf{S}(x)_{j_{0}}
\rangle-\langle\mathsf{S}(x)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle\cdot\langle
\mathsf{L}(y)_{i_{0}},\mathsf{S}(x)_{j_{0}}\rangle)\cdot\mathsf{V}(x)_{j_{0},i
_{0}}$

###### Proof.

Proof of Part 1.
We have

$$ $\displaystyle\frac{\mathrm{d}x}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\frac{\mathrm{d}[x_{1},x_{2},\dots,x_{d^{3}}]^{\top}}{\mathrm{ d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}e_{i}$ $$

where the first step follows from $x$ is a vector, and the second step follows from all coordinates are independent to each other.

Proof of Part 2.
We have

$$ $\displaystyle\frac{\mathrm{d}\mathsf{A}_{j_{0}}x}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\underbrace{\mathsf{A}_{j_{0}}}_{n^{2}\times d^{3}}\underbrace {\frac{\mathrm{d}x}{\mathrm{d}x_{i}}}_{d^{3}\times 1}$ $\displaystyle=$ $\displaystyle~{}\underbrace{\mathsf{A}_{j_{0}}}_{n^{2}\times d^{3}}\cdot \underbrace{e_{i}}_{d^{3}\times 1}$ $\displaystyle=$ $\displaystyle~{}\underbrace{\mathsf{A}_{j_{0},i}}_{n^{2}\times 1}$ $$

where the second step follows from Part 1.

Proof of Part 3.

It’s easy to show that

$$ $\displaystyle\underbrace{\frac{\mathrm{d}\mathsf{K}(x)_{j_{0}}}{\mathrm{d}x_{i }}}_{n^{2}\times 1}=$ $\displaystyle~{}\frac{\mathrm{d}\exp(\mathsf{A}_{j_{0}}x)}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\exp(\mathsf{A}_{j_{0}}x)\circ\frac{\mathrm{d}\mathsf{A}_{j_{0 }}x}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\exp(\mathsf{A}_{j_{0}}x)\circ\mathsf{A}_{j_{0},i}$ $\displaystyle=$ $\displaystyle~{}\underbrace{\mathsf{K}(x)_{j_{0}}}_{n^{2}\times 1}\circ \underbrace{\mathsf{A}_{j_{0},i}}_{n^{2}\times 1}$ $$

where the third step is because of Part 2, the last step follows from definition of $\mathsf{K}(x)_{j_{0}}$.

Proof of Part 4.

To further simplify the writing of proofs, we represent $(x)$ as $(\cdot)$.

It’s easy to see that

$$ $\displaystyle\frac{\mathrm{d}\alpha(\cdot)_{j_{0}}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\frac{\mathrm{d}\langle\mathsf{K}(\cdot)_{j_{0}},{\bf 1}_{n^{2 }}\rangle}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\langle\mathsf{K}(\cdot)_{j_{0}}\circ\mathsf{A}_{j_{0},i},{\bf 1 }_{n^{2}}\rangle$ $\displaystyle=$ $\displaystyle~{}\langle\mathsf{K}(\cdot)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle$ $$

where the first step is due to definition of $\alpha(\cdot)$, the second step is because of Part 3, the third step comes from $\langle a\circ b,{\bf 1}_{n^{2}}\rangle=\langle a,b\rangle$.

Proof of Part 5.

To further simplify the writing of proofs, we represent $(x)$ as $(\cdot)$.

It’s easy to see that

$$ $\displaystyle\frac{\mathrm{d}\mathsf{S}(\cdot)_{j_{0}}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\frac{\mathrm{d}\alpha(\cdot)_{j_{0}}^{-1}\mathsf{K}(\cdot)_{j _{0}}}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\alpha(\cdot)_{j_{0}}^{-1}\frac{\mathrm{d}\mathsf{K}(\cdot)_{j _{0}}}{\mathrm{d}x_{i}}+(\frac{\mathrm{d}\alpha(\cdot)_{j_{0}}^{-1}}{\mathrm{d }x_{i}})\mathsf{K}(\cdot)_{j_{0}}$ $$

For the first term, we have

$$ $\displaystyle\alpha(\cdot)_{j_{0}}^{-1}\frac{\mathrm{d}\mathsf{K}(\cdot)_{j_{0 }}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\alpha(\cdot)_{j_{0}}^{-1}\mathsf{K}(\cdot)_{j_{0}}\circ \mathsf{A}_{j_{0},i}$ $\displaystyle=$ $\displaystyle~{}\mathsf{S}(\cdot)_{j_{0}}\circ\mathsf{A}_{j_{0},i}$ $$

where the first step is due to Part 3, the second step is because of definition of $\mathsf{S}(\cdot)$.

For the second term, we have

$$ $\displaystyle(\frac{\mathrm{d}\alpha(\cdot)_{j_{0}}^{-1}}{\mathrm{d}x_{i}}) \mathsf{K}(\cdot)_{j_{0}}=$ $\displaystyle~{}-\alpha(\cdot)_{j_{0}}^{-2}\frac{\mathrm{d}\alpha(\cdot)_{j_{0 }}}{\mathrm{d}x_{i}}\mathsf{K}(\cdot)_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}-\alpha(\cdot)_{j_{0}}^{-2}\cdot\langle\mathsf{K}(\cdot)_{j_{0 }},\mathsf{A}_{j_{0},i}\rangle\cdot\mathsf{K}(\cdot)_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}-\mathsf{S}(\cdot)_{j_{0}}\cdot\langle\mathsf{S}(\cdot)_{j_{0} },\mathsf{A}_{j_{0},i}\rangle$ $$

where the first step is from simple calculus, the second step is from Part 4, and the third step is due to the definition of $\mathsf{S}(\cdot)_{j_{0}}$.

By applying all of the above, we have

$$ $\displaystyle\frac{\mathrm{d}\mathsf{S}(\cdot)_{j_{0}}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\mathsf{S}(\cdot)_{j_{0}}\circ\mathsf{A}_{j_{0},i}-\mathsf{S}( \cdot)_{j_{0}}\cdot\langle\mathsf{S}(\cdot)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle$ $$

Proof of Part 6.
From Part 5, clearly this holds.

Proof of Part 7.

To further simplify the writing of proofs, we represent $(x)$ as $(\cdot)$.

From definition of $\mathsf{V}$ in Definition D.7, it holds that

$$ $\displaystyle\mathsf{V}(\cdot)_{j_{0},i_{0}}:=\langle\mathsf{S}(\cdot)_{j_{0}} ,\mathsf{L}(y)_{i_{0}}\rangle-E_{j_{0},i_{0}}$ (2) $$

Thus it holds that

$$ $\displaystyle\frac{\mathrm{d}\mathsf{V}(\cdot)_{j_{0},i_{0}}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\frac{\mathrm{d}(\langle\mathsf{S}(\cdot)_{j_{0}},\mathsf{L}(y )_{i_{0}}\rangle-E_{j_{0},i_{0}})}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}\langle\mathsf{S}(\cdot)_{j_{0}},\mathsf{L}(y) _{i_{0}}\rangle}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\langle\mathsf{S}(\cdot)_{j_{0}}\circ\mathsf{A}_{j_{0},i}, \mathsf{L}(y)_{i_{0}}\rangle-\langle\mathsf{S}(\cdot)_{j_{0}},\mathsf{L}(y)_{i _{0}}\rangle\cdot\langle\mathsf{S}(\cdot)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle,$ $$

where the first step comes from Eq. (2), the second step follows from $\frac{\mathrm{d}E_{j_{0},i_{0}}}{\mathrm{d}x_{i}}=0$, and the last step is due to Part 6.

Proof of Part 8.

To further simplify the writing of proofs, we represent $(x)$ as $(\cdot)$.

From definition of $\mathsf{Loss}(\cdot)$ (see Definition D.8), it holds that

$$ $\displaystyle\mathsf{Loss}(\cdot)_{j_{0},i_{0}}=0.5\mathsf{V}(\cdot)_{j_{0},i_ {0}}^{2}$ (3) $$

Thus, we have

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(\cdot)_{j_{0},i_{0}}}{\mathrm{d}x_{ i}}=$ $\displaystyle~{}\frac{\mathrm{d}(0.5\mathsf{V}(\cdot)_{j_{0},i_{0}}^{2})}{ \mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\mathsf{V}(\cdot)_{j_{0},i_{0}}\frac{\mathrm{d}\mathsf{V}( \cdot)}{\mathrm{d}x_{i}}$ $\displaystyle=$ $\displaystyle~{}\mathsf{V}(\cdot)_{j_{0},i_{0}}\cdot(\langle\mathsf{S}(\cdot)_ {j_{0}}\circ\mathsf{A}_{j_{0},i},\mathsf{L}(y)_{i_{0}}\rangle-\langle\mathsf{S }(\cdot)_{j_{0}},\mathsf{L}(y)_{i_{0}}\rangle\cdot\langle\mathsf{S}(\cdot)_{j_ {0}},\mathsf{A}_{j_{0},i}\rangle),$ $$

where the 1st step comes from the Eq. (3), the second step follows from the chain rule, and the last step is because of Part 7.

∎

###### Remark D.9 .

###### Lemma D.10 (The gradient computation for various functions w.r.t. x i subscript 𝑥 𝑖 x_{i} italic_x start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) .

###### Proof.

## Appendix E Tensor Attention Exact Gradient Computation Time Complexity

Section E.1 demonstrates how to calculate $\mathsf{S}$ ($1/d$ factor is still ignored) and $\mathsf{L}$.
Section E.2 explains the straightforward method for calculating $\mathsf{V}$.
Section E.3 and Section E.4 define $\mathsf{F}$ and $\mathsf{W}$, and demonstrate their computations.
Section E.5 presents a more elegant way to express the gradient.
Finally, Section E.6 combines all these elements and determine the overall time complexity of our algorithm.

### E.1 Time complexity to get 𝖲 𝖲 \mathsf{S} sansserif_S and 𝖫 𝖫 \mathsf{L} sansserif_L

###### Remark E.1 .

Note that ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})\geq\Omega(n^{3})$.

Now we will show the time complexity for computing $\mathsf{S}$ and $\mathsf{L}$.

###### Lemma E.2 (Computing $\mathsf{S}$ and $\mathsf{L}$ ) .

If the following conditions hold

- •
Let $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definition D.3)
- •
Let $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ (see Definition D.4)

Then, we have

- •
the time complexity of $\mathsf{S}(x)$ is ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})+{\cal T}_{\mathrm{mat}}(n,d,d^{2})$
- •
the time complexity of $\mathsf{L}(y)$ is ${\cal T}_{\mathrm{mat}}(n^{2},d^{2},d)$

###### Proof.

Note that

$$ $\displaystyle\mathsf{S}(x)=\underbrace{D^{-1}}_{n\times n}\exp(\underbrace{A_{ 1}}_{n\times d}\underbrace{X}_{d\times d^{2}}\underbrace{(A_{2}\otimes A_{3})^ {\top}}_{d^{2}\times n^{2}})$ $$

and

$$ $\displaystyle D=\operatorname{diag}(\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}){ \bf 1}_{n^{2}})$ $$

We firstly compute $\exp(A_{1}X(A_{2}\otimes A_{3})^{\top})$, this takes time of

- •
$\underbrace{A_{1}}_{n\times d}\underbrace{X}_{d\times d^{2}}$ takes ${\cal T}_{\mathrm{mat}}(n,d,d^{2})$
- •
Computing $A_{2}\otimes A_{3}$ takes $O(n^{2}d^{2})$ time
- •
Computing $A_{1}X\cdot(A_{2}\otimes A_{3})^{\top}$ takes ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})$ time

The overall time complexity of above three parts is dominated by

$$ $\displaystyle{\cal T}_{\mathrm{mat}}(n,d,d^{2})+O(d^{2}n^{2})+{\cal T}_{ \mathrm{mat}}(n,d^{2},n^{2})={\cal T}_{\mathrm{mat}}(n,d,d^{2})+{\cal T}_{ \mathrm{mat}}(n,d^{2},n^{2})$ $$

Therefore, computing $D$ takes $O(n^{3})$ time.

Computing $D^{-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top})$ requires $O(n^{3})$ time.

Therefore, the overall time complexity is

$$ $\displaystyle~{}{\cal T}_{\mathrm{mat}}(n,d,d^{2})+{\cal T}_{\mathrm{mat}}(n,d ^{2},n^{2})$ $$

It is noted that computing $\mathsf{L}(y)=\underbrace{\mathsf{A}_{3}}_{n^{2}\times d^{2}}\underbrace{Y}_{d
^{2}\times d}$ takes time of ${\cal T}_{\mathrm{mat}}(n^{2},d^{2},d)$.

Thus, we complete the proof.
∎

###### Remark E.1 .

###### Lemma E.2 (Computing 𝖲 𝖲 \mathsf{S} sansserif_S and 𝖫 𝖫 \mathsf{L} sansserif_L ) .

###### Proof.

### E.2 Time complexity to get 𝖵 𝖵 \mathsf{V} sansserif_V

We will explain the calculation of $\mathsf{V}$.

###### Lemma E.3 (Computing $\mathsf{V}$ ) .

If the following conditions hold

- •
Let $E\in\mathbb{R}^{n\times d}$
- •
Let $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$.
- •
Let $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$.

Then one can get $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$ in $O({\cal T}_{\mathrm{mat}}(n,n^{2},d))$ time.

###### Proof.

Based on the definition of $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$ which is

$$ $\displaystyle\mathsf{V}(x)=\underbrace{\mathsf{S}(x)}_{n\times n^{2}} \underbrace{\mathsf{L}(y)}_{n^{2}\times d}-\underbrace{E}_{n\times d}$ $$

It is easy to see that we can compute $\mathsf{S}(x)\mathsf{L}(y)$ in time ${\cal T}_{\mathrm{mat}}(n,n^{2},d)$, and $\mathsf{S}(x)\mathsf{L}(y)-E$ in time $O(nd)$.

Therefore, overall running time is

$$ $\displaystyle{\cal T}_{\mathrm{mat}}(n,n^{2},d)+O(nd)=O({\cal T}_{\mathrm{mat} }(n,n^{2},d)).$ $$

∎

###### Lemma E.3 (Computing 𝖵 𝖵 \mathsf{V} sansserif_V ) .

###### Proof.

### E.3 Time complexity to get 𝖶 𝖶 \mathsf{W} sansserif_W

We will explain how to calculate $\mathsf{W}$.

###### Lemma E.4 .

If the below holds that

- •
Let $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$
- •
Let $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$

Then, computing $\mathsf{W}(x)$ takes time of $O({\cal T}_{\mathrm{mat}}(n,d,n^{2}))$.

###### Proof.

Let use recall that $\mathsf{W}(x)=\mathsf{V}(x)\mathsf{L}(y)^{\top}$. This need time of ${\cal T}_{\mathrm{mat}}(n,d,n^{2})$ to compute.
∎

###### Lemma E.4 .

###### Proof.

### E.4 Time complexity to get 𝖥 𝖥 \mathsf{F} sansserif_F

We can show how to construct $\mathsf{F}$.

###### Lemma E.5 .

If the following conditions hold

- •
Let $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$
- •
Let $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$

Then, computing takes time of $\mathsf{F}(x)$ in $O(n^{3})$.

###### Proof.

For every $j_{0}\in[n]$, it follows that $\mathsf{F}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ can be computed in $O(n^{2})$, given that $\operatorname{diag}(\mathsf{S}(x)_{j_{0}})$ is a diagonal matrix and $\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top}$ is a rank-one matrix.
Consequently, constructing the matrix $\mathsf{F}(x)\in\mathbb{R}^{n\times n^{2}}$ takes a total time of $n\times O(n^{2})=O(n^{3})$.
∎

###### Lemma E.5 .

###### Proof.

### E.5 Closed form of gradient

We will give the closed form the gradient of the loss function.

###### Lemma E.6 (Closed form of gradient, formal version of Lemma 4.1 ) .

Let us define functions $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$, $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$, $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$, $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ and $\mathsf{F}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definitions D.3, D.7, D.4, D.5 and D.6 respectively). Suppose three matrices $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ are given. We define $\mathsf{A}=A_{1}\otimes A_{2}\otimes A_{3}$. Let $\mathsf{Loss}(x)$ and $\mathsf{Loss}(x)_{j_{0},i_{0}}$ be defined as Definition 3.8 and D.8.
Then, we can show that

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x)}{\mathrm{d}x}=\operatorname{vec} (A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3}))\in\mathbb{R}^{d^{3}}.$ $$

###### Proof.

From the Lemma statement and Lemma D.10 Part 8, we have

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x,y)_{j_{0},i_{0}}}{\mathrm{d}x_{i} }=\mathsf{V}(x,y)_{j_{0},i_{0}}\cdot(\langle\mathsf{S}(x)_{j_{0}}\circ\mathsf{ A}_{j_{0},i},\mathsf{L}(y)_{i_{0}}\rangle-\langle\mathsf{S}(x)_{j_{0}},\mathsf {L}(y)_{i_{0}}\rangle\cdot\langle\mathsf{S}(x)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle)$ (4) $$

We know that for all $a,b\in\mathbb{R}^{n}$, we have $\mathrm{diag}(a)\cdot b=\mathrm{diag}(b)\cdot a=a\circ b=b\circ a$.
Then, we have

$$ $\displaystyle\langle\mathsf{S}(x)_{j_{0}}\circ\mathsf{A}_{j_{0},i},\mathsf{L}( y)_{i_{0}}\rangle=(\operatorname{diag}(\mathsf{S}(x)_{j_{0}})\mathsf{A}_{j_{0} ,i})^{\top}\mathsf{L}(y)_{i_{0}}=\mathsf{A}_{j_{0},i}^{\top}\operatorname{diag }(\mathsf{S}(x)_{j_{0}})\mathsf{L}(y)_{i_{0}}$ $$

and

$$ $\displaystyle\langle\mathsf{S}(x)_{j_{0}},\mathsf{L}(y)_{i_{0}}\rangle\cdot \langle\mathsf{S}(x)_{j_{0}},\mathsf{A}_{j_{0},i}\rangle=\mathsf{A}_{j_{0},i}^ {\top}\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top}\mathsf{L}(y)_{i_{0}}$ $$

Therefore, Eq. (4) becomes

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x)_{j_{0},i_{0}}}{\mathrm{d}x_{i}}=$ $\displaystyle~{}\mathsf{V}(x,y)_{j_{0},i_{0}}\cdot(\mathsf{A}_{j_{0},i}^{\top} \operatorname{diag}(\mathsf{S}(x)_{j_{0}})\mathsf{L}(y)_{i_{0}}-\mathsf{A}_{j_ {0},i}^{\top}\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top}\mathsf{L}(y)_{i _{0}})$ $\displaystyle=$ $\displaystyle~{}\mathsf{V}(x,y)_{j_{0},i_{0}}\cdot\mathsf{A}_{j_{0},i}^{\top}( \operatorname{diag}(\mathsf{S}(x)_{j_{0}})-\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_ {j_{0}}^{\top})\mathsf{L}(y)_{i_{0}},$ (5) $$

where the second step is due to basic algebra.

Note that we defined $\mathsf{W}(x)_{j_{0}}$ in Definition D.5.

$$ $\displaystyle\mathsf{W}(x)_{j_{0}}:=\sum_{i_{0}=1}^{d}\mathsf{V}(x)_{j_{0},i_{ 0}}\mathsf{L}(y)_{i_{0}}.$ (6) $$

Also, we defined $\mathsf{F}(x)_{j_{0}}\in\mathbb{R}^{n^{2}}$ in Definition D.6,

$$ $\displaystyle\mathsf{F}(x)_{j_{0}}:=(\operatorname{diag}(\mathsf{S}(x)_{j_{0}} )-\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top})\mathsf{W}(x)_{j_{0}}.$ (7) $$

We can show

$$ $\displaystyle~{}\frac{\mathrm{d}\mathsf{Loss}(x)}{\mathrm{d}x}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\sum_{i_{0}=1}^{d}\frac{\mathrm{d}\mathsf{ Loss}(x)_{j_{0},i_{0}}}{\mathrm{d}x}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\sum_{i_{0}=1}^{d}\underbrace{\mathsf{V}(x)_ {j_{0},i_{0}}}_{\mathrm{scalar}}\cdot\underbrace{\mathsf{A}_{j_{0}}^{\top}}_{d ^{3}\times n^{2}}\underbrace{(\operatorname{diag}(\mathsf{S}(x)_{j_{0}})- \mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top})}_{n^{2}\times n^{2}} \underbrace{\mathsf{L}(y)_{i_{0}}}_{n^{2}\times 1}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\mathsf{A}_{j_{0}}^{\top}(\operatorname{diag }(\mathsf{S}(x)_{j_{0}})-\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top}) \mathsf{W}(x)_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}\sum_{j_{0}=1}^{n}\mathsf{A}_{j_{0}}^{\top}\mathsf{F}(x)_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}\mathsf{A}^{\top}\operatorname{vec}(\mathsf{F}(x))$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3 }))\in\mathbb{R}^{d^{3}}$ $$

where the first step comes from Definition 3.8, the second step is due to Eq. (E.5), the third step is because of Eq. (6), the fourth step is due to Eq. (7), the fifth step utilize the notation of $\operatorname{vec}(\cdot)$, and the last step follows from Fact C.16.

∎

###### Lemma E.6 (Closed form of gradient, formal version of Lemma 4.1 ) .

###### Proof.

### E.6 Putting all together

We now show the overall running time of computing the gradient.

###### Theorem E.7 (Tensor attention gradient computation, formal version of Theorem 4.3 ) .

If we have the following conditions

- •
Suppose that we have input fixed matrices $A_{1},A_{2},A_{3},A_{4},A_{5},E\in\mathbb{R}^{n\times d}$.
- •
We denote $X\in\mathbb{R}^{d\times d^{2}}$ and $Y\in\mathbb{R}^{d^{2}\times d}$ as matrix variables (gradient is computed w.r.t. $X$ )
–
For simplicity of calculation, we utilize vector variables $x\in\mathbb{R}^{d^{3}\times 1}$ and $y\in\mathbb{R}^{d^{3}\times 1}$, i.e., $\operatorname{vec}(X)=x$.
–
For simplicity of calculation, we use tensor variables $\mathsf{X}\in\mathbb{R}^{d\times d\times d}$ and $\mathsf{Y}\in\mathbb{R}^{d\times d\times d}$
- •
Let $g=\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}\in\mathbb{R}^{d\times d^{2}}$ (see $\mathsf{Loss}(X)$ in Definition 3.8)

Then it’s plain to see that we can compute gradient $g\in\mathbb{R}^{d\times d^{2}}$ in ${\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})$ time.

###### Proof.

Step 1. We compute $\mathsf{S}(x)$ and $\mathsf{L}(y)$. According to Lemma E.2, this takes $O({\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})+{\cal T}_{\mathrm{mat}}(n,d,d^{2}))$ time.

Step 2. We compute $\mathsf{V}(x)$. According to Lemma E.3, this takes $O({\cal T}_{\mathrm{mat}}(n,n^{2},d))$ time.

Step 3. We compute $\mathsf{W}(x)$. According to Lemma E.4, this takes $O({\cal T}_{\mathrm{mat}}(n,d,n^{2}))$ time.

Step 4. We compute $\mathsf{F}(x)$. According to Lemma E.5, this takes $O(n^{3})$ time.

Step 5. From Lemma E.6, the gradient is give by $\operatorname{vec}(A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3}))$. We know that $A_{1}^{\top}\in\mathbb{R}^{d\times n}$, $\mathsf{F}(x)\in\mathbb{R}^{n\times n^{2}}$, and $A_{2}\otimes A_{3}\in\mathbb{R}^{n^{2}\times d^{2}}$, it can be calculated in $O({\cal T}_{\mathrm{mat}}(d,n,d^{2})+{\cal T}_{\mathrm{mat}}(n,n^{2},d^{2}))$ time.

Thus, the overall running time complexity for computing the gradient is $O({\cal T}_{\mathrm{mat}}(n,d^{2},n^{2})+{\cal T}_{\mathrm{mat}}(n,d,d^{2}))$.
∎

###### Theorem E.7 (Tensor attention gradient computation, formal version of Theorem 4.3 ) .

###### Proof.

## Appendix F Running Acceleration via Polynomial Method

Remember that in the preceding section, for simplicity in the computations of the gradient, we didn’t consider the $d$ factor in $\mathsf{S}$. This factor does not affect the time complexity in our algorithms as it merely acts as a rescaling factor. We will now retake the $1/d$ in $\mathsf{S}$ factor into consideration to utilize the tools from previous work [].

In Section F.1, we demonstrate how to create a low-rank representation for $\mathsf{S}$ efficiently and explicitly.
In Section F.2, we show how to make a low-rank construction for $\mathsf{V}(x)$.
In Sections F.3, F.4, and F.5, we present low-rank representations for $\mathsf{W}(x)$, $\mathsf{F}_{a}(x)$, and $\mathsf{F}_{b}(x)$, respectively.
Finally, in Section F.6, we will consolidate all these elements to prove our final algorithmic result.

### F.1 Fast computation of 𝖲 𝖲 \mathsf{S} sansserif_S

Using the polynomial method results in [, ], we have the following low-rank representation results.

###### Lemma F.1 .

For any $B=o(\sqrt[3]{\log n})$, we have $k_{1}=n^{o(1)}$ such that:
Let $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$, $X_{1},X_{2},X_{3}\in\mathbb{R}^{d\times d}$ and $X=X_{1}\cdot(X_{2}^{\top}\ominus X_{3}^{\top})\in\mathbb{R}^{d\times d^{2}}$. Assume that each number in $\mathsf{S}(x)$ can be written using $O(\log n)$ bits.
It holds that $\max\{\|A_{1}X_{1}\|_{\infty},\|A_{2}X_{2}\|_{\infty},\|A_{3}X_{3}\|_{\infty}
\}\leq B$, then there are three matrices $U_{1},V_{1},W_{1}\in\mathbb{R}^{n\times k_{1}}$ such that $\|U_{1}(V_{1}\oslash W_{1})^{\top}-\mathsf{S}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$. Here $\mathsf{S}(x)=D^{-1}\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d)\in\mathbb{R}^{n
\times n^{2}}$ and we define $D=\operatorname{diag}(\exp(A_{1}X(A_{2}\otimes A_{3})^{\top}/d){\bf 1}_{n^{2}})$.
Moreover, these matrices $U_{1},V_{1},W_{1}$ can be created explicitly in $n^{1+o(1)}$ time.

###### Proof.

We have

$$ $\displaystyle(X_{2}^{\top}\ominus X_{3}^{\top})\cdot(A_{2}\otimes A_{3})^{\top}=$ $\displaystyle~{}((A_{2}\otimes A_{3})\cdot(X_{2}^{\top}\ominus X_{3}^{\top})^{ \top})^{\top}$ $\displaystyle=$ $\displaystyle~{}((A_{2}\otimes A_{3})\cdot(X_{2}\oslash X_{3}))^{\top}$ $\displaystyle=$ $\displaystyle~{}((A_{2}\cdot X_{2})\oslash(A_{3}\cdot X_{3}))^{\top},$ $$

where the first step is due to simple algebra, the second step comes from Fact C.4, and the last step follows Fact C.7.

Thus, we can rewrite $\mathsf{S}(x)=D^{-1}\exp(Q(K_{1}\oslash K_{2})^{\top}/d)\in\mathbb{R}^{n\times
n
^{2}}$ and we define $D=\operatorname{diag}(\exp(Q(K_{1}\oslash K_{2})^{\top}/d){\bf 1}_{n^{2}})$, where $Q=A_{1}X_{1},K_{1}=A_{2}X_{2},K_{2}=A_{3}X_{3}$.

More explicitly, we have

$$ $\displaystyle Q(K_{1}\oslash K_{2})^{\top}=$ $\displaystyle~{}A_{1}X_{1}(A_{2}X_{2}\oslash A_{3}X_{3})^{\top}$ $\displaystyle=$ $\displaystyle~{}A_{1}X_{1}(X_{2}^{\top}\ominus X_{3}^{\top})\cdot(A_{2}\otimes A _{3})^{\top}$ $\displaystyle=$ $\displaystyle~{}A_{1}X(A_{2}\otimes A_{3})^{\top},$ $$

where the 1st step is due to $Q=A_{1}X_{1},K_{1}=A_{2}X_{2},K_{2}=A_{3}X_{3}$, the 2nd step is because of the identity in the beginning of the proof, and the 3rd step follows from $X=X_{1}(X_{2}^{\top}\ominus X_{3}^{\top})$.

Thus, we finish the proof by applying Lemma 5.1.
∎

###### Lemma F.1 .

###### Proof.

### F.2 Fast computation of 𝖵 𝖵 \mathsf{V} sansserif_V

We will explain how to obtain the low rank representation of $\mathsf{V}(x)$.

###### Lemma F.2 .

We assume conditions the same as Lemma F.1. Let $d=O(\log n)$ and $k_{1}=n^{o(1)}$. We also assume that we can write each number in $E\in\mathbb{R}^{n\times d}$ and $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ using $O(\log n)$ bits. Let $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$ (see Definition D.7). Then, there are three matrices $U_{1},V_{1},W_{1}\in\mathbb{R}^{n\times k_{1}}$ we have $\|U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E-\mathsf{V}(x)\|_{\infty}\leq
\epsilon/\operatorname{poly}(n)$, where $V_{1}\oslash W_{1}\in\mathbb{R}^{n^{2}\times k_{1}}$.
Moreover, we can construct these matrices $U_{1},V_{1},W_{1}$ in $n^{1+o(1)}$ time.

###### Proof.

Let $U_{1},V_{1},W_{1}$ be the matrices in Lemma F.1. We can show that

$$ $\displaystyle\|U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E-\mathsf{V}(x)\| _{\infty}=$ $\displaystyle~{}\|U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E-\mathsf{S}(x )\mathsf{L}(y)+E\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}(V_{1}\oslash W_{1})^{\top}-\mathsf{S}(x))\mathsf{L}(y )\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is due to $\mathsf{V}(x)=\mathsf{S}(x)\mathsf{L}(y)-E$, the 2nd step comes from basic algebra, and 3rd step is due to Lemma F.1 and each number in $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ can be written using $O(\log n)$.

∎

###### Lemma F.2 .

###### Proof.

### F.3 Fast computation of 𝖶 𝖶 \mathsf{W} sansserif_W

We will explain how to obtain the low rank representation of $\mathsf{W}(x)$.

###### Lemma F.3 .

Assume the same condition as Lemma F.2. Let $k_{2}=n^{o(1)}$.
We define $\mathsf{V}(x)\in\mathbb{R}^{n\times d}$ (see Definition D.7). We define $\mathsf{L}(y)\in\mathbb{R}^{n^{2}\times d}$ (see Definition D.4).
Let $\mathsf{W}(x):=\mathsf{V}(x)\mathsf{L}(y)^{\top}\in\mathbb{R}^{n\times n^{2}}$ be defined in Definition D.5. There are three matrices $U_{2},V_{2},W_{2}\in\mathbb{R}^{n\times k_{2}}$ such that $\|U_{2}(V_{2}\oslash W_{2})^{\top}-\mathsf{W}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$. We can construct the matrices $U_{2},V_{2},W_{2}$ in $n^{1+o(1)}$ time.

###### Proof.

For $\mathsf{W}(x)$, we define its approximation as $\widetilde{\mathsf{W}}(x)$.

According to Lemma F.2, we find a good approximation $U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E$ of $\mathsf{V}(x)$, where $k_{1}=n^{o(1)}$ and $U_{1},V_{1},W_{1}\in\mathbb{R}^{n\times k_{1}}$.

Now we turn $\widetilde{\mathsf{W}}(x)$ into low-rank representation

$$ $\displaystyle\widetilde{\mathsf{W}}(x)=$ $\displaystyle~{}\underbrace{(U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E)} _{n\times d}\underbrace{\mathsf{L}(y)^{\top}}_{d\times n^{2}}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E)} _{n\times d}\underbrace{((A_{4}\otimes A_{5})\cdot(Y_{1}\oslash Y_{2}))^{\top} }_{d\times n^{2}}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E)} _{n\times d}(\underbrace{(A_{4}\cdot Y_{1})}_{n\times d}\oslash\underbrace{(A_ {5}\cdot Y_{2})}_{n\times d})^{\top},$ $$

where the 1st step is because that $U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E$ is a good approximation to $\mathsf{V}(x)$, the 2nd step comes from definition of $\mathsf{L}(y)$ (see Definition D.4), the last step is due to Fact C.7.

Thus, we let $U_{2}=U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E$, $V_{2}=A_{4}\cdot Y_{1}$ and $W_{2}=A_{5}\cdot Y_{2}$, which only takes $n^{1+o(1)}$ time. (We remark that, if we use naive way to compute $U_{2}$ that it takes $\Omega(n^{2})$, however using Lemma C.13 can beat $O(n^{2})$ time.)
We can explicitly construct $U_{2},V_{2},W_{2}\in\mathbb{R}^{n\times k_{2}}$ where $k_{2}\leq\max\{d,k_{1}\}+d=n^{o(1)}$. (Here the reason is $k_{1}=n^{o(1)}$ and $d=n^{o(1)}$)

For controlling the error, we can show

$$ $\displaystyle\|\widetilde{\mathsf{W}}(x)-\mathsf{W}(x)\|_{\infty}=$ $\displaystyle~{}\|(U_{1}(V_{1}\oslash W_{1})^{\top}\mathsf{L}(y)-E)\mathsf{L}( y)^{\top}-\mathsf{V}(x)\mathsf{L}(y)^{\top}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}d\cdot\|\mathsf{L}(y)\|_{\infty}\cdot\|U_{1}(V_{1}\oslash W_{1 })^{\top}\mathsf{L}(y)-E-\mathsf{V}(x)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n),$ $$

where the first step follows from the definition of $\tilde{\mathsf{W}}(x),\mathsf{W}(x)$, the second step follows from $\|ab^{\top}\|_{\infty}\leq d\cdot\|a\|_{\infty}\cdot\|b\|_{\infty}$ for length $d$ vectors $a,b$, and the last step follows Lemma F.2.

Thus, we complete the proof.
∎

###### Lemma F.3 .

###### Proof.

### F.4 Fast computation of 𝖥 a subscript 𝖥 𝑎 \mathsf{F}_{a} sansserif_F start_POSTSUBSCRIPT italic_a end_POSTSUBSCRIPT : key step

###### Definition F.4 .

Let $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definition D.3).
Let $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definition D.5).
Then, we define

$$ $\displaystyle\mathsf{F}_{a}(x):=\mathsf{S}(x)\circ\mathsf{W}(x)\in\mathbb{R}^{ n\times n^{2}}.$ $$

We will explain how to obtain the low-rank representation of $\mathsf{F}_{a}(x)$.

###### Lemma F.5 .

Let $k_{1}=n^{o(1)}$, $k_{2}=n^{o(1)}$, $k_{3}=n^{o(1)}$.
We assume $U_{1},V_{1},W_{1}\in\mathbb{R}^{n\times k_{1}}$ approximates the $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ satisfying $\|U_{1}(V_{1}\oslash W_{1})^{\top}-\mathsf{S}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$.
Let us assume that $U_{2},V_{2},W_{2}\in\mathbb{R}^{n\times k_{2}}$ approximates the $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ satisfying $\|U_{2}(V_{2}\oslash W_{2})^{\top}-\mathsf{W}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$.
We assume that each number in $\mathsf{S}(x)$ and $\mathsf{W}(x)$ can be written using $O(\log n)$ bits.
Let $\mathsf{F}_{a}(x):=\mathsf{S}(x)\circ\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ be defined in Definition F.4.
Then there are matrices $U_{3},V_{3},W_{3}\in\mathbb{R}^{n\times k_{3}}$ such that $\|U_{3}(V_{3}\oslash W_{3})^{\top}-\mathsf{F}_{a}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$. We can construct the matrices $U_{3},V_{3},W_{3}$ in $n^{1+o(1)}$ time.

###### Proof.

If we choose $U_{3}=U_{1}\ominus U_{2}\in\mathbb{R}^{n\times k_{1}k_{2}}$ and $V_{3}=V_{1}\ominus V_{2}\in\mathbb{R}^{n\times k_{1}k_{2}}$, $W_{3}=W_{1}\ominus W_{2}\in\mathbb{R}^{n\times k_{1}k_{2}}$, this need $n^{1+o(1)}$ time to compute.

For further simplicity of proofs, we call $\widetilde{\mathsf{S}}(x)=U_{1}(V_{1}\oslash W_{1})^{\top}$ and $\widetilde{\mathsf{W}}(x)=U_{2}(V_{2}\oslash W_{2})^{\top}$.

According to Lemma C.13, we can show

$$ $\displaystyle\|U_{3}(V_{3}\oslash W_{3})^{\top}-\mathsf{F}_{a}(x)\|_{\infty}=$ $\displaystyle~{}\|U_{3}(V_{3}\oslash W_{3})^{\top}-\mathsf{S}(x)\circ\mathsf{W }(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}\ominus U_{2})((V_{1}\ominus V_{2})\oslash(W_{1} \ominus W_{2}))^{\top}-\mathsf{S}(x)\circ\mathsf{W}(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|(U_{1}(V_{1}\oslash W_{1})^{\top})\circ(U_{2}(V_{2}\oslash W _{2})^{\top})-\mathsf{S}(x)\circ\mathsf{W}(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|\widetilde{\mathsf{S}}(x)\circ\widetilde{\mathsf{W}}(x)- \mathsf{S}(x)\circ\mathsf{W}(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|\widetilde{\mathsf{S}}(x)\circ\widetilde{\mathsf{W}}(x)- \widetilde{\mathsf{S}}(x)\circ\mathsf{W}(x)+\widetilde{\mathsf{S}}(x)\circ \mathsf{W}(x)-\mathsf{S}(x)\circ\mathsf{W}(x)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|\widetilde{\mathsf{S}}(x)\circ\widetilde{\mathsf{W}}(x)- \widetilde{\mathsf{S}}(x)\circ\mathsf{W}(x)\|_{\infty}+\|\widetilde{\mathsf{S} }(x)\circ\mathsf{W}(x)-\mathsf{S}(x)\circ\mathsf{W}(x)\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the first step is due to the definition of $\mathsf{F}_{a}(x)$, the second step is because of the definition of $U_{3},V_{3},W_{3}$, the third step is due to Fact C.8, the fourth step follows from the definition of $\widetilde{\mathsf{S}}(x)$ and $\widetilde{\mathsf{W}}(x)$, the fifth step is because of basic algebra, the sixth step comes from triangle inequality, and the last step is because bounded entries (we can write each number in $\mathsf{S}(x)$ and $\mathsf{W}(x)$ using $O(\log n)$ bits) and Lemma assumptions that $\|\widetilde{\mathsf{S}}(x)-\mathsf{S}(x)\|_{\infty}\leq\epsilon/\operatorname
{poly}(n)$ and $\|\widetilde{\mathsf{W}}(x)-\mathsf{W}(x)\|_{\infty}\leq\epsilon/\operatorname
{poly}(n)$

∎

###### Definition F.4 .

###### Lemma F.5 .

###### Proof.

### F.5 Fast computation of 𝖥 b subscript 𝖥 𝑏 \mathsf{F}_{b} sansserif_F start_POSTSUBSCRIPT italic_b end_POSTSUBSCRIPT : key step

###### Definition F.6 .

Let $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definition D.3).
Let $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ (see Definition D.5).
Then, we define $\mathsf{F}_{b}(x)\in\mathbb{R}^{n\times n^{2}}$ whose $j_{0}$-th column

$$ $\displaystyle\mathsf{F}_{b}(x)_{j_{0}}=\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{ 0}}^{\top}\mathsf{W}(x)_{j_{0}},$ $$

for each $j_{0}\in[n]$.

We will explain how to obtain the low rank representation of $\mathsf{F}_{b}(x)$.

###### Lemma F.7 .

Let $k_{1}=n^{o(1)}$, $k_{2}=n^{o(1)}$, $k_{4}=n^{o(1)}$.
Let us assume that $U_{1},V_{1},W_{1}\in\mathbb{R}^{n\times k_{1}}$ approximates the $\mathsf{S}(x)\in\mathbb{R}^{n\times n^{2}}$ satisfying $\|U_{1}(V_{1}\oslash W_{1})^{\top}-\mathsf{S}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$.
We assume $U_{2},V_{2},W_{2}\in\mathbb{R}^{n\times k_{2}}$ approximates the $\mathsf{W}(x)\in\mathbb{R}^{n\times n^{2}}$ satisfying $\|U_{2}(V_{2}\oslash W_{2})^{\top}-\mathsf{W}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$.
Assume that we can write each number in $\mathsf{S}(x)$ and $\mathsf{W}(x)$ using $O(\log n)$ bits.
Let us assume that $\mathsf{F}_{b}(x)\in\mathbb{R}^{n\times n^{2}}$ whose $j_{0}$-th column $\mathsf{F}_{b}(x)_{j_{0}}=\mathsf{S}(x)_{j_{0}}\mathsf{S}(x)_{j_{0}}^{\top}
\mathsf{W}(x)_{j_{0}}$ for each $j_{0}\in[n]$ (see Definition F.6).
Then there are matrices $U_{4},V_{4},W_{4}\in\mathbb{R}^{n\times k_{4}}$ such that $\|U_{4}(V_{4}\oslash W_{4})^{\top}-\mathsf{F}_{b}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$. We can construct the matrices $U_{4},V_{4},W_{4}$ in $n^{1+o(1)}$ time.

###### Proof.

For further simplicity of proofs, we define $\mathsf{R}(x)\in\mathbb{R}^{n}$ to be a local vector function where $\mathsf{R}(x)_{j_{0}}$ is $\langle\mathsf{S}(x)_{j_{0}},\mathsf{W}(x)_{j_{0}}\rangle$. We denote the approximation of $\mathsf{R}(x)$ to be $\widetilde{\mathsf{R}}(x)$.

It is noted that a good approximation of $\mathsf{S}(x)_{j_{0}}$ is $(U_{1}(V_{1}\oslash W_{1})^{\top})_{j_{0},*}^{\top}$. We denote the approximation of $\mathsf{S}(x)$ to be $\widetilde{\mathsf{S}}(x)=U_{1}(V_{1}\oslash W_{1})^{\top}$.

It is noted that a good approximation of $\mathsf{W}(x)_{j_{0}}$ is $(U_{2}(V_{2}\oslash W_{2})^{\top})_{j_{0},*}^{\top}$.
Let denote the approximation of $\mathsf{W}(x)$ to be $\widetilde{\mathsf{W}}(x)=U_{2}(V_{2}\oslash W_{2})^{\top}$.

Suppose that $\widetilde{\mathsf{R}}(x)_{j_{0}}:=\langle\widetilde{\mathsf{S}}(x)_{j_{0}},
\widetilde{\mathsf{W}}(x)_{j_{0}}\rangle=(U_{1}(V_{1}\oslash W_{1})^{\top})_{j
_{0},*}\cdot(U_{2}(V_{2}\oslash W_{2})^{\top})_{j_{0},*}^{\top}$.

For the side of computation time, we compute $V_{1}^{\top}V_{2}$ first and this takes $n^{1+o(1)}$ time.
Then, we compute $W_{1}^{\top}W_{2}$ and this also takes $n^{1+o(1)}$ time.

Next, we have

$$ $\displaystyle\widetilde{\mathsf{R}}(x)_{j_{0}}=$ $\displaystyle~{}(U_{1}(V_{1}\oslash W_{1})^{\top})_{j_{0},*}\cdot(U_{2}(V_{2} \oslash W_{2})^{\top})_{j_{0},*}^{\top}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(U_{1})_{j_{0},*}}_{1\times k_{1}}\underbrace{(V_{ 1}\oslash W_{1})^{\top}}_{k_{1}\times n^{2}}\underbrace{(V_{2}\oslash W_{2})}_ {n^{2}\times k_{2}}\underbrace{((U_{2})_{j_{0},*})^{\top}}_{k_{2}\times 1}$ $\displaystyle=$ $\displaystyle~{}\underbrace{(U_{1})_{j_{0},*}}_{1\times k_{1}}(\underbrace{(V_ {1}^{\top}V_{2})}_{k_{1}\times k_{2}}\circ\underbrace{(W_{1}^{\top}W_{2})}_{k_ {1}\times k_{2}})\underbrace{((U_{2})_{j_{0},*})^{\top}}_{k_{2}\times 1}$ $$

where the first step follows from the definition of $\mathsf{R}(x)$, the second step follows from $(AB)_{j_{0},*}=e_{j_{0}}(AB)=(e_{j_{0}}A)B=A_{j_{0},*}B$ for any matrices $A$ and $B$, and the third step is due to Lemma C.13.

Once we have pre-computed $V_{1}^{\top}V_{2}\in\mathbb{R}^{k_{1}\times k_{2}}$ and $W_{1}^{\top}W_{2}\in\mathbb{R}^{k_{1}\times k_{2}}$, the above step only takes $O(k_{1}k_{2})$ time. Since there $n$ coordinates, so the overall time complexity is still $O(nk_{1}k_{2})=n^{1+o(1)}$.

We can use $\widetilde{\mathsf{S}}(x)$ and $\widetilde{\mathsf{R}}(x)$ to approximate $\mathsf{F}_{b}(x)$. Let $\widetilde{\mathsf{F}}_{b}(x)=\underbrace{\operatorname{diag}(\widetilde{
\mathsf{R}}(x))}_{n\times n}\underbrace{\widetilde{\mathsf{S}}(x)}_{n\times n^
{2}}$. Because $\operatorname{diag}(\widetilde{\mathsf{R}}(x))$ is a diagonal matrix and $\widetilde{\mathsf{S}}(x)$ has low-rank representation, then obviously we know how to construct $U_{4},V_{4},W_{4}$. Basically $U_{4}=\operatorname{diag}(\widetilde{\mathsf{R}}(x))U_{1}$ and $V_{4}=V_{1}$, $W_{4}=W_{1}$.

Now, we need to control the error, and we have

$$ $\displaystyle~{}\|U_{4}(V_{4}\oslash W_{4})^{\top}-\mathsf{F}_{b}(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|\widetilde{\mathsf{F}}_{b}(x)-\mathsf{F}_{b}(x)\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}} \widetilde{\mathsf{R}}(x)_{j_{0}}-\mathsf{S}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}}\| _{\infty}$ $\displaystyle=$ $\displaystyle~{}\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}} \widetilde{\mathsf{R}}(x)_{j_{0}}-\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}( x)_{j_{0}}+\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}}-\mathsf{S}(x )_{j_{0}}\mathsf{R}(x)_{j_{0}}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}} \widetilde{\mathsf{R}}(x)_{j_{0}}-\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}( x)_{j_{0}}\|_{\infty}+\|\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}} -\mathsf{S}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}}\|_{\infty}$ $$

where the first step is due to the definition of $\widetilde{\mathsf{F}}_{b}(x)$, the second step follows from the definition of $\mathsf{F}_{b}(x)$ and $\widetilde{\mathsf{F}}_{b}(x)$, the third step follows from simple algebra, and the last step follows from triangle inequality.

For the 1st term, we have

$$ $\displaystyle\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}}\widetilde{ \mathsf{R}}(x)_{j_{0}}-\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}} \|_{\infty}\leq$ $\displaystyle~{}\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}}\|_{ \infty}\cdot|\widetilde{\mathsf{R}}(x)_{j_{0}}-\mathsf{R}(x)_{j_{0}}|$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

For the 2nd term, we have

$$ $\displaystyle\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}}\mathsf{R}(x )_{j_{0}}-\mathsf{S}(x)_{j_{0}}\mathsf{R}(x)_{j_{0}}\|_{\infty}\leq$ $\displaystyle~{}\max_{j_{0}\in[n]}\|\widetilde{\mathsf{S}}(x)_{j_{0}}-\mathsf{ S}(x)_{j_{0}}\|_{\infty}\cdot|\mathsf{R}(x)_{j_{0}}|$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

We complete the proof, by using all three equations we derived above.
∎

###### Definition F.6 .

###### Lemma F.7 .

###### Proof.

### F.6 Gradient computation in almost linear time by low rank tensor approximation

We now present our main result regarding the time complexity of our algorithm.

###### Theorem F.8 (Main result for fast gradient computation, Restatement of Theorem 5.2 ) .

Assuming the entries of $A_{1},A_{2},$ $A_{3},$ $A_{4},A_{5},E\in\mathbb{R}^{n\times d}$ and $X_{1},X_{2},X_{3},Y_{1},Y_{2}\in\mathbb{R}^{d\times d}$ are represented using $O(\log n)$ bits.
Then, there exist an algorithm that runs in $n^{1+o(1)}$ time to solve $\mathsf{ATAttLGC}(n,d=O(\log n),B=o(\sqrt[3]{\log n}),\epsilon=1/\operatorname
{poly}(n))$ (see Definition 3.9), i.e., our algorithm computes a gradient matrix $\widetilde{g}\in\mathbb{R}^{d\times d^{2}}$ satisfying $\|\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}-\widetilde{g}\|_{\infty}\leq 1
/\operatorname{poly}(n)$.

###### Proof of Theorem 5.2 .

Given size $n\times n^{2}$ matrices $\mathsf{F}(x)$ (see Definition D.6), $\mathsf{F}_{a}(x)$ (see Lemma F.7) and $\mathsf{F}_{b}(x)$ (see Lemma F.5), obviously we know

$$ $\displaystyle\mathsf{F}(x)=\mathsf{F}_{a}(x)-\mathsf{F}_{b}(x).$ $$

By applying Lemma F.1, Lemma F.2, and Lemma F.3, we confirm that the assumptions in Lemma F.5 and Lemma F.7 hold true. Therefore, we can utilize Lemma F.5 and Lemma F.7 to conclude that

- •
Let $k_{3}=n^{o(1)}$. We know that $\mathsf{F}_{a}(x)$ has approximate low rank representation $U_{3},V_{3},W_{3}\in\mathbb{R}^{n\times k_{3}}$, let $\widetilde{\mathsf{F}}_{a}(x)$ denote $U_{3}(V_{3}\oslash W_{3})^{\top}$.
- •
Let $k_{4}=n^{o(1)}$. We know that $\mathsf{F}_{b}(x)$ has approximate low rank representation $U_{4},V_{4},W_{4}\in\mathbb{R}^{n\times k_{4}}$, let $\widetilde{\mathsf{F}}_{b}(x)$ denote $U_{4}(V_{4}\oslash W_{4})^{\top}$.
- •
Let $U_{5},V_{5},W_{5}\in\mathbb{R}^{n\times k_{5}}$ denote the approximate low rank representation for $\mathsf{F}(x)$, call it $\widetilde{\mathsf{F}}(x)=U_{5}(V_{5}\oslash W_{5})^{\top}$. We have $k_{5}\leq k_{3}+k_{4}=n^{o(1)}$.

Thus, Lemmas F.1, F.2, F.3, F.5 and F.7 all are taking $n^{1+o(1)}$ time to compute.

From the Lemma E.6, we know that

$$ $\displaystyle\frac{\mathrm{d}\mathsf{Loss}(x)}{\mathrm{d}x}=\operatorname{vec} (A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3}))$ $$

We use $\operatorname{vec}(A_{1}^{\top}\widetilde{\mathsf{F}}(x)(A_{2}\otimes A_{3}))$ to do approximation, then

$$ $\displaystyle\operatorname{vec}(\underbrace{A_{1}^{\top}}_{{d\times n}} \underbrace{\widetilde{\mathsf{F}}(x)}_{{n\times n^{2}}}\underbrace{(A_{2} \otimes A_{3})}_{{n^{2}\times d^{2}}})=$ $\displaystyle~{}\operatorname{vec}(\underbrace{A_{1}^{\top}}_{{d\times n}} \underbrace{\widetilde{\mathsf{F}}(x)}_{{n\times n^{2}}}\underbrace{(A_{2}^{ \top}\otimes A_{3}^{\top})^{\top}}_{{n^{2}\times d^{2}}})$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(\underbrace{[U_{5}\odot V_{5}\odot W_{5}]}_ {{n\times n\times n}}(A_{1}^{\top},A_{2}^{\top},A_{3}^{\top}))$ $\displaystyle=$ $\displaystyle~{}\operatorname{vec}(((A_{1}^{\top}U_{5})\odot(A_{2}^{\top}V_{5} )\odot(A_{3}^{\top}W_{5}))),$ $$

where the first step is due to Fact C.4, the second step is because of Claim C.20 and Fact C.11, and the last step follows Fact C.12.

The above computation takes $n^{1+o(1)}d+d^{3}n^{o(1)}$ time.
So, overall time complexity is still $n^{1+o(1)}$.

Recall that $\widetilde{g}\in\mathbb{R}^{d\times d^{2}}$ and $\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}\in\mathbb{R}^{d\times d^{2}}$ .

We have

$$ $\displaystyle\|\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}-\widetilde{g}\|_ {\infty}=$ $\displaystyle~{}\|\operatorname{vec}(A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_ {3}))-\operatorname{vec}(A_{1}^{\top}\widetilde{\mathsf{F}}(x)(A_{2}\otimes A_ {3}))\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|A_{1}^{\top}\mathsf{F}(x)(A_{2}\otimes A_{3})-A_{1}^{\top} \widetilde{\mathsf{F}}(x)(A_{2}\otimes A_{3})\|_{\infty}$ $\displaystyle=$ $\displaystyle~{}\|A_{1}^{\top}(\mathsf{F}_{a}(x)-\mathsf{F}_{b}(x))(A_{2} \otimes A_{3})-A_{1}^{\top}(\widetilde{\mathsf{F}}_{a}(x)-\widetilde{\mathsf{F }}_{b}(x))(A_{2}\otimes A_{3})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|A_{1}^{\top}(\mathsf{F}_{a}(x)-\widetilde{\mathsf{F}}_{a}(x) )(A_{2}\otimes A_{3})\|_{\infty}+\|A_{1}^{\top}(\mathsf{F}_{b}(x)-\widetilde{ \mathsf{F}}_{b}(x))(A_{2}\otimes A_{3})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|A_{1}\|_{\infty}\|A_{2}\|_{\infty}\|A_{3}\|_{\infty}\cdot n^ {3}\cdot(\|\mathsf{F}_{a}(x)-\widetilde{\mathsf{F}}_{a}(x)\|_{\infty}+\| \mathsf{F}_{b}(x)-\widetilde{\mathsf{F}}_{b}(x)\|_{\infty})$ $\displaystyle\leq$ $\displaystyle~{}\epsilon/\operatorname{poly}(n)$ $$

where the 1st step is due to definition of $\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}$ in the above, the 2nd step follows from the definition of $\operatorname{vec}(\cdot)$, the 3rd step follows from simple algebra, the 4th step follows from triangle inequality, the 5th step follows from $\|\mathsf{T}(A_{1},A_{2},A_{3})\|_{\infty}\leq n^{3}\cdot\|\mathsf{T}\|_{
\infty}\cdot\|A_{1}\|_{\infty}\cdot\|A_{2}\|_{\infty}\cdot\|A_{3}\|_{\infty}$, where $\mathsf{T}$ is a tensor, and the last step follows from entries in $A_{1},A_{2},A_{3}$ are bounded, and $\|\mathsf{F}_{a}(x)-\widetilde{\mathsf{F}}_{a}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$, $\|\mathsf{F}_{b}(x)-\widetilde{\mathsf{F}}_{b}(x)\|_{\infty}\leq\epsilon/
\operatorname{poly}(n)$.

By picking $\epsilon=1/\operatorname{poly}(n)$, we complete the proof.
∎

###### Theorem F.8 (Main result for fast gradient computation, Restatement of Theorem 5.2 ) .

###### Proof of Theorem 5.2 .

## Appendix G Hardness

In this section, we will show the hardness of our algorithm. In Section G.1, we provide some useful tools for our results. In Section G.2,we present our main hardness results.

### G.1 Tools for backward complexity

Next, we demonstrate that the tensor attention optimization problem (see Definition 3.8) exhibits favorable behavior when applied to matrices constrained as described in Lemma 6.2:

###### Lemma G.1 .

Suppose that a fixed matrix $H\in\mathbb{R}^{n\times n^{2}}$ with entries in the interval $[1,B_{a}]$ satisfying that more than half entries of $H$ in each row are equal to $B_{a}$. Let a matrix $V\in\mathbb{R}^{n^{2}\times d}$ with entries in $\{0,1\}$. For $\lambda\in\mathbb{R}$, let us define $M_{\lambda}:=\exp(\lambda H)\in\mathbb{R}^{n\times n^{2}}$. We denote the function $f:\mathbb{R}\to\mathbb{R}$ as

$$ $\displaystyle f(\lambda):=\|\underbrace{\operatorname{diag}(M_{\lambda}{\bf 1} _{n^{2}})^{-1}}_{n\times n}\underbrace{M_{\lambda}}_{n\times n^{2}}\underbrace {V}_{n^{2}\times d}\|_{F}^{2},$ $$

Then, for every $\lambda\in\mathbb{R}$ we get

- •
$|f^{\prime}(\lambda)|\leq O(B_{a}n{{}{}{}d})$,
- •
$|f^{\prime\prime}(\lambda)|\leq O(B_{a}^{2}n{{}{}{}d})$.

###### Proof.

Let $G$ denote the $n\times n^{2}$ matrix $G=\operatorname{diag}(M_{\lambda}{\bf 1}_{n})^{-1}M_{\lambda}$. For $i\in[n],j\in[n^{2}]$, we calculate that ${M_{\lambda}}_{i,j}=e^{\lambda H_{i,j}}$ and so

$$ $\displaystyle G_{i,j}=\frac{e^{\lambda H_{i,j}}}{\sum_{k=1}^{n^{2}}e^{\lambda H _{i,k}}}.$ $$

For $\ell\in[d]$, let $S_{\ell}\subseteq[n^{2}]$ represent the set of $1$s in column $\ell$ of $V$, defined as $S_{\ell}=\{j\in[n^{2}]\mid V_{j,\ell}=1\}$.
Therefore, for each $i\in[n],\ell\in[d]$, the $(i,\ell)$ entry of the matrix $\operatorname{diag}(M_{\lambda}{\bf 1}_{n})^{-1}M_{\lambda}V$ can be shown that

$$ $\displaystyle(\operatorname{diag}(M_{\lambda}{\bf 1}_{n})^{-1}M_{\lambda}V)_{i ,\ell}$ $\displaystyle=(GV)_{i,\ell}$ $\displaystyle=\sum_{j=1}^{n^{2}}G_{i,j}V_{j,\ell}$ $\displaystyle=\sum_{j\in S_{\ell}}G_{i,j}$ $\displaystyle=\frac{\sum_{j\in S_{\ell}}e^{\lambda H_{i,j}}}{\sum_{k=1}^{n^{2} }e^{\lambda H_{i,k}}}.$ $$

where the 1st step comes from definition, the 2nd step is due to simple algebra, the 3rd step is because of definition of $S_{\ell}$, and the last step comes from definition of $G$.

Thus, we obtain:

$$ $\displaystyle f(\lambda)$ $\displaystyle=\sum_{i=1}^{n}\frac{\sum_{\ell=1}^{d}\left(\sum_{j\in S_{\ell}}e ^{\lambda H_{i,j}}\right)^{2}}{\left(\sum_{k=1}^{n^{2}}e^{\lambda H_{i,k}} \right)^{2}}$ $\displaystyle=\sum_{i=1}^{n}\frac{\sum_{\ell=1}^{d}\sum_{j_{1}\in S_{\ell}} \sum_{j_{2}\in S_{\ell}}e^{\lambda(H_{i,j_{1}}+H_{i,j_{2}})}}{\sum_{k_{1}=1}^{ n^{2}}\sum_{k_{2}=1}^{n^{2}}e^{\lambda(H_{i,k_{1}}+H_{i,k_{2}})}}.$ $$

We define

$$ $\displaystyle g(\lambda,i):=\sum_{\ell=1}^{d}\sum_{j_{1}\in S_{\ell}}\sum_{j_{ 2}\in S_{\ell}}e^{\lambda(H_{i,j_{1}}+H_{i,j_{2}})}.$ $$

We also define

$$ $\displaystyle h(\lambda,i):=\sum_{k_{1}=1}^{n^{2}}\sum_{k_{2}=1}^{n^{2}}e^{ \lambda(H_{i,k_{1}}+H_{i,k_{2}})}$ $$

By the previous three equations, we have:

$$ $\displaystyle f(\lambda)=\sum_{i=1}^{n}g(\lambda,i)/h(\lambda,i).$ $$

As at least half of the entries in each row of $H$ are equal to $B_{a}$ and all entries lie within the interval $[1,B_{a}]$, we can bound:

$$ $\displaystyle\left(\frac{n^{2}}{2}\right)^{2}\cdot e^{2B_{a}\lambda}\leq h( \lambda,i)\leq\left(n\right)^{4}\cdot e^{2B_{a}\lambda}.$ (8) $$

Furthermore, since the derivative of $e^{\lambda(H_{i,k_{1}}+H_{i,k_{2}})}$ with respect to $\lambda$ is $(H_{i,k_{1}}+H_{i,k_{2}})\cdot e^{\lambda(H_{i,k_{1}}+H_{i,k_{2}})}$, we can bound

$$ $\displaystyle 2\cdot h(\lambda,i)\leq\frac{\mathrm{d}h(\lambda,i)}{\mathrm{d} \lambda}\leq 2B_{a}\cdot h(\lambda,i).$ (9) $$

We may similarly bound

$$ $\displaystyle 0\leq g(\lambda,i)\leq{{}{}{}d\cdot}n^{4}\cdot e^{2B_{a}\lambda},$ (10) $$

and

$$ $\displaystyle 2\cdot g(\lambda,i)\leq\frac{\mathrm{d}g(\lambda,i)}{\mathrm{d} \lambda}\leq 2B_{a}\cdot g(\lambda,i).$ (11) $$

The derivative of $f$ can be bounded by (where the ^′ notation denotes the derivative w.r.t. $\lambda$):

$$ $\displaystyle f^{\prime}(\lambda)$ $\displaystyle=\sum_{i=1}^{n}\frac{g^{\prime}(\lambda,i)\cdot h(\lambda,i)-g( \lambda,i)\cdot h^{\prime}(\lambda,i)}{(h(\lambda,i))^{2}}$ $\displaystyle\leq\sum_{i=1}^{n}\frac{g^{\prime}(\lambda,i)\cdot h(\lambda,i)}{ (h(\lambda,i))^{2}}$ $\displaystyle=\sum_{i=1}^{n}\frac{g^{\prime}(\lambda,i)}{h(\lambda,i)}$ $\displaystyle\leq\sum_{i=1}^{n}\frac{2B_{a}{{}{}{}d}\cdot n^{4}e^{2B_{a} \lambda}}{(n^{2}/2)^{2}\cdot e^{2B_{a}\lambda}}$ $\displaystyle=\sum_{i=1}^{n}8B_{a}{{}{}{}d}$ $\displaystyle=8B_{a}\cdot n{{}{}{}d}.$ $$

where the first step is due to the calculation of derivative, the second step is due to basic algebra, the third step is because of cancelling $h(\lambda,i)$, the fourth step is by Eq. (8) ($h(\lambda,i)$ term) and Eq. (11) ( $g^{\prime}(\lambda,i)$ term), the fifth step is due to basic algebra, and the last step is due to basic algebra.

In a similar manner, a lower bound for $f^{\prime}(\lambda)$ can be,

$$ $\displaystyle f^{\prime}(\lambda)$ $\displaystyle=\sum_{i=1}^{n}\frac{g^{\prime}(\lambda,i)\cdot h(\lambda,i)-g( \lambda,i)\cdot h^{\prime}(\lambda,i)}{(h(\lambda,i))^{2}}$ $\displaystyle\geq-\sum_{i=1}^{n}\frac{g(\lambda,i)\cdot h^{\prime}(\lambda,i)} {(h(\lambda,i))^{2}}$ $\displaystyle\geq-\sum_{i=1}^{n}\frac{({{}{}{}d}n^{4}\cdot e^{2B_{a}\lambda}) \cdot(2B_{a}\cdot h(\lambda,i))}{((n^{2}/2)^{2}\cdot e^{2B_{a}\lambda})\cdot(h (\lambda,i))}$ $\displaystyle=-\sum_{i=1}^{n}8B_{a}{{}{}{}d}$ $\displaystyle=-8B_{a}\cdot n{{}{}{}d}.$ $$

where the first step is due to the definition, the second step is due to basic algebra, the third step comes from Eq. (8) ($h(\lambda,i)$ term), Eq. (9) ($h^{\prime}(\lambda,i)$term), and Eq. (10) ($g(\lambda,i)$ term), the fourth step is due to basic algebra, and the final step comes from basic algebra.

Finally, we let $f(\lambda,i):=\frac{g(\lambda,i)}{h(\lambda,i)}$, and we can have $f^{\prime\prime}(\lambda)$ is equal to the following using the quotient rule:

$$ $\displaystyle\sum_{i=1}^{n}\frac{g^{\prime\prime}(\lambda,i)-h^{\prime\prime}( \lambda,i)\cdot f(\lambda,i)-2\cdot h^{\prime}(\lambda,i)\cdot f^{\prime}( \lambda,i)}{h(\lambda,i)},$ $$

which we can likewise bound in magnitude by $O(B_{a}^{2}n{{}{}{}d})$.
∎

We have the following tool from previous work.

###### Lemma G.2 (Lemma 5.4 in [ 14 ] ) .

Suppose that $f:[0,1]\to\mathbb{R}$ is a twice-differentiable function that satisfy $|f^{\prime\prime}(\lambda)|\leq b$ for all $\lambda\in[0,1]$. And for any positive integer $t$, we define

$$ $\displaystyle s_{t}:=\sum_{i=0}^{t-1}\frac{f^{\prime}(i/t)}{t}$ $$

Then, we have

$$ $\displaystyle|s_{t}-(f(1)-f(0))|\leq b/t.$ $$

###### Lemma G.1 .

###### Proof.

###### Lemma G.2 (Lemma 5.4 in [ 14 ] ) .

### G.2 Main result for lower bound

Finally, we are prepared to present our main result:

###### Theorem G.3 (Main result for hardness, Restatement of Theorem 6.3 ) .

Let $\gamma:\mathbb{N}\to\mathbb{N}$ be any function with $\gamma(n)=o(\log n)$ and $\gamma(n)=\omega(1)$.
Assuming $\mathsf{SETH}$, for any constant $\delta>0$, it is impossible to solve $\mathsf{ATAttLGC}(n,d=\Theta(\log n),B=\Theta(\sqrt[3]{\gamma(n)\cdot\log n}),
\epsilon=O(1/(\log n)^{4}))$ (Definition 3.9) in time $O(n^{{{}{}{}3}-\delta})$
when $E=0$, $\mathsf{Y}=\mathsf{I}_{d}$, $\mathsf{X}=\lambda\mathsf{I}_{d}$ for some scalar $\lambda\in[0,1]$.

###### Proof of Theorem 6.3 .

Let us assume that such an algorithm do exist. Then we can call it $O((\log n)^{11})$ times to refute Lemma 6.2 using parameter $\gamma=\gamma(n)$, i.e., we can get $f(1)$ by solving $\mathsf{ATAttLGC}$ with $O((\log n)^{11})$ times.

Suppose that $\mathsf{I}_{d}\in\mathbb{R}^{d\times d\times d}$ is an identity tensor.
Also suppose that the input matrices to Lemma 6.2 are $Q,K_{1},K_{2},V_{1},V_{2}$. And we set $A_{1}=Q$, $A_{2}=K_{1}$,$A_{3}=K_{2}$, $A_{4}=V_{1}$,$A_{5}=V_{2}$, $Y=I$, and $X=\lambda\cdot\underbrace{\mathrm{mat}(\mathsf{I}_{d})}_{d\times d^{2}}$, with some $\lambda\in[0,1]$.
Let $f:[0,1]\to\mathbb{R}$ be defined in Lemma G.1 where $H$ is the matrix $A_{1}(A_{2}\oslash A_{3})^{\top}$, so that $M_{\lambda}$ is the matrix $\exp(A_{1}X(A_{2}\otimes A_{3})^{\top})$ by Fact C.19. It follows from Lemma G.1 and $d=\Theta(\log n)$ that

$$ $\displaystyle|f^{\prime\prime}(\lambda)|\leq O(n{{}{}{}}\log^{{}{}{}5}n\cdot( \gamma(n))^{2}),$ $$

where $B_{a}=O(\gamma(n)\log^{2}n)$ in Lemma G.1 by the second bullet point of Lemma 6.2.

It is worth noting that $f(0)$ can be computed in $\widetilde{O}(n)$ time because of the all-1s matrix $M_{f}$. Our final target is to calculate $f(1)$.

From Lemma G.2, $f^{\prime}(\lambda)$ can be computed on
$O(\log^{9}(n)(\gamma(n))^{2})=O(\log^{11}n)$
points up to error $O(1/(\log n)^{4})$, and give back their average.
Because we have already chosen $X=\lambda I$, $f^{\prime}(\lambda)$ can be calculated from the gradient $\frac{\mathrm{d}\mathsf{Loss}(X)}{\mathrm{d}X}$ in (see Definition 3.9), by our assumed approximated algorithm.
∎

###### Theorem G.3 (Main result for hardness, Restatement of Theorem 6.3 ) .

###### Proof of Theorem 6.3 .

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: - Line 260-261 are not clear at all: Z has not been defined, it is impossible to guess what A1 to A5 are. A posteriori, after reading the next page, the reader can retrospectively guess what was meant by this statement but it is, in my opinion, not reasonable to expect the reader to not stumble on these lines and be confused (as I was). 

- Similarly, when reading Definition 3.8 it is not clear at all why the minimization is only w.r.t. X. Only after the definition is explained that minimization w.r.t. the other variables is ignored because the minimization w.r.t. X is the computational bottleneck. In my opinion, the authors should explained this *before* the definition. 

- Line 3.9: unless I am mistaken the log(n) bit model has never been introduced.

- Introducing the column-wise Kronecker product as "one kind of tensor operation" (line 59 and in other places) can be confusing. I would suggest directly calling it by its name along with a forward reference to the definition. 

- The column-wise Kronecker product is known as the Kathri-Rao product. The row-wise Kronecker product is sometimes referred to as the face-splitting product (https://en.wikipedia.org/wiki/Khatri%E2%80%93Rao_product), this should be mentioned. 

- Most of the proofs of the basic facts in appendix can be found in the seminal survey (Kolda and Bader, 2009), e.g. C.7, C.11, C.12. Some are trivial, e.g. part 1 of C.10. While I understand the desire to be self-contained, this long list of basic facts intertwined with less obvious identities  sometimes hinders the readability and understanding of the overall proof. 

**Typos / imprecisions**

line 170: w.r.t -> w.r.t.

line 173: "to denote a length nd vector" is confusing. Directly gives the definition: "to denote the length nd vector obtained by stacking the rows of A into a column vector". 

line 176: "denote an identity matrix" -> "denote the identity matrix" (there is only one). Idem for the following one when defining *the* identity tensor (not "an" identity tensor).

line 181: "And we have X = mat(X)" is not a proper way to introduce a mathematical statement, in my opinion. 

line 180: tesnorization

def 3.2 and 3.3. : "define matrix" -> "define the matrix"

[...]
_I stopped keeping track of such small typos / grammatical errors at this point. The paper needs a thorough proof read._
soundness: 2
strengths: - The problem of efficient computation of higher-order attention mechanism is relevant

- A hardness analysis is provided to strengthen the result.
summary: This paper derives an efficient algorithm to approximate the gradients of high-order transformers. Such models extend the classical self attention mechanism, where only pairwise interactions are modelled, to higher order interactions where e.g. each token can attend to every possible pairs of tokens in the sentence. This model can also be used to model interactions between different modalities or views of a same object. 

The attention model itself was already introduced previously in 2023 and 2024, along with an efficient algorithm to approximate the forward pass computation. The contributions of the submission is to design an analogous efficient algorithm for fast approximation of the backward pass computation. 

The main theoretical result shows that the proposed approximation algorithm can achieve epsilon = 1 / poly(n) approximation guarantee in almost linear time to compute gradients over a sequence of length n. The authors also provide a hardness analysis showing that their assumption are tight, in some sense. 

The contribution is only theoretical, no experiments are provided.
weaknesses: - The paper is poorly written and hard to follow. In my opinion, it is not ready for publication.

- No experiments are provided to demonstrate the effectiveness (and correctness) of the proposed analysis / algorithm.

- Some key aspects of the results are (at least to me) very difficult to get from the paper. E.g., what are the "assumptions" referred to in the statement "_we proved the necessity and tightness of our assumption_"? What were the key novel techniques needed to extend the previous approach to the gradient computation? 

- The proof (in appendix) is very difficult to follow and the results hard to interpret because they are presented in a very convoluted manner. E.g., Lemma 4.1 (which is central to the derivation and understanding of the main theorem) involves a quantity F(x) which is only defined in the Appendix. Furthermore, the definition given in the appendix itself relies on other quantities defined previously in the appendix (S and W), which are themselves  functions of other quantities previously defined (K,V,L...).... After processing backward through this series of definitions, one realize that F(x) is actually a function of many key variables of the problem (A1,A2,A3,Y,...). In the end it is thus very difficult to understand what is Lemma 4.1 suppose to highlight, since F(x) hides many dependencies on central variables of the attention computation. Note that not much context is given before the lemma to help/guide the reader. This is unfortunately not an isolated example.

- This is minor, but there are many grammatical imprecisions / errors that need to be fixed.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: See above
soundness: 3
strengths: 1. It is important to study the gradient approximation of tensor attention, which could make this structure practically usable.

2. For proving the results, the authors establish some results of tensor computation, which may be useful for related fields.
summary: This paper studies an approximation algorithm for computing gradients of high-order tensor attention. While the exact computation has cubic complexity, the authors prove that the proposed approximation algorithm has almost linear complexity. Also, the authors prove that the assumptions cannot be further weakened to achieve sub-cubic complexity.
weaknesses: 1. Since this paper  aims to propose a fast training algorithm for tensor attention, the lack of experiments could be a weakness. I am wondering about the practicality of the proposed algorithm. And are there still large gaps to implement or have experiments on these algorithms?

2. I am wondering whether the approximation error would have a major issue when training large models in practice. For example, for large models with deep architectures, will these errors accumulate and cause training difficulties?

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: -
soundness: 3
strengths: 1. The work is neatly structured and the presentation is consistent.
2. The formulations seem correct and the evidence rigorous.
3. If successfully implemented in practice, tensor attention could be in demand in ML community as it has the potential to account for more complex correlations in data.
summary: Authors consider the tensor attention within the framework of transformer models and propose an algorithm that allows calculating gradients for the tensor attention with linear complexity. While standard attention captures pairwise interactions between tokens, the tensor attention (third-order tensor attention variant) was suggested by [Sanford et al; NeurIPS-2023] in the context of capturing triplet interactions. In the basic implementation, such a mechanism has cubic complexity on the forward pass, and in a subsequent work [Alman et al; ICLR-2024] a new algorithm based on polynomial expansions was proposed that allows (theoretically) to have linear complexity for forward passes. The authors, continuing this line of work, constructed an algorithm for calculating derivatives with linear complexity (theoretically), which potentially allows the method to be practically implemented and overcome the theoretical cubic time complexity barrier both in inference and training.
weaknesses: 1. If in [Alman et al; ICLR-2024] work the lack of experimental confirmation of the proposed approach could be attributed to the lack of an effective method for calculating gradients, then in this work, as stated, the last ingredient for the successful implementation of tensor attention has been prepared. In this regard, it is confusing that no practical implementation or numerical experiments are provided.

2. In the context of the previous point, the degree of scientific novelty raises certain doubts and it seems that it is only a certain development of the previous theoretical works [Sanford et al; NeurIPS-2023], [Alman et al; ICLR-2024].

3. The list of references takes up 9 full pages. The reasons for such extensive citations are not entirely clear, and the usefulness of these references for the reader is questionable. For example, after the phrase "Moreover, tensors are crucial in numerous machine learning applications..." there are about 9 references to various works of little relevance to the subject of these manuscript. Instead, the authors could have referred here to a 1-2 review works.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. The paper mentions that the "bounded entries assumption" needs to be met. Are
there any studies or experimental results that support that this assumption is correct
in actual transformer models? If this assumption is not met in some actual scenarios,
can your algorithm still remain efficient, or what adjustments need to be made?
soundness: 3
strengths: 1. The paper proposes the algorithm which can quickly compute the backward gradient 
of tensor attention training in almost linear time by utilizing polynomial 
approximation methods and tensor computation techniques, based on the closed-
form solution of the gradient computation of tensor attention provided in the 
paper.

2. Under hardness analysis, the paper proves that the assumption should be necessary
and tight, meaning that there is no algorithm which can achieve almost linear time in solving 
the tensor attention gradient computation, which is very important for future 
related research.
summary: The paper proposes a closed-form solution to compute gradients for Tensor Attention 
in higher-order transformer models efficiently. By utilizing polynomial approximation
methods and tensor algebraic techniques, the algorithm achieves nearly linear time
n1+O(1), the same complexity for both forward and backward computation under 
certain assumptions, which are proven to be necessary and tight in the paper. The 
theoretical results establish the feasibility of efficient higher-order transformer 
training and may promote practical applications of tensor attention architectures.
weaknesses: 1. The paper's assumptions are relatively strict and may be difficult to meet in real-
world situations. The paper lacks discussion on the practical applicability and 
limitations of the method.

2. The work should talk more about the benefits of using a tensor-based attention mechanism, because otherwise the motivation is quite weak, although I understand its theoretical novelty.

3. The authors should talk more about the tradeoffs of Algorithm 1. Namely what it sacrifices for faster computation, and how does it compares with matrix attention.

Meta-review: This paper gives an algorithm for approximating the gradient for tensor attention in n^{1+o(1)} time under certain bounds on the weights, and showed a lowerbound when then weights are not bounded. The reviewers have a large disagreement on the evaluation of this paper. Most reviewers agree that if the method can work in practice it would be interesting, however there are major concerns on how practical the algorithms are, even after the authors included some more empirical evidence in the response period (and this applies to both the most positive and most negative reviewers, except they decide to focus on different aspects of the paper).

Decision: Reject