多層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.