Training Tensor Attention Efficiently: From Cubic to Almost Linear Time

TL;DR

This paper demonstrates that the backward gradient computation for tensor attention can be achieved in almost linear time, significantly improving training efficiency and enabling practical higher-order transformer models.

Contribution

We prove the backward gradient of tensor attention can be computed in almost linear time and provide a closed-form solution with polynomial approximation techniques.

Findings

Backward gradient computation is almost linear in sequence length.

Provided a closed-form solution for efficient gradient calculation.

Proved the tightness of assumptions necessary for subcubic complexity.

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)$ time complexity of tensor attention poses a significant obstacle to its utilization in transformers, where $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)}$, 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.