本文介绍了一种用于长文本检索-增强生成（RAG）的新方法--UncertaintyRAG，它利用基于信噪比（SNR）的跨度不确定性来估计文本块之间的相似性。

基于跨度不确定性可增强模型校准、提高鲁棒性并减轻随机分块带来的语义不一致的洞察力，它还提出了一种高效的无监督学习技术来训练检索模型，以及一种有效的数据采样和缩放策略。

主要贡献：

• 提出一种新颖的信噪比不确定性测量技术，通过解决随机分块带来的预测误差来实现更好的校准，从而提高分块间相似性估计的性能。
• 提出了一种无监督学习方法并训练了一个检索模型，该模型在分布偏移设置下的长上下文 RAG 任务中表现优于强大的开源嵌入模型。
• 设计了一种高效的数据采样策略来扩展数据，从而增强了我们的检索模型训练并显著提高了性能
• 对检索模型进行深入分析，证明在两个关键指标上的持续改进

方法论

• 介绍一种基于信噪比的跨度不确定性方法，以获得数据块之间的置信度分数。
• 然后利用这些置信度分数构建正样本和负样本，用于训练我们的检索模型。
• 这一过程通常不涉及使用查询数据进行训练，因此我们进一步开发了两种方法来扩展块数据，以增强模型的分布转移泛化能力。

i) 跨度不确定性

• 将两个数据块输入 LLM，使用模型输出的自信息作为不确定性的度量。
• 从 [2] 中用于估计泛化误差的样本梯度信噪比统计中汲取灵感，通过计算模型样本输出概率的信噪比，利用这一统计来量化不确定性。
• 首先，我们将两个数据块连接起来并输入模型，以获得概率。
• 然后，我们利用这一信噪比来测量跨度的不确定性，最后将跨度的不确定性转换为相似度的测量值。具体来说，不确定性可以用自我信息来表示：

其中，xi 是当前标记，p(xi|xi-1, ... , x0) 是给定序列中前面标记 xi-1, xi-2, ... , x0 的标记 xi 的条件概率。, x0 的条件概率。

• 使用长度为 20 的滑动窗口计算窗口内的信噪比，滑动步长为 10，直到信噪比最终稳定为止
• 将该窗口内的跨度概率定义为不确定性度量，并将该置信度分数视为一种特殊的蕴含关系[3]，它代表了一种语义等价形式。
• 基于上述观察和 I(x) 的加法属性，我们提出了以下定义：

• 跨度不确定性（表示为 SU(x)）被定义为所选窗口中所有连接词的平均自信息，其中重叠词只计算一次。计算公式为：

其中 1SNRj<σ 是一个指示函数，当第 j 个窗口的 SNR 低于阈值 σ 时，该函数等于 1，否则等于 0。nj 表示第 j 个窗口中的标记数量，I(xij |x(i−1)j , . . . , x0j) 表示给定其前一上下文的标记 xij 的自信息。

ii) 训练策略

a) 构建正样本和负样本

• 统一设置分块大小为 300 以分割数据。
• 该矩阵中的每个元素 Sij 代表将 chi 和 chj 依次组合并输入 LLM 时估计出的跨度不确定性，表示它们的相似程度。
• 选择每个锚 chi 的 BM25 分数最高的 M 个样本来估计 SU(x)，从而得到最终的稀疏矩阵 ˆS。用公式表示，让 sBM25(chi, chj) 表示 i 和 j 块组合的 BM25 分数：

其中，TopM(sBM25(chi, chM))表示锚点 chi 的第 M 个最高 BM25 分数 chM

• 在 M 中设置两个窗口，每个窗口包含 m 个样本。其中，顶部的 m 个样本为正样本，底部的 m 个样本为负样本。

b) 数据缩放策略

(1) 锚点样本缩放策略

• 每个数据块称为一个锚样本
• 最初合并五个数据集： 最初合并五个数据集：HotpotQA、MultiFieldQA、Qasper、NarrativeQA 和 QMSum。
• 每个数据集根据空格进行拆分，每 300 个单词组成一个词块。因此，我们分别获得了 37,799 个、14,547 个、15,865 个、72,146 个和 38,398 个块
• 首先使用 k-nearest neighbors（KNN）将每个数据集聚成 k 个群组。
• 然后，从每个聚类中随机抽取 c 个数据块，每个数据集抽取 10*c 个数据块。
• 通过混合五个数据集，我们总共得到 50*c 个数据块。通过增加 c 的倍数，我们可以扩大锚点样本的总数。

(2) 正样本和负样本缩放策略

• 选定一个锚点样本后，我们需要将其与其他数据块连接起来，然后输入 LLM 进行评分。由于计算成本较高，我们采用了一种类似于锚样本缩放的方法
• 首先，使用 KNN 方法，我们从 k 个群组中的每个群组随机抽取 n 个样本 chj，其中群组由前面提到的 50∗c 块组成。这样，我们就能在一次采样过程中获得 10∗n 个样本。
• 然后，我们使用 BM25 对锚点和每个样本进行评分，即 sBM25(chi,chj)，并选出最终的 M 个样本。使用 LLM 对这些样本进行评分，得到最终的 chi+ 和 chi-。
• 最后，我们得到一个三元组（chi、chi+、chi-）。通过重复这一采样过程，我们可以缩放锚点 chi 的正负样本数量。

c) 对比学习

• 设计一种对比学习策略来训练我们的检索模型：

其中，chi 表示锚块，即用于比较的块。术语 chi+ 表示正块，即与 chi 相似的块。chi - 表示硬负块，即与 chi 不相似的序列。

• 锚语块 chx 和另一个语块 chy 之间的相似度 f（chx，chy）是根据它们的语块嵌入 BERT 的内积计算得出的

其中，be(chx) 和 be(chy) 分别表示块 chx 和块 chy 的嵌入度

iii) 推论

分为以下几个步骤：

(1) 我们直接根据固定长度对输入文档进行分块，这样效率高，无需使用复杂的分块机制。

(2）将查询输入检索模型，检索出最相似的 m 个分块。

(3) 将检索到的 m 个块和查询输入模型，生成最终答案。

实验

i) 主要结果

• 下表列出了长语境检索增强语言生成的实验结果。All Chunking 表示在对比学习中使用两个连接块中所有标记的平均自信息作为选择正负样本的相似度得分。相比之下，精确分块（Precise Chunking）指的是只利用第二个分块的平均自信息对分块进行精确分割。

根据上表，我们得出了以下结论：

• 强大的 RAG 检索模型可以显著提高 4K 上下文窗口 LLM 的性能；但是，这在很大程度上取决于检索模型的性能。
• 强大的检索模型可以提高 LLM 的性能。Vicuna-7B 和 LLaMA-2-7B-Chat-HF 在其 4K 上下文窗口内的性能类似，在没有检索增强的情况下，平均性能差异约为 1.93%。如果配备鲁棒检索模型，平均性能差距可增至 6.14%。
• 我们的方法取得了最高的平均性能，超过了一些在大型数据集上训练的开源嵌入模型。但是，由于需要对上下文中的示例进行精确分割，因此在 TREC 等少量学习任务中的性能受到了限制。

ii) 表征分析

a) 表征相似性分析

• 表征相似性分析是一种测量给定刺激集的两个不同表征空间之间相似性的技术。
• 下图展示了不同表征层之间相似性的变化

基于上图的观察结果：

• 在第八层之前，不同数据集和评分方法之间的模型表征变化极小。
• 在第八层之后，模型在 GSM8K 数据集上的变化较小，分布偏移明显；在 2WikiMultihopQA 数据集上的变化较大，分布偏移较小；而在训练集上的变化最大。这表明 RSA 是评估分布偏移的有效指标。
• 在带有跨度不确定性（Span Uncertainty）的增强型长语境建模下，模型在不同的分布变化中表现出更大的表征变化，始终优于精确分块，从而提高了性能。

b) 表征属性分析

• 介绍对比学习在检索模型中改进的统一性和对齐性，这对于评估机器学习中检索模型的质量至关重要。
• 对齐性是通过成对正片块的嵌入之间的预期距离来量化的，假设这些嵌入是标准化的。对齐公式如下：

• 均匀性评估嵌入式在嵌入空间中的均匀分布程度。其计算公式为

其中，ppos 表示正对（相似句块），pdata 指数据集中数据点的分布，样本独立抽取，这与对比学习的目标一致。

• 下图显示了不同句子嵌入模型的统一性和一致性，以及它们的平均 STS 结果

根据上图，我们可以看出以下几点：

• 当我们缩放正负样本时
• 和负样本时，检索模型的统一性会提高，这与 InfoNCE 的理念一致，即增加噪声样本可以提高表征的统一性。
• 当我们扩展锚点时，检索的对齐度也会提高，这意味着我们可以通过对模型的锚点样本进行数据选择来对齐检索模型。
• 提高模型的对齐程度并不能提高检索模型的性能，而提高模型的均匀性则可以提高检索模型的性能。

iii) 不确定性建模

• 下表显示了不同不确定性建模方法的比较结果

• 发现我们方法的性能远远超过简单的标记级不确定性估计和句子级不确定性估计
• 对于不确定性校准评估，下图显示了不确定性度量的 AUROCs。横轴表示阈值 τ

• 上图显示，我们的不确定性测量校准效果明显优于其他方法。这两个结果都是在使用 TriviaQA 的 Llama-2-7b-chat-hf 模型上进行评估的。

结论

• UncertaintyRAG 引入了一种新颖的基于信噪比的跨度不确定性方法，以改进长语境检索增强生成（RAG）中的校准。该方法的性能优于 BGE-M3 等功能强大的开源嵌入模型，但只使用了 4% 的训练数据。
• 通过利用跨度的不确定性，UncertaintyRAG 促进了无监督学习，减少了对大型标签数据集的需求，使其具有可扩展性和高效性。
• 结果表明，即使使用最少的数据，UncertaintyRAG 也能在分布偏移的情况下实现卓越的性能和泛化。
• 这种轻量级解决方案无需微调，可无缝集成到任何 LLM 中，因此适用于各种长语境任务。