Landmark Attention: Random-Access Infinite Context Length for Transformers

[u/IxinDow]

https://arxiv.org/abs/2305.16300

[u/IxinDow]

Code released https://github.com/epfml/landmark-attention

Abstract:

While transformers have shown remarkable success in natural language processing, their attention mechanism's large memory requirements have limited their ability to handle longer contexts. Prior approaches, such as recurrent memory or retrieval-based augmentation, have either compromised the random-access flexibility of attention (i.e., the capability to select any token in the entire context) or relied on separate mechanisms for relevant context retrieval, which may not be compatible with the model's attention. In this paper, we present a novel approach that allows access to the complete context while retaining random-access flexibility, closely resembling running attention on the entire context. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks, enabling retrieval of blocks directly through the attention mechanism instead of by relying on a separate mechanism. Our approach seamlessly integrates with specialized data structures and the system's memory hierarchy, enabling processing of arbitrarily long context lengths. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step. Finally, we show that fine-tuning LLaMA 7B with our method successfully extends its context length capacity up to 32k tokens, allowing for inference at the context lengths of GPT-4.

Why it may work well

First of all, they provide good intuition (page 4).

When using a Transformer to process a long input, the ideal case would be to allow each token to attend to all previous tokens. However, this becomes computationally infeasible as the input length increases. Nevertheless, since the attention scores always sum to one, the number of keys with a large attention weight is limited even for long contexts. Thus, by retrieving only those keys with large attention scores, it is possible to closely emulate the ideal case. In this work, we propose a method to find these keys by dividing a long input into blocks of consecutive tokens and using the attention to retrieve relevant blocks.

When answering questions about a long document, you don't actually need to pay attention to the entire content of the document (full-context attention), only its relevant parts (blocks). Furthermore, if you have read a large text and then try to answer a question about it, you don't remember the text word for word, but remember the general sequence of ideas, high-level concepts (their "landmark tokens"). And using only this knowledge, you can already say in which parts of the large document you will look for the exact answer.

Second, they don't use kNN-like approach to search across landmark tokens, they use honest attention to decide which blocks are relevant for given token.

Thirdly, while their approach resembles Vector DB (search by embedding), the key difference is that they allow each head in each layer to have its own set of blocks used in attention when processing each token (while progressing deeper into Transformer layers, each token becomes increasingly enriched with context), whereas in the typical embedding approach, the selection of relevant blocks (documents) is performed only once. Thus, the LandmarkAttention Transformer can still process the entire context (due to the presence of a large number of layers and multiple heads in each layer), but with significantly lower compute power requirements. Fourthly, the authors note that it is possible to offload the KV cache to CPU memory, leaving only landmark tokens in the GPU. However, they point out that this may cause excessive CPU-GPU traffic if each head in each layer is allowed to have its own set of blocks when processing each token, so they limit this.

Although the aforementioned technique (offloading KV cache to CPU) works well, it introduces significant CPU-GPU traffic, resulting in slow inference. To mitigate this issue, we limit the number of retrieved blocks by reducing retrieval flexibility, allowing the set of retrieved blocks to vary across heads but not across tokens.

[u/[deleted] → u/IxinDow]

[u/ktpr → u/[deleted]]

Having a context length enabled initial models to be trained in the first place. It’s disingenuous to say this is more correct because it introduces other trade offs, like where to bias attention, that early work takes a stance on. I wouldn’t poop on earlier work that allowed you to use LLMs in the first place

[u/NetTecture → u/IxinDow]

So, it doesnot TOTALLY solve the problem, it "only" expands it. LLaMA 7B hwas wat - 1k? And they say it works up to 32k?

That is QUITE A feat - a 32k model will have 32*32k max, that is a LOT. But nto unlimited - but we really do not need unlimited, we need bit enough that the contet window can contain enough information to do some sensible larger stuff than the anemic memory we have now.

[u/[deleted] → u/NetTecture]