让大模型实时接入大规模知识库，推理更快更省显存。
SR-KI: Scalable and Real-Time Knowledge Integration into LLMs via Supervised Attention
- 用预训练编码器将知识库转为键值对，注入模型缓存中。
- 在模型内完成检索，支持单卡运行4万条知识，召回率超98%。
- 适合需要频繁更新知识的大模型应用，压缩率高达99.75%。
本文提出SR-KI，一种将实时、大规模结构化知识库（KB）高效集成到大语言模型（LLM）中的新方法。SR-KI首先使用预训练编码器将知识库编码为键值对，并注入到LLM的键值缓存中。在此基础上，采用两阶段训练：先定位模型内的专用检索层，再在该层施加基于注意力的损失函数，显式监督模型关注相关知识条目。与依赖外部检索器和多阶段流水线的传统检索增强生成方法不同，SR-KI通过模型隐空间内完成全部检索，实现端到端推理。该设计可有效压缩注入知识并支持动态更新。大量实验表明，SR-KI可在单张A100 40GB GPU上将最多4万条知识库注入7B参数量的LLM，最佳任务上保持超过98%的Recall@10，所有任务平均超过88%。问答与知识库条目生成任务表现优异，同时实现高达99.75%的知识压缩率。
原文摘要 · Abstract (English)
This paper proposes SR-KI, a novel approach for integrating real-time and large-scale structured knowledge bases (KBs) into large language models (LLMs). SR-KI begins by encoding KBs into key-value pairs using a pretrained encoder, and injects them into LLMs' KV cache. Building on this representation, we employ a two-stage training paradigm: first locating a dedicated retrieval layer within the LLM, and then applying an attention-based loss at this layer to explicitly supervise attention toward relevant KB entries. Unlike traditional retrieval-augmented generation methods that rely heavily on the performance of external retrievers and multi-stage pipelines, SR-KI supports end-to-end inference by performing retrieval entirely within the models latent space. This design enables efficient compression of injected knowledge and facilitates dynamic knowledge updates. Comprehensive experiments demonstrate that SR-KI enables the integration of up to 40K KBs into a 7B LLM on a single A100 40GB GPU, and achieves strong retrieval performance, maintaining over 98% Recall@10 on the best-performing task and exceeding 88% on average across all tasks. Task performance on question answering and KB ID generation also demonstrates that SR-KI maintains strong performance while achieving up to 99.75% compression of the injected KBs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。