Random-Access Infinite Context Length for Transformers
Amirkeivan Mohtashami EPFL、Martin Jaggi EPFL
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 Transforme
中文速览
大语言模型（LLM）的注意力机制（attention）虽然强大，却因显存消耗随上下文长度二次增长而难以处理长文本。作者提出"地标注意力"（landmark attention）：在输入文本每隔固定长度的块之后插入一个特殊的"地标token"，让模型学会用它来代表整个块，推理时先对各块的地标token打分、只取得分高的块加载进来做精细注意力计算，从而把实际参与计算的token数量压缩到原来的约1/50。这一设计完全基于原有注意力机制完成检索，无需额外的独立检索器，也保留了随机访问任意历史token的灵活性，并可与近似最近邻索引（如FAISS）结合进一步加速。实验表明，用该方法对LLaMA 7B进行微调后，模型上下文处理能力从原本的几千token扩展到32k token以上，达到GPT-4的上下文长度，而推理时间和显存占用却大幅低于直接扩展上下文窗口的方案，为低成本地赋予现有大模型"长记忆"能力提供了一条实用路径。
原文 arXiv:2305.16300;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2305.16300v2