The models that we have discussed so far have a limited context length of 512 or 1024(GPT-2)

What about other sequences like audio and images (viewing pixels as sequences)?

The sequence length of an audio signal of 10s duration is about 80000 (assuming sampling rate of 8000 Hz)

The sequence length of a colour image of size 100 by 100 is 30000. 

So, enabling the models to attend to longer sequences is crucial. 

What if we want to summarize a book that contains, say 64000 words?

What if we want to generate a story that contains at least 5000 words?

What if we want to find an answer to a question from an instruction manual that contains a thousand pages?

The context lengths of some of the recent models are listed below

Why is it difficult to scale the context length?

One of the primary factors is the computational complexity of the self-attention mechanism.

Let us first try to understand the computational complexity and memory requirements to train Transformer based models.

Context window size:   \(T\) 

The size of the (projected) embeddings:   \(d = d_q=d_k\)

 \(T \cdot d\) multiplications and \(T \cdot (d-1)\) additions

The number of tokens \(T\) is typically greater than embedding dimension \(d\).

Therefore, to compute \(QK^T\)  we require \(\mathcal{O}(T^2d)\) operations

What about computing attention score by normalizing the unnormalized matrix \(Z\) ?

There are \(T\) rows in \(Z\), therefore computational complexity for \(softmax\) is  \(\mathcal{O}(T^2)\).

\(\therefore\) The computational complexity of \(A\) is \(\mathcal{O}(T^2d)\).

We need to normalize each row of \(Z\), this requires \(O(T)\) operations.

Finally, the attention score matrix \(A\) is multiplied with the value matrix \(V\)

Therefore, the entire self-attention block is quadratic \(\mathcal{O}(T^2d)\) in the length of the input sequence.

Therefore, the entire self-attention block is quadratic \(\mathcal{O}(T^2d)\) in the length of the input sequence.

\(B\): Batch Size

\(T\) : Context (sequence) Length 

\(n_h\) : Number of attention heads

\(d_{model}\) : Dimension of input embedding, \(d_{att}=512=\frac{d_{ff}}{4}\)

\(N\) : Number of parameters in the model

\(|\mathcal{V}|\) : Number of tokens in the vocabulary

\(d_{ff}\) : Dimension of feed-forward layer

Storing Activation values 

Storing Parameters 

Storing Gradients 

Which of these do you think takes the most memory?

Memory for Attention

Memory scales up linearly with respect to batch size \(B\) and quadratically with respect to the length of the context window \(T\).

For a single sample

For a batch of \(B\) samples

Memory scales up linearly with respect to batch size \(B\) and quadratically with respect to the length of the context window \(T\).

So, the next question is, How do we reduce the time complexity to enable  models to handle long sequences (during both training and inference)?

What we have is the full attention mechanism with \(O(T^2d)\) complexity. What we want is sub-quadratic complexity \(O(cTd)\). How do we achieve that?

Well, we can localize the attention in multiple ways.

This gives rise to a plethora of approaches that approximate the full attention.

A study on BERT [Ref]  in NLP tasks concluded that the neighbouring inner products are extremely important in self-attention and not all heads attend to all tokens.

How do we reduce the computational complexity?

One way is to make it sparse by attending to a subset of words instead of all words

\(O(c \times T \times d)\) where \(c < T\)

Another way is to divide the sequence into \(n\) blocks and do the computations independently as in [Ref] 

Typically, \(n=2,3,4\). Therefore, for a context length of 512, the block sizes will be 256, 128, 64.

These are resonable context sizes for majority of NLP tasks.

We could also use different patterns in each attention head/layer.

Suppose we have \(n=4\) blocks. Then we can permute these blocks in 16 possible ways (how? will see soon)

Therefore, we can permute these blocks in each attention head so that it could effectively capture long range dependencies.

Block Attention is given by,

For example, suppose context length \(T=15\) and the number of blocks \(n=3\). 

Write \(Q\),\(K^T\), and \(V\) into \(n=3\) block matrices (following the blocks in \(M\))

Then we would have 3! ways of permuting the blocks.

Consider the permutation,

\(\pi=\text{perm}(0,1,2,\cdots, n-1)\)

Let the \(k\)-th element of \(\pi\) be \(\pi(k)\), We define the masking  matrix \(M\)  of size \(T \times T\) 

Using these we can compute Block Attention as follows

We can extend the same concept to Multi-head attention by allowing each head to use a different masking matrix (derived from the permutation function).

This would be called Blockwise Multi-Head Attention.

50% non-zero values

6.25% non-zero values

(50% sparse)

(93% sparse)

Using blockwise sparse attention reduces the training time significantly for longer sequences.

Note, however, that the perplexity score increases as the number of blocks \(n\) increases. 

This is reasonable given that local attention is just an approximation to full attention. (This is the case for all local attention schemes.)

Let us pick up two rows in the  attention matrix show on the left

For the given query \(q_i\), we know that the attention score is higher for some of the keys  \(k_j\)

Can we find those \(k_j\)'s without resorting to computing pair-wise dot product with all the keys?

It is similar to approximating attention using the K-Nearest Neighbours (KNN) search. For example, if K is of length 64K, for each \(q_i\) we could only consider a small subset of, say, the 32 or 64 closest keys

One of many possible ways of doing this is to use Locality Sensitive Hashing (LSH) (Yes, the same one we used for deduplication of contents!)

The attention matrix is mostly sparse !

For this to work, we need \(Q=K\).This requires us to share the weight matrices \(W_Q=W_K\). 

We look for a hashing scheme that assigns each vector \(x\) (query) to a hash bucket such that nearby vectors falls in the same hash bucket and distant ones do not. This is called locality-sensitive hashing

Let's see how we construct hash buckets using a hash function with a help of an example (in 2D case)

For each query, it takes \(log(T)\) time to find the approximate nearest neighbours [ref], therefore, its computational complexity is \(O(T \log(T))\)

Consider a query (and a key) \((x=q_1,y=k_2)\) that are away from each other and first projected onto a sphere

Define the number of hash buckets, \((b=4)\)

Let the random rotation matrix \(R\) be of size \([d_k,b/2]\). 

for illustration,

Let us define the hash function 

\(hash(q_1)= \textit{argmax}(concat(0.5,0.25,-0.5,-0.25))=0\)

\(hash(k_2)= \textit{argmax}(concat(-0.25,-0.5,0.25,0.5))=3\)

Recompute the hash by changing the rotation matrix a few times

The probability that the given query and the key share the same bucket is \(\frac{1}{3}\).  Therefore, they are not close to each other

Let's consider a query and a key that are close to each other

The probability that the queries share the same bucket is \(\frac{3}{3}\).  Therefore, they are close to each other

Therefore, all the points that are close to each other would share the same bucket.

Essentially compute the hash for all the queries (repeat \(n_R\) times if required)

This can be computed in parallel by multiplying all the queries with the random matrix \(R\)

Now, sort the queries by bucket number, 

Compute the attention scores for queries within each bucket. For example, in ClM, \(q_3 \) attends to \(q_1\),  and \(q_{16}\) attends to \(q_3\) and \(q_{1}\) and so on

 and within each bucket, by sequence position

The methods that we have seen so far localize or sparsify the attention matrix in different ways

The theoretical time complexity is sub-quadratic \(O(cTd)\), with \(c\) being a constant (manually set, or function of \(T\) such as \(log(T)\))

Can we reduce it to \(O(Td)\)? Yes, as follows :-)

Well, this is not helpful to capture the context.

However, we can see from empirical observation that the self-attention matrix is essentially low-rank (especially for the top layers)!

Apply SVD on the matrix \(A\) across different layers and different heads of the model, and plot the normalized cumulative singular value averaged over 10k sentences.

Suppose we plot the distribution of all 512 singular values . Then, for it to be a low-rank it should follow a ____ distribution.

They observed a long tail distribution. Hence they concluded the matrix A is a low-rank matrix.

Using this observation, they show that there exists a low-rank approximation of \(QK^T\)  and we can let the network to find (learn) it! How?

Introduce two learnable linear projection matrices

The computational complexity is \(O(kTd)\), where \(k\) is the rank and can be set according to the error \(\epsilon\) 

Fundamentally, the exponential function in \(A\) takes in the query and key vectors and outputs a positive number.

where,

\(\mathbf{1}\) is a vector of all ones of length \(T\)

So we can generalize this

\(\mathcal{K}\) is called a Kernel function, and defined as follows

Where \(\phi(x)\) is some randomized feature mapping and \(\mathbb{E}\) is an expectation operator.

Note that \(r>0\), and it is not necessary that \(r <d\).

For example, \(T=64K,d=4K\), then \(r=8k\) is also reasonable 

The generalized form of a kernel transformation is given by,

where \(w_i\) is sampled from \(\mathcal{N}(\mathbf{0}_d,\mathbf{ I}_d)\)

\(w_i^Tx\) is a random projection, that is, the vector \(x\) is projected in a random direction given by \(w_i\) vector.

We can construct a random feature vector of size \(r\) by setting \(l=1\), considering \(f_1=exp\) and \(h(x)=exp(\frac{-||x||^2}{2})\) (Note, this particular setting is theoretically motivated among many possible choices.)

where,

where \(\mathbf{W} \in \mathbb{R}^{d \times r}\) ,\(X \in \mathbb{R}^{d \times T}\) and \(||X||^2\) is the \(L2-\)norm of column vectors in \(X\).

\(w_i \sim \mathcal{N}(\mathbf{0}_d,\mathbf{ I}_d)\)

We can collect all \(w_i\)s in a matrix \(\mathbf{W}\)

For all the queries \(q_i\)s in \(Q\), we can compute  \(Q'\) as follows

here, \(\mathbf{W} \in \mathbb{R}^{r \times d}\) ,\(Q \in \mathbb{R}^{T \times d}\) and \(||Q||^2\) is the \(L2-\)norm of row vectors in \(Q\).

Similary, we can compute \(K' \in \mathbf{R}^{T \times r}\).

The  \(\mathbf{W}\) matrix is orthogonalized using Gram-Schmidt orthogonalization procedure.

Orthogonalization of W helps reduce the variance of the softmax estimator for any dimensionality of \(d\). This requires \( r \leq d\).

The time complexity is \(O(rTd)\).

All these approaches that we have discussed so far approximated the full attention.

However, it does not translate to wall-clock speed-up!

That is, suppose we train a model with the compute requirement \(C \propto T^2\) FLOPS for a day.

Reducing the compute requirement \(C \propto \frac{T^2}{2}\) FLOPS will not reduce the training time to half a day

So, the problem is not in the algorithm but..

All Transformer based language models are trained on GPUs. 

However, it was observed that the implementation of MHA did not efficiently utilize GPUs.

In fact, the training time is now memory bound not compute bound!