Skill-RAG：通过隐层状态探测和技能路由的故障感知检索增强

# Skill-RAG: 通过隐状态探测和技能路由的故障感知检索增强

来源: https://arxiv.org/html/2604.15771

Raymond Li（英属哥伦比亚大学，温哥华，加拿大，[email protected]），Xi Zhu（罗格斯大学，新不伦瑞克，新泽西，美国，[email protected]），Zhaoqian Xue（宾夕法尼亚大学佩雷尔曼医学院，费城，宾夕法尼亚，美国，[email protected]），Jiaojiao Han（新泽西理工学院，纽瓦克，新泽西，美国，[email protected]），Jingcheng Niu（达姆施塔特工业大学，达姆施塔特，德国，[email protected]），Fan Yang（维克森林大学，温斯顿-塞勒姆，北卡罗来纳，美国，[email protected]）

###### 摘要

检索增强生成（RAG）已成为大语言模型（LLM）基于外部知识的基础范式。虽然自适应检索机制提高了检索效率，但现有方法将检索失败作为重试信号，而非诊断信号，未能解决查询-证据错位的结构性原因。我们观察到，相当一部分持久性检索失败并非源于缺少相关证据，而是源于查询与证据空间之间的对齐差距。我们提出 Skill-RAG，一个故障感知的 RAG 框架，它将轻量级隐状态探针与基于提示的技能路由器相结合。探针在管道的两个阶段进行门控；当检测到故障状态时，技能路由器诊断根本原因，并从四种检索技能中选择——查询改写、问题分解、证据聚焦和用于真正不可约情况的退出技能——在下一次生成尝试前纠正错位。在多个开放领域问答和复杂推理基准上的实验表明，Skill-RAG 在多轮检索后持久的硬案例上显著提高准确率，在分布外数据集上收益尤为显著。表示空间分析进一步揭示，提议的技能在故障状态空间中占据结构化、可分离的区域，支持将查询-证据错位视为有类型而非整体现象的观点。

参见图1：Skill-RAG 管道概览。给定输入查询，隐状态探针在两个阶段进行检索决策门控；当检测到故障状态时，基于提示的技能路由器在四种检索技能中选择，以在下一次生成尝试前纠正查询-证据错位。

## 1. 引言

检索增强生成（RAG）已成为大语言模型（LLM）基于外部知识的基础范式，大幅提升了知识密集型任务的事实准确性。基于此，自适应和迭代检索机制已被提出来动态确定何时以及多久执行检索。然而，现有方法主要将检索控制视为粗粒度决策——专注于是否检索以及检索多少次——同时忽视了检索失败的结构性原因及其所需的纠正策略。因此，相当一部分硬案例表现出持久性失败，仅靠重复检索无法解决。

仔细检查表明，这些失败的很大一部分并非源于相关证据的缺失，而是源于结构性错位：查询相对于证据空间的表述不当，导致连续检索返回的文档在话题上相邻但推理上不足。这类失败在模型的内部表示中表现出结构化模式——过于宽泛的查询需要证据聚焦、纠缠的前提需要分解、不同的表面形式需要改写——如我们在第 4.3 节分析的失败表示的几何结构所证明的那样。我们通过引入**故障状态**来体现这一洞察：从模型隐层派生的潜在表示，用于表示检索何时停滞，使技能路由器能够选择有针对性的检索行为以替代通用的重新检索。

我们提出 Skill-RAG，一个故障感知的 RAG 框架，采用轻量级隐状态探针检测检索何时停滞并进行技能路由的门控。当检测到故障状态时，基于提示的技能路由器诊断根本原因，并从四个**检索技能**中选择——查询改写、问题分解、证据聚焦和用于识别真正不可约情况并优雅终止检索的退出技能。与优化检索触发或迭代深度的先前工作不同，Skill-RAG 将检索后恢复重新表述为**条件技能选择问题**，对 LLM 如何获取外部知识提供细粒度、故障条件化的控制。

我们做了三项贡献。（1）我们提出首个集成隐状态探针门控与基于提示的技能路由的框架，用于检索后故障恢复，产生统一的探测和路由管道，既不需要额外的 LLM 调用来做任何决策。（2）我们引入可转移的四种检索技能词汇，以观察到的失败模式为基础，在多个模型和数据集上产生一致改进，建立用于查询-证据对齐纠正的可复用分类法。（3）跨多个基准的实验表明，Skill-RAG 实现了有竞争力或最先进的性能，同时在分布外数据集上大幅超越仅探针基线，凸显了故障条件化技能路由相比简单门控的好处。

## 2. 相关工作

**自适应和迭代检索**。早期 RAG 系统在生成前检索一次，后续工作探索了自适应和迭代机制以改进检索覆盖率和效率。IRCoT 交错链式思维推理与检索，使用每个推理步骤指导下一个查询；Iter-RetGen 将模型的上一输出作为下一轮检索的上下文。FLARE 在令牌级生成置信度低于阈值时触发检索；DRAGIN 使用注意力信号确定检索时机；Self-RAG 训练模型发出控制检索和自我批评的特殊令牌；Adaptive-RAG 对查询复杂度分类以在不同深度的检索策略之间路由。Probing-RAG 利用隐状态表示进行检索决策门控。尽管取得这些进展，所有现有方法都将检索后故障视为重试信号而非诊断信号，未能解决查询-证据错位的结构性原因。

**查询改写和纠正性检索**。平行的工作通过查询侧和证据侧干预改进检索质量。查询改写方法重新表述原始查询以更好地匹配语料库索引约定；问题分解方法将复杂的多跳查询分解为顺序子查询；CRAG 评估检索文档质量，在检索置信度低时触发纠正行为——包括网络搜索和证据过滤。虽然 CRAG 在文档级别运行，评估检索段落是否相关，但 Skill-RAG 在故障状态级别运行，诊断模型为何未能生成正确答案，并根据诊断的故障状态路由到有针对性的对齐纠正。Skill-RAG 在单一故障条件化路由框架中统一查询改写、分解和证据聚焦，基于诊断的故障状态而非无条件地应用任何单一策略。

## 3. 方法

图1 展示了 Skill-RAG 管道。给定输入查询，隐状态探针首先评估模型的参数知识是否足以在无需检索的情况下回答；如果是，答案直接返回。否则，执行标准检索步骤，并对增强生成重新应用探针。如果检索证据充分，答案最终确定；如果不充分，基于提示的技能路由器接收失败的推理、答案和检索证据，诊断错位的根本原因，并从四个检索技能中选择以改写查询或重新聚焦证据。修订的查询触发新的检索轮次，探针进行下一次迭代的门控。此过程重复进行，直到探针判断模型状态充分或达到最大检索轮次。

### 3.1. 探针训练

为了训练探针，我们对我们的在域数据集的训练集（HotpotQA、NQ 和 TriviaQA）应用两种检索策略——无检索和单步检索——并提示模型产生链式思维推理轨迹，随后是最终答案。对于每个例子，我们从模型后三分之二层对应推理和答案令牌的隐状态中提取，并通过比较生成答案与金标答案分配二元标签，产生隐状态表示与正确性信号配对的标记数据集。探针实现为单隐层的前馈网络和二元分类头。为了利用跨深度的信息，我们为每层训练一个探针，并在推理时通过对预测概率求平均聚合输出，产生单一门控信号，反映跨表示级别的答案就绪性。

### 3.2. 技能路由器

当探针检测到故障状态时，基于提示的技能路由器被调用。路由器接收原始问题、模型的失败推理和答案以及当前检索的证据，诊断错位的原因，并从四个检索技能中选择。

**查询改写**针对查询的表面形式与语料库索引约定不同的情况，产生更好地与可检索证据对齐的改写查询。

**问题分解**处理具有纠缠前提的多跳查询，生成一系列子查询在发出最终检索查询前隔离每个推理步骤。

**证据聚焦**处理语义上宽泛的查询，从当前上下文中提取缺失的证据槽位并发出针对特定信息缺口的有根据的查询。

**退出**识别错位不可约的情况（由于缺失知识或模型容量限制），并终止检索以避免不必要的推理开销。

### 3.3. 迭代技能检索和终止

执行技能后，改写的查询被发送到检索器，模型基于更新的证据生成新答案。探针随后进行下一次迭代的门控。此循环继续进行直到满足三个终止条件之一：技能路由器选择退出、探针判断模型状态充分，或达到预定的最大检索轮次。

## 4. 实验

### 4.1. 设置

我们在跨越单跳和多跳推理的五个开放领域问答基准上进行评估。三个数据集——NQ、TriviaQA 和 HotpotQA——作为在域基准，我们从中采样 3,000 个例子用于探针训练，500 个用于开发。两个多跳数据集——MuSiQue 和 2WikiMultiHopQA——保留为分布外（OOD）测试集，各在 500 个例子上评估。所有方法使用 BM25 作为检索器。

我们将 Skill-RAG 与六个基线进行比较：无检索（仅从参数知识生成答案）；单步 RAG（在生成前执行一轮检索）；FLARE（基于令牌级生成不确定性触发检索）；DRAGIN（通过基于注意力的相关性信号确定检索时机）；Adaptive-RAG（通过训练分类器将查询路由到不同复杂度的检索策略）；和 Probing-RAG（使用隐状态探测进行检索决策门控）。我们使用 Gemma2-9B 作为骨干模型进行实验；跨其他模型族的结果将在未来工作中报告。所有方法使用 4-shot 提示法，并在精确匹配（EM）和准确率（ACC）上评估。

表1

### 4.2. 主要结果

表1 报告了 Gemma2-9B 在五个基准上的结果。Skill-RAG 在在域数据集上实现了最先进或有竞争力的性能，在 HotpotQA、NQ 和 TriviaQA 上的 EM 和 ACC 都匹配或超越 Probing-RAG。最显著的收益出现在 OOD