论文链接
Abstract
密集检索需要学习区分性文本嵌入来表示查询和文档之间的语义关系。考虑到大型语言模型在语义理解方面的强大能力，它可能受益于大型语言模型的使用。然而，LLM是由文本生成任务预先训练的，其工作模式与将文本表示为嵌入完全不同。因此，必须研究如何正确地调整LLM，以便它们能够有效地初始化为密集检索的骨干编码器。
在本文中，我们提出了一种新的方法，称为LLaRA（适用于密集检索的LLM），它作为LLM的后验自适应，用于密集检索应用。LLaRA由两个前置任务组成：基于嵌入的自动编码（EBAE）和基于嵌入的自回归（EBAR），其中来自LLM的文本嵌入分别用于重构输入句子的标记和预测下一个句子的标记。LLaRA被证明是简单、轻便和高效的。它被应用于维基百科语料库上的LLaMA-2-7B（基础），在那里它大大提高了模型在各种密集检索基准上的微调性能，如MSMARCO和BEIR。我们的模型和代码将在BGE存储库中公开。
1. Introduction
密集检索是深度神经网络提出的一种新的检索范式。与传统的IR方法不同，密集检索学习将查询和文档表示为同一潜在空间内的嵌入，其中查询和文档之间的语义关系可以通过嵌入相似性来反映。如今，密集检索已经是许多现实应用中的重要组成部分，比如网页搜索和开放域问答。
密集检索的质量受到其骨干编码器容量的严重影响。在过去的几年里，预训练的语言模型，例如BERT、RoBERTa、T5，被广泛应用于查询和文档的表示。实验发现，模型大小和训练规模的扩展可以显著提高密集检索的准确性和通用性。
最近，大型语言模型（LLM）被作为许多经典NLP任务的通用解决方案进行了微调。考虑到LLM在语义理解方面的卓越能力，利用这种强大的模型进行密集检索也很有希望。事实上，在这个方向上已经有了开创性的努力，LLM被提示或微调以生成判别嵌入，从而促进密集检索。
尽管取得了初步进展，但要充分释放LLM在密集检索应用中的潜力仍然很困难。特别地，LLM通过文本生成进行预训练，学习文本生成的嵌入来预测下一个token。因此，LLM的输出嵌入将主要集中于捕获上下文的局部和未来的语义。然而，密集检索需要嵌入来表示整个上下文的全局语义。如此大的差异将严重限制LLM在密集检索中的直接应用。
为了解决上面提到的问题，我们提出了一个新的方法，名叫LLaRA（如Figure 1），它作为LLM的后验自适应，以提高它们的密集检索能力。LLaRA可以看作是无监督生成预训练的一个扩展训练阶段。通过适当设计的文本前置任务，它旨在增强LLM，使其能够生成用于全局上下文语义表示的文本嵌入。
特别是LLaRA引入了两个前置训练任务：EBAE（Embedding Based Auto Encoding）和EBAR（Embedding-Based Auto Regression）。在EBAE中，LLM被提示生成文本嵌入，该文本嵌入可用于预测输入句子本身的token。而使用EBAR时，LLM会被提示生成文本嵌入，该嵌入可用于预测下一个句子的token。通过从上述前置任务中学习，可以将LLM的文本嵌入从局部语义表示（即对下一个标记的预测）调整为全局语义表示（例如对句子级特征的预测）。有了这两个不同的提示模板，LLM的嵌入能力可以区分开来，以处理不同的语义匹配场景，例如相似性搜索（使用EBAE的提示）和问答（使用EBAR的提示）。
在LLaRA中，句子级特征的预测是通过LLM的输出嵌入的线性投影进行的，其中不需要额外的解码组件。因此，LLaRA可以直接在现有的生成预训练管道之上实现，并且它导致了极具竞争力的训练效率。此外，不需要收集任何被标记的数据，因为LLaRA完全基于普通语料库进行。
我们使用Wikipedia语料库，应用LLaRA对LLaMA-2-7B进行适应，从而显著提高了LLM的检索质量。在进行了常见的微调操作后，适应良好的模型在各种评估场景中恢复了最先进的性能，例如在MSMARCO上的文章和文档检索，以及在BEIR基准上的零样本检索。
总之，我们在这项工作中做出了以下技术贡献。1） 我们提出了LLaRA，这是第一项将LLM应用于密集检索应用的研究工作。2） LLaRA设计简单但有效。通过对未标记数据执行EBAE和EBAR这两个前置任务，LLaRA大大提高了LLM的检索能力。3） LLM的预训练和微调将花费巨大的成本。为了便于未来在这一领域的研究，我们的模型和源代码将公开。
2. Related Works
密集检索是将查询和文档表示为同一潜在空间内的嵌入，其中可以基于嵌入相似性为查询检索相关文档。如今，它被广泛应用于许多重要的现实世界应用中，如网页搜索、问答和会话系统。密集检索的准确性由其嵌入的质量决定，其中骨干编码器是学习有判别性嵌入的决定性因素。在过去的几年里，预训练语言模型（PLM），如BERT、RoBERTa和T5，被广泛用于查询和文档的编码。得益于大规模的预训练和基于transformer的架构，PLM能够为输入文本生成细粒度的语义表示。此外，研究还发现，随着模型和训练规模的扩大，以及预训练方法的改进，可以进一步提高基于PLM的密集检索的准确性和通用性。
遵循同样的精神，利用LLM不断扩大骨干编码器是一个自然的举措。LLM的使用在许多方面都是有前景的。值得注意的是，考虑到LLM强大的语义理解能力，它可以大大有助于复杂查询和文档的建模。此外，考虑到LLM的上下文长度大大扩展，它为构建文档级检索器提供了直接的基础。由于LLM具有前所未有的通用性和指令跟随能力，它也有利于多任务嵌入模型的学习。最近，已经有几项工作对应用LLM作为密集检索的骨干编码器进行了初步的努力。然而，现有的方法只是直接使用LLM。由于文本生成和文本嵌入任务之间的巨大差异，LLM可能还有许多未被发掘的潜力。事实上，如何将LLM作为密集检索应用的更好的基础模型，还有待研究。
3. Methodology
3.1 初步的
密集检索利用文本嵌入模型来生成查询和文档的嵌入： e q e_q eq和 e d e_d ed。查询和文档的相关性通过它们的嵌入相似性来反映： < e q , e d > <e_q, e_d> <eq,ed>。因此，可以通过embedding空间内的ANN搜索来检索查询（ D q D_q Dq）的相关文档：
D q ← T o p − k ( { d : < e q , e d > ∣ D } ) D_q \gets Top-k(\{d:<e_q, e_d>|D\}) Dq←Top−k({d:<eq,ed>∣D})
预训练语言模型曾经是嵌入模型的骨干编码器。以BERT为例。输入文本被标记为序列T: [ C L S ] ， t 1 , … , t N , [ E O S ] [CLS]，t1, \dots, tN,[EOS] [CLS]，t1,…,tN,[EOS]。然后，通过BERT对标记化序列（tokenized sequence）进行编码，其中输出嵌入被集成为文本嵌入。执行集成有两个常用选项：[CLS]或平均池化：
e t ← B E R T ( T ) [ C L S ] e_t \gets BERT(T)[CLS] et←BERT(T)[CLS]
e t ← A V G ( B E R T ( T ) ) e_t \gets AVG(BERT(T)) et←AVG(BERT(T))
当使用大型语言模型（LLM）作为骨干编码器时，文本嵌入需要以不同的方式生成。考虑到现有的LLM主要采用仅解码器的架构，全局上下文只能通过输入序列末尾的token来获取。因此，特殊标记 ⟨ ∖ s ⟩ ⟨\setminus s⟩ ⟨∖s⟩或[EOS]的输出嵌入被用作文本嵌入。以LLaMA为例，我们有以下更新的文本嵌入形式：
e t ← L L a M A ( T ) [ < ∖ s > ] e_t \gets LLaMA(T)[<\setminus s>] et←LLaMA(T)[<∖s>]
3.2 LLaRA
尽管LLM中的最后一个token可以在给定仅解码器的架构的情况下处理整个上下文，但其输出嵌入并不是输入文本的合适表示。这是因为LLM是通过文本生成进行预训练的，其中每个token的嵌入用于预测其下一个token。换句话说，LLM的输出嵌入侧重于捕获局部和不久的将来（near-future）的语义，而不是全局上下文的语义。
- 目标： 为了解决上述问题，我们提出了LLaRA（适用于检索的LLM），用于LLM的面向检索的自适应。LLM的文本嵌入有望通过自适应过程实现两个特性。
- 文本嵌入需要表示全局上下文的语义。
- 全局上下文表示应该便于说明查询和文档之间的关联。
- 前置文本任务： 基于以上两个目标，我们为LLaRA引入了两个前置任务。第一个叫做EBAE（Embedding-based Auto-Encoding)，文本嵌入 e t e_t et 用于预测它自己的输入文本。特别地，如果原始输入文本可以由 e t e_t et 预测，那么输入文本的全局语义必须由 e t e_t et 完全编码。第二个是EBAR（基于嵌入的自回归），其中文本嵌入 e t e_t et 用于预测输入文本的下一个句子。已知相关文档是查询的合理的下一个句子，例如，问题和答案、对对话上下文的回应，可以通过对这种语义进行表示来建立查询和文档之间的关联。
- 文本嵌入： LLM由两个不同的模板提示生成EBAE和EBAR的文本嵌入（图1）。对于EBAE，LLM通过这个模版生成prompt：
"[Place-holder for input] The original sentence: ⟨\s⟩"
文本embedding的生成方式为：
e t α ← L L a M A ( T , S E L F , ⟨ ∖ s ⟩ ) [ − 1 ] e_t^{\alpha} \gets LLaMA(T, SELF, \left\langle \setminus s \right\rangle)[-1] etα←LLaMA(T,SELF,⟨∖s⟩)[−1]
其中，“SELF” 表示EBAE的prompt：“The original sentence: ⟨\s⟩”。对于EBAR，LLM的prompt模版为：
"[Placeholder for input] The next sentence: ⟨\s⟩"
基于此，文本嵌入被生成为：
e t β ← L L a M A ( T , N E X T , ⟨ ∖ s ⟩ ) [ − 1 ] e_t^{\beta} \gets LLaMA(T, NEXT, \left\langle \setminus s \right\rangle)[-1] etβ←LLaMA(T,NEXT,⟨∖s⟩)[−1]
其中"NEXT"表示EBAR的prompt：“The next sentence:”。
如果直接计算 e t α e_t^{\alpha} etα 和 e t β e_t^{\beta} etβ 将会造成很大的运算开销的浪费，因为输入文本 T T T 将会被处理两次。为了解决这个问题，我们提出一次性计算 e t α e_t^{\alpha} etα 和 e t β e_t^{\beta} etβ。特别地，EBAE和EBAR的prompts被合并到了LLM的一个prompt里面：
"[Placeholder for input] SELF ⟨\s⟩ NEXT ⟨\s⟩"
因为这两个文本嵌入需要被单独计算，我们修改了传统casual language modeling的注意力掩码，其中 “SELF
⟨
∖
s
⟩
\left\langle \setminus s \right\rangle
⟨∖s⟩” 和 “NEXT
⟨
∖
s
⟩
\left\langle \setminus s \right\rangle
⟨∖s⟩” 是相互不可见的（Figure2）。
现在，第一个和第二个 <\s> 的输出embedding分别被用作
e
t
α
e_t^{\alpha}
etα 和
e
t
β
e_t^{\beta}
etβ。考虑到输入文本T将占据联合提示（joint prompt）的大部分长度，与直接计算相比，这种处理将节省几乎50%的成本。
-
训练目标： 如前所述，LLaRA的文本嵌入是为了捕捉输入文本本身的全局语义和输入文本的下一个句子的语义。在这里，我们提出了一个简单但有效的训练目标，将文本嵌入转换为全局语义表示。理论上，我们认为，如果一个嵌入本身能够准确预测特定上下文的标记，那么该嵌入一定是相应上下文的全局语义的有力表示。
基于这一基本原理，文本嵌入的训练被公式化为一个多类分类问题，其中文本嵌入是线性投影的，用于预测目标上下文中的tokens。上述问题的目标函数推导为：
在这个地方，
W
∈
R
∣
V
∣
×
d
W \in R^{|V| \times d}
W∈R∣V∣×d 是线性投影矩阵，
V
V
V 是词汇空间。
T
\mathcal{T}
T代表输入文本本身或下一个句子的标记的集合，分别取决于
e
t
α
e_t^{\alpha}
etα 和
e
t
β
e_t^{\beta}
etβ的处理。事实证明，上述训练目标简单而有效。它可以很容易地在现有的语言建模训练管道之上实现。
4. Experimental Study
4.1 设置
进行实验研究是为了验证LLaRA的有效性，特别是它对微调后检索准确性的影响，以及它在不同场景下的通用性。为了实现这些目标，我们使用MS MARCO（Nguyen et al.，2016）作为我们的微调数据集，对段落检索和文档检索任务进行评估。为了评估模型的通用性，我们还利用了BEIR基准涵盖了多种检索场景，如问答、事实验证、实体检索、重复检测等。将MS MARCO的微调模型直接移植到BEIR的zero-shot评价中。
- 训练： LLaRA应用于LLaMA-2-7B（基础）模型。这是基于DPR策划的维基百科的未标记语料库进行的（Karpukhin等人，2020）。我们总共执行了10000个LLaRA自适应步骤，批量大小为256，序列长度为1024，学习率为1e-5。LLaRA按照RepLLaMA提出的过程进行微调：它利用LoRA进行LLM的参数有效训练，并简单地利用ANN hard negatives进行嵌入模型的对比学习。
4.2 分析
分别用Table 1、2和3给出了在MS MARCO上的文章和文献检索以及在BEIR基准上的零样本检索的评价结果。我们与各种各样的基线方法进行了比较，包括基于预训练的语言模型的具有代表性的密集检索器，例如 ANCE，RocketQA，GTR，RetroMAE ，SimLM以及传统的基于BM25的稀疏检索。我们还介绍了利用LLM作为骨干编码器的最新方法，包括CPT，SGPT和RepLLaMA。
主要观察结果如下。首先，LLaRA在每个评估场景中都实现了最高的重新评估性能。值得注意的是，LLaRA在MS MARCO 段落检索中MRR@10指标获得了43.1分，在文档检索中MRR@100指标获得了47.5分，且NDCG@10平均为55.1。这样的性能甚至高于交叉编码器的大多数re-ranking结果。此外与跟它最接近的baseline RepLLaMA相比较，LLaRA在MS MARCO段落检索上的MRR@10指标高出1.9%，在MS MARCO文档检索上的MRR@100指标也要高出1.9%，在BEIR上zero-shot检索的NDGC@10指标要高出1.0%。这种显著且一致的经验增益验证了LLM的文本嵌入能力由于LLaRA的自适应而得到了显著提高。
我们对每个特定场景都有以下观察结果。首先，MS MARCO段落检索（Table 1）曾经是信息检索中引用最广泛的基准。近几年，由于对预训练的语言模型和微调方法的重大改进，一系列具有竞争力的基线极大地提高了SOTA。随着LLM的使用，新提出的稠密检索又向前迈出了一大步。需要注意的是，LLaRA只是通过hard negative采样进行了微调。如果未来可以使用更先进的微调方法，那么报告的性能很可能会得到进一步改进。与基于BERT的模型（比如RetroMAE和SimLM）相比，骨干编码器的变化在MRR@10指标上带来了将近4%的增益。这种巨大的进步表明LLM在密集检索中的巨大潜力。
对MS MARCO的文档检索任务也进行了同样的观察（Table 2）。基于LLM的检索器带来了优越的经验性能，其中LLaRA的结果相较于以前基于BERT的方法在MRR@100上有了超过+5%的改进。事实上，考虑到LLM的文本长度显著扩展，例如LLaMA-2的4K，文档检索是使用基于LLM的作为骨干编码器的直接优势。
根据BEIR基准零样本评估结果，检索器的通用性是使用基于LLM的骨干编码器的另一个明显优势。此前，人们普遍观察到，密集型检索很难处理zero-shot场景，尽管它们在监督数据集中表现出了竞争性。对于BEIR基准中的许多评估任务，基于BERT的方法甚至比基于BM25的稀疏检索器差得多。然而，通过切换到基于LLM的骨干编码器，密集检索器的zero-shot性能可以得到很大的改善。值得注意的是，随着模型尺寸的大幅扩大，所有大型基线，如GTR-XXL、CPT-XL、Ada-002、SGPT、RepLLaMA，在大多数情况下都能够优于BM25。从侧面来看，与BERT基线相比，LLaRA在每个单独任务中都取得了更好的性能，最终NDCG@10指标平均提高了+16%。
5. Conclusion
在本文中，我们提出了LLaRA，这是一种新的方法，旨在通过提高LLM的文本嵌入能力，使其成为密集检索的更好基础。LLaRA由两个前置任务组成，即EBAE和EBAR。这两个任务协作将LLM的文本嵌入转换为全局上下文的表示，从而促进查询与其相关目标之间的语义匹配。作为对预训练良好的LLM的后验自适应，LLaRA不仅有效而且实现简单。它可以完全基于未标记的语料库进行，并与现有的语言建模训练管道完全兼容。基于全面的实验研究验证了LLaRA的有效性，其中基于LLM的检索器的准确性和通用性可以显著提高，从而在有监督和zero-shot评估场景中实现最先进的性能。
