原文: English original · Anthropic/OpenAI 官方

Contextual Retrieval 简介

要让 AI 模型在特定上下文中发挥作用,它通常需要访问背景知识。

要让 AI 模型在特定上下文中发挥作用,它通常需要访问背景知识。例如,客户支持聊天机器人需要了解它所服务的具体业务,法律分析机器人则需要掌握大量过往案例。

开发者通常使用 Retrieval-Augmented Generation(RAG,检索增强生成)来增强 AI 模型的知识。RAG 会从知识库中检索相关信息,并将这些信息追加到用户的 prompt 中,从而显著改善模型回复。问题在于,传统 RAG 方案在编码信息时会移除上下文,这往往导致系统无法从知识库中检索到相关信息。

本文将介绍一种显著改善 RAG 检索步骤的方法。这个方法叫做 “Contextual Retrieval”,包含两项子技术:Contextual Embeddings 和 Contextual BM25。该方法可以将检索失败次数降低 49%;与 reranking 结合后,可以降低 67%。这些都是检索准确率上的显著提升,会直接转化为下游任务中更好的表现。

你可以通过我们的 cookbook,用 Claude 轻松部署自己的 Contextual Retrieval 方案。

关于直接使用更长 prompt 的说明

有时,最简单的方案就是最好的。如果你的知识库少于 200,000 tokens(约 500 页材料),你可以直接把整个知识库放入给模型的 prompt 中,而不需要 RAG 或类似方法。

几周前,我们为 Claude 发布了 prompt caching,让这种方法明显更快,也更具成本效益。开发者现在可以在多次 API 调用之间缓存常用 prompt,将延迟降低 2 倍以上,并将成本最多降低 90%(你可以阅读我们的 prompt caching cookbook 了解它的工作方式)。

不过,随着知识库增长,你会需要更可扩展的方案。这就是 Contextual Retrieval 的用武之地。

RAG 入门:扩展到更大的知识库

对于无法放入 context window 的更大知识库,RAG 是典型方案。RAG 会通过以下步骤预处理知识库:

  1. 将知识库(文档“语料库”)拆分成更小的文本块,通常每块不超过几百个 tokens;
  2. 使用 embedding 模型将这些文本块转换为编码语义的向量 embeddings;
  3. 将这些 embeddings 存入支持按语义相似度搜索的向量数据库。

运行时,当用户向模型输入查询,系统会使用向量数据库,根据与查询的语义相似度找到最相关的文本块。随后,最相关的文本块会被加入发送给生成式模型的 prompt 中。

虽然 embedding 模型擅长捕捉语义关系,但它们可能错过关键的精确匹配。幸运的是,有一种更早出现的技术可以在这些场景下提供帮助。BM25(Best Matching 25)是一种排名函数,使用词法匹配来查找精确的词或短语匹配。它对包含唯一标识符或技术术语的查询尤其有效。

BM25 建立在 TF-IDF(Term Frequency-Inverse Document Frequency,词频-逆文档频率)概念之上。TF-IDF 衡量一个词对于集合中某篇文档的重要性。BM25 在此基础上进行了改进:它会考虑文档长度,并对词频应用饱和函数,从而避免常见词主导结果。

下面这个例子说明了 BM25 如何在语义 embeddings 失败的地方成功:假设用户在技术支持数据库中查询 “Error code TS-999”。embedding 模型可能会找到关于错误码的一般内容,却可能错过精确的 “TS-999” 匹配。BM25 会查找这个特定文本字符串,以识别相关文档。

RAG 方案可以通过以下步骤结合 embeddings 和 BM25 技术,更准确地检索最适用的文本块:

  1. 将知识库(文档“语料库”)拆分成更小的文本块,通常每块不超过几百个 tokens;
  2. 为这些文本块创建 TF-IDF 编码和语义 embeddings;
  3. 使用 BM25 根据精确匹配找到排名靠前的文本块;
  4. 使用 embeddings 根据语义相似度找到排名靠前的文本块;
  5. 使用 rank fusion 技术合并并去重步骤(3)和(4)的结果;
  6. 将 top-K 文本块加入 prompt,用于生成回复。

通过同时利用 BM25 和 embedding 模型,传统 RAG 系统可以提供更全面、更准确的结果,在精确术语匹配与更广泛的语义理解之间取得平衡。

一个标准 Retrieval-Augmented Generation(RAG)系统,同时使用 embeddings 和 Best Match 25(BM25)来检索信息。TF-IDF(term frequency-inverse document frequency)衡量词语重要性,并构成 BM25 的基础。

这种方法让你能够以具备成本效益的方式扩展到巨大知识库,远远超过单个 prompt 能容纳的范围。但这些传统 RAG 系统有一个重大限制:它们经常破坏上下文。

传统 RAG 中的上下文难题

在传统 RAG 中,文档通常会被拆分成更小的文本块,以便高效检索。虽然这种方法适用于许多应用,但当单个文本块缺乏足够上下文时,它可能导致问题。

例如,假设你的知识库中嵌入了一组财务信息(比如美国 SEC 文件),并收到如下问题:“ACME Corp 在 2023 年第二季度的收入增长是多少?”

相关文本块可能包含这段文字:“该公司的收入较上一季度增长了 3%。”然而,这个文本块本身没有说明它指的是哪家公司,也没有说明相关时间段,因此很难检索到正确信息,或者有效使用这些信息。

Contextual Retrieval 简介

Contextual Retrieval 通过在 embedding 之前、以及创建 BM25 索引之前,给每个文本块前置特定于该文本块的解释性上下文,来解决这个问题。这两项分别称为 “Contextual Embeddings” 和 “Contextual BM25”。

让我们回到 SEC 文件集合的例子。下面展示一个文本块可能如何被转换:

original_chunk = "The company's revenue grew by 3% over the previous quarter."
contextualized_chunk = "This chunk is from an SEC filing on ACME corp's performance in Q2 2023; the previous quarter's revenue was $314 million. The company's revenue grew by 3% over the previous quarter."

值得注意的是,过去也有人提出过其他利用上下文改善检索的方法。其他方案包括:给文本块添加通用文档摘要(我们做过实验,只看到非常有限的提升)、hypothetical document embedding,以及基于摘要的索引(我们评估后发现性能较低)。这些方法都不同于本文提出的方法。

实现 Contextual Retrieval

当然,要手动标注知识库中成千上万甚至数百万个文本块,工作量会大得离谱。为了实现 Contextual Retrieval,我们转向 Claude。我们写了一个 prompt,指示模型基于整篇文档的上下文,提供简洁且特定于文本块的上下文说明。我们使用下面这个 Claude 3 Haiku prompt 为每个文本块生成上下文:

<document>
{{WHOLE_DOCUMENT}}
</document>
Here is the chunk we want to situate within the whole document
<chunk>
{{CHUNK_CONTENT}}
</chunk>
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.

生成的上下文文本通常为 50-100 tokens,会在创建 embedding 之前,以及创建 BM25 索引之前,被前置到文本块前面。

下面是实践中的预处理流程:

Contextual Retrieval 是一种提升检索准确率的预处理技术。

如果你有兴趣使用 Contextual Retrieval,可以从我们的 cookbook 开始。

使用 prompt caching 降低 Contextual Retrieval 的成本

得益于我们上面提到的特殊 prompt caching 功能,Contextual Retrieval 可以用 Claude 以低成本实现。使用 prompt caching 时,你不需要为每个文本块都传入参考文档。你只需将文档加载到缓存中一次,然后引用之前缓存过的内容。

假设文本块为 800 tokens,文档为 8k tokens,上下文指令为 50 tokens,并且每个文本块生成 100 tokens 的上下文,那么生成 contextualized chunks 的一次性成本是每百万文档 tokens 1.02 美元。

方法

我们在多个知识领域(代码库、小说、ArXiv 论文、科学论文)、embedding 模型、检索策略和评测指标上进行了实验。我们在 Appendix II 中放入了每个领域所用问题和答案的若干示例。

下面的图展示了在所有知识领域上的平均表现,使用表现最佳的 embedding 配置(Gemini Text 004),并检索 top-20 chunks。我们使用 1 减 recall@20 作为评测指标,它衡量的是相关文档未能在前 20 个文本块内被检索到的比例。你可以在附录中看到完整结果;在我们评估的每一种 embedding 与来源组合中,上下文化都会提升性能。

性能提升

我们的实验显示:

  • Contextual Embeddings 将 top-20-chunk 检索失败率降低了 35%(5.7% → 3.7%)。
  • 结合 Contextual Embeddings 和 Contextual BM25 后,top-20-chunk 检索失败率降低了 49%(5.7% → 2.9%)。

结合 Contextual Embedding 和 Contextual BM25 将 top-20-chunk 检索失败率降低了 49%。

实现注意事项

实现 Contextual Retrieval 时,需要注意以下几点:

  1. 文本块边界:思考如何将文档拆分成文本块。文本块大小、文本块边界和文本块重叠的选择都可能影响检索性能^{1}。
  2. Embedding 模型:虽然 Contextual Retrieval 会提升我们测试过的所有 embedding 模型的性能,但有些模型可能获益更多。我们发现 GeminiVoyage embeddings 尤其有效。
  3. 自定义 contextualizer prompts:虽然我们提供的通用 prompt 效果不错,但你也许可以通过针对特定领域或使用场景定制 prompts 来取得更好的结果(例如,加入关键术语表,其中某些术语可能只在知识库的其他文档中定义)。
  4. 文本块数量:向 context window 中加入更多文本块,会提高包含相关信息的概率。不过,更多信息可能会分散模型注意力,因此这里存在上限。我们尝试传入 5、10 和 20 个文本块,发现使用 20 个在这些选项中表现最好(参见附录中的对比),但仍然值得针对你的使用场景做实验。

务必运行评测:把 contextualized chunk 传给模型,并区分哪些内容是上下文、哪些内容是文本块,可能会改善回复生成。

用 Reranking 进一步提升性能

在最后一步,我们可以将 Contextual Retrieval 与另一项技术结合,以获得更大的性能提升。在传统 RAG 中,AI 系统会搜索知识库,寻找潜在相关的信息文本块。对于大型知识库,这一步初始检索经常会返回大量文本块,有时多达数百个,而且相关性和重要性各不相同。

Reranking 是一种常用过滤技术,用于确保只有最相关的文本块会被传给模型。Reranking 可以提供更好的回复,并降低成本和延迟,因为模型需要处理的信息更少。关键步骤如下:

  1. 执行初始检索,获得最可能相关的 top chunks(我们使用 top 150);
  2. 将 top-N 文本块连同用户查询一起传入 reranking 模型;
  3. 使用 reranking 模型,根据每个文本块与 prompt 的相关性和重要性给出分数,然后选择 top-K 文本块(我们使用 top 20);
  4. 将 top-K 文本块作为上下文传入模型,生成最终结果。

结合 Contextual Retrieval 和 Reranking,以最大化检索准确率。

性能提升

市场上有多种 reranking 模型。我们使用 Cohere reranker 进行了测试。Voyage 也提供 reranker,但我们没有时间测试。我们的实验显示,在多个领域中,加入 reranking 步骤会进一步优化检索。

具体来说,我们发现 Reranked Contextual Embedding 和 Contextual BM25 将 top-20-chunk 检索失败率降低了 67%(5.7% → 1.9%)。

Reranked Contextual Embedding 和 Contextual BM25 将 top-20-chunk 检索失败率降低了 67%。

成本与延迟注意事项

关于 reranking,一个重要考量是它对延迟和成本的影响,尤其是在对大量文本块进行 reranking 时。因为 reranking 会在运行时增加一个额外步骤,所以它不可避免地会增加少量延迟,即使 reranker 会并行给所有文本块打分也是如此。在对更多文本块做 reranking 以获得更好性能,和对更少文本块做 reranking 以降低延迟与成本之间,存在内在权衡。

我们建议在你的具体使用场景中尝试不同设置,找到合适平衡。

结论

我们进行了大量测试,对上文描述的所有技术的不同组合进行了比较(embedding 模型、是否使用 BM25、是否使用 contextual retrieval、是否使用 reranker,以及检索得到的 top-K 结果总数),并覆盖多种不同数据集类型。以下是我们的发现摘要:

  1. Embeddings+BM25 优于单独使用 embeddings;
  2. 在我们测试的 embeddings 中,Voyage 和 Gemini 表现最好;
  3. 向模型传入 top-20 文本块,比只传入 top-10 或 top-5 更有效;
  4. 给文本块添加上下文会大幅提升检索准确率;
  5. Reranking 优于不使用 reranking;
  6. 所有这些收益可以叠加:为了最大化性能提升,我们可以将 contextual embeddings(来自 Voyage 或 Gemini)与 contextual BM25 结合,再加上一个 reranking 步骤,并把 20 个文本块加入 prompt。

我们鼓励所有使用知识库的开发者使用我们的 cookbook 来实验这些方法,以解锁新的性能水平。

附录 I

下面是按数据集、embedding 提供商、是否在 embeddings 之外使用 BM25、是否使用 contextual retrieval、以及是否使用 reranking 划分的 Retrievals @ 20 结果明细。

参见附录 II,了解 Retrievals @ 10 和 @ 5 的明细,以及每个数据集的示例问题和答案。

跨数据集和 embedding providers 的 1 minus recall @ 20 结果。

致谢

研究与写作:Daniel Ford。感谢 Orowa Sikder、Gautam Mittal 和 Kenneth Lien 提供关键反馈,感谢 Samuel Flamini 实现 cookbooks,感谢 Lauren Polansky 进行项目协调,也感谢 Alex Albert、Susan Payne、Stuart Ritchie 和 Brad Abrams 帮助打磨这篇博客文章。