[Paper Review] Landmark Attention: Random-Access Infinite Context Length for Transformers

Amirkeivan Mohtashami, Martin Jaggi|arXiv (Cornell University)|May 25, 2023

Topic Modeling4 citations

TL;DR

This paper proposes Landmark Attention, a method that enables Transformers to achieve random-access, infinite context length by using landmark tokens to gate retrieval of fixed-length input blocks. By training attention to retrieve relevant blocks via landmark similarity, the approach reduces computation and memory by a factor equal to block size (e.g., 50x), enabling LLaMA 7B to perform inference at over 32k tokens—matching GPT-4’s context length—while maintaining performance comparable to Transformer-XL.

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 to over 32k tokens, allowing for inference at the context lengths of GPT-4. We release the implementation of landmark attention and the code to reproduce our experiments at https://github.com/epfml/landmark-attention/.

Motivation & Objective

• To overcome the quadratic memory and computation costs of self-attention in long-context Transformers.
• To preserve random-access flexibility of attention while enabling retrieval of long contexts without recurrence or external retrievers.
• To allow inference at arbitrary context lengths, even beyond those seen during training, without retraining.
• To integrate seamlessly with memory hierarchy and data structures like FAISS for efficient large-context processing.
• To extend pre-trained models like LLaMA 7B to ultra-long contexts (e.g., 32k tokens) via fine-tuning.

Proposed method

• The input is divided into fixed-length blocks, each represented by a landmark token that acts as a gate for attention to that block.
• During inference, attention scores to landmark tokens determine which blocks are retrieved and attended to, enabling random access to any prior block.
• The model uses standard attention mechanisms but conditions retrieval on similarity between query and landmark key vectors.
• A cache of previous tokens is maintained (KV cache), with non-landmark tokens in each block swapped out and loaded only when the corresponding landmark is activated.
• The method integrates with efficient data structures like FAISS for k-nearest neighbor retrieval of relevant blocks.
• Fine-tuning is applied to pre-trained models (e.g., LLaMA 7B) using this mechanism to extend context length without retraining from scratch.

Experimental results

Research questions

• RQ1 Can Transformers achieve random-access, infinite context length without relying on recurrence or external retrieval models?
• RQ2 Can landmark-based retrieval maintain attention flexibility while drastically reducing computation and memory?
• RQ3 Can this method be applied to fine-tune existing large models to handle ultra-long contexts (e.g., 32k tokens) without performance degradation?
• RQ4 How does the performance of landmark-attentioned models compare to recurrent models like Transformer-XL at long context lengths?
• RQ5 Can the method be integrated with advanced data structures to further reduce resource usage in long-context inference?

Key findings

• Landmark Attention enables inference at arbitrary context lengths, even when the model was trained on much shorter contexts, without additional training costs.
• The method reduces computation and memory usage by a factor equal to the block size—e.g., 50x reduction with 50-token blocks—making long-context inference efficient.
• Fine-tuning LLaMA 7B with Landmark Attention allows it to achieve high accuracy in retrieving hidden passphrases from contexts up to 32,000 tokens long, matching GPT-4’s context length.
• The model maintains high performance on long-context tasks, with near-perfect accuracy in phrase recovery even beyond 32k tokens, while the base model fails.
• The approach is compatible with external data structures like FAISS, enabling further optimization of retrieval efficiency and memory usage.
• The method preserves interpretability, as attention scores to landmarks reveal which blocks are being accessed, enabling insight into model reasoning.