arXiv · 2405.16411
Training Tensor Attention Efficiently: From Cubic to Almost Linear Time
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.
Explore related subjects
Keep this discovery
Yang Cao, Yingyu Liang, Zhenmei Shi, Zhao Song. 2024-05-26. Training Tensor Attention Efficiently: From Cubic to Almost Linear Time. https://arxiv.org/abs/2405.16411
Cite the original work for its findings. Save a collection to share your selection of sources.