摘要

本文介绍了一种名为Neurocache的新方法，旨在通过使用外部向量缓存来扩展大型语言模型（LLMs）的有效上下文大小。Neurocache利用高效的k-最近邻（kNN）算法从缓存中检索相关过去状态，并将其整合到注意力过程中。该方法通过存储压缩状态、每个令牌执行单次检索操作以及扩展检索窗口至相邻状态，提高了推理速度和下游任务的准确性。实验表明，无论是从头开始训练的模型还是如Llama27B和Mistral-7B这样的预训练模型，Neurocache都能有效提升其性能。此外，Neurocache在与文本检索方法的比较中，在单文档问答和少样本学习任务中显示出改进。

原理

Neurocache的核心在于其高效的k-最近邻（kNN）检索策略，该策略从外部向量缓存中检索压缩的过去状态。这种方法通过以下几个关键步骤实现其先进性：

1. 状态压缩：文本段落通过Transformer解码器堆栈处理，并在中间层将隐藏状态压缩成更紧凑的形式，以提高后续kNN检索的效率。
2. 状态检索：对于每个压缩状态，从缓存中识别出最相似的top-k状态，基于这些状态与缓存中状态的L2距离。
3. 缓存更新：缓存随着新压缩状态的加入而更新，同时保持固定大小，丢弃最旧的条目。
4. 缓存增强层：从检索到的状态开始，缓存增强层通过专门的注意力机制整合这些状态，从而优化语言建模质量和推理速度。

流程

Neurocache的工作流程包括以下几个步骤：

1. 文档分段：将长文本序列分割成较小的段落，每个段落包含n个令牌，适合模型的注意力窗口大小。
2. 状态压缩：通过Transformer解码器堆栈顺序处理每个文本段落，并在中间层将隐藏状态压缩。
3. 状态检索：对于每个压缩状态，从缓存中检索最相似的top-k状态。
4. 缓存更新：更新缓存以包含新的压缩状态，同时保持固定大小。
5. 缓存增强层：使用检索到的状态，缓存增强层通过专门的注意力机制生成键和值，用于缓存注意力。

应用

Neurocache的应用前景广泛，特别是在需要处理长文档的任务中，如文档摘要和学术文献回顾。通过扩展模型的上下文长度至128K令牌，Neurocache显著提升了长文档处理的能力，适用于各种自然语言处理任务，包括单文档问答和少样本学习。