UNREAL:用单一模型统一检索与长上下文 UNREAL: Unifying Retrieval and Long-Context with a Single Model
冻结LLM内部读出检索信号,仅训不到50万参数,同时胜任语料检索与长上下文证据选择
前置知识
RAG(检索增强生成)
先用检索系统(BM25、BGE、Qwen3-Embedding 等稠密模型,常再配 cross-encoder reranker)从语料中选出与问题相关的段落,再把这些段落与问题拼进 LLM 的 prompt 生成答案。检索器与生成器是两个独立训练、独立部署的模型,生产上需要保持语义对齐。
UNREAL 要替代的正是这套外部流水线;理解 RAG 的组成才能明白“单模型内置检索”省掉了哪些训练、部署与对齐成本。
长上下文推理与干扰项(distractor)
把所有可能相关的文本全部放进上下文窗口,让模型靠注意力自己定位信息。缺点是无关内容会干扰生成(已有研究证明 LLM 极易被 distractor 带偏),且标准 softmax 注意力的预填充开销随长度二次增长,扩大窗口还需专门的长上下文训练与效率机制。
论文把长上下文任务重新定义为“检索问题”,其动机正是全上下文精度随长度崩塌——NoLiMa 在 128K token 处全上下文准确率仅 1.0%。
残差流(residual stream)
Transformer 每层输出以加法方式写入一条贯穿全网的向量通道,各层注意力与 FFN 都在读写它;中间层的残差状态常被认为比最后一层蕴含更丰富的语义信息。
UNREAL 的块编码与检索 query 都直接从残差流读出($k_i=\mathrm{LLM}_{\ell_c}(c_i)$、$q^\ell$),这是它不训练主干、且跨注意力/SSM 架构通用的关键。
MaxSim 后期交互打分
ColBERT 提出的多向量打分:query 与文档都保留 token 级向量,对 query 每个向量取文档侧最大相似度再求和,$\mathrm{MaxSim}(u,v)=\sum_a\max_b\langle u_a,v_b\rangle$。比单向量点积更细粒度,代价是存储与计算更大。
UNREAL 的打分函数就是 MaxSim,只是多向量来自冻结 LLM 的内部 states,并新增了跨层 $\alpha$ 加权这一自由度。
InfoNCE 对比学习
把正样本对得分拉高、负样本对压低的训练目标 $\mathcal{L}=-\log\frac{e^{s_+/\tau}}{\sum_i e^{s_i/\tau}}$,多正例版本对所有正例求平均。是稠密检索模型的标准训练法,通常配合 hard negatives 提升判别力。
UNREAL 仅用这一个损失训练全部可学习参数(检索 token、层混合权重、soft prompt),是全文唯一的训练信号来源。
decoder-only 架构与先行工作 INTRA
当今主流 LLM 只有带因果掩码的解码器,没有独立 encoder 与 cross-attention。先行工作 INTRA(Hoffer et al., 2026)证明 encoder–decoder 模型能从自身编码表征中做内在检索,但其机制依赖 encoder 表征与 cross-attention,无法直接搬到 decoder-only 模型。
UNREAL 是把“内在检索”推广到 decoder-only 的尝试;理解该背景才能看清其核心动作——用残差流同时替代 encoder 表征和 cross-attention query。
线性注意力与状态空间模型(Mamba)混合架构
用线性注意力或选择性 SSM(如 Mamba)替代部分 softmax 注意力层,把长序列的二次成本降为线性;现代长上下文模型常将少量全局注意力层与这些层交错混合。
UNREAL 声称机制架构无关:四个测试骨干覆盖 dense、线性注意力混合、Mamba–注意力混合三类,验证完全依赖“残差流读出”不绑定具体序列混合器这一性质。
complete-evidence recall@n
严格检索指标:只有当标注的全部证据 chunk 都出现在 top-$n$ 召回结果中,该样例才算检索成功。对多跳问题尤其苛刻——缺任何一块证据即计失败。
论文主实验全部采用该口径:文中 recall 数字(如 HotpotQA 49.1%→73.2%)都是“证据全中”的召回率,误当成普通 recall 会误判基线水平。
TTFT 与 FLOPs break-even
TTFT(time-to-first-token)是从请求发出到生成第一个 token 的延迟,由预填充主导;FLOPs break-even 指新方法与基线计算量相等时的上下文长度,超过它新方法开始省算力。
效率部分的核心结论即:UNREAL 的 FLOPs 在约 1.8 万 token、实测 TTFT 在约 3.2 万 token 之后优于全上下文推理,且优势随长度扩大。
研究动机
当代 LLM 处理的上下文里,与问题真正相关的证据往往极其稀疏,证据选择因此成为核心环节,但现有技术把它拆成两套系统。长上下文一侧是隐式选择:模型把整个 prompt 塞进注意力自己“扛”,既容易被无关内容干扰(Shi et al. 2023、NoLiMa 等均证明 distractor 显著拉低精度),又受制于标准注意力随长度二次增长的预填充开销,扩窗还需要专门的长上下文训练与效率机制;论文自建的 8K–100M token 基准显示全上下文 EM 随长度急剧崩塌,NoLiMa 在 128K token 处准确率仅剩 1.0%。语料检索一侧是显式外置:独立 retriever(BGE、Qwen3-Embedding 等)加 reranker 先选后读,生产上要训练、部署、维护多套模型并保持语义对齐。即便如此,在 30 亿 token、2100 万 chunk 的 Wikipedia-2018 全库上,“全部证据都被召回才算对”的 complete-evidence recall@10 也只有 HotpotQA 49.1%、2WikiMultiHopQA 31.7%、MuSiQue 8.8%——多跳问题缺一块证据即全盘皆输。
本文的目标是论文想回答一个统一性问题:能否让一个预训练 LLM 自己完成跨尺度的证据选择,把语料级检索和长上下文推理变成同一机制的两种规模,而非两套竞争系统。具体目标有三:其一,在冻结的 decoder-only LLM 上实现内在检索,不需要外部 encoder 或检索模型,新增可训练参数不超过 50 万(不足骨干参数的 0.005%);其二,该机制要在 2100 万 chunk 的全库检索上击败专门训练的 retriever–reranker 系统;其三,同一模块不做任何长上下文数据的额外训练,就能直接用于长上下文场景,在生成前移除 distractor,并同时降低 FLOPs 与首 token 延迟,且对 dense attention、线性注意力、状态空间(Mamba)混合架构一律通用——同一个选择器覆盖从 8K token 上下文到 30 亿 token 语料约六个数量级的范围。
与已有工作不同的是,独特切入点是把“读哪里”转移到模型内部。最接近的先行工作 INTRA(Hoffer et al., 2026)证明 encoder–decoder 模型内部存在检索能力,但它依赖独立 encoder 的块编码 $k_i=\mathrm{Enc}(c_i)$ 和解码器 cross-attention,无法迁移到主流 decoder-only LLM。UNREAL 的关键观察是:decoder 的残差流本身就是可读出的表征总线——用它同时替代 INTRA 的 encoder 表征与 cross-attention query,就能在完全冻结的骨干上做检索。这带来两个别人没有的性质:一是架构无关性,残差流在 softmax 注意力、线性注意力、SSM 层中都存在(论文借 selective SSM 的因果注意力解释为跨层型聚合提供理论动机),一个读出头即可覆盖所有混合架构;二是统一性,检索训练学到的选择器原封不动拿去做长上下文的 chunk 过滤,把长期被反复比较的 RAG 与 long-context 两条路线收敛为同一操作“query 条件化的证据选择”的两个尺度。
核心方法
直觉上,冻结 LLM 的中间层已经隐含“哪些 chunk 与问题相关”的信号:作者用 Nemotron-3-Nano 做了无训练实验,把每层注意力的 key 按 chunk 平均、query 按问题 token 平均后做点积打分,未训练表征的召回率已明显高于随机基线(Fig. 6 右)。UNREAL 只需用极少量可训练参数把这个信号放大成可靠检索器。技术路线分三块:其一,块编码——直接用冻结 LLM 第 $\ell_c$ 层的残差状态当块表征 $k_i=\mathrm{LLM}_{\ell_c}(c_i)\in\mathbb{R}^{T_i\times d}$,无需独立 encoder;其二,查询侧——在问题 token 后拼接 $R=64$ 个可学习的检索 token $\rho$,前面再拼 BM25 召回的 top-5 chunk 作为初始上下文 $C_0(x)$ 和一个可学习 soft prompt,做一次前向,读取 $\rho$ 位置所有层的 states $q^\ell$,用可学习的层混合系数 $\alpha_\ell$ 加权后,以 ColBERT 式 MaxSim 对每个 chunk 打分 $s_i=\mathrm{MaxSim}(\sum_\ell\alpha_\ell q^\ell,\,k_i)$;其三,生成——把得分 top-$n$ 的 chunk 原文与问题拼接,让同一个 LLM 再做一次普通生成前向。训练只用 multi-positive InfoNCE 对比损失,骨干全程冻结。
核心创新是“模型内证据选择”的统一框架:块编码与检索 query 都取自同一个冻结 LLM 的残差流,检索与生成共享表征空间,匹配发生在“最终消费者”自己的表征里,天然规避了外部 retriever 与 reader 之间的语义错位。与 INTRA 的本质区别在于彻底去掉 encoder 与 cross-attention:INTRA 检索时把选中记忆经 cross-attention 注入解码器,UNREAL 则把选中 chunk 以纯文本重新拼接生成,因此适用于任何 decoder-only 模型。与外部检索器(RAG)的本质区别是无需维护第二个模型——BM25 只提供廉价初始上下文(消融显示完全去掉会掉 10–16 个百分点,说明它是重要词法先验而非主体)。与稀疏注意力方法(Landmark Attention、Quest、NSA、MoBA 等)的本质区别是:那些方法在注意力内部选择性读取 KV 块,UNREAL 在生成前就把未选中文本丢掉,LLM 只读干净证据,从机理上消除 distractor 干扰(Fig. 11 显示检索到的干扰项比随机干扰项更伤精度)。
方法步骤详情
完整流程分离线索引、查询打分、生成、训练四段。离线建库:每个 chunk $c_i$(中位约 141 token)过一遍冻结 LLM 至选定层 $\ell_c$,取残差状态 $k_i=\mathrm{LLM}_{\ell_c}(c_i)\in\mathbb{R}^{T_i\times d}$,切成 $L_p=7$ 个连续组均值池化成 $7\times d$ 存入索引;全库约 21M chunk、30 亿 token。查询打分:构造 $x_{\mathrm{ret}}=[\,C_0(x),\,x_1..x_{T_q},\,\rho_1..\rho_R\,]$(最前置 1 个 soft prompt;$C_0(x)$ 为 BM25 top-5;$R=64$),一次全层前向后读出检索 token 位置的逐层 states $q^\ell(x_{\mathrm{ret}})$,按 $G=4$ 组均值池化,以 $\alpha_\ell$ 跨层加权得 $\sum_\ell\alpha_\ell q^\ell$,与每个 chunk 的 $7\times d$ 向量做 $s_i=\mathrm{MaxSim}(\sum_\ell\alpha_\ell q^\ell,k_i)$,取 top-$n$(长上下文实验 $n=10$)。生成:选中 chunk 按原文档顺序拼回 $y=\mathrm{LLM}([\,\mathcal{C}_{\mathrm{UNREAL}}(x),\,x\,])$ 常规生成,推理即“一次检索前向 + 一次生成前向”。训练:oracle 集 $O(x)$ 为包含标注证据的 chunk,负例以 oracle 文本为 BM25 种子挖掘最多 500 个 hard negatives(丢弃 top-10 近重复),在 $B(x)=O(x)\cup N(x)$ 上最小化 $\mathcal{L}=-\frac{1}{|O(x)|}\sum_{j\in O(x)}\log\frac{\exp(s_j(x)/\tau)}{\sum_{i\in B(x)}\exp(s_i(x)/\tau)}$;共 15K 步、全局 batch 256、AdamW($\beta=(0.9,0.95)$、weight decay 0.1、梯度裁剪 1.0、100 步 warmup 后线性衰减),学习率 $5\times10^{-3}$ 至 $10^{-2}$。
技术新颖性
技术新颖性有四点。第一,极低的可训练预算:只训 $(R+1)D+|\alpha|$ 个参数(检索 token、soft prompt、层混合权重),四个骨干上均不足 50 万、占比小于 $2\times10^{-5}$,却超过了在高质量检索数据上微调、参数量数亿到数十亿的 Qwen3-Embedding 加 reranker 的组合。第二,跨架构通用:读出只依赖残差流的存在,论文在 dense(Qwen3.5-4B、Muse-Glimmer-30B)、线性注意力混合(Qwen3.5-35B-A3B)、Mamba–注意力混合(Nemotron-3.5-Lightning-30B-A3B)上全部验证成功。第三,一次训练两个用途:检索数据训练出的同一模块直接用于 NoLiMa、HELMET、LV-Eval、LOFT 等长上下文基准,无任何长上下文训练数据,支撑“检索监督即长上下文能力”的论点。第四,系统级分析完整:推导了 UNREAL 与全上下文推理的 FLOPs 差项分解(式 8)与 break-even 长度 $N^*$ 的闭式估计(式 10),并在 vLLM/H100 上实测 TTFT 与吞吐,而非只报理论复杂度。
实验结果
全库检索(Fig. 3、8–10):在 21M chunk 的 Wiki-2018 上四个 UNREAL 骨干全面领先,最强骨干把 complete-evidence recall@10 从最强基线的 49.1% 提到 73.2%(HotpotQA)、31.7%→60.1%(2WikiMultiHopQA)、8.8%→14.4%(MuSiQue),增益集中在多跳任务。端到端生成(Table 1、2):固定 Nemotron-3.5-Lightning 为生成器、各检索器只给 top-5 上下文,UNREAL 在三个多跳数据集上 EM/F1 全部最高——HotpotQA 52.6/65.8(最强基线 Qwen3-Emb-4B+reranker 43.3/55.3,oracle 上限 64.3/78.9)、2Wiki 42.0/50.1、MuSiQue 16.9/25.4;换 Qwen3.5-35B-A3B 生成器结论一致(HotpotQA EM 48.2 vs 41.3)。长上下文(Fig. 1、4、5、12):8K–100M token 新基准上全上下文 EM 崩塌且约 256K 后不可行,UNREAL 全程最高并能在 100M token 处选题作答;NoLiMa 在 128K 处由 1.0% 升至 24.83%,LV-Eval 在 256K 词处 F1 由 49.97% 提到 54.66%,HELMET RAG 子集在两个生成器上 substring EM 均第一,标准 LOFT 上 k=5/10/20 预算全部优于全文阅读。消融(Table 3–5):跨多层读出与 BM25 初始上下文最关键(多层→单层掉 8.26–11.92 点,去掉 $C_0(x)$ 掉 10.20–16.20 点),块编码取中后层最佳(第 42/52 层),召回随 $L_p$ 单调上升。效率(Fig. 7、13、14、Table 6):FLOPs 约在 1.8 万 token 以上占优(各骨干 break-even 4.1K–17.4K),实测 TTFT 从约 32K token 起反超全上下文,100M token 处加速达数百至上千倍(图中最高标注 1455×)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| HotpotQA 全库检索(21M Wiki-2018 chunks) | Complete-evidence recall@10 | 73.2%(最佳 UNREAL 骨干) | 49.1%(最强 retriever–reranker 基线) | +24.1 个百分点 |
| 2WikiMultiHopQA 全库检索 | Complete-evidence recall@10 | 60.1% | 31.7% | +28.4 个百分点 |
| MuSiQue 全库检索 | Complete-evidence recall@10 | 14.4% | 8.8% | +5.6 个百分点 |
| 端到端多跳 QA(HotpotQA,Nemotron-3.5-Lightning 生成器,top-5 检索上下文) | EM | 52.6 | 43.3(Qwen3-Emb-4B + reranker) | +9.3(oracle 上限 64.3) |
| 端到端多跳 QA(2Wiki,同上设置) | EM | 42.0 | 32.5(Hybrid RAG,最强基线) | +9.5(oracle 上限 57.5) |
| 端到端多跳 QA(MuSiQue,同上设置) | EM | 16.9 | 13.2(Qwen3-Emb-4B + reranker) | +3.7(oracle 上限 47.5) |
| NoLiMa 长上下文(128K token 最大长度) | 准确率 | 24.83%(UNREAL) | 1.0%(全上下文推理) | +23.83 个百分点 |
| LV-Eval 长上下文(256K 词) | F1 | 54.66% | 49.97%(全上下文推理) | +4.69 个百分点 |
| 超长上下文推理效率(约 100M token) | 首 token 时间(TTFT)加速比 | 数百至上千倍(图中最高 1455×,随骨干而异) | 全上下文推理(该长度下单次前向已不可行) | 约 32K token 起即开始更快,长度越大加速越大 |
局限与改进
作者承认的局限:训练数据全部来自 Wikipedia QA,向异构领域与跨语言的迁移留待未来;长上下文实验集中在稀疏证据任务;索引每 chunk 存 $L_p$ 个向量,构建与存储成本高于单向量稠密索引;依赖 BM25 初始上下文;用通用 LLM 而非 SpanBERT 类任务专用抽取器生成答案(后者在抽取式 QA 上 EM 更高);检索路径引入了额外的 chunk 编码与搜索步骤;聚焦单轮检索,未覆盖多轮 agentic RAG。我的补充观察:其一,评测是对全部 21M chunk 穷举 MaxSim 打分、不用任何近似索引,生产系统必须换 ANN,而多向量($7\times d$)索引的内存开销与召回损失未被评估;其二,训练与评测高度同分布——八个 QA 数据集的训练 split 训练、相同数据集的验证 split 评测,对未见任务的 zero-shot 泛化缺少证据;其三,hard negatives 全部由 BM25 挖掘,对词法不重叠的语义难负例覆盖不足,而 NoLiMa 恰以“无关键词重叠”著称,选择器在更开放分布上的鲁棒性存疑;其四,长上下文场景每次都要对全部输入 chunk 重新编码,短上下文(约 32K 以下)反而比全上下文慢;其五,检索到的干扰项比随机干扰项更伤生成(Fig. 11),选择器的误差会集中传导为语义近似干扰,$n$、$\ell_c$、$L_p$ 等超参迁移到新域都需重调。
独立分析的弱点
具体场景化的弱点与改进方向。其一,域迁移风险:所有训练与评测都在英文 Wikipedia 上,面对代码库问答、法律合同、多语言场景,最优编码层 $\ell_c$ 与检索信号可能漂移——可先做跨域 probing(复现 Fig. 6 右的无训练内在检索评估)判断信号是否仍在,再决定是否引入领域对比数据重训。其二,稠密证据任务失效:摘要、聚合类任务需要整合大部分输入,top-$n$ 选择必然丢信息——可与全上下文或层级式读出自适应融合,按估计的证据密度动态决定“选择还是全读”。其三,对 BM25 的隐性依赖:去掉 $C_0(x)$ 掉 10–16 个百分点,说明所谓“单模型”仍内嵌词法检索器,在同义改写、跨语言等词法失配场景会退化——可用可学习稀疏检索头或迭代查询扩展替代。其四,索引经济性:30 亿 token 需要 30B 级模型全库前向一遍建索引,换模型版本即需重建,$7\times d$ 多向量的存储与 ANN 成本高——可引入量化(标量/乘积量化)、两级检索(单向量粗排 + MaxSim 精排)或蒸馏小编码器专管建库。其五,两遍前向的编排复杂且短上下文无延迟优势,论文自认 per-stage 基准未计入编排开销,生产栈需要把编码、选择、生成拼成低开销流水线并做好 KV cache 管理。
未来方向
作者提出三个方向:一是“训练模型去回忆”——把检索监督放进预训练或后训练,使模型内部检索能力从被动涌现变为显式训练目标,联合优化检索与生成;二是把稀疏证据型长上下文任务在训练与评测中都形式化为检索问题,评测同时报告证据召回与答案准确率;三是基准应标注每题所需证据,区分稀疏证据任务与需要跨大半输入整合信息的稠密证据任务。基于本文成果可延伸的方向:把 UNREAL 选择器作为 agentic RAG(Search-o1、Self-RAG 等)中的单步检索组件,支持多轮迭代检索与自适应 $n$;用强化学习以下游答案 EM 直接奖励选择质量,替代或增强 InfoNCE;扩展到多模态(长视频、文档图像)的块级证据选择;研究检索训练对生成主干本身的影响——是否零成本提升模型的注意力定位与抗干扰能力;以及在流式、持续更新语料上的增量索引与缓存复用。
复现评估
复现条件较为透明。模型:四个骨干均为公开权重(Qwen3.5-4B、Qwen3.5-35B-A3B、Muse-Glimmer-30B、Nemotron-3.5-Lightning-30B-A3B)。数据:训练只用公开 QA 数据集(NQ、SQuAD v2、HotpotQA、2WikiMultiHopQA、MuSiQue、IIRC、FEVER、HoVer)加公开的 DPR Wikipedia-2018 语料,KILT 式证据映射、质量阈值 0.6 的过滤、oracle 种子 BM25 挖掘 hard negatives 的流程在附录 B.1 写得很细。超参完整给出(附录 B.3:$R=64$、$G=4$、$L_p=7$、$|C_0(x)|=5$、各骨干的 $\ell_c$ 与学习率)。算力:重头在建索引——21M chunk 过 30B 级模型取中间层,需要可观的批量推理资源;对比训练本身只更新不足 50 万参数,成本不高;效率基准用随机权重 + vLLM 单张 H100(bf16)即可复现。文本未明确给出代码发布链接,MaxSim 穷举打分、池化压缩、长上下文基准(8K–100M、固定相对深度、固定随机种子)的工程实现需自行完成,复现 95% BCa 置信区间也有工作量。总体难度:中——材料公开、配方详尽,但需要较大的推理算力与工程能力。
论文图表
三个骨干(Nemotron-3.5-Lightning-30B-A3B、Qwen3.5-35B-A3B、Muse-Glimmer-30B)在上下文从 8K 增至 100M token 时的 QA 精确匹配曲线:全上下文只画到计算不可行为止;BM25、BGE+reranker、Qwen3-Emb-0.6B/4B+reranker 等外置检索方法退化较缓但更低;UNREAL 用冻结 LLM 内部表征检索,在所有长度上保持最高 EM,并能在 100M token 规模完成检索作答。
一图点题:全上下文长到一定程度既不可行也不准确,把检索与长上下文统一为一个机制后在两个尺度上都占优,是全文动机与主张的浓缩。