【论文速读】LLM-Augmented Retrieval:EnhancingRetrievalModels Through LanguageModels and DocLevel Embedding

论文链接：https://arxiv.org/html/2404.05825v1
文章标题：LLM-Augmented Retrieval: Enhancing Retrieval Models Through Language Models and Doc-Level Embedding
这篇文章提出了一种与检索模型无关的框架框架，通过大型语言模型来丰富文档的嵌入，显著提高了现有检索模型的性能。

摘要：

文章介绍了一种通过大型语言模型（LLM）增强的文档级别嵌入框架，用于提升基于嵌入的检索（embedding-based retrieval）或密集检索（dense retrieval）的效果。这种方法与传统的基于稀疏向量或词袋模型的方法相比，展示了最先进的结果。作者通过实施这种LLM增强的检索框架，显著提高了广泛使用的检索模型（如Bi-encoders和late-interaction models）的有效性，并在LoTTE数据集和BEIR数据集上取得了最先进的结果。

解决的主要问题：

1. 信息检索的改进：传统的信息检索方法依赖于稀疏或词袋（bag of words）模型，这些方法在处理语义相关性方面存在局限性。
2. 检索模型训练的改进：包括负采样（negative sampling）和损失函数等重要组件的改进。

主要方法：

1. LLM增强检索：提出了一个与模型无关的框架，通过大型语言模型增强来丰富文档的上下文信息，从而提高现有检索器的质量和鲁棒性。
2. 文档级别嵌入：提出了一种结合更多上下文信息的文档级别嵌入方法。
3. 跨不同模型和数据集的评估：在不同的模型和广泛的数据集上评估了这个框架，确立了超越原始模型的最新质量标准。
4. 检索模型训练的关键组件改进：对检索模型训练中的一些关键组件进行了改进，例如负采样和损失函数等。

1 Introduction的重点内容

1. Bi-encoder架构：Bi-encoder这种神经网络架构，它在信息检索中被广泛使用。Bi-encoder由两个编码器组成，通常采用Transformer模型，分别对用户查询和潜在文档或段落进行编码，生成向量表示。这两个编码器可以共享或使用两个独立的模型。通过计算这两个嵌入向量之间的相似度（通常使用点积或余弦相似度），确定文档或段落与用户查询的相关性。

2. Cross-encoders：与Bi-encoder不同，Cross-encoder在早期阶段就合并输入，允许用户查询和文档之间更复杂的交互。用户查询和文档被连接起来，基于此计算联合嵌入向量，然后使用该向量进行预测，例如文档对查询的相关性。

3. Late-interaction模型：如ColBERT、ColBERTv2或SPALDE++，这些模型架构混合了Cross-encoder模型和Bi-encoder模型的特点。查询和文档被独立编码成token级别的向量表示，然后通过计算这些表示之间的余弦相似度或点积分数来进行“后期交互”。

4. 文档嵌入的丰富性：文章提出，虽然在检索任务中我们无法控制用户查询，但作者假设通过丰富文档的嵌入可以提高文本检索的质量和鲁棒性。

5. 贡献：作者列出了他们在这项工作中的贡献，包括提出了LLM增强检索框架、文档级别嵌入、在不同模型和数据集上的评估以及改进检索模型训练的关键组件。
   
   图1提供了LLM增强检索框架的总体视图。在这个框架中：

• 利用大型语言模型（LLM）生成合成的相关查询和合成标题。
• 将这些合成的查询和标题与原始文档中分块（passages）的内容一起组合成文档级别的嵌入。
• 最终的检索是基于用户查询和文档级别嵌入之间的相似度来进行的。

图1展示了如何通过LLM增强来丰富文档嵌入，从而提高检索系统的性能。这个过程包括了从原始文档中提取信息，并结合LLM生成的内容来创建一个更全面、更相关的文档表示，以便更准确地响应用户的查询。

2 Background & Related Work的重点内容

1. 数据增强在信息检索中的应用：讨论了数据增强技术在信息检索训练中的广泛使用，包括对比学习、逆完形任务、独立裁剪、随机单词删除、替换或掩码等方法来增加训练数据的多样性。

2. 大型预训练语言模型的文本生成能力：强调了大型预训练语言模型（如GPT-2、BERT等）在生成高质量文本数据方面的能力，并提到了以往研究中如何利用这些模型的生成能力来为检索模型创建合成训练数据。

3. 伪查询生成：探讨了预生成的伪查询在提高检索性能方面的效果，包括通过BM25或BERT模型计算伪查询和用户查询之间的相似度，以及通过K-means聚类算法生成伪查询嵌入。

4. 检索增强生成（RAG）系统：介绍了RAG系统，该系统通过检索模块根据输入查询从语料库中检索相关文档，并将检索到的文档作为上下文整合到语言模型中，以优化最终的响应生成。

5. 相似度分数：讨论了如何计算输入查询和目标文档的嵌入向量之间的相似度分数以确定它们之间的相关性，包括点积和余弦相似度的比较及其在不同检索系统中的应用。

6. 训练方法：介绍了在密集检索训练中广泛采用的对比InfoNCE损失函数和其他流行的损失函数，如点式二元交叉熵、列表交叉熵、RankNet和LambdaLoss等。

3 LLM-augmented Retrieval

3.1 LLM增强检索框架

在本节中，首先讨论LLM增强检索框架的组成部分。之后，将解释如何将此框架适应于不同的检索模型架构。特别是，为Bi-encoders和late-interaction编码器提出了文档级别嵌入，并展示了如何将其应用于提高端到端检索质量。

3.1.1 合成相关查询

这个概念的灵感来自网页搜索技术。为了说明这个想法，让我们考虑一个用户查询“MIT”的例子。没有先验知识的情况下，很难弄清楚“麻省理工学院”和“MIT”是等价的。然而，在网页搜索中，我们可以观察到“麻省理工学院”的主页收到了来自查询“MIT”的许多点击，这使我们能够推断“麻省理工学院”的主页必须与查询“MIT”密切相关。另一方面，我们通常在上下文检索的场景中没有每个用户查询的点击数据。然而，大型语言模型擅长生成合成查询，因此我们可以使用合成查询作为代理的“点击数据”，以将用户查询引导到相关文档。

一个重要的点是，在传统检索任务中，我们使用相似度来表达相关性。相似度分数在数学上定义为用户查询和文档编码向量的点积或余弦。然而，有时这种相似度分数可能无法反映语义相关性。例如，“谁是美国的第一位总统？”在相似度分数上可能与“谁成为了美国的第一位总统？”非常接近。但我们的目标答案可能是关于“乔治·华盛顿”的维基页面或自传，其与查询的相似度分数并不高。而如果我们使用华盛顿的自传来创建合成查询，“谁成为了美国的第一位总统？”可能就是其中之一。用户查询“谁是美国的第一位总统？”可以通过相似度分数轻松匹配到相关查询。后者指向目标文档（华盛顿的自传）。因此，生成的合成查询从不同角度表达了原始文档的语义，这有助于匹配相关查询。

图2：通过合成相关查询，相关关系现在不仅通过相似度表达，而且还通过大型语言模型的增强步骤来表达。

3.1.2 标题

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

3.1.3 块（段落）

块化是指将大型文档或文本分割成更小、更易于管理的单元。这些单元被称为“块”或“段落”，通常通过将相关信息组合在一起来创建。由于检索模型的上下文窗口限制（换句话说，模型输入的最大长度），我们通常将长文档分割成几个token数量低于上下文窗口限制的块。块数据来自原始文档，而不是来自LLM增强。对于不同的检索模型，最优的块化大小是不同的。对于像Contriever和DRAGON这样的Bi-encoders，我们通过实证研究发现64是最优的块化大小。对于像ColBERT和ColBERTv2这样的token级别的late-interaction模型，由于它已经在token级别计算相似度分数，所以除非达到上下文窗口限制，否则没有必要对原始文档进行分块化。

3.2 文档级别嵌入

在本节中，首先介绍文档级别嵌入在信息检索中的高级概念，然后使用Bi-encoders和token级别的late-interaction模型来说明文档级别嵌入如何适应不同的检索模型结构。

文档字段。为了方便起见，将上述信息源，即合成查询、标题和块，称为文档的字段。这些字段从不同角度表达了原始文档的语义，并将组合成一个静态的、可以预先计算和缓存的文档级别嵌入，用于信息检索。可以预先构建嵌入索引以加快检索推理速度，每个文档级别嵌入指向原始文档。

图3：文档级别嵌入中“相关性”的图形表示

3.2.1 对于Bi-encoders

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

右侧的第一项计算查询和块嵌入对的最大相似度分数，其中s是相似度分数函数，q是输入查询的嵌入向量， c_i 是文档中第 i 个块的嵌入向量。这一项通常用于当前基于嵌入的检索中，它基于查询和文档中最相关的块来确定查询和文档之间的相似性。

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

现在让我们考虑如何为每个字段计算查询字段嵌入e_f 。对于标题字段，它只包含一个短语或句子，这很直接。我们只应用文档编码器来计算标题的嵌入向量作为标题字段嵌入。

对于块字段，可能包含多个文档块，我们可以使用文档编码器计算每个块的嵌入向量。问题是，如何组合这些嵌入向量来表示整个文档？实际上，人们已经考虑了句子嵌入的类似问题：拥有一个句子中所有标记的嵌入向量后，如何得到整个句子的表示？在(Arora et al., 2017)中提出了一个简单但有效的想法来解决这个问题。在这里，我们将这个想法适应于文档嵌入问题。即我们计算所有块嵌入向量的平均值作为块字段嵌入。同样，对于合成查询字段，我们使用查询编码器计算每个查询的嵌入向量，然后计算这些嵌入向量的平均值作为查询字段嵌入。这种简单的方法在我们的实验中效果很好，尽管未来可以探索更先进的方法。

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

因此，我们可以将c_i + sum（w_f * e_f)视为原始文档中每个块c_i的块嵌入向量，并仍然应用近似最近邻等算法来检索最相关的文档。

3.2.2 对于Token级别的Late-Interaction模型

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

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

4 Experiments的要点内容

1. 数据集：实验使用了两个数据集：BEIR（Benchmark for Evaluating Information Retrieval）和LoTTE（Long-Tail Topic-stratified Evaluation）。BEIR是一个全面的基准测试，用于评估和比较多个信息检索模型，特别是在领域外测试的背景下。LoTTE数据集专注于长尾主题的自然用户查询。

2. 模型：实验涉及了三种模型：

   • Contriever：使用基于Roberta-base的模型架构，通过对比学习在Wiki和CC100数据上训练。
   • DRAGON：同样使用Roberta-base模型架构，但查询编码器和文档编码器是分开的。
   • ColBERTv2：使用bert-base-uncased模型架构，与原始论文中的设置相同。

3. 开源工具：选择了开源的Llama-70B模型用于合成查询和标题生成。

4. 超参数设置：对于Bi-encoders，作者实现了文档级别嵌入，并为Contriever和DRAGON模型选择了特定的字段权重。例如，对于Contriever模型，查询字段的权重为1.0，标题字段的权重为0.5，块字段的权重为0.1。这些权重是基于在单个LoTTE-lifestyle-forum数据集上的性能选择的，并在其他所有数据集的评估中固定。

5. 块大小：对于Bi-encoders，作者通过实证研究发现块大小为64时，在检索结果上通常表现最佳。

6. ColBERTv2的实现：与Bi-encoders不同，ColBERTv2将标题和所有合成查询连接到每个文档，并将其作为原始文档的附加“段落”。因此，在这个实验中没有字段权重超参数。作者还提到，可能有其他更好的组装方法来在late-interaction模型架构下组成文档级别嵌入。

7. 索引构建：对于ColBERTv2，作者在构建索引时设置了index_bits=8。

8. 提示模板：文章还分享了用于生成合成查询和标题的提示模板，这些模板用于与Llama-70B模型交互。