百万令牌上下文检索：语言模型能否真正召回文档
速览
本文对大规模上下文检索进行了全面探索，首次在百万令牌级语料库和远超训练长度的长度泛化任务上进行了系统性分析。研究者开发了BlockSearch检索模型，通过注意力机制优化实现了长度泛化，显著优于现有基线模型。结果表明，大型语言模型在多样化任务中展现出强大的召回能力，为构建高效的AI检索系统开辟了新路径，强调了在极端上下文增长下控制注意力稀疏性的关键挑战。
AI 深度解读
背景
随着大语言模型（LLMs）的快速发展，它们在自然语言处理任务中展现出强大的能力，其中一种新兴研究方向是将检索作为一种原上下文（in-context）方式来实现：即直接让模型在输入的文档集合中进行检索，而非依赖向量数据库或传统检索系统。这类方法通过将相关文档与查询一起作为上下文提供给模型，让模型自行生成答案。然而，现有研究大多集中在专有系统或小型的重新排序任务上，对于百万级文档规模的原上下文检索仍缺乏系统性探索。尤其在上下文长度远超训练阶段的情况时，性能表现如何仍未充分研究。
核心内容
该论文提出首个针对百万级文档规模和长度泛化能力的原上下文检索系统研究。作者首先引入了 BlockSearch，这一 0.6B 参数的语言模型检索器。通过对架构和训练进行修改，它在原有的语言模型基线之上取得了改进，并在长度泛化方面能够达到训练阶段大小的 10 倍。然而，当上下文扩展到更极端程度时，检索性能仍会发生崩溃。论文通过分析指出，这一失败源于注意力稀释效应：随着文档集合增大，不相关文档会主导 softmax 分母，导致即使金标准文档的预 softmax 分数保持较高，归一化后的注意力质量也会下降。
受此分析驱动，论文提出了两种改进方法：一是长度感知的注意力 softmax 调整，二是文档级稀疏注意力机制。在百万级文档规模下，改进后的模型在广泛研究的基准任务上（如 MS MARCO 和 NQ）与密集检索（dense retrieval）性能相当，同时在参数量上仅为并发工作的 MSA 模型的 1/7，却仍优于该模型。此外，在需要全新相似度概念的任务（如 LIMIT）上，该模型显著超越密集检索，获得 3 倍更高的分数。这些结果表明，原上下文检索作为经典检索方法的替代方案具有潜力，同时凸显在极端上下文增长下进行注意力控制成为一项新的挑战。
关键要点
- 首次系统研究百万级文档规模下的原上下文检索，涵盖长度泛化远超训练大小的情况。
- 提出 BlockSearch：0.6B 参数模型，通过架构与训练修改实现改进，并能泛化至训练阶段长度的 10 倍。
- 检索性能在极端外推时崩溃，主要原因是注意力稀释效应，不相关文档主导 softmax 分母导致金标准质量下降。
- 提出长度感知注意力 softmax 调整和文档级稀疏注意力机制。
- 改进模型在 MS MARCO 和 NQ 等基准上与密集检索相当，同时在 LIMIT 任务上获得 3 倍更高分数。
意义与影响
该研究为大语言模型在原上下文下的检索能力提供实证基础，推动传统检索范式向模型原生支持的方向演进。在百万级文档规模和长度泛化挑战上取得突破，意味着未来搜索系统可以更高效地整合大量文档，而无需依赖外部向量索引。这不仅降低了开发门槛，也为处理超长上下文提供了新路径。然而，注意力控制在极端增长下的难题仍需进一步解决，为后续研究指明了方向。该工作强调了模型在复杂环境中维持检索精度的能力，有望影响多模态搜索和长文档分析等应用场景的未来发展。