论文阅读：Making Large Language Models A Better Foundation For Dense Retrieval
Abstract
在密集检索任务中，学习区分查询与文档之间的语义关联是必要的前提。尽管大型语言模型（LLM）在语义理解方面具有强大的能力，但其预训练任务侧重于文本生成，其工作模式与基于文本的嵌入表示不同。因此，必须研究如何调整LLM使其能够有效地初始化为密集检索的骨干编码器。在本文中，我们提出了一种名为LLaRA（LLM适配密集检索）的新方法。LLaRA作为一种后验自适应模型，旨在支持密集检索应用。该方法由两个自适应任务组成：基于嵌入的自动编码（EBAE）和基于嵌入的自回归（EBAR）。EBAE利用LLM的文本嵌入来重构输入句子的标记，而EBAR则利用这些嵌入来预测下一个句子的标记。通过实验结果表明，LLaRA是一种简单、高效且实用的方法。具体而言，LLaRA在应用于维基百科语料库上的LLM基础模型（LLaMA-2-7B）时，显著提升了模型在多种密集检索基准上的微调性能，包括MSMARCO和BEIR指标。我们的方法和实现代码已将在BGE存储库中公开。
1. Introduction
密集检索是一种由深度神经网络提出的新型检索范式。与传统信息检索方法不同的是，密集检索学习将查询和文档表示为同一潜在空间内的嵌入形式，其中查询与文档的语义关联可通过嵌入间的相似度来表征。如今，密集检索已成为多个现实应用场景中的关键技术基础，其中包括网页搜索和开放信息检索。
在密集检索系统中，其性能主要取决于关键编码器的容量。近年来，预训练语言模型的广泛应用为信息检索领域带来了显著的进展，其中包括BERT、RoBERTa、T5等模型。研究表明，模型规模和训练强度的扩大不仅显著提升了密集检索的准确率，还显著增强了其通用性。
近年来，大型语言模型（LLM）被广泛应用于多种经典的自然语言处理任务，作为统一的解决方案进行微调。鉴于LLM在语义理解方面的卓越能力，利用这种强大的模型进行密集检索具有巨大潜力。在此领域，已取得了一系列开创性成果，LLM通过被提示或微调以生成判别嵌入，从而显著提升了密集检索的效果。
尽管已取得一定进展，但全面释放LLM在密集检索应用中的潜力仍面临诸多挑战。值得注意的是，LLM通过预训练进行文本生成，并学习文本生成的嵌入以预测下一个token。其输出嵌入主要关注局部语义和未来语义的捕获。然而，密集检索需要嵌入表示整个上下文的全局语义。这种显著的差异将严重限制LLM在密集检索任务中的直接应用。
为了解决上文提及的问题，我们提出了一种新方法，命名为LLaRA（如图1所示）。该方法可被视为LLM的后验自适应机制，旨在增强其密集检索能力。LLaRA可被视为无监督生成预训练任务的一个扩展训练阶段。通过精心设计的文本预训练任务，该方法旨在提升LLM的能力，使其能够生成用于全局上下文语义表示的文本嵌入。
LLaRA方法引入了两个重要的前置训练任务，即EBAE（基于嵌入的自编码器）和EBAR（基于嵌入的自回归）。在EBAE任务中，LLM被引导生成文本嵌入，这些嵌入可以用于预测输入句子的token。而在EBAR任务中，LLM同样被引导生成嵌入，但这些嵌入被用于预测下一个句子的token。通过从这两个不同的提示模板中进行学习，LLM的文本嵌入能力得以从局部语义表示（即仅针对下一个标记的预测）被调整为全局语义表示（例如，对句子级别的特征进行预测）。这种区分使得LLM在处理不同的语义匹配场景时展现出不同的能力，例如在相似性搜索任务中（使用EBAE的提示），LLM能够基于嵌入进行高效的相似文本检索；而在问答任务中（使用EBAR的提示），LLM则能够通过嵌入理解并回答与上下文相关的复杂问题。
在LLaRA方法中，句子级特征的预测过程基于LLM输出嵌入的线性投影实现，其中无需额外的解码组件辅助。由此可知，LLaRA能够直接嵌入到现有的生成预训练管道中，并显著提升了训练效率。此外，LLaRA完全依赖于普通语料库，无需收集任何标记数据，这使其在实际应用中具有较高的灵活性和实用性。
我们基于Wikipedia语料库，通过LLaRA对LLaMA-2-7B进行适应性训练，从而显著提升了LLM的检索质量。经过常规微调后，适应性良好的模型在多种评估场景中恢复了最先进水平的性能，包括在MSMARCO上的文章和文档检索任务，以及在BEIR基准上的零样本检索任务中表现优异。
在本研究工作中，我们完成了以下技术贡献。第一项将LLM应用于密集检索应用的研究工作，我们提出LLaRA。LLaRA设计简洁且效能突出。通过执行EBAE和EBAR两个关键预处理任务，LLaRA显著提升了LLM的检索性能。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_d。查询与文档之间的相关性则通过它们的内积
D_q \gets Top-k(\{d:
预训练语言模型曾作为嵌入模型的关键编码器。以BERT为例说明，输入文本被标记为序列T，其中[CLS]，t1, \dots, tN，以及[EOS]。BERT对标记化序列进行编码，其中输出嵌入被集成为文本嵌入。在集成操作中，通常采用两种常见方式：一种是使用[CLS]$，另一种是进行平均池化。
e_t \gets BERT(T)[CLS]
e_t \gets AVG(BERT(T))
当大型语言模型（LLM）被用作骨干编码器时，文本嵌入的生成方式需采取不同的策略。由于现有的LLM主要依赖于仅解码器的架构，全局上下文的获取仅能通过输入序列末尾的token来实现。因此，为了获取文本嵌入，我们采用⟨\setminus s⟩或[EOS]作为标记。以LLaMA为例，这一过程具体而言，我们采用了一种更新的文本嵌入形式。
e_t \gets LLaMA(T)[<\setminus s>]
3.2 LLaRA
在仅解码器架构的LLM设计中，最后一个token能够捕获整个输入序列的语义信息，然而，其输出嵌入并不适合作为输入文本的表示方式，因为LLM主要通过生成任务进行训练阶段，其中每个token的嵌入被用于估计其后续token的出现概率。换句话说，LLM的输出嵌入更倾向于反映局部和近未来的语义特征，而非全局语义信息。
目标：为了解决上述问题，我们提出了LLaRA（适用于检索的LLM），用于LLM的自适应。文本嵌入需具备表示全局语义上下文的能力。全局语义表示需便于阐述查询与文档之间的关联。
- 前置文本任务： 基于以上两个目标，我们为LLaRA引入了两个前置任务。第一个叫做EBAE（Embedding-based Auto-Encoding)，文本嵌入 e_t 用于预测它自己的输入文本。特别地，如果原始输入文本可以由 e_t 预测，那么输入文本的全局语义必须由 e_t 完全编码。第二个是EBAR（基于嵌入的自回归），其中文本嵌入 e_t 用于预测输入文本的下一个句子。已知相关文档是查询的合理的下一个句子，例如，问题和答案、对对话上下文的回应，可以通过对这种语义进行表示来建立查询和文档之间的关联。
- 文本嵌入： LLM由两个不同的模板提示生成EBAE和EBAR的文本嵌入（图1）。对于EBAE，LLM通过这个模版生成prompt：
"[Place-holder for input] The original sentence: ⟨\s⟩"
AI助手文本embedding的生成方式为：
e_t^{\alpha} \gets LLaMA(T, SELF, \left\langle \setminus s \right\rangle)[-1]
其中，"SELF" 代表 EBAE 的 prompt，具体表述为：The original sentence: ⟨\s⟩。对于 EBAR，LLM 的 prompt 模版为：
"[Placeholder for input] The next sentence: ⟨\s⟩"
AI助手基于此，文本嵌入被生成为：
e_t^{\beta} \gets LLaMA(T, NEXT, \left\langle \setminus s \right\rangle)[-1]
其中"NEXT"表示EBAR的prompt：“The next sentence:”。
在实际应用中，单独计算e_t^{\alpha}和e_t^{\beta}将带来显著的计算开销，因为输入文本T会被处理两次。针对这一问题，我们提出同时计算e_t^{\alpha}和e_t^{\beta}。值得注意的是，通过整合EBAE和EBAR的prompts，我们可以将两个prompt整合到LLM的一个prompt中：
"[Placeholder for input] SELF ⟨\s⟩ NEXT ⟨\s⟩"
AI助手在对这两个文本嵌入进行单独计算的需求下，我们对传统casual language modeling中的注意力掩码进行了调整，其中，SELF \left\langle \setminus s \right\rangle和NEXT \left\langle \setminus s \right\rangle之间是不可见的（Figure2）。
在对这两个文本嵌入进行单独计算的需求下，我们对传统casual language modeling中的注意力掩码进行了调整，其中，SELF \left\langle \setminus s \right\rangle和NEXT \left\langle \setminus s \right\rangle之间是不可见的（Figure2）。
在当前阶段，第一和第二个 <\s> 的输出嵌入表示分别用于生成 e_t^{\alpha} 和 e_t^{\beta}。鉴于输入文本T占据了联合提示（joint prompt）的 majority portion，采用这种处理方式相比直接计算，将节省约50%的计算成本。
训练目标： 如前所述，LLaRA的文本嵌入是为了捕捉输入文本本身的全局语义和输入文本的下一个句子的语义。在这里，我们提出了一个简单但有效的训练目标，将文本嵌入转换为全局语义表示。理论上，我们认为，如果一个嵌入本身能够准确预测特定上下文的标记，那么该嵌入一定是相应上下文的全局语义的有力表示。
基于这一理论，文本嵌入的训练被建模为一个多标签分类问题，其中，文本嵌入被设计为通过线性变换识别目标上下文中的tokens。针对该问题，目标函数的推导结果为：
在该处，线性投影矩阵W \in R^{|V| \times d}代表词汇空间V，而标记集合\mathcal{T}则根据e_t^{\alpha}和e_t^{\beta}的处理方式来确定，具体取决于模型的处理策略。研究表明，这一训练目标简洁且具有显著效果。该方法可以很容易地在现有的语言建模训练管道之上进行部署。
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等，同时也包含了传统BM_25框架下的稀疏检索方法。此外，我们还探讨了以LLM为骨干编码器的最新检索技术，如CPT、SGPT和RepLLaMA等。
主要观察结果表明，LLaRA在各个评估场景中均展现了卓越的重新评估性能。具体而言，在MS MARCO段落检索任务中，MRR@10的值达到了43.1%，而在文档检索任务中，MRR@100的值则为47.5%。此外，NDCG@10的平均值为55.1%。这些指标显示，LLaRA的性能不仅超越了交叉编码器的大多数重新排序结果，而且在多个评估维度上均表现优异。与最接近的基准模型RepLLaMA相比，LLaRA在MS MARCO段落检索任务中，MRR@10提升了1.9个百分点；在文档检索任务中，MRR@100的提升幅度同样达到了1.9个百分点；而在BEIR平台上的zero-shot检索任务中，NDCG@10的提升幅度为1.0个百分点。这些显著且一致的性能提升进一步验证了LLM在文本嵌入能力方面得益于LLaRA的自适应机制的优化。
在各个特定场景下，我们进行了深入观察。首先，MS MARCO段落检索（见表1）曾是信息检索领域应用最广泛的基准。近年来，预训练语言模型和微调方法的显著改进推动了一系列具有竞争力的基准模型，显著提升了SOTA水平。值得注意的是，LLaRA仅通过硬负样本采样进行微调。随着LLM技术的引入，稀疏检索方法已取得了显著进展。与基于BERT的模型（如RetroMAE和SimLM）相比，骨干编码器的改进在MRR@10指标上带来了约4%的显著提升。这种改进表明LLM在密集检索领域展现出巨大潜力。
在对MS MARCO的文档检索任务进行观察时，参考了Table 2的数据。LLM驱动的检索器展现出卓越的经验性能，其中LLaRA在MRR@100指标上较基于BERT的方法实现了超过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。从性能对比来看，LLaMA在各个单独任务中均表现优异，最终的NDCG@10指标较BERT基线提升了16%以上。
5. Conclusion
在本文中，我们提出了一种名为LLaRA的方法，旨在通过提升LLM在文本嵌入方面的性能，为其构建更高效密集检索的基础。LLaRA由两个关键组件构成，即EBAE和EBAR。这两个组件协同作用，将LLM的文本嵌入转换为全局上下文表示，从而实现查询与目标语义匹配的优化。作为对经过良好预训练的LLM的后验自适应方法，LLaRA在效果上表现出色且实现过程简便。它完全基于未标注的语料库构建，并与现有的语言建模训练流程无缝衔接。基于全面的实验研究，LLaRA的有效性得到了充分验证，基于LLM构建的检索器在准确性和通用性方面均有显著提升，从而在监督学习和zero-shot评估场景中达到了当前最先进的性能。