Landmark Attention: Random-Access Infinite Context Length for Transformers

By Amirkeivan Mohtashami, Martin Jaggi. 2023

222 citations recorded by the source corpus.

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/.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

Mixtral of Experts (2024, methodology): “To assess the capabilities of Mixtral to tackle long context, we evaluate it on the passkey retrieval task introduced in [23], a synthetic task designed to measure the ability of the model to retrieve a passkey inserted randomly in a long prompt.”
Extending Context Window of Large Language Models via Positional Interpolation (2023, methodology): “To measure this, we follow a synthetic evaluation task of passkey retrieval proposed by Mohtashami & Jaggi (2023).”
LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding (2023, unknown): “…include extending the context window (Press et al., 2022; Chen et al., 2023), utilizing recurrent memory (Dai et al., 2019; Bulatov et al., 2023), using sparsed attention (Ding et al., 2023; Mohtashami and Jaggi, 2023), and augmenting with an external memory (Liang et al., 2023; Zhou et al., 2023).”
YaRN: Efficient Context Window Extension of Large Language Models (2023, unknown): “The passkey retrieval task as defined in [25] measures a model’s ability to retrieve a simple passkey (i.e., a five-digit number) from amongst a large amount of otherwise meaningless text.”
Long-context LLMs Struggle with Long In-context Learning (2024, unknown): “2. passkey retrieval (Mohtashami & Jaggi, 2023; Chen et al., 2023a; Li et al., 2023a) or needle-in-a-haystack (Team et al., 2023; Fu et al., 2024), which requires reciting a randomly inserted information in a long sequence.”
KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (2024, methodology): “We used the passkey evaluation framework from [37] (which is based on the methodology from [24]) to evaluate retrieval performance.”
RULER: What's the Real Context Size of Your Long-Context Language Models? (2024, methodology): “Recent works have primarily focused on retrieval-based synthetic tasks (Kamradt, 2023; Mohtashami & Jaggi, 2023; Li et al., 2023a; Liu et al., 2024d; Lee et al., 2024), with a few investigate other aspects, including fact reasoning (Kuratov et al., 2024; Karpinska et al., 2024), long-range…”
MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (2024, unknown): “Vertical-Slash (VS) pattern The attention weights are concentrated on specific tokens (vertical lines) [MJ23] and tokens at fixed intervals (slash lines).”
∞Bench: Extending Long Context Evaluation Beyond 100K Tokens (2024, unknown): “Extending Context Length Transformers, typically trained on text sequences under 8K tokens due to self-attention’s quadratic complexity, face challenges in longer downstream tasks.”
LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens (2024, methodology): “We follow a synthetic evaluation task of passkey retrieval proposed by (Mohtashami & Jaggi, 2023).”
Effective Long-Context Scaling of Foundation Models (2023, background): “Moreover, existing open-sourced long-context models (Tworkowski et al., 2023a; Chen et al., 2023; Mohtashami and Jaggi, 2023; MosaicML, 2023b) often fall short on evaluations and primarily measure long-context capabilities with the language modeling loss and synthetic tasks, which do not…”
Improving Text Embeddings with Large Language Models (2023, methodology): “…… Figure 4: Illustration of the personalized passkey retrieval task adapted from Mohtashami and Jaggi [25].”
LongEmbed: Extending Embedding Models for Long Context Retrieval (2024, background): “Passkey retrieval (Mohtashami and Jaggi, 2023) requires LLMs to recover a random passkey hidden within a long document comprising garbage information.”