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

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.

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.

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.

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.

Deepest thanks to all the thoughtful and valuable reviews.

We appreciate reviewers tbYL, egs4, kn9h, and ZjZd for recognizing the relevance and impact of our work on the efficient computation of tensor attention gradients. tbYL and ZjZd commend the importance of the hardness analysis, emphasizing its role in establishing the necessity and tightness of our assumptions. egs4 highlights the practical potential of gradient approximation for tensor attention and the broader applicability of our tensor computation results to related fields. kn9h values the structured and rigorous presentation, noting the potential demand for tensor attention in the ML community due to its ability to capture complex data correlations. ZjZd specifically acknowledges our almost linear time algorithm for tensor attention gradient computation, achieved through polynomial approximation and tensor techniques, as a significant contribution to advancing the field.

We have updated a revision for our draft and put all major updates in brown. We also update the code about experiments in the supplemental material. Here, we summarize all major updates in the revision. Note that all line numbers in the rebuttal are corresponding to the revised version.

• Line 167-168: Give a definition of function mat.
• Line 321-323: Give explanation $\log(n)$ bits model.
• Line 270-288: Provide a figure to illustrate the proof flow. And also mention the figure in Line 310
• Line 258-259: Provide the reference of variable $Z$.
• Line 260-265: Reorganize the paper to provide intuition first.
• Line 184 and Line 191: We give the other name of column/row-wise Kronecker product.
• Line 945: Add citation [9].
• Line 51-53 and 57: Add more motivation about high-order attention.

Here, we address common concerns. We will address other questions and concerns in individual rebuttals.

Empirical validation

We admit that the practical implementation of high-order attention to real-world tasks poses additional significant challenges. From the authors’ perspective, these challenges include but are not limited to (1) as the kernel module, i.e., attention, changes, numerous other modules, such as layer normalization, residual connections, positional embeddings, and many others, may require corresponding updates; (2) integrating this module directly into existing Large Language Models or Large Vision-Language Models to save training time is also a challenge. Despite these challenges, these uncertainties may open a broad space for algorithm design. We hope our work inspires new model architectures and further algorithmic design.

We have carefully checked the current deep learning package, PyTorch, to see how our algorithm could be practically implemented on GPU for numerical verification. Due to the time limit, we provide approximation results of our Lemma E.1 below, where the approximation in Lemma E.1 is the key to our whole approximation process.

All input matrix elements are drawn from standard Gaussian. We compare the difference between $\mathsf{S}(x)$ and $U_1 (V_1 \oslash W_1)^\top$ (see detail in Lemma E.1).

In the above Table, we use $n=10 \times 2^d$ as $d = O(\log(n))$ in our assumptions. We can show that our error is getting small when $n$ becomes larger, which is consistent with our theorem that the error bound is $O(\epsilon / \mathrm{poly} (n))$. This well support the correctness of our theoretical bounds. This approximation may inspire a new training paradigm of large multimodal models.

We update the code about the above experiments in the supplemental material. We also put the kernel code below. The key idea is to use torch.einsum() to implement tensor operations.

import torch
import math
CHAR_LIST = 'abcdefghijklmnopqrstuvw'
def poly_order_approx(x, order=1):
    shape = list(x.shape)
    shape[-1] = -1
    denorm = (1 / math.factorial(order)) ** 0.5

    operation_in = [f'xyz{CHAR_LIST[i]}' for i in range(order)]
    operation_in = ','.join(operation_in)
    operation_out = [f'{CHAR_LIST[i]}' for i in range(order)]
    operation_out = 'xyz' + ''.join(operation_out)
    operation = operation_in+'->'+operation_out # xyza,xyzb,xyzc,xyzd,xyze->xyzabcde
    x_command = ', '.join(['x'] * order)
    # torch.einsum('xyza,xyzb,xyzc,xyzd,xyze->xyzabcde', x, x, x, x, x).view(shape)
    torch_command = f"torch.einsum('{operation}', {x_command}).view(shape)" 

    x_order = eval(torch_command)
    # (q.T k)^d / d!
    return x_order * denorm 

We will check other bounds further and include them in the revision. We hope the above experiment fixes your concern, and we are glad to discuss it further per the reviewer's request.

On ther other hand, we would like to clarify that our work is purely a theoretical study, which does have a large place in the top machine learning conferences, like NeurIPS, ICLR, and ICML. For instance, ICLR accepted papers [1,2,3,4], all of which are entirely theoretical and do not contain any experiments. Additionally, the papers [4,6,7,8] that focus on designing efficient algorithms for low-rank approximation and attention approximation also do not contain experimental results. Our work is similar to theoretical studies like [4,5,6], concentrating on the theoretical aspects of attention computation.

Moreover, 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. We believe that our extensive theoretical analysis can outweigh the lack of experimental results.

[1] Zhan, W., Uehara, M., Kallus, N., Lee, J. D., & Sun, W. Provable Offline Preference-Based Reinforcement Learning. ICLR’24.

[2] Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., & Zhang, A. R. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. ICLR’23.

[3] Wen, K., Ma, T., & Li, Z. How Sharpness-Aware Minimization Minimizes Sharpness?. ICLR’23.

[4] Alman, J., & Song, Z. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. ICLR’24.

[5] Han, I., Jayaram, R., Karbasi, A., Mirrokni, V., Woodruff, D. P., & Zandieh, A. Hyperattention: Long-context attention in near-linear time. ICLR’24.

[6] Alman, J., & Song, Z. Fast attention requires bounded entries. NeurIPS’23.

[7] Sarlos, T., Song, X., Woodruff, D., & Zhang, R. Hardness of low rank approximation of entrywise transformed matrix products. NeurIPS’23.

[8] Dexter, G., Drineas, P., Woodruff, D., & Yasuda, T. Sketching algorithms for sparse dictionary learning: PTAS and turnstile streaming. NeurIPS’23.

[9] Kolda, T. G., & Bader, B. W. (2009). Tensor decompositions and applications. SIAM review, 51(3), 455-500.

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

Reject