[论文解读] Landmark Attention: Random-Access Infinite Context Length for Transformers

Amirkeivan Mohtashami, Martin Jaggi|arXiv (Cornell University)|May 25, 2023

Topic Modeling被引用 4

一句话总结

本文提出 Landmark Attention，一种使 Transformer 模型能够实现随机访问、无限上下文长度的方法，通过使用地标标记来控制固定长度输入块的检索。通过训练注意力机制根据地标相似性检索相关块，该方法将计算和内存消耗降低至块大小的倒数（例如，50倍），使 LLaMA 7B 实现超过 32k token 的推理——与 GPT-4 的上下文长度相当——同时保持与 Transformer-XL 相当的性能。

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 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 中自注意力机制带来的二次方内存和计算成本问题。
• 在不使用循环或外部检索器的前提下，保持注意力机制的随机访问灵活性，同时实现长上下文的检索。
• 实现任意上下文长度的推理，即使该长度超过训练时所见的范围，且无需重新训练。
• 与内存层次结构及数据结构（如 FAISS）无缝集成，以实现高效的大上下文处理。
• 通过微调将预训练模型（如 LLaMA 7B）扩展至超长上下文（例如，32k token）

提出的方法

• 输入被划分为固定长度的块，每个块由一个地标标记表示，该标记作为注意力机制访问该块的门控信号。
• 在推理过程中，注意力分数对地标标记的计算决定了哪些块被检索并加以关注，从而实现对任意先前块的随机访问。
• 模型使用标准注意力机制，但将检索过程基于查询与地标键向量之间的相似性进行条件控制。
• 维护一个先前 token 的缓存（KV 缓存），其中每个块中的非地标 token 被暂存，仅在对应地标被激活时才加载。
• 该方法可与高效的数据结构（如 FAISS）集成，用于相关块的 k-最近邻检索。
• 通过在此机制下对预训练模型（如 LLaMA 7B）进行微调，实现无需从头开始训练即可扩展上下文长度。

实验结果

研究问题

• RQ1 Transformer 能否在不依赖循环或外部检索模型的前提下，实现随机访问、无限上下文长度？
• RQ2 基于地标检索的方法能否在大幅降低计算和内存消耗的同时，保持注意力机制的灵活性？
• RQ3 该方法能否应用于微调现有大模型，使其在不产生性能下降的情况下处理超长上下文（如 32k token）？
• RQ4 在长上下文长度下，地标注意力模型的性能与循环模型（如 Transformer-XL）相比如何？
• RQ5 该方法能否与先进数据结构集成，以进一步降低长上下文推理中的资源消耗？

主要发现

• Landmark Attention 允许在任意上下文长度下进行推理，即使模型在训练时仅接触过较短的上下文，且无需额外训练成本。
• 该方法将计算和内存使用量降低至块大小的倒数（例如，使用 50 个 token 的块时，降低 50 倍），使长上下文推理更加高效。
• 通过 Landmark Attention 对 LLaMA 7B 进行微调，使其能够在长达 32,000 个 token 的上下文中准确检索隐藏的密码短语，达到与 GPT-4 相当的上下文长度。
• 该模型在长上下文任务中保持高性能，即使在超过 32k token 的长度下，短语恢复的准确率也接近完美，而基础模型则完全失效。
• 该方法与外部数据结构（如 FAISS）兼容，可进一步优化检索效率和内存使用。
• 该方法保持可解释性，因为对地标注意力分数的分析可揭示模型正在访问哪些块，从而提供对模型推理过程的洞察。