LONGNET: Scaling Transformer to 1 billion tokens
In this article, we will discuss in detail a recently released advanced model - "LongNet". The model was developed by Microsoft Research Asia and was officially announced about two weeks ago. LongNet is built based on the Transformer model, and its core concept is to expand the application scale of Transformer. It is worth mentioning that the research team successfully scaled it to handle 1 billion tokens. For those who are familiar with language models, they will understand the impact of sequence length on model performance, because the sequence length determines the number of tokens that can be associated when executing the attention mechanism, thereby affecting the length of contextual information that the model can obtain. For example, we hope that models like GPT will have longer context, so that the model can refer to words further back to predict the next token. LongNet has successfully expanded this capability to 1 billion tokens. As an example in the figure below, it can be clearly seen that the sequence length of GPT is only 512, while the sequence length of Power Transformer can be extended to 12, 000, 64, 262, 000, or even 10 million tokens. However, LongNet extends the sequence length to an astonishing 1 billion tokens. Imagine that we could feed all of Wikipedia’s text information into the model, and the model could leverage all of those tokens for attention calculations. Next, let us first understand how LongNet works.
Summary
In the era of large language models, extending sequence length has become an important requirement. However, existing methods encounter difficulties in dealing with computational complexity or model expressiveness, resulting in limitations on the maximum sequence length. To solve this problem, we introduce LONGNET, a Transformer variant that can extend sequence length beyond 1 billion tokens without compromising performance on shorter sequences. Specifically, we propose dilated attention, which exponentially expands the attention span as distance increases. LONGNET has significant advantages: 1) It has linear computational complexity and there is a logarithmic dependence between any two markers in the sequence; 2) It can be used as a distributed trainer for very long sequences; 3) Its expanded attention can seamlessly replace standard attention and can be seamlessly integrated with existing Transformer-based optimization methods. Experimental results show that LONGNET exhibits strong performance in long sequence modeling and general language tasks. Our work opens new possibilities for modeling very long sequences, such as treating an entire corpus or even the entire Internet as a sequence.
Advantages of LongNet
LongNet has several advantages. First, its computational complexity is linearly related to the sequence length, and the reasons will be explained in detail later. Second, there is a logarithmic dependence between tokens, that is, the dependence between two tokens that are further apart is weaker, while the dependence between tokens that are closer is stronger. Furthermore, it can be trained in a distributed network, which means we can compute this attention mechanism using a distributed system, such as using multiple GPUs or multiple computers. At the same time, LongNet can be used as a replacement for standard attention, which means that if we already have a model that uses the attention mechanism, we only need to replace the attention mechanism with LongNet's mechanism without changing other parts of the model. The model is still able to run as before, but by using this improved attention mechanism, it can handle longer sequence lengths.
About the Transformer model
For the self-attention mechanism, we use a matrix called "Q, K, V". Among them, the "Q" matrix represents the query, and its scale is "sequence length times model size", where model size refers to the vector representation of each word embedding. When we compute the product of the query with the key (K), or the product of the query with the transpose of K, to produce this matrix, the number of operations required is "the sequence length squared times the model size" because we need to compute the dot product for each element in the matrix. This is why the complexity of self-attention is "the square of the sequence length times the model size". This comparison is also described in detail in related papers. The conventional attention complexity is "the square of the sequence length multiplied by the model size." However, the complexity of the attention mechanism of the new model LongNet is only "the sequence length multiplied by the model size." Below I will explain how to achieve this linear complexity.
LongNet’s attention allocation principle
The core principle of LongNet is that the distribution of attention between tokens decreases exponentially as the distance between them increases. Let us refer to the diagram to understand how it works. In the traditional way, we calculate the attention between all tokens and all other tokens, but LongNet does not do this. It uses a method of splitting the sequence into windows of different sizes. First, taking 4 as the window size as an example, where "N" is the number of tokens in the sequence, we divide it into four segments of size 4 and calculate the attention between all words within this small window and other words. We then do the same thing for all the words in these segments, but use a larger window, this time with a window size of 8. And so on until the entire sequence length is covered, then we proceed by increasing the window size, and we also increase the number of tokens skipped, which is "R". For example, we could first calculate a window of size 8, then set "R" to 2, which means we would skip one token, then calculate attention, skip another token, and continue calculating attention. In this way, as the window size and the number of skipped tokens increase, the computational complexity becomes smaller, since we are not computing attention between each token and all other tokens, but only within a limited range. In this way, LongNet's attention allocation follows the principle of logarithmic dependence, that is, the dependence between tokens that are farther away is weaker, while the dependence between tokens that are closer to each other is stronger. This is the key to LongNet being able to work on larger sequence lengths.
Extended attention consists of a series of attention patterns for modeling short- and long-range dependencies, and the number of attention patterns can be expanded according to the sequence length. In each attention mode, the dot product between the query vector and the key vector is decomposed into multiple sub-dot products, each involving only a small part of the key vectors. This decomposition method can reduce computational complexity and also allow the model to better handle long sequences. The details are shown in the figure below:
Expanded attention also introduces a "multi-head" mechanism that can calculate attention separately among different heads. Each head has its own offset, which allows attention to be calculated at different positions to better capture information in the sequence. In this way, expanded attention can better handle long sequences while maintaining performance on shorter sequences. The details are shown in the figure below:
Optimization of computational complexity
Now let's take a look at why LongNet's computational complexity is linear. In the traditional self-attention mechanism, we need to perform dot product operations of the square of the sequence length, but LongNet reduces the computational complexity to linear by using windows and skipping. If we assume that the window size is fixed, say 4, and then "R" is also fixed, say 2, then the computational complexity will be "O(N)". Of course, in actual operation, the window size and the number of skipped tokens may be adjusted according to the actual situation, but they are constants and do not increase as the sequence length increases. This is why the computational complexity of LongNet is linear.
Token expansion reaches 1 billion+
The distributed training method uses the linear computational complexity of LONGNET to train the sequence dimensions in a distributed manner. Specifically, the algorithm first splits the input sequence along the sequence dimension, and each sequence fragment is assigned to a different device for calculation. Each device then projects the sequence fragments into queries, keys and values, and uses local computation to derive local attention weights. For parts that exceed the sequence length on the local device, the keys and values are sent to other devices for calculation. Finally, all devices summarize the local attention weights to obtain the global attention weight, and use the global attention weight to calculate the representation of each mark. The details are shown in the figure below:
The algorithm is scalable across any number of devices, and the training process can be accelerated through parallel computing. Since LONGNET has linear computational complexity, the algorithm can efficiently handle very long sequences without sacrificing training speed and model performance. In addition, the algorithm also supports standard Transformer optimization technologies, such as kernel fusion, quantization, and distributed training, allowing LONGNET to be seamlessly integrated with existing deep learning frameworks.
Application prospects of LongNet
The release of LongNet has brought many potential application prospects to the field of natural language processing. First, it can be applied to longer text generation tasks, such as generating novels or long news reports. Second, it can be applied to more complex dialogue tasks, because in dialogue we often need to process a large amount of historical information and context. Additionally, it may play a greater role in translation tasks, which often involve dealing with long sentences or paragraphs. Overall, the release of LongNet gives us new tools and possibilities for processing longer texts.
Conclusion
In summary, LongNet is a Transformer-based model that successfully extends the self-attention mechanism to 1 billion tokens and achieves the ability to process longer texts. Its advantages include that the computational complexity is linear, follows the principle of logarithmic dependence, and can be trained on distributed systems. With LongNet, we can explore more natural language processing tasks and handle those that were intractable in the past due to sequence length limitations. The release of this new model brings us even more possibilities.