RAPTOR：树结构的索引和检索系统的递归抽象处理

论文地址：https://arxiv.org/pdf/2401.18059.pdf

摘要

增强型检索语言模型能够更好地适应世界状态的变化，并整合长尾知识，然而现有大多数方法仅能从检索语料库中检索到较短的连续文本片段，这限制了对整个文档上下文的整体理解。我们引入一种新颖的方法，即递归地嵌入、聚类及总结文本块，自下而上构建一个包含不同层次抽象的树状结构。在推理阶段，我们的RAPTOR模型会从这个树中检索信息，跨过冗长文档的不同抽象层次进行信息整合。受控实验表明，在多个任务上，采用递归摘要进行检索相较于传统的增强型检索语言模型有显著的性能提升。在涉及复杂、多步骤推理的问答任务上，我们展示了最先进的结果；例如，通过将RAPTOR检索与GPT-4相结合，我们在QuALITY基准测试上实现了绝对精度提升20%，达到了新的最佳性能水平。

概述

在构建检索增强生成（Retrieval-Augmented Generation, RAG）系统的过程中，文档分块是一个关键步骤，它影响着系统能否检索到足够的信息来回答用户的问题。目前，大多数系统采用滑动窗口方法进行分块，并通过调整窗口大小来优化分块效果。然而，这种方法存在一个主要缺点：检索到的往往是短且连续的文本块，这对于需要理解整个文档才能回答的问题来说，可能无法提供足够的信息。

论文指出，分块策略应该根据用户查询的具体需求来定制。有些查询可能只需要一个或几个特定的文本块，或者甚至只需要块中的某一部分。此外，文本通常包含多个主题，并且具有层次化的结构，这意味着分块策略需要能够捕捉到文本的这种复杂性。

为了解决这些问题，作者提出了RAPTOR（Recursive Abstractive Processing for Tree-Organized Retrieval）系统，这是一个树结构的索引和检索系统。RAPTOR的设计目的是捕捉文本的多尺度、不同层次的信息。系统通过对文本块进行总结，为大型语言模型（Large Language Models, LLMs）提供不同层次的信息，从而增强模型对复杂问题的理解和回答能力。

总结来说，RAPTOR系统通过其递归抽象处理和树形组织检索的设计，旨在改进现有RAG系统的分块策略，以更好地适应不同查询的需求，并提高对多层次、多主题文本的理解和检索效率。

RAPTOR方法介绍

RAPTOR系统通过在文本块的基础上构建一个递归的树结构来组织信息，树中的每个节点都代表了不同粒度的语义信息。

RAPTOR的基本流程

分块与嵌入
1. 文本首先按100个词的大小进行分块。如果某个句子超过100个词，则整个句子作为一个独立的文本块，以确保块内的语义连贯性。
2. 每个文本块都会被转换成嵌入向量（embedding）。
递归构建RAPTOR树
1. 文本块及其对应的嵌入向量作为树的叶子节点。
2. 通过聚类算法将相似的文本块聚集在一起。
生成总结与新的嵌入
1. 利用语言模型为每个簇内的文本生成总结。
2. 为每个生成的总结再次生成嵌入向量，并将这些总结作为树的新节点。
递归执行
1. 上述过程会递归进行，不断细化信息粒度，构建出具有多层次的RAPTOR树。
查询策略
1. Tree Traversal Retrieval：遍历RAPTOR树的每一层，剪枝并选择最相关的节点。
2. Collapsed Tree Retrieval：评估整个树中每个节点的相关性，找到最相关的几个节点。

总结来说，RAPTOR通过递归的方式构建了一个层次化的文本表示，每个层次都为不同粒度的信息提供了索引。这种结构不仅能够捕捉文本的细节，还能够抽象出更高层次的概念，从而提高检索系统的灵活性和效率。在查询时，RAPTOR提供了两种策略来检索相关信息，一种是逐层遍历，另一种是全局评估，这两种策略都旨在找到与查询最相关的文本块。

RAPTOR树构建流程

RAPTOR（Recursive Abstractive Processing for Tree-Organized Retrieval）树构建流程是一个旨在提高长文档检索效果和上下文理解能力的技术方案，具体步骤如下：

文本分块与初步嵌入
1. 首先，将检索语料库中的长文档分割成长度约为100个标记的连续文本块，若某句子超过这个限制，则将整句移到下一个块中，保证每个块内的语境和语义连贯性不受破坏。
2. 使用基于BERT的Sentence-BERT（SBERT）编码器（如multi-qa-mpnet-base-cos-v1版本）对这些文本块进行嵌入，将每个块转化为一个稠密向量，这些向量和对应的文本块共同构成树结构的基础——叶子节点。
层次聚类
1. 应用聚类算法对这些嵌入向量进行处理，将语义相似的文本块聚集在一起形成簇。
2. 选择合适的聚类方法（未指定具体类型，但通常可能包括K-means、层次聚类等）以确保相关的内容能够有效地组合在一起。
摘要生成与重新嵌入
1. 对每个形成的聚类，使用语言模型（如UnifiedQA等）生成一个对该簇内文本内容的简洁概括或摘要。
2. 生成的摘要随后再经过SBERT嵌入，形成新的向量表示。
3. 重复此嵌入、聚类、摘要生成的过程，直到无法进一步合理地进行聚类为止，由此形成一个多层、结构化的树状表示，其中包含了原始文档从详细信息到高级主题的递归抽象。
构建递归树结构
1. 经过上述循环，将低层次的摘要作为更高层次聚类的输入，逐级向上构建树的中间节点和根节点，使得树的每一层代表了不同程度的抽象级别。
两种不同的检索策略
1. 树遍历检索：从顶层节点开始，逐层向下遍历树结构，根据查询的相关性剪枝并选择每个层次的节点。
2. 折叠树检索：跨越所有层次一次性评估节点的相关性，找出最具相关性的节点。

通过这一系列流程，RAPTOR能够建立一个高度结构化、具备层次主题理解和细节保留特点的检索树，当面临复杂的问答和检索任务时，模型可以在树的不同层级中检索相关信息，从而大大增强了对长文档的理解和检索性能。此外，论文强调了RAPTOR在构建效率上的优势，无论是令牌消耗还是构建时间，都保持了与文档长度的线性关系，这意味着即使面对大规模、复杂的文本库，也能实现高效处理。

RAPTOR树的两种查询策略

RAPTOR树设计了两种不同的查询策略，分别是树遍历方法和折叠树方法，这两种策略都是用来在构建好的多层结构树中寻找与给定查询最为相关的文本信息。

树遍历方法
1. 在这一策略中，检索过程是从RAPTER树的顶层（即最抽象的层）开始，逐层向下搜索。
2. 对于每一层，计算查询向量与当前层节点（即该层的摘要或嵌入表示）之间的余弦相似度。
3. 选择相似度最高的前k个节点，并将它们标记为待进一步考虑的节点集合。
4. 在下一层中，仅针对上一步骤选出的节点的子节点重复上述相似度计算和选择过程，直到抵达叶子节点（最底层的原始文本块）。
5. 最后，将所有在遍历过程中选中的节点所对应的文本合并起来，形成针对查询请求的检索上下文。
折叠树方法
1. 折叠树方法不同于逐层遍历，它采取了一种更为综合的方式来评估树中所有节点的相关性。
2. 在这种方法中，不是分别考虑每一层，而是将整棵树的所有节点视为一个整体集合，一次性计算查询向量与所有节点（不论在哪一层）的相似度。
3. 基于这些相似度得分，不分层级地选择与查询最为匹配的节点。
4. 由于跳过了逐层筛选的过程，这种方法可能会更快地收敛到最有价值的信息，但也可能丧失了通过层次结构进行细粒度控制的优势。

通过这两种查询策略，RAPTOR可以根据应用场景和查询复杂度的需求，灵活地选择适合的检索路径，在RAPTER树的多层抽象表示中快速定位所需信息，从而在处理复杂问题和长文档检索时展现出卓越的性能。

整个流程很简洁，方法也很直观。其中比较重要的环节不言而喻：聚类。

RAPTOR系统在构建检索树的过程中采用了软聚类技术，让每个文本节点能够同时属于多个聚类，这充分考虑到实际文本数据中，单一文本片段常常含有与多种主题相关的多元信息，因此允许节点出现在多个摘要中。具体的聚类算法选用的是基于高斯混合模型（GMM）的方法，它具有灵活性，并且在一个概率框架下运行，假设数据点由多个高斯分布的混合生成。

为了克服高维文本嵌入空间中距离度量可能导致相似性估计不准确的问题，研究人员采用了统一流形近似和投影（UMAP），这是一种适用于高维数据降维的流形学习技术。通过调整UMAP中的最近邻数量参数，算法首先识别全局聚类，然后在这些全局聚类内部执行局部聚类，形成一个分层聚类结构，从而捕捉到文本数据之间从宽泛主题到具体细节的广泛联系。

在确定最优聚类数量时，使用了贝叶斯信息准则（BIC）。BIC不仅惩罚模型复杂度，还奖励拟合优度，依据BIC选择出来的聚类数目能够在平衡模型复杂性和数据拟合质量的基础上，有效地划分文本数据。一旦通过BIC得到最优簇的数量，接下来采用期望最大化（EM）算法来估计GMM的参数，包括各个高斯分布的均值、协方差矩阵以及混合权重。

尽管GMM中的高斯假设可能并不完美匹配文本数据的稀疏和偏斜分布特性，但实验观察显示，GMM在此情境下仍是一种有效的聚类工具。为进一步验证这一点，研究团队进行了消融研究，比较了GMM聚类与连续文本块直接摘要的效果，并在附录B中提供了详细对比结果。

总结

对于查询，在不考虑性能和其他问题的条件下，最理想的块应该是文本中与查询最相关的那段话、几句话甚至一句话，类似抽取式的文本摘要。

以下是RAPTOR模型构建和检索过程中各个步骤的具体做法：

文本分块与嵌入
- 将原始文档划分为若干个文本块。
- 使用预训练的BERT变种SBERT（multi-qa-mpnet-base-cos-v1）作为编码器，将每个文本块转换为固定维度的密集向量表示，这有助于保留文本块的语义和上下文一致性。
层次聚类与摘要生成
- 根据文本块的向量表示应用高斯混合模型（GMM）进行软聚类，允许单个文本块归属于多个聚类中心，体现了文本块内容通常与多个主题相关的特点。
- 对于每个形成的聚类簇，使用大型语言模型进行摘要生成，产生对簇内文本内容的概括性描述。
- 摘要文本再次被转换为向量表示，并进入下一个迭代循环，继续进行嵌入、聚类和摘要生成，直到无法进一步有效聚类为止。
构建RAPTOR树
- 这一系列嵌入、聚类和摘要操作形成了一个多层级的树状结构。其中，底层包含原始文本块及其向量表示，上层则是对下层节点内容的不断抽象和总结，直至达到最顶层根节点。
- 在树的构建过程中，通过对不同参数（如局部和全局邻域数量）的调整以及使用贝叶斯信息准则（BIC）来确定最优聚类数目，确保算法既能捕获文本数据间的广泛关系，又能关注具体细节。
检索策略
- 树遍历检索：针对查询，从树的顶层开始，按层次逐级向下检索，基于查询向量与各节点的余弦相似度选取每层的top-k节点。这个过程是递归进行的，形成多个集合S1到Sd，最终将这些集合拼接起来以提供与查询相关的上下文信息。
- 折叠树检索：另一种方式是将多层级的树结构“折叠”成单层，计算查询向量与所有节点的相似度，直接选择最高相似度的top-k节点。
效率评估与优化
- 对于RAPTOR树构建过程的计算效率，进行了实证研究，结果表明其时间和令牌消耗随着文档长度呈线性增长，证明了该系统在普通硬件上的可行性和可扩展性。