[论文笔记]RAPTOR: RECURSIVE ABSTRACTIVE PROCESSING FOR TREE-ORGANIZED RETRIEVAL

引言

今天带来又一篇RAG论文笔记：RAPTOR: RECURSIVE ABSTRACTIVE PROCESSING FOR TREE-ORGANIZED RETRIEVAL。

检索增强语言模型能够更好地适应世界状态的变化并融入长尾知识。然而，大多数现有方法只能从检索语料库中检索到短的连续文本片段，限制了对整个文档上下文的完整理解。作者引入了一种新颖的方法，即递归嵌入、聚类和总结文本片段，从底部开始构建具有不同摘要级别的树。在推理期间，RAPTOR模型从这棵树中检索，将跨不同抽象级别的庞大文档中的信息整合起来。

源码发布在： https://github.com/parthsarthi03/raptor 。

1. 总体介绍

大语言模型已经成为具有深远影响的变革性工具，在许多任务上表现出令人印象深刻的性能。然而，即使是一个大型语言模型也不包含特定任务所需的足够领域特定知识，并且世界持续变化，使LLM中的事实失效。通过额外的微调或编辑来更新这些模型的知识是困难的，特别是在处理庞大文本语料库时。在开放领域问答系统中开创性的另一种方法是将大量文本进行分块(段落)后在单独的信息检索系统中建立索引。所检索的信息随后与问题一起作为上下文呈现给LLM（RAG检索增强），这使得为系统提供特定领域当前知识变得容易，并且能够轻松解释和追溯来源。

然而，现有的检索增强方法也存在缺陷。大多数现有方法只检索到几个短小的相邻文本片段，这限制了它们表示和利用大规模话语结构的能力。这在需要整合来自文本多个部分知识的主题性问题中尤为重要，比如理解整本书的情况，就像NarrativeQA数据集中以灰姑娘的童话故事为例，针对问题“灰姑娘如何走向幸福结局？”，检索到的前K个相邻短文本不包含足够的上下文来回答问题。

为了解决这个问题，作者设计了一个索引和检索系统，使用树结构捕捉文本的高级和低级细节。如图1所示，RAPTOR对文本片段进行聚类，生成这些聚类的文本摘要，然后重复此过程，从底部生成树。这种结构使得RAPTOR能够加载不同级别代表文本的片段到LLM的上下文中，从而能够有效地回答不同级别的问题。

作者的主要贡献是利用文本摘要的想法，允许在不同尺度上进行检索增强上下文，并在长文集合的实验中展示其有效性。

2. 相关工作

为什么使用检索？ 硬件和算法的最新进展确实扩展了模型可以处理的上下文长度，引发了关于检索系统必要性的讨论。然而，模型往往未充分利用长距离上下文，并在上下文长度增加时性能下降，特别是当相关信息嵌入在冗长上下文中时。此外，在实际应用中，使用长上下文既昂贵又缓慢。这表明，为知识密集型任务选择最相关的信息仍然至关重要。

检索方法 检索增强语言模型（Retrieval-augmented language models, RALMs）在各个组成部分中取得了改进：检索器(retriever)、阅读器(reader)和端到端系统训练。检索方法已经从传统的基于单词(term)的技术(如TF-IDF和BM25)转变为基于深度学习的策略。一些最近的工作提出使用大型语言模型作为检索器，因为它们能够记忆大量知识。关于阅读器组件的研究包括Fusion-in-Decoder，它同时使用DPR和BM25进行检索，并在编码器中独立处理段落，以及RETRO，它利用跨块关注和分块检索生成基于检索上下文的文本。

端到端系统训练工作包括Atlas，它在与检索器一起微调编码器-解码器模型；REALM，一种为开放域问答微调的双向掩码语言模型；以及RAG，将预训练的序列到序列模型与神经网络检索器集成。尽管方法多样，模型的检索组件主要依赖于标准方法，即对语料库进行分块并使用基于BERT的检索器进行编码。

尽管这种方法被广泛采用，但Nair等强调了一个潜在的缺点：连续分割可能无法捕获文本的完整语义深度。从技术或科学文档中提取的片段阅读可能缺乏重要上下文，使其难以阅读甚至具有误导性。

递归摘要作为上下文 摘要(Summarization，总结)技术提供了对文档的简洁视图，使内容更加集中。Gao等人的摘要模型使用段落(passage)的摘要和片段，提高了大多数数据集的正确性，但有时可能是一种高失真的压缩手段。Wu等人的递归摘要生成模型采用任务分解来总结较小的文本块，然后将其集成以形成较大部分的摘要。虽然这种方法对捕获更广泛主题很有效，但可能会忽略细节。LlamaIndex通过类似摘要相邻文本块的方式来缓解这一问题，但同时保留中间节点，因此存储着不同层次的细节，保留了细粒度的详情。然而，由于这两种方法依赖于相邻节点的分组或摘要，可能仍会忽视文本内的远程相互依赖关系，而RAPTOR可以找到并组合这些关系。

3. Method

RAPTOR概览 建立在长文本通常呈现子主题和分层结构的思想基础上，RAPTOR通过构建一个递归树结构来解决阅读中的语义深度和连接问题，该结构平衡了更广泛的主题理解和细粒度细节，允许根据语义相似性而不仅仅是文本顺序对节点进行分组。

RAPTOR树的构建始于将检索语料库分割成长度为100的短连续文本，类似于传统的检索增强技术。如果一句话超过100个标记的限制，我们将整个句子移至下一个块(chunk)，而不是在句子中途截断。这保留了每个块内文本的上下文和语义连贯性。然后使用SBERT嵌入这些文本。这些块及其对应的SBERT嵌入形成了树结构的叶节点。

为了将相似的文本块分组，采用了聚类算法。一旦分组，将使用语言模型对分组文本进行总结。然后重新嵌入这些总结文本，并继续嵌入、聚类和总结循环，直到进一步聚类变得不可行为止，最终形成原始文档的结构化、多层树表示。RAPTOR的一个重要特点是其计算效率。该系统在构建时间和标记消耗方面呈线性缩放，适用于处理庞大和复杂的语料库。

对于在该树内查询，引入了两种不同的策略：树遍历和折叠树。树遍历方法逐层遍历树，在每个级别修剪和选择最相关的节点。折叠树方法评估所有级别上的节点，以找到最相关的节点。

聚类算法 聚类在构建RAPTOR树、将文本段组织成连贯簇(Group)方面发挥着关键作用。这一步将相关内容组合在一起，有助于后续的检索过程。

作者聚类方法的一个独特之处在于使用软聚类，其中节点可以属于多个簇，而不需要固定数量的簇。这种灵活性是必不可少的，因为单个文本段通常包含与各种主题相关的信息，因此需要将其纳入多个摘要中。

聚类算法基于高斯混合模型(GMMs)，这种方法提供了灵活性和概率框架。GMMs假设数据点是从多个高斯分布混合中生成的。

给定一个包含N个文本片段的集合，每个表示为一个d维稠密向量嵌入，假设文本向量 $x$ 属于第k个高斯分布的概率表示为 $\mathscr N(x;\mu_k,\Sigma_k)$ 。整体概率分布是加权和 $\sum_{k=1}^K \pi_k \mathscr N(x;\mu_k,\Sigma_k)$ ，其中 $\pi_k$ 表示第k个高斯分布的混合权重。

向量嵌入的高维度对传统GMMs构成挑战，因为在高维空间中使用距离度量可能表现不佳。为了缓解这一问题，作者采用均匀流形近似和投影(Uniform Manifold Approximation and Projection,UMAP)，一种用于降维的流形学习技术。UMAP中的最近邻参数n_neighbors决定了保留局部结构和全局结构之间的平衡。作者的算法通过变化n_neighbors来创建层次聚类结构：首先识别全局聚类，然后在这些全局聚类中执行局部聚类。这种两步聚类过程捕捉了文本数据之间的广泛关系，从广泛主题到具体细节。

如果局部聚类的组合上下文超过了总结模型的标记阈值，作者的算法会在该聚类内递归应用聚类，确保上下文保持在标记阈值内。

为了确定最佳聚类数，使用贝叶斯信息准则(Bayesian Information Criterion,BIC)进行模型选择。BIC不仅惩罚模型复杂性，还奖励拟合优度。对于给定的GMM，BIC为 $\ln(N)k - 2\ln(\hat L)$ ，其中 $N$ 为文本片段(或数据点)的数量， $k$ 为模型参数的数量， $\hat L$ 为模型似然函数的最大化值。在GMM的情境中，参数数量 $k$ 是输入向量的维度和聚类数量的函数。

得到由BIC确定的最佳聚类数量后，使用EM算法来估计GMM参数，即均值、协方差和混合权重。虽然GMM中的高斯假设可能与文本数据的性质不完全符合，后者通常表现出稀疏和偏斜的分布，但实证观察表明，它是一种有效的模型。

基于模型的摘要 在使用高斯混合模型对节点进行聚类后，将每个聚类中的节点发送到语言模型进行总结。这一步允许模型将大块文本转化为所选节点的简明、连贯摘要。作者使用gpt-3.5-turbo生成摘要。总结步骤将检索到的大量信息压缩为可管理的大小。

查询下面详细阐述了RAPTOR所采用的两种查询机制：树遍历和折叠树(collapsed tree，坍塌树)。这些方法提供了在多层RAPTOR树中获取相关信息的独特方式，每种方法都有其各自的优势和权衡。在附录F中提供了这两种方法的伪代码。使用SBERT对所有节点进行嵌入。

树遍历方法首先基于查询嵌入与根节点的余弦相似性选择top-k个最相关的根节点。这些所选节点的子节点在下一层被考虑，然后再次基于它们的子节点与查询向量的余弦相似性选择top-k个节点。该过程重复进行，直到达到叶节点。最后，将从所有选定节点中提取的文本连接起来形成检索到的上下文。算法步骤如下：

在RAPTOR树的根层开始。计算查询嵌入与该初始层中所有节点的嵌入之间的余弦相似性。
基于最高余弦相似性得分选择前k个节点，形成集合 $S_1$ 。
继续处理集合 $S_1$ 中元素的子节点。计算查询向量与这些子节点的向量嵌入之间的余弦相似性。
根据与查询的最高余弦相似性得分选择top-k个子节点，形成集合 $S_2$ 。
递归地继续这个过程达到d层，生成集合 $S_1、S_2...S_d$ 。
将集合 $S_1$ 至 $S_d$ 连接起来组装成与查询相关的上下文。

通过调整深度 $d$ 和每层选择的节点数 $k$ ，树遍历方法提供了对检索到的信息的特定性和广度的控制。该算法从树的顶层开始具有宽泛的视野，并随着向下穿越到较低层而逐渐聚焦于更细节的内容。

图2：树遍历和折叠树检索机制的示意图。树遍历从树的根层开始，并基于余弦相似性与查询向量检索top-k个(在此处为top-1个节点)。在每层，它从前一层top-k的子节点中检索top-k的节点。折叠树将树折叠成单层，并基于与查询向量的余弦相似性检索节点，直到达到标记数的阈值。这两种示意图中都突出显示了进行余弦相似性搜索的节点。

折叠树方法提供了一种更简单的搜索相关信息的方式，通过同时考虑树中的所有节点，如图2所示。这种方法不是逐层进行，而是将多层树折叠成单层，实质上将所有节点置于同一水平进行比较。该方法的步骤如下所述：

首先将整个RAPTOR树折叠为单层。这个新的节点集合，表示为 $C$ ，包含原始树的每一层中的节点。
接下来，计算查询嵌入与折叠集合 $C$ 中所有节点的嵌入之间的余弦相似性。
最后，选择与查询具有最高余弦相似性得分的top-k个节点。继续将节点添加到结果集，直到达到预定义的最大标记数，确保不超出模型的输入限制。

作者在QASPER数据集的20个故事上对这两种方法进行了测试。图3显示了树遍历在不同top-大小和折叠树在不同最大标记数下的性能。折叠树方法始终表现更好。因为它提供了比树遍历更大的灵活性；即，通过同时搜索所有节点，它检索出针对特定问题正确粒度的信息。相比之下，使用相同值的 $d$ 和 $k$ 进行树遍历时，树的每个级别的节点比例将保持恒定。因此，无论问题是什么，高阶主题信息与细节信息的比例将保持不变。

然而，折叠树方法的一个缺点是需要对树中的所有节点执行余弦相似性搜索。然而，使用诸如FAISS等快速k近邻库可以使这一过程更加高效。

定性研究 作者进行了定性研究，以比较RAPTOR的检索过程与Dense Passage Retrieval(DPR)方法的优劣。研究专注于使用一个1500字的灰姑娘童话故事的主题性、多跳问题。如图4所示，RAPTOR基于树的检索使其能够选择来自不同树层的节点，与问题的细节级别相匹配。这种方法通常比DPR为下游任务提供更相关和全面的信息。

4. 实验

在三个问答数据集：NarrativeQA、QASPER和QuALITY上评估。

使用UnifiedQA 3B作为阅读器，SBERT、BM25和DPR作为嵌入模型，在QASPER、NarrativeQA和QuALITY三个数据集上进行受控对比。如表1和表2所示，结果表明，当RAPTOR与任何检索器结合时，在所有数据集上均表现出优于各自检索器的一致性。RAPTOR与SBERT结合具有最佳性能。