Meerlaagse Transformers Gradienten Kunnen Benaderd Worden in Bijna Lineaire Tijd

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

August 23, 2024

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

cs.AI

Samenvatting

De kwadratische computationele complexiteit in het self-attention-mechanisme van populaire transformer-architecturen vormt aanzienlijke uitdagingen voor training en inferentie, met name wat betreft efficiëntie en geheugenvereisten. Om deze uitdagingen aan te pakken, introduceert dit artikel een nieuwe snelle berekeningsmethode voor gradientberekening in meerlaagse transformer-modellen. Onze aanpak maakt het mogelijk om de gradienten voor het volledige meerlaagse transformer-model te berekenen in bijna lineaire tijd n^{1+o(1)}, waarbij n de lengte van de invoerreeks is. Deze doorbraak vermindert de computationele bottleneck die gepaard gaat met de traditionele kwadratische tijdcomplexiteit aanzienlijk. Onze theorie geldt voor elke verliesfunctie en behoudt een begrensde benaderingsfout over het gehele model. Bovendien blijft onze analyse geldig wanneer het meerlaagse transformer-model veel praktische submodules bevat, zoals restverbindingen, causale maskers en multi-head attention. Door de efficiëntie van gradientberekening in grote taalmodelen te verbeteren, hopen we dat ons werk het effectiever trainen en implementeren van lang-context taalmodelle zal vergemakkelijken, gebaseerd op onze theoretische resultaten.

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.