Meta提出全新文档级嵌入框架，利用LLM来增强信息检索能力

科学

6月前

近年来，基于嵌入式检索（embedding-based search）或密集检索（dense retrieval）相比传统的稀疏检索（sparse retrieval）或基于词袋（bag of words）的方法，已经展示出了更先进的结果。

于此同时，Meta 的研究团队（Mingrui Wu 和 Sheng Cao）近日发表了一个通过大型语言模型（LLM）实现检索增强的的文档级嵌入框架，使得可以显著提高广泛使用的各类检索器模型（如 Bi-encoders (Contriever，DRAGON) 和 late-interaction模型 (ColBERTv2)）的检索有效性，并在 LoTTE 数据集和 BEIR 数据集上实现了 Sota（State of the art）的结果。

论文链接：

https://arxiv.org/abs/2404.05825

引言

Bi-encoder 是一种广泛用于信息检索的神经网络架构。它由两个编码器组成，通常以 Transformer 模型的形式，分别为用户查询（user query）和潜在文档（target document）或段落编码一个向量表示。这两个编码器可以共享或使用两个独立的模型。然后可以计算这两个嵌入向量之间的相似性，通常使用点积或余弦相似性，以确定文档或段落与用户查询条（user query）的相关性。

Cross-encoders 与 bi-encoders 不同，它们在早期阶段就将输入融合在一起，允许用户查询和文档之间进行更复杂的交互。在这里，用户查询和文档被连接在一起，基于此计算一个联合的潜在空间中的向量。然后使用此向量进行预测，例如在信息检索任务中确定文档对查询的相关性。在需要对输入之间的相互作用进行细致理解的任务中，Cross-encoders 通常优于 bi-encoders。

Late-interaction 模型，如 ColBERT 和 ColBERTv2，是混合了 cross-encoder 模型和 bi-encoder 模型的模型架构。查询和文档被独立地编码成 token 级别的向量表示。

因此，在某种意义上，这是一个向量袋（bag of embedding）的模型。这些向量表示之间的交互，即“后期交互”，涉及到在 token 级别的向量嵌入上计算余弦相似性或点积分数。所有的模型架构都需要用户查询和目标文档的信息嵌入向量。虽然我们不能在检索任务中控制用户查询，但我们假设通过丰富文档的嵌入式向量可以提高文本检索的质量和稳健性。

大语言模型增强检索及其框架

2.1 合成相关查询（synthetic relevant queries）

这个概念的灵感来自网络搜索技术。为了说明这个想法，让我们考虑一个用户查询 "MIT" 的例子。如果没有先验知识，很难弄清楚 "Massachusetts Institute of Technology" 和 "MIT" 是等价的。

然而，在网络搜索中，我们可以观察到 "Massachusetts Institute of Technology" 的主页从 "MIT" 的查询中收到了大量的点击，这使我们可以推断出 "Massachusetts Institute of Technology" 的主页必须与 "MIT" 的查询密切相关。

另一方面，我们通常没有每个用户查询的点击数据。然而，大型语言模型擅长生成合成查询，所以我们可以使用合成查询作为代理的“点击数据”来引导用户查询到相关的文档。

一个重要的点是，在传统的检索任务中，我们使用相似性来表示相关性。相似性得分在数学上定义为用户查询和文档的编码向量的点积或余弦。然而，有时这个相似性得分可能不能反映语义相关性。例如，“谁是美国的第一任总统？”在相似性得分上可能非常接近“谁成为了美国的第一任总统？”。但我们的目标答案可能是关于“乔治·华盛顿”的维基百科页面或自传，其与查询的相似性得分可能不那么高。

然而，如果我们使用华盛顿的自传来创建合成查询，“谁成为了美国的第一任总统？”可能是其中之一。用户查询“谁是美国的第一任总统？”可以通过相似性得分轻松地匹配到相关查询。后者指向目标文档（华盛顿的自传）。因此，生成的相关查询从不同的角度表达了原始文档的语义，这有助于匹配相关的查询。

2.2 标题（title）

文档的标题在确定其对用户查询的相关性和用途方面起着关键作用。标题通常是用户在搜索文档时看到的第一条信息，它可以极大地影响他们点击特定结果的决定。一个精心制作的标题可以提供重要的上下文和关键词，帮助用户快速理解文档的内容和目的。如果原始文档有标题，我们可以直接使用。如果没有，我们可以利用大型语言模型为该文档生成一个合成标题。

2.3 段落（chunk/passage）

分块是指将一大段文档或文本划分为更小、更易管理的单元的过程。这些单元，被称为"块"或"段落"，通常是通过将相关的信息组合在一起创建的。由于检索模型的上下文窗口的限制（换句话说，模型输入的最大长度），我们通常将一个长文档划分为几个块，其 token 数量低于上下文窗口限制。块数据来自原始文档，而不是来自 LLM 增强。最优的分块大小对于各种检索模型是不同

以下是对于文档级嵌入框架当中对于文档“相关性”搜索的表达。

2.4 大语言模型增强检索框架下的Bi-encoder

Bi-encoder 通常是“双塔”（two-tower）模型结构。给定一个查询和一个文档，分别应用查询编码器和文档编码器来计算查询和文档的嵌入向量。然后将这两个嵌入向量输入到点积（或余弦相似度）中，以计算查询和文档之间的相似度分数。由于我们希望通过注入合成查询和标题来丰富文档嵌入向量，作者建议计算相似度如下：

右手边的第一项计算查询块嵌入对的最大相似度分数，其中是相似度分数函数，是输入查询的嵌入向量，是文档中第块的嵌入向量。这一项在当前基于嵌入的检索中常用，它根据查询和文档中最相关的块确定查询和文档之间的相似度。

第二项是创新的，并考虑了更多的信息，其中是每个文档字段的嵌入向量。计算查询嵌入和每个字段嵌入之间的相似度分数，并将它们组合在一起，每个字段都有一个字段权重参数。如前所述，这些文档字段包括合成查询、标题和块字段。

现在让我们考虑如何计算每个字段的查询字段嵌入。对于只包含一个阶段或句子的标题字段，这是直接的。我们只需应用文档编码器来计算标题的嵌入向量作为标题字段嵌入。对于可能包含多个文档块的块字段，我们可以使用文档编码器计算每个块的嵌入向量。问题是如何组合这些嵌入向量来表示整个文档？作者建议计算所有块嵌入向量的平均值作为块字段嵌入。

同样，对于合成查询字段，作者使用查询编码器计算每个查询的嵌入向量，然后计算这些嵌入向量的平均值作为查询字段嵌入。这种简单的方法在实验中效果非常好，而在未来显然可以在这里探索更先进的方法。

此外，由于相似度函数是线性的（点积是线性的；如果我们将嵌入向量归一化到单位长度，余弦相似度也是线性的），上述方程可以简化为以下形式：

因此，我们可以将视为原始文档的每个块的块嵌入向量。这将与业内的嵌入式向量搜索数据库兼容。

2.5 大语言模型增强检索框架下的Token-Level Late-Interaction model

与使用单个嵌入向量表示查询和每个文档的单个嵌入向量不同，像 ColBERT 和 ColBERTv2 这样的后期交互模型 token 级嵌入向量交互，所有 token 的嵌入向量都被保留，并将参与计算查询和文档之间的相似度分数。

其中和分别是输入查询和文档的 token 级嵌入向量。因此，对于每个查询 token，都会识别出文档中最相似的 token，并记录他们的相似度分数。所有这些分数都会在所有查询 token 上求和，以获得查询和文档之间的总体相似度。由于相似度分数的计算是在 token 级别完成的，我们可以将合成查询和标题连接到原始文档段落。之后，如果 token 数量达到上下文窗口限制，我们将决定是否对连接的文档进行分块。

Experiments

作者选择开源的 Llama-70B 进行合成查询生成和标题生成。对于 Bi-encoders，作者按照上述方法实现了文档级嵌入，并为 Contriever 模型选择了=1.0，=0.5，=0.1，为 DRAGON 模型选择了=0.6，=0.3，=0.3。

作者在历经实验后选择了=64，并发现这样的 chunk size通常在 Bi-encoders 的检索结果中表现最好。注意，这些字段权重的超参数并没有完全优化。它们是根据 Bi-encoder 模型在单个数据集上的性能选择的，然后固定用于所有其他数据集的评估。

对于 ColBERTv2，如前所述，作者将标题与每个文档的所有合成查询连接起来，使其成为原始文档的额外“段落”。因此，在这些实验中没有字段权重超参数。在 late-interaction 模型架构下，可能有其他更好的组装方法来组成文档级嵌入。作者在构建 ColBERT 索引时设置 index_bits=8。

基于所有三种模型在 LoTTE 和 BEIR 数据上的结果，我们可以观察到，LLM 增强检索和文档级嵌入显著提高了Bi-encoders（Contriever 和 DRAGON）的recall@3 和 recall@10。

对于 token-level late-interaction 模型（ColBERTv2），在 LoTTE 和 BEIR 数据集上的增加仍然明显，尽管没有 Bi-encoders 那么多。这可能是因为 token-level late-interaction 模型的基线比 Bi-encoders 的基线要高。

此外，LLM 增强的 Contriver 的性能已经超过了大多数数据集中的 vanilla DRAGON。同样，LLM 增强的 DRAGON 甚至在 BEIR-ArguAna，BEIR-SciDocs 和 BEIR-CQADupstack-English数据集上超过了 vanilla ColBERTv2，并在其余数据集中大大减小了性能差距，尽管 ColBERTv2 引入了比 DRAGON 更复杂的 late-interaction 架构。

因此，我们可以看到，通过使用 LLM 增强来丰富文档的嵌入，可以在不进行进一步微调的情况下，大大提高检索器模型的召回性能。

Ablation Studies

作者进行了进一步研究来了解段落块（chunk）, 相关查询（relevant query）和标题（title）字段如何影响不同检索模型的检索质量。对于 Bi-encoders（Contriever 和 DRAGON），作者进一步控制了块、合成查询和标题的字段权重，以了解这些参数如何影响性能。

对于 token-level late-interaction 模型（ColBERTv2），作者只控制模型使用块、查询或标题字段中的一个，以了解它们如何影响端到端的检索质量。

对于 Contriever 模型，作者观察到大多数时候，合成查询在提升召回性能方面起着最关键的作用，与其他两个字段相比，而在 BEIR-SciDocs 和 BEIR-Scifact 中，合成查询的重要性较小。因此，文档级嵌入中多个字段的加权和在大多数情况下可以产生更好的性能。这些权重可以作为超参数进一步调整。

对于 DRAGON 模型，在文档级嵌入中哪个字段起着更重要的作用没有明显的模式。在 LoTTE 数据集中，它更多地受到标题字段的驱动。然而，在 BEIR-ArguAna、BEIR-Quora 和 BEIR-SciFact 的数据集中，块字段更重要。再次，文档级嵌入中多个文档字段的加权和在大多数情况下可以产生更好的性能。

对于作者在 DRAGON 和 Contriever 中观察到的不同模式，一个原因可能是 DRAGON 使用独立的查询和上下文编码器，而 Contriever 在作者的设置中使用共享的查询和上下文编码器。因此，Contriever 更擅长于识别相似性而不是相关性，这就是为什么合成查询字段在 Contriever 模型中有更大的影响，因为它更好地将相似性转化为相关性。

对于 ColBERTv2，作者从 LoTTE 数据集中观察到，进一步切分段落实际上会损害性能。这可能是因为相似性计算是在 token 级别进行的，所以切分段落（同时增加块的数量）不会帮助 ColBERTv2 模型消化粒度的上下文信息。因此，作者在 BEIR 数据集上没有评估仅块的情况。

在所有数据集中，合成查询对 ColBERTv2 的重要性都大于标题，而将它们全部组合在一起通常会产生更好的召回结果。再次注意，对于 token-level late-interaction 模型，没有字段权重超参数。

结论和局限性

本文提出了一种创新的框架，即大语言模型增强检索，通过大型语言模型来丰富文档的嵌入，显著提高了现有检索模型的性能。该框架包括一个文档级嵌入系统，它编码了来自合成查询、标题和段落块的上下文信息，可以适应各种主流检索模型架构。该方法在不同的模型和数据集上取得了 Sota 的结果，证明了其在提高神经信息检索的质量和稳健性方面的有效性。

未来的研究可以探索对大语言模型增强检索框架的进一步改进，例如将更多的上下文信息集成到文档级嵌入中，应用更先进的相似性得分度量，更复杂的方法来将多个块/查询的嵌入组合成一个块/查询字段嵌入等。

本研究的一个局限性是它在为原始文档增强相关查询和标题时需要额外的计算资源，有时增强文本的大小可以与原始文档的大小相当。这种计算量可能会限制这种方法在计算资源有限的地方的使用。另一个局限性或风险是，大型语言模型中的幻觉可能会对增强语料库对原始文档的准确性造成额外的影响。幻觉在大型语言模型的研究领域仍然是一个未解决的问题。

更多阅读