Abstract
While transformers have shown remarkable success in natural languageprocessing, their attention mechanism's large memory requirements have limitedtheir ability to handle longer contexts. Prior approaches, such as recurrentmemory or retrieval-based augmentation, have either compromised therandom-access flexibility of attention (i.e., the capability to select anytoken in the entire context) or relied on separate mechanisms for relevantcontext retrieval, which may not be compatible with the model's attention. Inthis paper, we present a novel approach that allows access to the completecontext while retaining random-access flexibility, closely resembling runningattention on the entire context. Our method uses a landmark token to representeach block of the input and trains the attention to use it for selectingrelevant blocks, enabling retrieval of blocks directly through the attentionmechanism instead of by relying on a separate mechanism. Our approachseamlessly integrates with specialized data structures and the system's memoryhierarchy, enabling processing of arbitrarily long context lengths. Wedemonstrate that our method can obtain comparable performance withTransformer-XL while significantly reducing the number of retrieved tokens ineach step. Finally, we show that fine-tuning LLaMA 7B with our methodsuccessfully extends its context length capacity up to 32k tokens, allowing forinference at the context lengths of GPT-4.