AI 系统中的上下文检索(Context Retrieval)

原文: Introducing Contextual Retrieval
日期:2024年9月19日

要让一个 AI 模型在特定情境下发挥作用,它通常需要访问背景知识。

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

开发者通常使用检索增强生成(Retrieval-Augmented Generation, RAG)来增强 AI 模型的知识。RAG 是一种从知识库中检索相关信息并将其附加到用户提示中的方法,从而显著提升模型的响应质量。问题在于,传统的 RAG 解决方案在编码信息时会移除上下文,这常常导致系统无法从知识库中检索到相关信息。

在这篇文章中,我们概述了一种能显著改进 RAG 中检索步骤的方法。该方法被称为“上下文检索”(Contextual Retrieval),并使用了两种子技术:上下文嵌入(Contextual Embeddings)和上下文 BM25(Contextual BM25)。这种方法可以将检索失败的次数减少 49%,当与重排序(reranking)结合使用时,可减少 67%。这代表了检索准确性的显著提升,并直接转化为下游任务中更好的性能。

您可以使用我们的 cookbook 轻松地在 Claude 中部署您自己的上下文检索解决方案。

关于简单使用更长提示的说明

有时候最简单的解决方案是最好的。如果您的知识库小于 200,000 个 token(约 500 页材料),您可以直接将整个知识库包含在给模型的提示中,无需使用 RAG 或类似方法。

几周前,我们为 Claude 发布了提示缓存(prompt caching)功能,这使得这种方法的速度显著加快且成本效益更高。开发者现在可以在 API 调用之间缓存频繁使用的提示,将延迟降低超过 2 倍,成本降低高达 90%(您可以通过阅读我们的提示缓存 cookbook 了解其工作原理)。

然而,随着您的知识库不断增长,您将需要一个更具扩展性的解决方案。这时,上下文检索就派上用场了。

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

对于无法容纳在上下文窗口中的更大型知识库,RAG 是典型的解决方案。RAG 的工作方式是通过以下步骤对知识库进行预处理:

  1. 将知识库(文档“语料库”)分解成更小的文本块,通常不超过几百个 token;
  2. 使用嵌入模型将这些文本块转换为编码了语义的向量嵌入;
  3. 将这些嵌入存储在向量数据库中,以便按语义相似性进行搜索。

在运行时,当用户向模型输入查询时,向量数据库会根据与查询的语义相似性找到最相关的文本块。然后,将最相关的文本块添加到发送给生成模型的提示中。

虽然嵌入模型擅长捕捉语义关系,但它们可能会错过关键的精确匹配。幸运的是,有一种更古老的技术可以在这些情况下提供帮助。BM25(Best Matching 25)是一种排序函数,它使用词法匹配来查找精确的单词或短语匹配。它对于包含唯一标识符或技术术语的查询特别有效。

BM25 的工作原理建立在 TF-IDF(词频-逆文档频率)概念之上。TF-IDF 衡量一个词在文档集合中对某个文档的重要性。BM25 对此进行了改进,考虑了文档长度并对词频应用了饱和函数,这有助于防止常用词在结果中占据主导地位。

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

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

  1. 将知识库(文档“语料库”)分解成更小的文本块,通常不超过几百个 token;
  2. 为这些文本块创建 TF-IDF 编码和语义嵌入;
  3. 使用 BM25 根据精确匹配找到排名靠前的文本块;
  4. 使用嵌入根据语义相似性找到排名靠前的文本块;
  5. 使用排名融合技术合并并去重来自 (3) 和 (4) 的结果;
  6. 将排名最高的 K 个文本块添加到提示中以生成响应。

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

一个标准的检索增强生成(RAG)系统,它同时使用嵌入和 BM25 来检索信息。TF-IDF(词频-逆文档频率)衡量词的重要性,并构成 BM25 的基础。

这种方法使您能够以经济高效的方式扩展到庞大的知识库,远远超出了单个提示所能容纳的范围。但这些传统的 RAG 系统有一个显著的局限性:它们常常会破坏上下文。

传统 RAG 中的上下文难题

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

例如,假设您的知识库中嵌入了一系列财务信息(比如美国证券交易委员会(SEC)的文件),然后您收到了以下问题:“ACME 公司在 2023 年第二季度的收入增长是多少?”

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

引入上下文检索

上下文检索通过在嵌入(“上下文嵌入”)和创建 BM25 索引(“上下文 BM25”)之前,为每个文本块前置特定于该块的解释性上下文来解决这个问题。

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

1
2
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."

值得注意的是,过去也曾提出过其他使用上下文来改进检索的方法。其他提案包括:向文本块添加通用的文档摘要(我们进行了实验,发现收益非常有限)、假设性文档嵌入基于摘要的索引(我们评估后发现性能较低)。这些方法与本文提出的方法不同。

实现上下文检索

当然,手动为知识库中成千上万甚至数百万个文本块添加注释的工作量太大了。为了实现上下文检索,我们求助于 Claude。我们编写了一个提示,指示模型提供简洁的、特定于文本块的上下文,利用整个文档的语境来解释该文本块。我们使用以下 Claude 3 Haiku 提示为每个文本块生成上下文:

1
2
3
4
5
6
7
8
<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 个 token,在嵌入和创建 BM25 索引之前被添加到文本块的前面。

以下是预处理流程在实践中的样子:

上下文检索是一种提高检索准确性的预处理技术。

如果您有兴趣使用上下文检索,可以从我们的 cookbook 开始。

使用提示缓存降低上下文检索的成本

借助我们上面提到的特殊提示缓存功能,上下文检索在 Claude 中能够以低成本实现,这是其独特之处。通过提示缓存,您无需为每个文本块都传入参考文档。您只需将文档加载到缓存中一次,然后引用之前缓存的内容。假设每个文本块 800 个 token,每个文档 8000 个 token,上下文指令 50 个 token,每个文本块的上下文 100 个 token,生成上下文化文本块的一次性成本为每百万文档 token 1.02 美元

方法论

我们跨越了各种知识领域(代码库、小说、ArXiv 论文、科学论文)、嵌入模型、检索策略和评估指标进行了实验。我们在附录 II 中为每个领域提供了一些我们使用的问题和答案示例。

下面的图表显示了在所有知识领域中使用性能最佳的嵌入配置(Gemini Text 004)并检索排名前 20 的文本块时的平均性能。我们使用 1 减去 recall@20 作为我们的评估指标,它衡量了在前 20 个文本块中未能检索到的相关文档的百分比。您可以在附录中看到完整的结果——在我们评估的每一种嵌入-来源组合中,上下文化都提高了性能。

性能提升

我们的实验表明:

  • 上下文嵌入将排名前 20 文本块的检索失败率降低了 35%(从 5.7% → 3.7%)
  • 结合上下文嵌入和上下文 BM25 将排名前 20 文本块的检索失败率降低了 49%(从 5.7% → 2.9%)

结合上下文嵌入和上下文 BM25 可将排名前 20 文本块的检索失败率降低 49%。

实现注意事项

在实现上下文检索时,需要考虑以下几点:

  1. 文本块边界(Chunk boundaries):考虑如何将文档分割成块。块的大小、边界和重叠的选择都会影响检索性能 ¹。
  2. 嵌入模型(Embedding model):尽管上下文检索在我们测试的所有嵌入模型中都提高了性能,但某些模型可能受益更多。我们发现 Gemini 和 Voyage 的嵌入特别有效。
  3. 自定义上下文生成提示(contextualizer prompts):虽然我们提供的通用提示效果很好,但您可能会通过针对特定领域或用例量身定制的提示(例如,包含可能只在知识库中其他文档中定义的关键术语词汇表)获得更好的结果。
  4. 文本块数量(Number of chunks):在上下文窗口中增加更多的文本块会增加包含相关信息的几率。然而,过多的信息可能会分散模型的注意力,因此这是有限度的。我们尝试了传递 5、10 和 20 个文本块,发现使用 20 个是这些选项中性能最好的(比较见附录),但这值得在您的具体用例中进行实验。

始终运行评估:通过将上下文化的文本块传递给模型,并区分什么是上下文、什么是文本块,可能会改善响应生成。

通过重排序进一步提升性能

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

重排序(Reranking)是一种常用的过滤技术,以确保只有最相关的文本块被传递给模型。重排序可以提供更好的响应,并降低成本和延迟,因为模型处理的信息更少。关键步骤是:

  1. 执行初始检索以获取排名靠前的潜在相关文本块(我们使用了前 150 个);
  2. 将排名前 N 的文本块连同用户的查询一起传递给重排序模型;
  3. 使用重排序模型,根据每个文本块与提示的相关性和重要性为其打分,然后选择排名前 K 的文本块(我们使用了前 20 个);
  4. 将排名前 K 的文本块作为上下文传递给模型以生成最终结果。

结合上下文检索和重排序以最大化检索准确性。

性能提升

市场上有几种重排序模型。我们使用 Cohere 的重排序器 进行了测试。Voyage 也提供重排序器,但我们没有时间进行测试。我们的实验表明,在各种领域中,添加重排序步骤可以进一步优化检索。

具体来说,我们发现重排序的上下文嵌入和上下文 BM25 将排名前 20 文本块的检索失败率降低了 67%(从 5.7% → 1.9%)。

重排序的上下文嵌入和上下文 BM25 可将排名前 20 文本块的检索失败率降低 67%。

成本和延迟的考虑

使用重排序时一个重要的考虑因素是对延迟和成本的影响,尤其是在对大量文本块进行重排序时。因为重排序在运行时增加了一个额外的步骤,它不可避免地会增加少量延迟,即使重排序器是并行地对所有文本块进行评分。在为了更好的性能而重排序更多文本块与为了更低的延迟和成本而重排序更少文本块之间存在固有的权衡。我们建议在您的特定用例中尝试不同的设置以找到合适的平衡点。

结论

我们进行了大量的测试,比较了上述所有技术(嵌入模型、BM25 的使用、上下文检索的使用、重排序器的使用以及检索的 top-K 结果总数)的不同组合,并涵盖了各种不同的数据集类型。以下是我们发现的总结:

  • 嵌入+BM25 优于单独使用嵌入;
  • 在我们测试的模型中,Voyage 和 Gemini 的嵌入效果最好;
  • 将排名前 20 的文本块传递给模型比仅传递前 10 或前 5 个更有效;
  • 向文本块添加上下文能极大地提高检索准确性;
  • 使用重排序优于不使用;
  • 所有这些好处都可以叠加:为了最大化性能提升,我们可以将上下文嵌入(来自 Voyage 或 Gemini)与上下文 BM25 相结合,再加上一个重排序步骤,并将 20 个文本块添加到提示中。

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

附录 I

以下是跨数据集、嵌入提供商、在嵌入基础上使用 BM25、使用上下文检索以及使用重排序的 Retrievals @ 20 结果细分。

有关 Retrievals @ 10 和 @ 5 的细分以及每个数据集的问题和答案示例,请参阅附录 II

跨数据集和嵌入提供商的 1 减去 recall @ 20 结果。