研究背景
- 研究问题:这篇文章要解决的问题是如何在信息检索(IR)中使用测试时计算(test-time compute)来改进重排模型的性能。具体来说,作者提出了RANK1,第一个利用测试时计算的重排模型,以快速提高较小模型的性能。
- 研究难点:该问题的研究难点包括:如何在没有指令数据的情况下,从较大的推理语言模型(如OpenAI的o1和Deepseek的R1)中蒸馏出较小的模型;如何使模型能够对用户输入的提示做出响应;如何在保持高性能的同时减少计算和内存使用。
- 相关工作:该问题的研究相关工作包括推理语言模型在数学、逻辑和编程任务中的改进(如OpenAI的o1模型),以及信息检索系统中复杂任务的处理(如指令遵循和推理能力的提升)。
研究方法
这篇论文提出了RANK1模型,用于解决信息检索中的重排问题。具体来说,
- 数据收集:首先,作者收集了超过600,000个R1推理轨迹的例子,这些例子来自MS MARCO查询和段落。
- 数据准备:然后,作者使用kluster.ai API服务从R1中生成数据,包括MS MARCO的正例、Tevatron2的负例、mT5-13B的rank 1-5硬负例和rank 5-10硬负例。经过多次生成和过滤,最终得到了386,336个高质量的训练样本。
- 模型训练:作者从Qwen 2.5系列模型中选择三个主要模型(7B、14B和32B参数模型)进行训练。使用LoRA和LLaMA-Factory进行微调,最多训练两个epoch或三天。
- 推理链生成:在推理过程中,RANK1模型会生成推理链,这些链可以帮助模型在给出最终答案之前进行“思考”。
实验设计
- 数据集:实验使用了MS MARCO数据集,并从中收集了635,264个R1推理轨迹的例子。这些例子被标记为相关或不相关。
- 基准模型:实验使用了BM25、mE5-base、MonoT5-3B、RankLLaMA-7B和RankLLaMA-13B等基准模型。
- 评估指标:实验使用了多个评估指标,包括nDCG@10、p-MRR等,用于评估模型在高级推理、指令遵循和语义理解任务上的性能。
- 实验设置:实验使用了mteb工具运行所有实验,除了DL19使用rankllm工具。推理由vllm驱动,显著提高了推理速度。
结果与分析
- 高级推理能力:在BRIGHT基准上,RANK1模型在所有类别中均表现出色,特别是在生物学和Stackoverflow类别中,nDCG@10得分分别达到48.8和37.8,接近两倍于其他模型。
- 语义理解:在NevIR基准上,RANK1模型在多个语言子集上的表现也非常出色,nDCG@20得分在Persian和Chinese子集上分别达到0.611和0.678,显著高于其他模型。
- 指令遵循:在mFollowIR数据集上,RANK1模型在跨语言和多语言设置中均表现出色,nDCG@20得分在Persian和Chinese子集上分别达到0.610和0.678,显著高于其他模型。
- 传统基准:在DL19和BEIR基准上,RANK1模型表现出色,尽管在这些基准上存在大量未判断的文档,但经过重新标注后,RANK1模型的性能仍然优于其他模型。
总体结论
这篇论文提出了RANK1,第一个利用测试时计算的重排模型。通过从R1中蒸馏推理链,RANK1模型在高级推理、指令遵循和语义理解任务上表现出色。尽管训练数据有限且仅使用英语数据,RANK1模型仍展示了强大的多语言推理能力。总体而言,RANK1引入了一种新的重排模型类别,使得更复杂的检索任务成为可能。
论文评价
优点与创新
- 首次引入测试时计算的重新排序模型:RANK1是第一个利用测试时计算进行重新排序的模型,展示了在信息检索中使用推理语言模型(如OpenAI的o1、Deepseek的R1等)进行蒸馏以快速提高较小模型性能的适用性。
- 数据集开放:收集了超过600,000个R1推理轨迹的示例,并开源这些数据集,促进了研究的复现和合作。
- 多语言和多任务能力:尽管仅使用英语MS MARCO数据进行训练,RANK1在多语言设置中仍表现出强大的推理和指令跟随能力。
- 可解释的推理链:RANK1生成的推理链可以向用户提供或用于RAG系统,增强了模型的可解释性和透明度。
- 量化模型的性能:量化版本的RANK1模型在保持强大性能的同时,减少了计算和内存使用。
- 广泛的应用场景:RANK1展示了测试时计算可以为搜索领域带来新的可解释且高性能的重新排序模型类型。
不足与反思
- 过度思考:RANK1可能会过度思考,特别是在交互式使用时,模型已经知道答案并在寻找非常具体的短语。通过特定数据集的训练来校准这种过度思考可能会减少这种情况。
- 推理速度:作为使用测试时计算的推理模型,RANK1比仅有分类头的模型(如RankLLaMA)慢。尽管现代分页注意力库(如vLLM)和量化技术可以缓解这一问题,但额外的计算使用是不可避免的。
- 未来工作:未来的工作包括使用强化学习进行微调、开发列表推理重新排序器以及训练多语言和指令微调的版本,以进一步提高在这些任务上的性能。
关键问题及回答
问题1:RANK1模型在数据收集和准备方面有哪些具体的步骤?
- 数据收集:作者从MS MARCO数据集中收集了超过600,000个R1推理轨迹的例子,这些例子包括查询和段落的推理过程。
- 数据生成:使用kluster.ai API服务从R1中生成数据,具体包括以下几类样本:
- MS MARCO的正例
- Tevatron2的负例(从BM25和CoCondenser中采样)
- mT5-13B的rank 1-5硬负例
- mT5-13B的rank 5-10硬负例
- 数据过滤:经过多次生成和过滤,最终得到了386,336个高质量的训练样本。过滤步骤包括:
- 使用标签最确定的数据(如MS MARCO的正例和Tevatron的负例)进行初步平衡
- 使用一个模型对数据进行自过滤,去除约10%的噪声数据
- 最终保留了所有MS MARCO的正例和自过滤后的负例
问题2:RANK1模型在实验中使用了哪些评估指标,这些指标分别反映了模型的哪些方面?
- nDCG@10:这是一个常用的排序评估指标,用于衡量模型在排序结果中的准确性。较高的nDCG@10得分表示模型能够更准确地将相关文档排在前面。
- p-MRR:这是Precision@10的归一化版本,用于衡量模型在前10个结果中的精度。较高的p-MRR得分表示模型在前10个结果中有更多的正确答案。
- nDCG:全称为Normalized Discounted Cumulative Gain,用于衡量模型在所有文档中的排序质量。它考虑了文档的排名和相关性。
- p-MRR:这是Precision@10的归一化版本,用于衡量模型在前10个结果中的精度。较高的p-MRR得分表示模型在前10个结果中有更多的正确答案。
这些指标反映了RANK1模型在高级推理、指令遵循和语义理解任务上的性能,以及在传统信息检索基准上的表现。
问题3:RANK1模型在传统基准(如DL19和BEIR)上的表现如何,为什么这些基准可能存在问题?
- DL19:在DL19基准上,RANK1模型表现出色,尽管发现这些基准可能存在模型饱和问题。具体表现为RANK1模型找到的未判断文档比例较高(364%以上),这导致其性能受到影响。作者建议重新评估旧的TREC集合,因为这些基准可能无法有效区分最佳重排器。
- BEIR:在BEIR基准上,RANK1模型的表现也较为出色,但在某些数据集上存在噪声和不准确性。例如,SciFact数据集中的NEI查询被错误地包含在内,导致噪声增加。此外,部分数据集(如FiQA2018和DBPedia)的部分相关标签也会增加噪声。这些问题使得这些传统基准在区分最佳重排器方面可能不够有效。
Links
https://github.com/orionw/rank1
Models
Resource | Description |
---|---|
rank1-7b | Trained from Qwen2.5-7B base |
rank1-14b | Trained from Qwen2.5-14B base |
rank1-32b | Trained from Qwen2.5-32B base |
rank1-mistral-2501-24b | Trained from Mistral-Small 2501 24B base |
rank1-llama3-8b | Trained from Llama 3.1 8B base |
Quantized Models (fits in 24GB GPUs)
Resource | Description |
---|---|
rank1-7b-awq | Quantized version of rank1-7b |
rank1-14b-awq | Quantized version of rank1-14b |
rank1-32b-awq | Quantized version of rank1-32b |
rank1-mistral-2501-24b-awq | Quantized version of rank1-mistral-24b |
rank1-llama3-8b-awq | Quantized version of rank1-llama3-8b |
Datasets
Resource | Description |
---|---|
rank1-r1-msmarco | All R1 output examples from MS MARCO |
rank1-training-data | Training data used for rank1 models |
rank1-run-files | Pre-computed run files for use in top 100 doc reranking |
Using Rank1 in Your Own Code
You can integrate rank1 into your code:
from rank1 import rank1
# Initialize the model
model = rank1(
model_name_or_path="jhu-clsp/rank1-7B",
num_gpus=1,
device="cuda",
context_size=16000,
max_output_tokens=8192,
fp_options="float16"
)
# Rerank documents
results = model.predict({
"query": ["Your query/prompt here", "Same number as docs"],
"corpus": ["Document 1 content", "Document 2 content", ...],
})