提出BlockRank方法,让大模型高效处理长文档列表的排序任务。
Scalable In-context Ranking with Generative Models
- 通过分块稀疏注意力机制,将复杂度从二次方降至线性。
- 在MSMarco等数据集上达到顶尖性能,100文档推理快4.7倍。
- 适合需要快速处理长上下文排序的检索系统应用。
上下文排序(ICR)是信息检索的新范式,利用大语言模型的上下文理解能力,将查询、候选文档与任务描述一同输入模型以识别相关文档。然而随着候选文档增多,注意力计算量呈二次或超线性增长,效率成为瓶颈。本文首次发现微调后的ICR模型中存在两类可利用结构:(1)文档块间注意力稀疏,即同一文档内注意力密集,跨文档则稀疏;(2)查询到文档块的注意力分数在中间层与真实相关性高度相关。基于此,提出BlockRank方法,通过架构强制实现文档块间稀疏性,使注意力复杂度由二次降至线性,且在微调中引入辅助对比学习目标优化相关性预测。在BEIR、MSMarco和NQ数据集上使用Mistral-7B实验表明,BlockRank在保持或超越现有最先进列表排序器表现的同时,推理效率显著提升(100文档时快4.7倍),并能以秒级延迟支持约500文档(约10万上下文长度)的长列表排序,为ICR提供了一种高效可扩展的解决方案。
原文摘要 · Abstract (English)
In-context Ranking (ICR) is an emerging paradigm for Information Retrieval (IR), which leverages contextual understanding of LLMs by directly incorporating the task description, candidate documents, and the query into the model's input prompt and tasking the LLM to identify relevant document(s). While it is effective, efficiency is a significant challenge in this paradigm, especially as the candidate list grows due to quadratic/super-linear scaling of attention operation with context length. To this end, this paper first identifies inherent and exploitable structures in the attention of LLMs finetuned for ICR: (1) inter-document block sparsity: attention is dense within each document block but sparse across different documents in the context; and (2) query-document block relevance: the attention scores from certain query tokens to a document block in middle layers strongly correlate with that document's actual relevance. Motivated by these observations, we introduce BlockRank (Blockwise In-context Ranking), a novel method that adapts the attention operation in an LLM by (a) architecturally enforcing the observed inter-document block sparsity, reducing attention complexity from quadratic to linear without loss in performance, and (b) optimizing query-document block relevance for true relevant documents during fine-tuning using an auxiliary contrastive training objective, improving retrieval in attention. Experiments on BEIR, MSMarco and NQ with Mistral-7B demonstrate that BlockRank Mistral matches or outperforms existing SOTA listwise rankers and controlled fine-tuned baseline while being significantly more efficient at inference (4.7x for 100 MSMarco documents in context) and scaling gracefully to long-context shortlists, around 500 documents in-context (approximately 100K context length) within a second, presenting a scalable and effective solution for ICR.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。