Los gradientes de los Transformers de Múltiples Capas pueden ser aproximados en casi tiempo lineal.

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

August 23, 2024

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

cs.AI

Resumen

La complejidad computacional cuadrática en el mecanismo de autoatención de las arquitecturas de transformadores populares plantea desafíos significativos para el entrenamiento y la inferencia, especialmente en términos de eficiencia y requisitos de memoria. Para abordar estos desafíos, este artículo introduce un nuevo método de cálculo rápido de gradientes en modelos de transformadores de múltiples capas. Nuestro enfoque permite el cálculo de gradientes para todo el modelo de transformador de múltiples capas en casi tiempo lineal n^{1+o(1)}, donde n es la longitud de la secuencia de entrada. Este avance reduce significativamente el cuello de botella computacional asociado con la complejidad temporal cuadrática tradicional. Nuestra teoría es válida para cualquier función de pérdida y mantiene un error de aproximación acotado en todo el modelo. Además, nuestro análisis puede aplicarse cuando el modelo de transformador de múltiples capas contiene muchos submódulos prácticos, como conexiones residuales, máscaras casuales y atención multi-cabeza. Al mejorar la eficiencia del cálculo de gradientes en modelos de lenguaje grandes, esperamos que nuestro trabajo facilite el entrenamiento y despliegue más efectivos de modelos de lenguaje de largo contexto basados en nuestros resultados teóricos.

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.