In the current field of NLP, Attention is everywhere. Of course, it is not limited to NLP; Attention also holds a place in the CV field (e.g., Non-Local, SAGAN, etc.). In early 2018, in the article "A Brief Reading of <Attention is All You Need> (Introduction + Code)", we discussed the Attention mechanism. The core of Attention lies in the interaction and fusion of three vector sequences: \boldsymbol{Q}, \boldsymbol{K}, and \boldsymbol{V}. The interaction between \boldsymbol{Q} and \boldsymbol{K} provides a certain degree of correlation (weight) between pairs of vectors, and the final output sequence is obtained by summing \boldsymbol{V} according to these weights.
Clearly, numerous achievements in NLP & CV have fully affirmed the effectiveness of Attention. In this article, we will introduce some variants of Attention. The common characteristic of these variants is that they are "born for efficiency"— saving both time and video memory (VRAM).
Background Overview
The Attention discussed in "Attention is All You Need" is what we call "Multiplicative Attention," which is currently the most widely used form: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = \text{softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{V}
Additionally, there is Additive Attention, but it is difficult to implement in parallel (or consumes a lot of memory when implemented), so it is generally only used to encode variable-length vector sequences into fixed-length vectors (replacing simple Pooling) and is rarely used for sequence-to-sequence encoding. Among Multiplicative Attentions, the most widely used is Self-Attention. In this case, \boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V} are all results of the same \boldsymbol{X} after linear transformations. Consequently, the output is a vector sequence of the same length as \boldsymbol{X}, and it can directly capture the correlation between any two vectors in X and is easy to parallelize, which are the advantages of Self-Attention.
However, theoretically, the computation time and memory usage of Self-Attention are both of the order \mathcal{O}(n^2) (where n is the sequence length). This means that if the sequence length doubles, the memory usage and computation time both quadruple. While computation time might not strictly quadruple if there are enough parallel cores, the fourfold increase in memory usage is real and unavoidable. This is also the reason why Out-of-Memory (OOM) errors frequently occur when fine-tuning BERT.
Sparse Attention
We say Self-Attention is \mathcal{O}(n^2) because it calculates the correlation between any two vectors in the sequence, resulting in a correlation matrix of size n^2:
In the figure above, the left side shows the attention matrix, and the right side shows the connectivity. This indicates that every element is associated with every other element in the sequence.
Therefore, to save memory and speed up computation, a basic idea is to reduce the calculation of correlations, assuming that each element is only related to a subset of elements in the sequence. This is the basic principle of Sparse Attention. The Sparse Attention introduced in this article originates from OpenAI’s paper "Generating Long Sequences with Sparse Transformers", but it is presented here in a way that the author finds more natural rather than strictly following the original paper’s structure.
Atrous Self-Attention
The first concept to introduce is Atrous Self-Attention (also known as Dilated Self-Attention). This name, like the subsequent Local Self-Attention, was coined by the author based on its characteristics. These terms do not appear in the original paper, but I believe it is meaningful to introduce them separately.
Obviously, Atrous Self-Attention is inspired by "Atrous Convolution." As shown in the right figure below, it constrains the correlation, forcing each element to only associate with elements at relative distances of k, 2k, 3k, \dots, where k > 1 is a preset hyperparameter. From the attention matrix on the left, it forces the attention to be zero (white represents zero) for relative distances that are not multiples of k:
Since the attention calculation is now "jumping," each element only calculates correlation with approximately n/k elements. Ideally, the operational efficiency and memory usage become \mathcal{O}(n^2/k), effectively reducing them to 1/k of the original.
Local Self-Attention
Another transitional concept is Local Self-Attention. While self-attention mechanisms in the CV field are generally called "Non-Local," Local Self-Attention abandons global correlation and reintroduces local correlation. Specifically, it constrains each element to only associate with itself and the k elements before and after it, as shown below:
From the attention matrix perspective, any attention where the relative distance exceeds k is set to 0.
In fact, Local Self-Attention is very similar to ordinary convolution. Both maintain a window of size 2k+1 and perform operations within it. The difference is that ordinary convolution flattens the window and passes it through a fully connected layer, whereas here, the output is a weighted average within the window via attention. For Local Self-Attention, each element only calculates correlation with 2k+1 elements, making the efficiency and memory usage \mathcal{O}((2k+1)n) \sim \mathcal{O}(kn). This grows linearly with n, which is an ideal property—though it directly sacrifices long-range correlations.
Sparse Self-Attention
Now, we can naturally introduce OpenAI’s Sparse Self-Attention. We notice that Atrous Self-Attention has "holes," and Local Self-Attention happens to fill these holes. A simple way is to alternate between Local Self-Attention and Atrous Self-Attention. Accumulating both can theoretically learn global correlations while saving memory.
(A simple sketch shows that if the first layer uses Local Self-Attention, each output vector fuses several local input vectors. Then, if the second layer uses Atrous Self-Attention, even though it jumps, because the first layer’s output already fused local inputs, the second layer’s output can theoretically relate to any input vector, thus achieving long-range correlation.)
However, OpenAI did not do this. They directly merged Atrous Self-Attention and Local Self-Attention into one, as shown below:
It is easy to understand from the attention matrix: attention is set to zero except for relative distances not exceeding k and relative distances that are multiples of k (k, 2k, 3k, \dots). In this way, the Attention possesses the characteristic of "dense local correlation and sparse remote correlation." This might be a good prior for many tasks, as tasks truly requiring dense long-range correlation are actually rare.
Code Implementation
The Atrous, Local, and Sparse Self-Attentions mentioned above are all types of Sparse Attention. Intuitively, the attention matrix becomes very sparse. How do we implement them? If we simply apply a mask to the zero parts of the attention matrix, it is mathematically correct but does not speed up the process or save memory.
Official Implementation
OpenAI has open-sourced their implementation at: https://github.com/openai/sparse_attention
This is based on TensorFlow and uses their own sparse matrix library, blocksparse. However, the encapsulation seems quite strange; I am not sure how to migrate it to Keras, and it uses many Python 3 features, making it incompatible with Python 2. Friends using Python 3 and pure TensorFlow can give it a try.
Another issue is that OpenAI’s original paper mainly uses Sparse Attention to generate ultra-long sequences, so they masked all upper triangular parts of the attention matrix in both the paper and code (to avoid using future information). However, not all scenarios using Sparse Attention are generative, and for an introduction to basic concepts, this is unnecessary. This is one reason why I did not follow the original paper’s presentation logic.
Keras Implementation
For Keras, I have implemented the three types of Sparse Attention
based on my own logic and standardized them with my previous Attention
code. They are available at:
https://github.com/bojone/attention/blob/master/attention_keras.py
After experimentation, I found that in my implementation, these three
types of Sparse Attention do save some memory compared to full
Attention. Unfortunately, except for Atrous Self-Attention, the other
two implementations do not speed up the process and actually slow it
down slightly. This is because the implementation does not fully exploit
the sparsity. OpenAI’s blocksparse, on the other hand, is
highly optimized and written directly in CUDA, which is incomparable.
Regardless of speed, the three sparse attention variants should be
functionally sound.
Article Summary
There isn’t much more to summarize; we have introduced and implemented three types of sparse attention. Besides saving memory, sparse attention should better adapt to certain tasks, as most task correlations are primarily local and build from local to global. In particular, the "dense local and sparse remote" characteristic of Sparse Self-Attention should meet the needs of most tasks. Readers with relevant tasks are encouraged to try it out.
When reposting, please include the original link: https://kexue.fm/archives/6853
For more detailed information on reposting, please refer to: "Scientific Space FAQ"