颠覆传统检索：RAPTOR检索树提升检索准确率20%！

今天分享的是斯坦福大学在ICLR 2024上发表的一篇工作

论文题目：RAPTOR: RECURSIVE ABSTRACTIVE PROCESSING FOR TREE-ORGANIZED RETRIEVAL

论文链接：https://arxiv.org/abs/2401.18059

代码地址：https://github.com/parthsarthi03/raptor

论文概述

检索增强语言模型能够更好地适应世界状态的变化，并融入长尾知识。然而，现有的检索增强方只能检索几个简短的、连续的文本块，这对于需要整合文本多个部分的知识的问题是不够的，限制了它们表示和利用大规模语义结构的能力。这篇文章提出了一种新颖的方法——检索树，即考虑了广泛的主题理解，也考虑了细粒度的细节信息。通过递归地嵌入、聚类和总结文本片段，从底部向上构建一个具有不同总结层次的树，来解决阅读中的语义深度和连接性问题。在推理时，使用RAPTOR模型从这棵树中进行检索，在不同抽象层次上整合信息，以跨越较长文档进行理解。

论文核心

RAPTOR核心主要分为三个步骤：

• 文本切分与嵌入：首先将长文档切分成短小的文本块，然后使用预训练的语言模型(如SBERT)对每个文本块进行嵌入，得到其语义向量表示。这些文本块及其嵌入向量构成了树的叶子节点。
• 聚类与总结：使用高斯混合模型(GMM)对叶子节点进行软聚类，同一簇内的节点在语义上相近。接着利用LLM(如GPT-3.5)对每个簇的文本块进行总结，生成一个摘要节点，将大量的文本块转化为简明、连贯的摘要，涵盖所选节点的核心内容。这些摘要节点及其新的嵌入向量构成了树的中间层。生成摘要的提示词如下所示：

• 递归构建：重复步骤2，对上一层的摘要节点进行聚类和总结，直到最后只剩下一个根节点。这样就得到了一棵自底向上的树形结构，不同层次的节点代表了不同粒度的语义概括。

聚类算法

聚类算法采用的是高斯混合模型（Gaussian Mixture Models, GMMs），同时由于单个文本可能包含与多个主题相关的信息，所以这篇文章采用了软聚类，即节点可以同时属于多个聚类，而不需要事先设定聚类的固定数量，将它们包含在多个摘要中。

查询机制

• 树遍历查询

树遍历查询可以控制查询的深度和每层查询的节点数，并在树的深度下降时关注更精准的细节，该算法的步骤如下：

1. 从树的根层开始。计算查询嵌入与此初始层上存在的所有节点的嵌入之间的余弦相似度；
2. 根据最高的余弦相似度分数选择前 k 个节点，形成集合 S1；
3. 继续处理集合 S1 中元素的子节点。计算查询向量与这些子节点的向量嵌入之间的余弦相似度；
4. 选择与查询余弦相似度分数最高的前 k 个子节点，形成集合 S2；
5. 对 d 层递归地继续此过程，生成集合 S1， S2， . . . ， Sd；
6. Concatenate 将 S1 设置为 Sd，以将相关上下文组合到查询中。

• 压缩树查询

压缩树查询同时考虑树中的所有节点，将所有节点带到同一级别进行比较，该算法的步骤如下：

1. 首先，将整个 RAPTOR 树折叠成一个图层。这组新的节点（表示为 C）包含来自原始树的每一层的节点；
2. 接下来，计算查询嵌入与折叠集合 C 中存在的所有节点的嵌入之间的余弦相似性；
3. 选择与查询具有最高余弦相似度分数的前 k 个节点，继续向结果集添加节点，直到达到预定义的最大令牌数，确保不超过模型的输入限制。

• 树遍历 VS 压缩树 性能

本文做实验对比了树遍历方法在不同 k 值下，以及压缩树方法在不同最大令牌数下的性能表现。结果显示，压缩树方法的表现始终优于树遍历方法。这说明压缩树查询提供了更大的灵活性，同时搜索所有节点，压缩树方法能够根据具体问题检索到适当粒度的信息。

论文总结

论文介绍了 RAPTOR，这是一种新颖的基于树的检索系统，它通过各种抽象级别的上下文信息来增强大型语言模型的参数知识。采用递归聚类和汇总技术，RAPTOR创建了一个分层树结构，能够跨检索语料库的各个部分综合信息。在查询阶段，RAPTOR 利用此树结构进行更有效的检索。实验表明，使用递归总结的检索方法在多个任务上相较于传统的检索增强语言模型提供了显著的改进。在涉及复杂、多步骤推理的问题解答任务中，展示了最优的结果。例如，通过将RAPTOR检索与GPT-4结合使用，能够将QuALITY基准上的最佳性能提高20%的准确率。