PaperlayerView exact revision on arXiv

Source-linked AI summary

Landmark Attention: Random-Access Infinite Context Length for Transformers

Amirkeivan Mohtashami, Martin Jaggi

arXiv:2305.16300v2cs.CLcs.LG

TL;DR

Transformers’ quadratic attention cost and memory footprint limit long-context processing, while prior recurrent or retrieval-based approaches do not fully preserve attention’s random access. The paper trains landmark tokens to select relevant blocks directly through attention, achieving comparable Transformer-XL performance and enabling LLaMA 7B inference beyond 32k tokens. The approach also reduces inference resources and supports arbitrary context lengths at inference.

  • Problem

    Transformer attention’s quadratic computational cost and memory footprint limit context length, while prior approaches can compromise random access or rely on separate retrieval mechanisms.

  • Method

    The method divides inputs into blocks, trains landmark tokens to represent them, and uses attention over landmarks to retrieve relevant blocks for standard attention.

  • Results

    The method achieves comparable performance with Transformer-XL and enables fine-tuned LLaMA 7B to retrieve information from contexts exceeding 32k tokens.

  • Takeaways & Limitations

    Landmark attention provides direct, random-access retrieval of earlier blocks while reducing inference computation and memory by the block-size factor.

  • Takeaways & Limitations

    The method’s positional indexing limits attention to tokens that are too far away based on semantic rather than positional relevance, and exact indexing could improve performance.

Abstract

from arXiv · show

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

1 Introduction

The paper addresses Transformers’ context-length limits by integrating block retrieval into attention through landmark tokens, preserving random access while reducing inference resources. Experiments show comparable performance to Transformer-XL and retrieval from contexts exceeding 32k tokens with fine-tuned LLaMA 7B.

  • Transformer attention provides flexible access to prior tokens but incurs quadratic computation and a problematic memory footprint that limits context length.
  • Recurrent-memory methods can sacrifice attention’s random-access flexibility, while retrieval-based methods require separate retrievers and may not mimic attention over long documents.
  • Landmark attention divides inputs into blocks and uses a landmark token as an attention-controlled gate for retrieving relevant earlier blocks.
  • Sequential chunk processing with landmark-based block selection reduces attention computation by the block length; experiments use blocks of 50 tokens.
  • The method supports arbitrary inference context lengths, achieves comparable performance with Transformer-XL, and retrieves relevant information from over 32k-token contexts using fine-tuned LLaMA 7B.
  • The stated advantages include inference at any context length, substantial reductions in inference time and memory, and compatibility with advanced data structures.

2 Related Work

Related work addresses long-context modeling through retrieval augmentation and positional-encoding extrapolation, but these approaches retain limitations involving separate retrieval, long-range attention, or extrapolation beyond training contexts.

  • Very long context windows such as 32k tokens for GPT-4 and 100k for Claude motivate research into efficient and accurate long-context models.
  • Retrieval-augmented language models use a separate retriever to find relevant knowledge-base documents before the main reader processes the augmented input.
  • Transformers struggle to extrapolate to contexts longer than those seen during training, and existing solutions can weaken long-range attention or require windowed attention.

3 Methodology

Landmark Attention represents input blocks with landmark tokens whose attention scores select relevant blocks, preserving random access while reducing long-context computation. Grouped Softmax gates access to other blocks, and inference retrieves selected blocks from a cached memory.

  • 3.1 Training Landmark Tokens: Landmark tokens represent input blocks and let attention scores retrieve relevant earlier blocks directly through the attention mechanism.Each block receives a special landmark token whose key serves as a representative vector for block retrieval.
  • 3.1 Training Landmark Tokens: During training, a landmark token is inserted after every ℓblock tokens and trained to represent the preceding block.The landmark corresponding to each token’s block is identified by its position pi.
  • 3.1 Training Landmark Tokens: Grouped Softmax places regular tokens within each block in separate groups and shares groups between current tokens and landmarks for other blocks.The current block’s landmark is ignored, while other landmarks compete with current-block tokens for attention.
  • 3.1 Training Landmark Tokens: Final attention weights are gated by landmark scores, while landmark tokens themselves receive zero attention weight.This forces the model to trade off attending to current tokens against retrieving information from other blocks.
  • 3.2 Inference: At inference, landmark scores select the highest-scoring blocks from a cache, and attention is then computed over the retrieved blocks together with the local context.The retrieved attention matrix is combined with local attention before GroupedSoftmax produces token representations.
  • 3.2.1 Positional Encoding: The stingy position mapping scheme places the latest k retrieved blocks at their corresponding positions and maps older retrieved blocks to an earlier prefix position.For k=2, the scheme separates the latest two blocks from older blocks with an empty block.
  • 3.3 Memory & Computation: The method decouples training context length from inference context length, giving constant training time O(1) with respect to operational inference context length.Standard Transformer training time scales quadratically with the operational inference context length.
  • 3.3 Memory & Computation: With ℓblock = 50, retrieving landmarks and selected blocks can reduce attention operations and memory by a factor of 50 compared with standard full-context attention.The retrieved blocks have constant size, while landmark search grows linearly with context length at one landmark per ℓblock + 1 tokens.

4 Experiments

The experiments evaluate landmark retrieval for long-context language modeling and LLaMA fine-tuning. Landmark models achieve Transformer-XL-comparable performance while supporting flexible retrieval and substantially longer inference contexts.

  • Language Modeling: Landmark models retrieve relevant blocks on PG-19 and arXiv, achieving comparable perplexity to Transformer-XL while reducing FLOPs.The evaluation covers English books and mathematical papers with billions of tokens.
  • Language Modeling: Each landmark block contains 50 normal tokens plus one landmark token, while Table 1 compares training and inference settings using language-modeling perplexity.Transformer-XL uses an XL cache; landmark settings report the maximum number of stored blocks.
  • Language Modeling: Using a local context of 250 tokens and retrieving the top 2 relevant blocks achieves a competitive perplexity setting.The passage introduces this result as a Table 1 finding across inference settings.
  • Language Modeling: Landmarks enable larger inference contexts than training contexts while preserving random access to past tokens and retaining fine-grained details.The method is compared with Transformer-XL trained on segments of length 2048.
  • Language Modeling: At context lengths 2048 and 4096, the model still outperforms the baseline with only 2 retrieved blocks, although reducing retrieval generally harms performance.With only the last 40 blocks stored, performance is better at evaluation length 4096 than 2048, suggesting learned recurrent behavior.
  • Language Modeling: Restricting retrieval to the same blocks across tokens reduces flexibility but still improves over the baseline, with a 0.23-point perplexity cost.Table 2 compares retrieval granularity and counts theoretically accessible blocks.
  • Fine-Tuning Pre-Trained Models: Fine-tuning LLaMA 7B with landmarks retrieves hidden pass phrases with high accuracy at context lengths far beyond the base model’s reliable range.At very large inputs such as 32K, KV-cache offloading to CPU reduces memory usage.

5 Future Work

The paper identifies positional extrapolation, hierarchical retrieval, and cache-aware training as directions for extending landmark attention. It also describes the evaluation setup used to test retrieval at long contexts.

  • Extrapolating Positional Encoding: The proposed positional indexing lets landmark models extrapolate to large contexts, but distant tokens can be selected only by semantic rather than positional information.The authors expect exact positional indexing could further improve performance and leave its investigation for future work.
  • Long-Context Retrieval Evaluation: Figure 3 tests pass-key retrieval by varying total prompt length, randomizing the key’s position, and retrieving the top 4 landmarked blocks.Prompts are processed in 250-token windows, and results average 50 random pass-key generations.
  • Hierarchical Landmarks: Hierarchical landmark tokens could gate lower-level landmarks, decide whether retrieval is needed, and model cache misses across memory levels.The authors leave systematic exploration of hierarchical landmarks as future work.
  • Training with Cache: Training with the cache may provide additional benefits beyond the standard softmax training procedure, but this variant remains unexplored.The authors expect standard training to resemble inference-time retrieval under their indexing scheme.

6 Conclusion

The method trains attention to retrieve relevant blocks from memory, providing direct access to previous tokens without recurrent forgetting. It achieves comparable performance to recurrent methods with fewer computational resources and supports arbitrary context lengths, including through fine-tuning LLaMA 7B.

  • The method trains attention to retrieve relevant blocks from memory rather than relying on recurrence.
  • Direct access to previous tokens is intended to avoid the slowly forgetting past data associated with recurrence.
  • The approach achieves comparable performance to recurrent methods such as Transformer-XL while using fewer computational resources.
  • Attention-based retrieval provides tracking and interpretability into the information used to generate outputs.
  • The method supports efficient inference with arbitrary context lengths and was incorporated into LLaMA 7B through fine-tuning.

A Grouped Softmax Example

The grouped-softmax example shows how landmark tokens gate attention to their corresponding blocks under a causal mask. Landmark attention weights are redistributed to normal tokens, leaving the landmarks themselves with zero final attention weight.

  • Tokens at indices 2, 5, and 8 serve as landmarks in the example.The landmarks satisfy the condition p_i = i.
  • For the token at index 6, the causal mask ignores the landmark at position 8.
  • The other two landmarks gate attention to tokens in their corresponding blocks.
  • Attention weight assigned to landmarks is distributed among normal tokens in their blocks, making final landmark attention weights zero.

B Dataset Description

The paper uses cleaned arXiv mathematics submissions and PG-19 books as large-scale datasets. The arXiv math subset contains around 5.6B tokens, while PG-19 contains 3.7B training tokens and supports evaluation of long-range token interactions.

  • The cleaned arXiv math subset of Proof-Pile 2 contains around 5.6B training tokens.It consists of UTF-encoded, heuristically filtered TEX files from arXiv mathematics submissions.
  • PG-19 contains 3.7B tokens in its training dataset and consists of English books published before 1919.
  • PG-19 is widely used to evaluate models’ ability to utilize long-range token interactions.

C Number of Unique Retrieved Blocks

The PG-19 analysis measures how many unique cache blocks are retrieved across tokens and layers. Retrieval patterns can substantially reduce distinct blocks loaded, with trade-offs between bandwidth savings and perplexity depending on whether variation occurs across heads or tokens.

  • The evaluation uses 2048-token PG-19 inputs, processes them in 250-token chunks, and retrieves the top 4 cache blocks.Unique blocks are counted in the penultimate chunk for each batch element and layer.
  • Figure 5 shows the distribution of unique retrieved blocks for 255 tokens, including landmarks, in the penultimate chunk.
  • A flexible-retrieval spike corresponds to very few retrieved blocks, attributed to the initial layer’s retrieval pattern.The finding is consistent with locality of attention in earlier layers.
  • Typically fewer than 10 unique blocks are observed when 8 heads each retrieve 4 blocks but the retrieved set varies only across heads.This setting theoretically permits 32 distinct blocks and can improve offloading performance, at a perplexity cost.
  • Allowing retrieved blocks to vary across tokens but not heads can reduce cache-loading bandwidth while achieving better perplexities.

D Context Miss Token

The context miss token (CMT) adds a hierarchy above landmark tokens to signal when the available in-context landmarks lack sufficient information. During inference, it can suppress unnecessary retrievals based on a learned attention score.

  • D Context Miss Token: CMT is trained as a higher-level gateway that regulates selected landmark tokens through a changed grouping scheme.The grouping separates CMT-controlled landmarks from free landmarks and places CMT at the beginning of the input.
  • D Context Miss Token: Randomly selecting CMT-controlled landmarks with probability PCMT = 0.5 teaches the model to route some retrieval decisions through CMT.The formulation preserves an attention-weight sum of one when the controlled-landmark set is nonempty.
  • D Context Miss Token: Because CMT cannot attend to input tokens, its representation is input-independent and serves as a beacon for insufficient context information.The model should attend to CMT when free landmarks do not provide enough information.
  • D Context Miss Token: At inference, setting LCMT to memory landmarks lets CMT indicate whether retrieval is necessary, with low scores emulating retrieval omission.The implementation always performs retrieval but sets the CMT score to zero below a cutoff threshold.

E Positional Augmentation

The positional-augmentation scheme introduces random jumps after landmark tokens to help models extrapolate beyond training contexts. Its evaluation examines extrapolation behavior, context-length perplexity, and implementation trade-offs.

  • E Positional Augmentation: Stingy mapping is required because the model cannot extrapolate to positions unseen during training, and a more thorough positional-encoding assessment remains future work.Existing techniques that dampen or limit attention to distant tokens do not fit the goal of attending to any token.
  • E Positional Augmentation: Random positional jumps increase all subsequent token positions by 1 to pjump after each landmark, with pjump = 1 recovering standard positions.The augmentation is intended to support extrapolation to longer contexts.
  • E Positional Augmentation: With pjump = 100 and 512-token training contexts, the model is theoretically expected to extrapolate to 1612 tokens.Each input contains 10 to 11 landmark tokens under this configuration.
  • E Positional Augmentation: Perplexity decreased with longer contexts through 1400 tokens under augmentation, whereas the standard model stopped improving before 1024 tokens.The augmented model’s observed range was close to its theoretical extrapolation estimate.
  • E Positional Augmentation: Landmark Attention can combine with FlashAttention by matching landmark frequency to FlashAttention’s block size, making GroupedSoftmax overhead negligible.The evaluation can reduce retrieval flexibility so all tokens in a chunk share retrieved blocks.
  • E Positional Augmentation: Larger blocks reduce memory usage and block-search time, but with fixed retrieved-block count they require shorter chunks.This creates a trade-off between block size, retrieved blocks, and chunk length.

G Offloading KV Cache to CPU

Landmark Attention supports offloading the KV cache to CPU while retaining landmarks on GPU and loading only retrieved blocks per step. Limiting retrieval to five blocks reduces CPU-GPU traffic for long-context inference.

  • G Offloading KV Cache to CPU: Offloading the KV cache to CPU allows only retrieved blocks to be loaded back into GPU memory while all landmarks remain on GPU.This reduces memory usage for very high context lengths.
  • G Offloading KV Cache to CPU: Significant CPU-GPU traffic can slow inference, so retrieval flexibility is reduced by sharing retrieved blocks across tokens within each head.The method computes each block’s maximum score across tokens and selects the top k = 5 blocks.
  • G Offloading KV Cache to CPU: 98% pass-key retrieval accuracy was achieved across 50 randomly generated prompts at a 32070-token context length after fine-tuning LLaMA.The authors describe this context length as comparable to those supported by GPT-4.
Loading 2305.16300v2…