多层Transformer的梯度可以在几乎线性时间内进行近似计算。

Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time

August 23, 2024

作者: Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, Yufa Zhou

cs.AI

摘要

流行的Transformer架构中自注意机制中的二次计算复杂度对训练和推断提出了重大挑战，特别是在效率和内存需求方面。为了解决这些挑战，本文介绍了一种新颖的快速计算方法，用于多层Transformer模型中的梯度计算。我们的方法使得几乎可以在线性时间n^{1+o(1)}内计算整个多层Transformer模型的梯度，其中n为输入序列长度。这一突破显著降低了与传统二次时间复杂度相关的计算瓶颈。我们的理论适用于任何损失函数，并在整个模型中保持有界的近似误差。此外，我们的分析可以应用于包含许多实用子模块的多层Transformer模型，例如残差连接、自回归掩码和多头注意力。通过提高大型语言模型中梯度计算的效率，我们希望我们的工作能基于我们的理论结果促进更有效的长上下文语言模型的训练和部署。

English

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.