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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Liang, Yingyu, Sha, Zhizhou, Shi, Zhenmei, Song, Zhao, Zhou, Yufa
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866929542754467840
author Liang, Yingyu
Sha, Zhizhou
Shi, Zhenmei
Song, Zhao
Zhou, Yufa
author_facet Liang, Yingyu
Sha, Zhizhou
Shi, Zhenmei
Song, Zhao
Zhou, Yufa
contents The computational complexity of the self-attention mechanism in popular transformer architectures poses significant challenges for training and inference, and becomes the bottleneck for long inputs. Is it possible to significantly reduce the quadratic time complexity of computing the gradients in multi-layer transformer models? This paper proves that a novel fast approximation method can calculate the gradients in almost linear time $n^{1+o(1)}$ where $n$ is the input sequence length, while it maintains a polynomially small approximation error $1 / \mathrm{poly}(n)$ across the entire model. Our theory holds for general loss functions and 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, we hope that this work will facilitate more effective training and deployment of long-context language models based on our theoretical results.
format Preprint
id arxiv_https___arxiv_org_abs_2408_13233
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time
Liang, Yingyu
Sha, Zhizhou
Shi, Zhenmei
Song, Zhao
Zhou, Yufa
Machine Learning
Artificial Intelligence
Computation and Language
The computational complexity of the self-attention mechanism in popular transformer architectures poses significant challenges for training and inference, and becomes the bottleneck for long inputs. Is it possible to significantly reduce the quadratic time complexity of computing the gradients in multi-layer transformer models? This paper proves that a novel fast approximation method can calculate the gradients in almost linear time $n^{1+o(1)}$ where $n$ is the input sequence length, while it maintains a polynomially small approximation error $1 / \mathrm{poly}(n)$ across the entire model. Our theory holds for general loss functions and 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, we hope that this work will facilitate more effective training and deployment of long-context language models based on our theoretical results.
title Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2408.13233