[论文解读] Fine-Tuning LLaMA for Multi-Stage Text Retrieval

Xueguang Ma, Liang Wang+3 more|arXiv (Cornell University)|Oct 12, 2023

Topic Modeling被引用 7

一句话总结

该论文提出微调 LLaMA-2 大型语言模型作为密集检索器（RepLLaMA）和点对点重排序模型（RankLLaMA），用于多阶段文本检索，在 MS MARCO 和 BEIR 基准测试中达到最先进性能。该方法利用大语言模型的长上下文能力，整体表示整个文档，消除了传统分割和池化操作，并展示了使用开源模型检查点的强零样本有效性。

ABSTRACT

The effectiveness of multi-stage text retrieval has been solidly demonstrated since before the era of pre-trained language models. However, most existing studies utilize models that predate recent advances in large language models (LLMs). This study seeks to explore potential improvements that state-of-the-art LLMs can bring. We conduct a comprehensive study, fine-tuning the latest LLaMA model both as a dense retriever (RepLLaMA) and as a pointwise reranker (RankLLaMA) for both passage retrieval and document retrieval using the MS MARCO datasets. Our findings demonstrate that the effectiveness of large language models indeed surpasses that of smaller models. Additionally, since LLMs can inherently handle longer contexts, they can represent entire documents holistically, obviating the need for traditional segmenting and pooling strategies. Furthermore, evaluations on BEIR demonstrate that our RepLLaMA-RankLLaMA pipeline exhibits strong zero-shot effectiveness. Model checkpoints from this study are available on HuggingFace.

研究动机与目标

• 探究像 LLaMA-2 这类最先进大语言模型（LLMs）是否能在多阶段文本检索中超越较小的预训练模型。
• 通过微调大语言模型作为端到端检索器和重排序器，解决基于提示的 LLM 重排序方法存在的非并行推理和缺乏标注数据利用等问题。
• 探索大语言模型凭借其长上下文能力整体表示完整文档的潜力，从而消除对文档分割和池化策略的需求。
• 建立一个完全开源、优化的多阶段检索流水线，以 LLaMA-2 作为骨干模型，提升检索效果和推理效率。
• 评估所提出的 RepLLaMA–RankLLaMA 流水线在多样化检索基准上的零样本泛化能力。

提出的方法

• 使用对比损失微调 LLaMA-2 作为双编码器密集检索器（RepLLaMA），以在密集向量空间中学习查询-文档相似性。
• 使用标注的相关性判断微调 LLaMA-2 作为点对点重排序器（RankLLaMA），直接预测查询-文档对的相关性得分。
• 使用 MS MARCO 数据集在段落和文档检索任务上进行域内微调，并采用标准评估指标。
• 利用 LLaMA-2 的长上下文注意力机制对整个文档进行编码，无需分割，实现整体文档表示。
• 采用标准监督微调方案分别训练两个模型，损失函数分别为对比损失和点对点排序损失。
• 在 HuggingFace 上发布模型检查点，以实现开源访问和可复现性。

Figure 1: Cumulative distribution function of document lengths in the MS MARCO document corpus, showing the proportion of documents that has a length less than a specific value (determined by the LLaMA tokenizer). For clarity, we exclude 3% of documents with a length exceeding 10,000 tokens.

实验结果

研究问题

• RQ1 微调后的 LLaMA-2 模型是否能在多阶段文本检索中超越较小的预训练模型？
• RQ2 当使用标注数据进行微调时，LLaMA-2 是否能作为有效的密集检索器和点对点重排序器？
• RQ3 LLaMA-2 的长上下文能力是否消除了文档检索中对文档分割和池化策略的需求？
• RQ4 RepLLaMA–RankLLaMA 流水线在多样化 BEIR 基准上的零样本检索中效果如何？
• RQ5 与基于提示的生成方法相比，大语言模型的端到端微调是否能提升检索效果？

主要发现

• RepLLaMA 和 RankLLaMA 在使用 MS MARCO 数据集的段落和文档检索任务中均达到最先进性能，超越了以往较小的模型。
• RepLLaMA–RankLLaMA 流水线在 BEIR 基准上展现出强大的零样本有效性，表明其在多样化领域中具有稳健的泛化能力。
• LLaMA-2 的长上下文注意力机制支持整体文档编码，消除了传统文档分割和池化策略在文档检索中的需求。
• 将 LLaMA-2 微调为检索器和重排序器可实现高效、并行的打分，克服了基于提示的 LLM 重排序中非并行、多轮解码的局限性。
• 该方法有效利用了标注数据（例如来自 MS MARCO 的数据），性能优于忽略此类标注信息的零样本提示方法。
• 在 HuggingFace 上开源的模型检查点支持可复现性，并推动基于大语言模型的检索系统进一步研究。

Figure 2: Comparison of document ranking MRR@100 scores for RankLLaMA trained with different maximum input lengths and evaluated using different maximum input lengths. Each line represents a model trained with a specific maximum length, while points along the line indicate the effectiveness when var