原理、代码解释与洞察 许多重要的现实任务,包括科学文献综述、法律案件简报和医学诊断,都需要跨片段或文档的知识理解。
现有的RAG方法无法帮助LLMs完成需要跨片段边界理解信息的任务,因为每个片段都是独立编码的。
本文将介绍四种创新方法,以增强文档或语料库的全局理解,并分享从中获得的洞察和思考。
这四种方法如下:
RAPTOR :这是一种基于树的检索系统,递归地嵌入、聚类和总结文本片段。
Graph RAG :该方法结合了知识图谱生成、社区检测、RAG和查询聚焦摘要(QFS),以促进对整个文本语料库的全面理解。
HippoRAG :这一检索框架受到人类长期记忆的海马索引理论启发,与LLMs、知识图谱和个性化PageRank算法协同工作。
spRAG :该方法通过两种关键技术,即AutoContext和相关段落提取(RSE),来提升标准RAG系统的性能。
RAPTOR:递归抽象化处理树形组织检索系统 RAPTOR 是一种新颖的基于树的检索系统,旨在递归地嵌入、聚类和总结文本段落。它自底向上构建树结构,提供不同层次的总结。
在推理过程中,RAPTOR 从该树中检索信息,结合长文档在不同抽象层次的数据。
关键概念 RAPTOR 采用递归方法,根据文本块的嵌入将其组织成簇。它为每个簇生成摘要,自底向上构建树结构。这一过程如图1所示。
以下我们将深入探讨与图1相关的具体主题:
构建 RAPTOR 树 文本分块
将检索语料库划分为每块包含100个词元的连续块。如果一个块超过100个词元 ,RAPTOR会将整个句子移至下一个块,以保持上下文和语义的一致性。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 def split_text ( text: str , tokenizer: tiktoken.get_encoding("cl100k_base" ), max_tokens: int , overlap: int = 0 ): """ 根据分词器和允许的最大词元数将输入文本分割成更小的块。 参数: text (str): 待分割的文本。 tokenizer (CustomTokenizer): 用于分割文本的分词器。 max_tokens (int): 允许的最大词元数。 overlap (int, optional): 块之间的重叠词元数。默认为0。 返回: List[str]: 文本块列表。 """ ... ... elif current_length + token_count > max_tokens: chunks.append(" " .join(current_chunk)) current_chunk = current_chunk[-overlap:] if overlap > 0 else [] current_length = sum (n_tokens[max (0 , len (current_chunk) - overlap):len (current_chunk)]) current_chunk.append(sentence) current_length += token_count ... ...
嵌入
使用Sentence-BERT 生成这些块的密集向量表示。
这些块及其对应的嵌入构成了RAPTOR树结构的叶节点 。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 class TreeBuilder : """ TreeBuilder类负责构建一个层次化的文本抽象结构,即“树”,使用摘要模型和嵌入模型。 """ ... ... def build_from_text (self, text: str , use_multithreading: bool = True ) -> Tree: """从输入文本构建一个黄金树结构,可选是否使用多线程。 参数: text (str): 输入文本。 use_multithreading (bool, optional): 创建叶节点时是否使用多线程。默认: True。 返回: Tree: 黄金树结构。 """ chunks = split_text(text, self .tokenizer, self .max_tokens) logging.info("创建叶节点" ) if use_multithreading: leaf_nodes = self .multithreaded_create_leaf_nodes(chunks) else : leaf_nodes = {} for index, text in enumerate (chunks): __, node = self .create_node(index, text) leaf_nodes[index] = node layer_to_nodes = {0 : list (leaf_nodes.values())} logging.info(f"创建了{len (leaf_nodes)} 个叶节点嵌入" ) ... ...
聚类方法
聚类在构建RAPTOR树中至关重要,它将文本段落组织成连贯的组。通过将相关内容聚集在一起,它增强了后续的检索过程。
RAPTOR的聚类方法 具有以下特点:
使用高斯混合模型(GMMs)和UMAP降维进行软聚类。
UMAP参数可调整以识别全局和局部聚类。
使用贝叶斯信息准则(BIC)进行模型选择,以确定最佳聚类数。
该聚类方法的核心是一个节点可以属于多个聚类。这消除了固定类别数量的需求,因为单个文本段通常包含多个主题的信息,从而确保其被包含在多个摘要中。
使用GMM对节点进行聚类后,每个聚类内的节点由LLM进行摘要。这一过程将大块内容转化为所选节点的简洁连贯摘要。
在实现中,使用gpt-3.5 turbo生成摘要。相应的提示如图2所示。
构建算法
至此,我们已经获得了整个树的叶节点并确定了聚类算法。
如图1中部所示,被分组在一起的节点形成兄弟节点,而父节点则包含该特定聚类的摘要。生成的摘要包含树的非叶节点。
摘要节点被重新嵌入,嵌入、聚类和摘要的过程持续进行,直到无法进一步聚类为止。这最终形成了一个结构化、多层次的原始文档树形表示。
相应的代码 如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class ClusterTreeConfig (TreeBuilderConfig ): ... ... def construct_tree ( self, current_level_nodes: Dict [int , Node], all_tree_nodes: Dict [int , Node], layer_to_nodes: Dict [int , List [Node]], use_multithreading: bool = False , ) -> Dict [int , Node]: ... ... for layer in range (self .num_layers): new_level_nodes = {} logging.info(f"构建第{layer} 层" ) node_list_current_layer = get_node_list(current_level_nodes) if len (node_list_current_layer) <= self .reduction_dimension + 1 : self .num_layers = layer logging.info( f"停止层构建: 无法创建更多层。树中总层数: {layer} " ) break clusters = self .clustering_algorithm.perform_clustering( node_list_current_layer, self .cluster_embedding_model, reduction_dimension=self .reduction_dimension, **self .clustering_params, ) lock = Lock() summarization_length = self .summarization_length logging.info(f"摘要长度: {summarization_length} " ) ... ...
检索过程 拥有了一棵 RAPTOR 树后,应如何利用它进行查询?
查询方式有两种:基于树遍历和基于压缩树,如图 3 所示。
树遍历从根层开始,根据与查询向量的余弦相似度检索前 k 个节点(此处为前 1 个)。在每一层,从前一层的前 k 个节点的子节点中检索前 k 个节点,相应代码 如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 class TreeRetriever (BaseRetriever ): ... ... def retrieve_information ( self, current_nodes: List [Node], query: str , num_layers: int ) -> str : """ 根据查询从树中检索最相关的信息。 参数: current_nodes (List[Node]): 当前节点的列表。 query (str): 查询文本。 num_layers (int): 遍历的层数。 返回: str: 使用最相关节点创建的上下文。 """ query_embedding = self .create_embedding(query) selected_nodes = [] node_list = current_nodes for layer in range (num_layers): embeddings = get_embeddings(node_list, self .context_embedding_model) distances = distances_from_embeddings(query_embedding, embeddings) indices = indices_of_nearest_neighbors_from_distances(distances) if self .selection_mode == "threshold" : best_indices = [ index for index in indices if distances[index] > self .threshold ] elif self .selection_mode == "top_k" : best_indices = indices[: self .top_k] nodes_to_add = [node_list[idx] for idx in best_indices] selected_nodes.extend(nodes_to_add) if layer != num_layers - 1 : child_nodes = [] for index in best_indices: child_nodes.extend(node_list[index].children) child_nodes = list (dict .fromkeys(child_nodes)) node_list = [self .tree.all_nodes[i] for i in child_nodes] context = get_text(selected_nodes) return selected_nodes, context
相比之下,压缩树将树压缩为单层,并根据余弦相似度 检索节点,直到达到阈值数量的令牌,相应代码 如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 class TreeRetriever (BaseRetriever ): ... ... def retrieve_information_collapse_tree (self, query: str , top_k: int , max_tokens: int ) -> str : """ 根据查询从树中检索最相关的信息。 参数: query (str): 查询文本。 max_tokens (int): 最大令牌数。 返回: str: 使用最相关节点创建的上下文。 """ query_embedding = self .create_embedding(query) selected_nodes = [] node_list = get_node_list(self .tree.all_nodes) embeddings = get_embeddings(node_list, self .context_embedding_model) distances = distances_from_embeddings(query_embedding, embeddings) indices = indices_of_nearest_neighbors_from_distances(distances) total_tokens = 0 for idx in indices[:top_k]: node = node_list[idx] node_tokens = len (self .tokenizer.encode(node.text)) if total_tokens + node_tokens > max_tokens: break selected_nodes.append(node) total_tokens += node_tokens context = get_text(selected_nodes) return selected_nodes, context
那么哪种方法更好呢?
RAPTOR 进行了比较,如图 4 所示。
如图 4 所示,使用 2000 个令牌的压缩树效果最佳。这是因为与树遍历相比,它提供了更多的灵活性。具体而言,通过同时搜索所有节点,它能够以适合给定问题的适当粒度级别检索信息。
图 5 展示了 RAPTOR 如何检索与灰姑娘故事相关的两个查询的信息:“故事的中心主题是什么?”和“灰姑娘是如何找到幸福结局的?”。
高亮节点表示 RAPTOR 的选择,而箭头指向 DPR(密集段落检索)的叶节点。重要的是,RAPTOR 提供的上下文通常包括 DPR 检索到的信息,无论是直接还是通过更高层次的摘要。
Graph RAG Graph RAG 利用LLM分两阶段构建基于图的文本索引:
首先,从源文档中提取知识图谱。
随后,为所有紧密关联的实体组生成社区摘要。
针对查询,每个社区摘要贡献部分回答。这些部分回答随后被聚合,形成最终的全局答案。
概述 图6展示了Graph RAG的流程。紫色框表示索引操作,绿色框表示查询操作。
Graph RAG利用特定于数据集领域的LLM提示来检测、提取和总结节点(如实体)、边(如关系)和协变量(如声明)。
社区检测用于将图划分为LLM可以在索引和查询时总结的元素(节点、边、协变量)组。
特定查询的全局答案是通过对与该查询相关的所有社区摘要进行最后一轮查询聚焦总结生成的。
图6中每一步的实现将在下文解释。值得注意的是,截至2024年6月12日,Graph RAG尚未开源 ,因此无法与源代码相关讨论。
第一步:源文档 → 文本块 块大小权衡是RAG长期存在的问题。
如果块太长,LLM调用次数减少。然而,由于上下文窗口的限制,理解和处理大量信息变得困难,这可能导致召回率下降。
如图7所示,对于HotPotQA 数据集,600个token的块大小提取的有效实体数量是2400个token块大小的两倍。
步骤2:文本块 → 元素实例(实体与关系) 该方法通过从每个文本块中提取实体及其关系来构建知识图谱,这一过程结合了大型语言模型(LLMs)和提示工程技术。
同时,Graph RAG采用多阶段迭代流程。这一流程要求LLM判断是否已提取所有实体,类似于二元分类问题。
步骤3:元素实例 → 元素摘要 → 图社区 → 社区摘要 在前一步骤中,提取实体、关系和声明实际上是一种抽象摘要的形式。
然而,Graph RAG 认为这还不够,需要使用 LLM 对这些“元素”进行进一步的摘要。
一个潜在的担忧是,LLM 可能不会总是以相同的文本格式提取对同一实体的引用。这可能导致重复的实体元素,从而在图中生成重复的节点。
这种担忧很快就会消失。
Graph RAG 采用社区检测算法来识别图中的社区结构,将紧密关联的实体纳入同一社区。图8展示了在 MultiHop-RAG 数据集中使用 Leiden 算法 识别的图社区。
在这种情况下,即使 LLM 在提取过程中未能一致地识别实体的所有变体,社区检测也能帮助建立这些变体之间的联系。一旦被归入一个社区,就意味着这些变体指的是同一实体内涵,只是表达方式或同义词不同。这类似于知识图谱领域的实体消歧。
在识别社区之后,我们可以在 Leiden 层次结构中为每个社区生成类似报告的摘要。这些摘要在理解数据集的全局结构和语义方面具有独立的价值。它们也可以用来理解语料库,没有任何问题。
图9展示了社区摘要的生成方法。
第四步:社区总结 → 社区答案 → 全局答案 我们现在来到了最后一步:基于上一步的社区总结生成最终答案。
由于社区结构的层级性,不同层级的总结可以回答各种问题。
然而,这引出了另一个问题:在多个层级的社区总结中,哪个层级能在细节和覆盖范围之间取得平衡?
Graph RAG 在进一步评估(Graph RAG 论文中的第3节)后,选择了最合适的抽象层级。
对于给定的社区层级,任何用户查询的全局答案都会生成,如图10所示。
HippoRAG HippoRAG 是一种新颖的检索框架,灵感源自人类长期记忆的海马索引理论。它与大型语言模型(LLMs)、知识图谱及个性化PageRank算法协同工作,模拟了人类记忆中大脑新皮层和海马体的不同角色。
关键理念 图11展示了人类大脑如何相对容易地处理知识整合的复杂任务。
海马体记忆索引理论 ,一种著名的人类长期记忆理论,为这种卓越能力提供了一种可能的解释。
具体而言,基于环境的、持续更新的记忆依赖于新皮层与C形海马体之间的相互作用。新皮层处理并存储实际的记忆表征,而海马体则维护海马体索引。这一索引是一组相互连接的索引,指向新皮层中的记忆单元并存储它们的关联。
在图11中,我们的目标是识别出一位参与阿尔茨海默病研究的斯坦福大学教授,从众多可能描述成千上万斯坦福教授和阿尔茨海默病研究者的段落中。
传统的RAG,独立编码段落,除非一个段落同时提及这两个特征,否则难以识别托马斯教授。
相比之下,熟悉这位教授的人可以迅速记住他,这得益于我们大脑的关联记忆能力,据信是由图11中蓝色所示的C形海马体索引结构驱动的。
受此机制启发,HippoRAG使LLMs能够构建和利用类似的关联图来管理知识整合任务。
概述 受图11启发,HippoRAG的每个组成部分对应于人类长期记忆的三个组成部分之一,如图12所示。
HippoRAG模拟人类长期记忆的三个组成部分,以模拟其模式分离和完成功能。
对于离线索引,LLM将段落处理成开放知识图谱(KG)三元组。这些随后被添加到人工海马索引中,同时合成旁海马区域(PHR)检测同义词。在上例中,HippoRAG提取涉及托马斯教授的三元组,并将其整合到KG中。
对于在线检索,LLM大脑皮层从查询中提取命名实体。旁海马检索编码器随后将它们与海马索引关联。HippoRAG利用个性化PageRank算法进行基于上下文的检索,并提取与托马斯教授相关的信息。
整体流程演示 以下是一个实际示例,介绍HippoRAG的流程。
图13展示了问题、答案及其支持段落和干扰段落。
图14描绘了索引阶段,包括OpenIE过程和相关知识图谱子图。
最后,图15展示了检索阶段,展示了查询命名实体识别(NER)、查询节点检索、个性化PageRank(PPR)算法对节点概率的影响以及顶级检索结果的计算。
接下来,结合源代码,我们将具体讨论HippoRAG如何构建长期记忆以及如何在两个方面进行检索。
构建长期记忆的方法 构建长期记忆的过程主要包含以下三个步骤。
首先,利用LLM通过OpenIE从检索语料库的每个段落中提取一组命名实体 ,如图16所示。
接下来,将这些命名实体添加到OpenIE提示中,以提取最终的三元组 ,如图17所示。
最后,使用经过微调的现成密集编码器来创建知识图谱 ,该图谱也将用于检索。
如何检索 首先,使用LLM从用户查询中提取一组命名实体 ,如图18所示。
然后,根据检索编码器确定的相似度,将这些命名实体链接到知识图谱中的节点。我们将这些选定的节点称为查询节点。
在海马体中,海马索引元素之间的神经通路使得相关邻域能够被激活并向上游回忆。
为了模仿这一高效的图搜索过程,HippoRAG利用了个性化PageRank(PPR)算法 ,这是一种仅通过一组用户定义的源节点在图上分布概率的PageRank版本。相应的代码 如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 def rank_docs (self, query: str , top_k=10 ): """ Rank documents based on the query @param query: the input phrase @param top_k: the number of documents to return @return: the ranked document ids and their scores """ ... ... if len (query_ner_list) > 0 : combined_vector = np.max ([top_phrase_vectors], axis=0 ) if self .graph_alg == 'ppr' : ppr_phrase_probs = self .run_pagerank_igraph_chunk([top_phrase_vectors])[0 ] elif self .graph_alg == 'none' : ppr_phrase_probs = combined_vector elif self .graph_alg == 'neighbor_2' : ppr_phrase_probs = self .get_neighbors(combined_vector, 2 ) elif self .graph_alg == 'neighbor_3' : ppr_phrase_probs = self .get_neighbors(combined_vector, 3 ) elif self .graph_alg == 'paths' : ppr_phrase_probs = self .get_neighbors(combined_vector, 3 ) else : assert False , f'Graph Algorithm {self.graph_alg} Not Implemented' fact_prob = self .facts_to_phrases_mat.dot(ppr_phrase_probs) ppr_doc_prob = self .docs_to_facts_mat.dot(fact_prob) ppr_doc_prob = min_max_normalize(ppr_doc_prob) else : ppr_doc_prob = np.ones(len (self .extracted_triples)) / len (self .extracted_triples) ... ...
最后,正如海马体信号向上游发送时所做的那样,HippoRAG聚合输出PPR节点概率 ,并将其用于对先前索引的段落进行排序以进行检索。
spRAG spRAG 是一种用于管理复杂查询的方法。它通过以下两种关键技术提升了标准 RAG 的性能:
AutoContext
相关段落提取(Relevant Segment Extraction,简称 RSE)
我们重点探讨 spRAG 如何跨块处理复杂查询。值得注意的是,目前关于 spRAG 的资料仅限于分析结合代码,尚无相关论文。
AutoContext:自动注入文档级上下文 在传统的RAG中,文档通常被分割成固定长度的块进行嵌入。这种简单的方法往往忽略了文档级的上下文信息,导致上下文嵌入不够准确和全面。
为了解决这一问题,开发了AutoContext。其核心思想是在嵌入每个块之前,自动将文档级的上下文信息融入其中。
具体来说,它会生成一个1-2句的文档摘要 ,并连同文件名一起添加到每个块的开头。这样一来,每个块不再是孤立的,而是携带了整个文档的上下文信息。获取文档摘要的代码如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 def get_document_context (auto_context_model: LLM, text: str , document_title: str , auto_context_guidance: str = "" ): max_content_tokens = 6000 text, num_tokens = truncate_content(text, max_content_tokens) if num_tokens < max_content_tokens: truncation_message = "" else : truncation_message = TRUNCATION_MESSAGE prompt = PROMPT.format (auto_context_guidance=auto_context_guidance, document=text, document_title=document_title, truncation_message=truncation_message) chat_messages = [{"role" : "user" , "content" : prompt}] document_context = auto_context_model.make_llm_call(chat_messages) return document_context
相关片段提取:智能组合相关文本块 RSE 是一个后处理步骤。其目标是通过智能识别和组合能够提供最相关信息的文本块,从而形成更长的片段。
具体来说,RSE 首先将检索到的内容相似或语义相关的文本块进行分组。然后,根据查询需求,智能地选择和组合这些文本块,形成最佳片段。相应的代码 如下所示。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 def get_best_segments (all_relevance_values: list [list ], document_splits: list [int ], max_length: int , overall_max_length: int , minimum_value: float ) -> list [tuple ]: """ 该函数接收文本块的相关性值,然后运行优化算法以找到最佳片段。 - all_relevance_values: 每个元文档的每个文本块的相关性值列表,每个外层列表代表一个查询 - document_splits: 表示每个文档开始的索引列表 - 最佳片段不会与这些索引重叠 返回 - best_segments: 表示元文档中最佳片段索引的元组列表(结束索引不包括在内) """ best_segments = [] total_length = 0 rv_index = 0 bad_rv_indices = [] while total_length < overall_max_length: if rv_index >= len (all_relevance_values): rv_index = 0 if len (bad_rv_indices) >= len (all_relevance_values): break if rv_index in bad_rv_indices: rv_index += 1 continue relevance_values = all_relevance_values[rv_index] best_segment = None best_value = -1000 for start in range (len (relevance_values)): if relevance_values[start] < 0 : continue for end in range (start+1 , min (start+max_length+1 , len (relevance_values)+1 )): if relevance_values[end-1 ] < 0 : continue if any (start < seg_end and end > seg_start for seg_start, seg_end in best_segments): continue if any (start < split and end > split for split in document_splits): continue if total_length + end - start > overall_max_length: continue segment_value = sum (relevance_values[start:end]) if segment_value > best_value: best_value = segment_value best_segment = (start, end) if best_segment is None or best_value < minimum_value: bad_rv_indices.append(rv_index) rv_index += 1 continue best_segments.append(best_segment) total_length += best_segment[1 ] - best_segment[0 ] rv_index += 1 return best_segments
洞察与思考 算法与数据结构的比较 RAPTOR 通过聚类构建树状数据结构,并基于此结构进行检索。
尽管 Graph RAG 和 HippoRAG 都采用知识图谱,但它们存在一些差异:
在数据结构方面,Graph RAG 通过汇总知识元素来整合信息。因此,每当新增数据时,汇总过程需要重复进行。RAPTOR 也是如此。相反,HippoRAG 只需在知识图谱中添加边即可无缝集成新知识。
在检索算法方面,Graph RAG 依赖于社区检测,而 HippoRAG 则使用个性化 PageRank(PPR)算法。
与其他算法不同,spRAG 不使用复杂的数据结构。它仅将文档摘要和文件名添加到每个块中,然后基于相关性值进行检索。这也意味着 spRAG 的索引和查询速度应该是最快的。
关于性能 HippoRAG 以 RAPTOR 为基准进行了实验,展示了一些超越 RAPTOR 的结果,如图 19 所示。
在 Graph RAG 论文 中,并未包含性能对比实验。
此外,目前尚无关于 spRAG 的论文可供参考。
关于增强范围 四种方法——RAPTOR、Graph RAG、HippoRAG 和 spRAG——旨在提升对整个语料库的理解。
它们各自基于整个语料库构建数据结构。
关于可定制性 在此背景下,HippoRAG 因其所有组件均为现成产品而表现更优,无需额外训练,如图 20 所示。
因此,通过微调特定组件,存在显著的改进潜力。
结论 本文介绍了四种新方法,以提升传统RAG在文档或语料库上的全球理解能力,并辅以代码解释。同时,也分享了我的见解和思考。
若您对RAG感兴趣,欢迎查阅我的其他相关文章。
最后,如有任何错误或遗漏,或您有任何想法愿意分享,请随时在评论区讨论。