Algorithm and Hardness for Dynamic Attention Maintenance in Large Language Models
Jan van den Brand, Zhao Song, Tianyi Zhou
Introduction
Large language models (LLMs) such as Transformer , BERT , GPT-3 , PaLM , and OPT offer better results when processing natural language compared to smaller models or traditional techniques. These models possess the capability to understand and produce complex language, which is beneficial for a wide range of applications like language translation, sentiment analysis, and question answering. LLMs can be adjusted to multiple purposes without requiring them to be built from scratch. A prime example of this is ChatGPT, a chat software developed by OpenAI utilizing GPT-3’s potential to its fullest. GPT-4 , the latest iteration, has the potential to surpass the already impressive abilities of GPT-3, including tasks such as language translation, question answering, and text generation. As such, the impact of GPT-4 on NLP could be significant, with new applications potentially arising in areas like virtual assistants, chatbots, and automated content creation.
The primary technical foundation behind LLMs is the attention matrix . Essentially, an attention matrix is a square matrix with corresponding rows and columns representing individual words or “tokens,” and entries indicating their correlations within a given text. This matrix is then utilized to gauge the essentiality of each token in a sequence, relative to the desired output. As part of the attention mechanism, each input token is assigned a score or weight based on its significance or relevance to the current output, which is determined by comparing the current output state and input states through a similarity function.
Furthermore, the static attention computation and approximation has been studied by from both algorithmic and hardness perspectives. However, in practice, the attention matrix needs to be trained and keeps changing. In this work, we study the dynamic version of the attention computation problem. By using a dynamic approach, the attention weights can be updated on-the-fly as new information is introduced, enabling the model to adapt more effectively to changes in the input. This is particularly beneficial in cases where the input data is highly dynamic and subject to frequent changes, such as in natural language processing applications where the meaning and context of words and phrases can be influenced by the surrounding text.
Following the prior work , we formally define the standard attention computation problem as follows. To distinguish their standard model with the dynamic version studied in this paper, we call the problem defined in “static” version of attention multiplication. Another major difference between previous work and our work is that they studied an approximate version, whereas we study the exact version.
In applied LLMs training, the model parameters are changing slowly during training . Thus, it is worth considering the dynamic version of Attention multiplication problem. Next, we formally define the “dynamic” or “online” version of attention multiplication problem, we call it The name of our problem is inspired by a well-known problem in theoretical computer science which is called Online Matrix Vector multiplication problem () .. For consistency of the discussion, we will use the word “online” in the rest of the paper.
The goal of Online Diagonal-based normalized Attention Matrix Vector multiplication problem is to design a data-structure that satisfies the following operations:
Init: Initialize on three matrices , , .
Query: For any given , , return .
Here denotes the set .
In this paper, we first propose a data-structure that efficiently solves the problem (Definition 1.2) by using lazy update techniques. When then complement our result by a conditional lower bound. On the positive side, we use lazy update technique in the area of dynamic algorithms to provide an upper bound. In the area of theoretical computer science, it is very common to assume some conjecture in complexity when proving a lower bound. For example, , (strong) exponential time hypothesis, orthogonal vector and so on. To prove our conditional lower bound, we use a conjecture which is called Hinted Matrix Vector multiplication () conjecture . On the negative side, we show a lower bound of computing solving assuming the conjecture holds.
We first show our upper bound result making use of the lazy update strategy.
For any constant . Let . There is a dynamic data structure that uses space and supports the following operations:
Query. This operation outputs and takes worst-case time.
Our second result makes use of a variation of the popular online matrix vector multiplication () conjecture which is called hinted matrix vector multiplication conjecture (see Definition 5.2 and ). Next, we present a lower bound for the problem of dynamically maintaining the attention computation .
Assuming conjecture is true. For every constant , there is no algorithm that solve problem (see formal version in Definition 5.8) with
worst query time .
2 Related Work
A recent work by Zandieh, Han, Daliri, and Karbasi was the first to give an algorithm with provable guarantees for approximating the attention computation. Their algorithm makes use of locality sensitive hashing (LSH) techniques . They show that the computation of partition functions in the denominator of softmax function can be reduced to a variant of the kernel density estimation (KDE) problem, and an efficient KDE solver can be employed through subsampling-based swift matrix products. They propose the KDEformer which can approximate the attention within sub-quadratic time and substantiated with provable spectral norm bounds. In contrast, earlier findings only procure entry-wise error bounds. Based on empirical evidence, it was confirmed that KDEformer outperforms other attention approximations in different pre-trained models, in accuracy, memory, and runtime.
In another recent work , they focus on the long-sequence setting with . The authors established that the existence of a fast algorithm for approximating the attention computation is dependent on the value of , given the guarantees of , , and . They derived their lower bound proof by building upon a different line of work that dealt with the fine-grained complexity of KDE problems, which was previously studied in . Their proof was based on a fine-grained reduction from the Approximate Nearest Neighbor search problem . Additionally, their findings explained how LLM computations can be made faster by assuming that matrix entries are bounded or can be well-approximated by a small number of bits, as previously discussed in , Section 2 and , Section 3.2.1. Specifically, they showed a lower bound stating that when , there is no algorithm that can approximate the computation in subquadratic time. However, when , they proposed an algorithm that can approximate the attention computation almost linearly.
Transformer Theory
Although the achievements of transformers in various fields are undeniable, there is still a significant gap in our precise comprehension of their learning mechanisms. Although these models have been examined on benchmarks incorporating numerous structured and reasoning activities, comprehending the mathematical aspects of transformers still considerably lags behind. Prior studies have posited that the success of transformer-based models, such as BERT , can be attributed to the information contained within its components, specifically the attention heads. These components have been found to hold a significant amount of information that can aid in solving various probing tasks related to syntax and semantics, as noted by empirical evidence found in several studies .
Various recent studies have delved into the representational power of transformers and have attempted to provide substantial evidence to justify their expressive capabilities. These studies have employed both theoretical as well as controlled experimental methodologies through the lens of Turing completeness , function approximation , formal language representation , abstract algebraic operation learning , and statistical sample complexity aspects. According to the research conducted by , transformers possess the capability of functioning as universal approximators for sequence-to-sequence operations. Similarly, the studies carried out by have demonstrated that attention models may effectively imitate Turing machines. In addition to these recent works, there have been several previous studies that aimed to assess the capacity of neural network models by testing their learning abilities on simplistic data models . Furthermore, conducted a formal analysis of the training dynamics to further understand the type of knowledge that the model learns from such data models. According to findings from a recent study , moderately sized masked language models have demonstrated the ability to parse with satisfactory results. Additionally, the study utilized BERT-like models that were pre-trained using the masked language modeling loss function on the synthetic text generated with probabilistic context-free grammar. The researchers empirically validated that these models can recognize syntactic information that aids in partially reconstructing a parse tree. studied the computation of regularized version of exponential regression problem (without normalization factor).
Dynamic Maintenance
In recent years, projection maintenance has emerged as a crucial data structure problem. The effectiveness and efficiency of several cutting-edge convex programming algorithms greatly hinge upon a sturdy and streamlined projection maintenance data structure . There are two major differences between the problem in the dynamic data structure for optimization and our dynamic attention matrix maintenance problem. The first notable difference is that, in the optimization task, the inverse of a full rank square matrix is typically computed, whereas, in the attention problem, we care about the inverse of a positive diagonal matrix which behaves the normalization role in LLMs. The second major difference is, in the standard optimization task, all the matrix matrix operations are linear operations. However, in LLMs, non-linearity such as softmax/exp function is required to make the model achieve good performance. Therefore, we need to apply an entry-wise nonlinear function to the corresponding matrix. In particular, to compute when is linear function, we can pre-compute . However when is function, we are not allowed to compute directly.
Roadmap
The rest of the paper is organized as follows. In Section 2, we give some preliminaries. In Section 3, we explain the techniques used to show our upper bound and lower bound results. In Section 4, we present our dynamic data-structure. Our algorithm shows the upper bound results. In Section 5, we give our conditional lower bound result by assuming the Hinted MV conjecture.
Preliminary
In many TCS/ML literature, denotes the matrix exponential, i.e., . However, in this paper, we use to denote the entry-wise exponential, i.e.,
We use to denote the length- vector where all the entries are ones. We use to denote the length- vector where all entries are zeros.
We give a standard fact that is used in our proof.
We define a standard notation for describing the running time of matrix multiplication, see literature for examples.
Technique Overview
For the algorithmic result in , they make use of the “polynomial method in algorithm design”. The polynomial method is a technique for finding low-rank approximations of the attention matrix , which can be computed efficiently if the entries are bounded. For the hardness result in , they assume the strong exponential time hypothesis and use nearest neighbor search hardness result in the reduction.
For each update, we receive as input and update one entry in either matrix or . In the query function, we take index as input, and return the -th element in the target matrix .
Let denote . Let denote the updated target matrix . We notice that the computation of the attention can be written as
Let denote the change in the -th iteration. In a lazy-update fashion, we write in the implicit form
Lazy Update
Re-compute
Fast Query
Note that we maintain in our re-compute function. Hence, computing the first part takes time. As each column of and row of is 1-sparse, computing the second part takes time. The total running time needed for the query function is (Lemma 4.7, Lemma 4.6).
2 Hardness
We now turn to our lower bound result, which is inspired by the conjecture . Let us firstly define the problem (see formal definition in Definition 5.2).
Let the computation be performed over the boolean semi-ring and let . The problem has the following three phases
Phase 1. Input two matrices and
Phase 2. Input an matrix with at most non-zero entries
According to , the above problem is conjectured to be hard in the following sense,
For every constant no algorithm for the hinted Mv problem (Definition 5.2) can simultaneously satisfy
time complexity in Phase 2. and
Specifically, let us take an instance for the problem (Definition 5.2)
Let denote two matrices from Phase 1. from .
We create a new instance where
In Claim 5.6 and Claim 5.7, by making use of our construction of and , we show that for each and ,
Main Upper Bound
In Section 4.1, we show the running time of initializing our data structure. In Section 4.2, we show the running time of updating and . In Section 4.3, we show the correctness and the running time of querying the target matrix. In Section 4.4, we show the correctness and the running time of recomputing the variables in our data-structure.
We propose our upper bound result as the following:
For any constant . Let . There is a dynamic data structure that uses space and supports the following operations:
Query. This operation outputs operation takes in worst case time.
The amortized time in UpdateK and UpdateV can be made into worst case time by using standard techniques, e.g. see Section B of .
We first give the running time of the initialization procedure.
It is trivially from applying fast matrix multiplication. ∎
2 Update
Next, we give the running time of updating .
The procedure UpdateK (Algorithm 2) takes
Now, we give the running time of updating .
The procedure UpdateV (Algorithm 3) takes
3 Query
We show the correctness of our Query that queries only one element in the target matrix.
The procedure Query (Algorithm 4) outputs
For the -th element, by using simple algebra, we have
By summing up and , we have
The running time of procedure Query (Algorithm 4) is .
Computing takes time as is -sparse in columns and is -sparse in rows.
Hence, the total running time needed is ∎
4 Re-compute
We show the correctness of our re-compute function.
The procedure Recompute (Algorithm 5) correctly re-compute .
By computing , we correctly get the updated . By computing the inverse of a diagonal matrix we get .
By using Fact 2.1, we have .
Similar to the proof of re-computing .
By using Fact 2.1, we have .
By using the definition of , we can update by using .
Computing takes time as .
Main Lower Bound
In Section 5.1, we give the definition of Online Matrix Vector () problem. In Section 5.2, we introduce the definition of Hinted MV and its conjecture (from previous work ). In Section 5.3, we show the hardness of computing the target matrix without the normalization factor. In Section 5.4, we show the hardness of computing the target matrix with the normalization factor.
Before studying the hardness of our problem, we first review a famous problem in theoretical computer science which is called online matrix vector multiplication problem. Here is the definition of online matrix vector multiplication, which has been a crucial task in many fundamental optimization problems.
Given a matrix , let , there is an online sequence of vectors . The goal is to design a structure that whenever receives a new vector and output .
Such a problem is widely believed in the community that there is no algorithm to solve it in truly subquadratic time per vector and there is no algorithm to solve it in truly subcubic time over all vectors.
2 Hardness from Previous Work
We define the hinted Mv problem from previous work .
Let the computations be performed over the boolean semi-ring and let , . The hinted problem consists of the following phases:
Input two matrices and
Input an matrix with at most non-zero entries
We give the hinted Mv conjecture which is from prior work .
For every constant no algorithm for the hinted Mv problem (Definition 5.2) can simultaneously satisfy
time complexity in phase 2 and
3 Online Attention Matrix Vector Multiplication
We define the dynamic attention matrix vector problem here. For the following definition, we ignore the effect by the normalization factor. We will handle it in the later section.
The goal of the Online Attention Matrix Vector Multiplication problem is to design a data structure that satisfies the following operations:
Init: Initialize on matrices , , .
Update: Change any entry of , , or .
Query: For any given , , return .
Next, we present our lower bound result ignoring the normalization factor.
Assuming the hinted conjecture (5.3): For every constant , there is no dynamic algorithm for problem (Definition 5.4) with
worst query time .
Let us take an instance for the -hinted Mv problem (Definition 5.2) with . We create a new instance where
During phase 1, we give this input to the dynamic algorithm for the problem (Definition 5.4). During phase 2, when we receive the matrix with non-zero entries, we perform updates to the data structure to set . This takes
At last, in phase 3, we perform queries to obtain the column in time.
Using Claim 5.6, and Claim 5.7, we know that is enough to reconstruct for the hinted problem.
For each and , if is , then ,
We defined , so we can rewrite it as
Using the definition of matrix multiplication, and the fact that for all , we have some with
We can conclude that for each , there is at least one such that
Therefore, by using the definition of boolean semi-ring, we can conclude that
For each and , if is then .
where the first step follows from the definition of matrix multiplication and the second step follows from the definition of and .
By using the above equation, if , we have
Eq. (1) implies that, for all such that , we have , which also implies that .
Now, we can conclude that for each and . ∎
4 Online Diagonal-normalized Attention Matrix Vector Multiplication
Next, we consider the normalization factor and defined the problem as the following.
The goal of Online Diagonal-based normalized Attention Matrix Vector Multiplication problem is to design a data structure that satisfies the following operations:
Init: Initialize on matrices , , .
Update: Change any entry of , , or .
Query: For any given , , return , where .
Next, we present our lower bound result with the normalization factor.
Assuming the hinted conjecture (5.3): For every constant , there is no algorithm that solve problem (Definition 5.8) with
worst query time .
Let us take an instance for the -hinted Mv problem (Definition 5.2) with
We can construct matrix and as follows
where is a matrix that .
Note that , for each .
Based on the above construction, we will create a new instance , where
During phase 1, we give this input to the dynamic algorithm for the problem (Definition 5.8).
Let denote a diagonal matrix, where
During phase 2, we receive the diagonal matrix , where
and .
We perform updates to the data structure to set . This takes
, for each .
, for each .
By using the definition of , we know that, for each
Hence, we don’t need to update .
At last, in phase 3, we perform queries to obtain the column in time.
Using Claim 5.11 and Claim 5.10, we know that, for any and for any , if there is an algorithm that can find , then using is enough to reconstruct . Here can be computed in just time via Eq. (2). Thus, we can know the for the hinted problem in time, contradicting the hinted conjecture.
For each and , if is , then ,
By using the fact that and , we have
For , as , we know .
Using the definition of matrix multiplication, and the fact that for all , we have some with
We can conclude that for each , there is at least one such that
Therefore, by using the definition of boolean semi-ring, we can conclude that
For each and , if is then .
By using the fact that and , we have
For , as , we know .
For all such that , we have , which also implies that .
Now, we can conclude that for each and . ∎