让大模型在超长文本中精准检索,无需额外训练或工具。
Infinite Retrieval: Attention Enhanced LLMs in Long-Context Processing
- 利用模型自身注意力分布实现无限长度文本的精准检索
- 0.5B模型在100万词长文本中达100%准确率,超越更大模型
- 适配任意Transformer模型,显著降低推理延迟和计算开销
受限于大语言模型(LLM)的上下文窗口大小,处理输入超过上限的各类任务(如简单直接检索或复杂多跳推理)始终是挑战。尽管已有多种方法增强长文本处理能力,但或需高昂后训练成本,或依赖额外工具(如RAG),或在真实任务中未体现显著提升。本文观察到各层注意力分布与生成答案的相关性,实验验证注意力分配与检索增强能力一致。基于此,提出新方法InfiniRetri,利用模型自身注意力信息实现任意长度输入下的精准检索。评估显示,使用0.5B参数模型在100万词长的Needle-In-a-Haystack(NIH)测试中达到100%准确率,优于其他方法及更大模型,刷新当前最佳表现(SOTA)。同时,在真实世界基准上最大实现288%性能提升。InfiniRetri可无缝应用于任意Transformer-based LLM,无需额外训练,大幅降低长文本推理延迟与计算开销。研究表明,该方法具有实际应用潜力,为利用模型自身能力实现无限长度检索提供了新范式。代码将公开。
原文摘要 · Abstract (English)
Limited by the context window size of Large Language Models(LLMs), handling various tasks with input tokens exceeding the upper limit has been challenging, whether it is a simple direct retrieval task or a complex multi-hop reasoning task. Although various methods have been proposed to enhance the long-context processing capabilities of LLMs, they either incur substantial post-training costs, or require additional tool modules(e.g.,RAG), or have not shown significant improvement in realistic tasks. Our work observes the correlation between the attention distribution and generated answers across each layer, and establishes the attention allocation aligns with retrieval-augmented capabilities through experiments. Drawing on the above insights, we propose a novel method InfiniRetri that leverages the LLMs's own attention information to enable accurate retrieval across inputs of infinitely length. Our evaluations indicate that InfiniRetri achieves 100% accuracy in the Needle-In-a-Haystack(NIH) test over 1M tokens using a 0.5B parameter model, surpassing other method or larger models and setting a new state-of-the-art(SOTA). Moreover, our method achieves significant performance improvements on real-world benchmarks, with a maximum 288% improvement. In addition, InfiniRetri can be applied to any Transformer-based LLMs without additional training and substantially reduces inference latency and compute overhead in long texts. In summary, our comprehensive studies show InfiniRetri's potential for practical applications and creates a paradigm for retrievaling information using LLMs own capabilities under infinite-length tokens. Code will be released in link.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。