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

翻译：虽然Transformer在自然语言处理领域取得了显著成功，但其注意力机制对内存的高需求限制了处理长上下文的能力。先前的方法（如循环记忆或基于检索的增强）要么牺牲了注意力的随机访问灵活性（即选择整个上下文中任意词元的能力），要么依赖独立的上下文检索机制，这可能与模型的注意力机制不兼容。本文提出了一种新方法，在保留随机访问灵活性的同时允许访问完整上下文，其效果接近对整个上下文运行注意力机制。该方法使用地标词元表示输入中的每个块，并训练注意力机制利用这些地标词元选择相关块，从而通过注意力机制直接检索块，而非依赖独立机制。我们的方法与专用数据结构及系统内存层次无缝集成，可处理任意长度的上下文。实验表明，该方法在显著减少每步检索词元数量的同时，可获得与Transformer-XL相当的性能。最后，我们展示了使用该方法微调LLaMA 7B模型，成功将其上下文长度扩展至超过32k词元，使其能够支持GPT-4级别的上下文长度推理。我们已在https://github.com/epfml/landmark-attention/ 开源地标注意力的实现及可复现实验的代码。