The quadratic computational complexity in the self-attention mechanism of popular transformer architectures poses significant challenges for training and inference, particularly in terms of efficiency and memory requirements. Towards addressing these challenges, this paper introduces a novel fast computation method for gradient calculation in multi-layer transformer models. Our approach enables the computation of gradients for the entire multi-layer transformer model in almost linear time $n^{1+o(1)}$, where $n$ is the input sequence length. This breakthrough significantly reduces the computational bottleneck associated with the traditional quadratic time complexity. Our theory holds for any loss function and maintains a bounded approximation error across the entire model. Furthermore, our analysis can hold when the multi-layer transformer model contains many practical sub-modules, such as residual connection, casual mask, and multi-head attention. By improving the efficiency of gradient computation in large language models, we hope that our work will facilitate the more effective training and deployment of long-context language models based on our theoretical results.

本文解决了热门变换器架构自注意力机制中存在的二次计算复杂度问题，提出了一种新的快速梯度计算方法，能够在几乎线性时间内计算多层变换器模型的梯度。研究结果表明，该方法显著降低了传统计算瓶颈，有助于改进大型语言模型的训练和部署效率。

多层变换器梯度的近似计算可在近线性时间内完成