开源 AI 原生数据库 Infinity 0.2 release 正式发布,提供了 2 种新数据类型:稀疏向量 Sparse Vector 和 张量 Tensor,在此前的全文搜索和向量搜索之外, Infinity 提供了更多的召回手段,如下图所示,用户可以采用任意 N 路召回(N ≥ 2)进行混合搜索,这是目前功能最强大的 RAG 专用数据库。
为什么需要混合搜索(多路召回)?
我们知道,仅仅依靠向量搜索(默认情况下,它用来特指稠密向量)并不总能提供令人满意的结果。当用户问题中的特定关键词与存储的数据不准确匹配时,这种问题尤为明显。这是因为向量本身不具备精确语义表征能力:一个词,一句话,乃至一篇文章,都可以只用一个向量来表示,这时向量本质上表达的是这段文字的“语义”,也就是这段文字跟其他文字在一个上下文窗口内共同出现概率的压缩表示 ,因此向量天然无法表示精确的查询。例如如果用户询问“2024 年 3 月我们公司财务计划包含哪些组合”,那么很可能得到的结果是其他时间段的数据,或者得到运营计划,营销管理等其他类型的数据。
因此,在一种好的解决方案是,利用基于关键词的全文搜索提供精确查询,它跟向量搜索共同工作,这就是全文搜索 + 向量搜索 的 2 路召回,又被称为混合搜索(hybrid search)。
多路召回,在 RAG 的使用场景中,有时候还被解释为其他选择:
一种是仍然用向量搜索,但是采用多种方式将改写查询,然后合并多个查询的返回结果,这其实解决的仍然是向量本身语义表征粒度难以控制的问题,并不能解决向量无法进行精确查询的问题。
另一种就是引入稀疏向量,跟稠密向量组合到一起提供混合搜索(hybrid search)。稀疏向量跟稠密向量并不是一回事,它并没有稠密向量那种针对语义的压缩表示,而是试图针对全文搜索的一种替代,它解决的是全文搜索过程中如何针对倒排索引的词典进一步对关键词进行裁剪、扩展和定义权重。这样,一篇原始文档,就可以用这些裁剪后的关键词组成的稀疏向量来表征。例如下边的例子,上边是稠密向量,下边则是稀疏向量,它的维度一般要远高于稠密向量,例如会有 3 万维,由于大多数维度并没有值,因此可以采用 (位置,值)的形式表达向量中每个存在权重的维度。
把文本转为稀疏向量最知名和有代表性的工作就是 SPLADE (参考文献 [1]),它利用一个标准的预训练数据,将文档中的冗余词删除,并且增加扩展词,从而形成一个标准的 3 万维的稀疏向量输出。这里的冗余词删除,其实就类似传统搜索引擎中分词过程中的“去停用词”(在构建索引的过程中,将英文中的 the、a、等频率很高但信息密度很低的词跳过,这并不影响整体召回的效果,中文也同样的道理);增加相应扩展,也类似传统搜索引擎的查询同义词等扩展技术。站在使用的角度,它可以把任何文档都表征为一个 3 万维的稀疏向量,向量每个维度表征这个单词的权重。在典型的信息检索 ( Information Retrieval) 评测任务中,采用 SPLADE 稀疏向量取得了比传统搜索引擎基于 BM25 排序方式更好的表现。而利用稀疏向量 + 稠密向量的混合搜索,其效果在近期的一篇采用 BGE M3 embedding 模型的论文中也得到了验证(参考文献 [2]),在典型评测中,取得了比 BM25 要好很多的效果:
那么看起来,似乎稠密向量 + 稀疏向量就是更好的多路召回方案,采用全文搜索 + 稠密向量,似乎没有必要?
为什么需要三路召回?
我们知道, RAG 在技术上总是会遇到很多挑战,尽管信息检索理论为 RAG 提供了很多支撑,包括信息检索的评测等,但在实际中仍然会遇到很多问题。稀疏向量通过预训练模型删除了很多无用词,并增加了许多用来做查询扩展的词,这在通用查询任务上必然会表现更好,然而在实际使用中,依然有大量用户提问的关键词,并不在生成稀疏向量的预训练模型中,例如各种机器型号,说明书,专用词汇等等,用 3 万维表达全部关键词,还是多语言,依然会有诸多信息损失,此外还有各类业务所必须的短语查询等等,这些都是必须采用全文搜索才能实现的功能。因此,近期 IBM 的研究文章(参考文献 [3])对比了各种召回方式的组合,包括 BM25,稠密向量,BM25 + 稠密向量, 稠密向量 + 稀疏向量,以及 BM25 + 稠密向量 + 稀疏向量,最终得出结论:采用 3 路召回是所有组合中最适合 RAG 的选择,而 Infinity 已经完全内置了这种混合搜索功能。
3 路召回表现好非常容易理解,因为稠密向量可以表征语义, 稀疏向量可以在训练数据类似的场景下提供更好的精确召回,而关键词全文搜索则在各种场景下提供更加鲁棒的精确召回选择。3 种查询选择,一种都不能少,这使得 RAG 的方案设计更加复杂化,如果这些查询方式不能在一个数据库内完成,就需要用户组合多个数据库的 pipeline 来完成这一个功能,从而引入更多的工程 tricks 和复杂度,也影响了 RAG 技术的推广。例如:采用向量数据库提供稠密向量搜索和稀疏向量搜索,采用 Elasticsearch 提供关键词全文搜索,因此确保两者的数据同步就带来了技术挑战,如果一个文档在一个存储中存在,但在另一个存储中却不存在,就会带来一些错误的查询结果,所以可能需要一些其他 workaround :例如再引入一个 OLTP 数据库如 PostgreSQL 用来存放元数据,然后用一些对象存储来存放原始的文档数据,两者结合向量数据库和 Elasticsearch 来提供最终的数据同步,这是一个非常复杂的后端架构。而如果采用 Infinity ,就无需依赖这种复杂架构,3 种格式的存储,连同原始数据一起,一次全部插入且保证数据 ACID ,3 路召回的混合搜索,一条语句就可以完成,方便且高效。
如何排序?
三路召回完成后,一个直接的问题就是如何进行融合排序。Infinity 内置了多种融合排序算法:
Reciprocal Rank Fusion (RRF) 算法,它是这样工作的:为每路召回的结果列表中的每个文档都根据其排序位置分配一个分数,通常,得分是其排名的倒数。例如,排名第一的文档得分为 1,排名第二的得分为 0.5,排名第三的得分为 0.33,以此类推。那么最终文档的得分就是各路召回结果的累加。RRF 算法的好处在于鲁棒性,它的简单使得这种排序不容易过拟合,针对各种用户的不同场景无需大量参数调整就可以适应。
简单权重加权融合,RRF 算法非常鲁棒,但它完全按照各路召回的排名进行打分,丢掉了原始召回中的相似度信息。在某些情况下,仍然需要进一步控制。例如这样一个问题“请问型号为 ADX-156 的机器不能工作该如何处理”,我们需要对关键词得分提升权重。
基于外部模型的重排序,Infinity 原生支持基于 ColBERT 的重排序功能,关于这部分细节,我们在下文进一步讲解。
Infinity 提供了强大的排序组合机制,给用户提供多样化选择,如下图的 2 个例子:
第一种是 3 路召回的结果直接用 RRF 融合排序后返回。使用方式极其简单:
第二种向量搜索用 ColBERT 重排序,然后结果跟关键词全文搜索做权重叠加再返回。
这些召回和融合排序手段,使得 Infinity 可以为 RAG 提供最强大易用的多路召回能力。在多路召回之外,Infinity 0.2 release 还引入了 引入了一种全新的数据类型——Tensor,在计算机科学中,Tensor 可以用来表达多个向量,多维数组,或者一个矩阵,为什么要支持这种类型呢?这要从 ColBERT 谈起。
ColBERT(参考文献 [4])是一种排序模型,距离今天已经有四年了,是信息检索领域近年来引用次数非常多的知名论文。目前排序模型的架构有这样几类范式:
1. 双编码器。以 BERT 模型为例,它针对查询和文档分别编码,最后再经过一个 Pooling 层,使得输出仅包含一个向量。在查询时的 Ranking 阶段,只需要计算两个向量相似度即可,如下图所示。双编码器既可以用于 Ranking 也可以用于 Reranking 阶段。由于双编码器针对查询和文档分别编码,因此无法捕获查询和文档的 Token 之间的复杂交互关系。
2. 交叉编码器(Cross Encoder)。Cross-Encoder 使用单编码器模型来同时编码查询和文档,它能够捕捉查询和文档之间的复杂交互关系,因此通常能够提供更精准的搜索排序结果。Cross-Encoder 并不输出查询和文档的 Token 所对应的向量,而是再添加一个分类器直接输出查询和文档的相似度得分。它的缺点在于,由于需要在查询时对每个文档和查询共同编码,这使得排序的速度非常慢,因此 Cross-Encoder 只能用于最终结果的重排序。
3. 延迟交互模型( Late Interaction Model ),就是以 ColBERT 为代表的工作。它具备一些显著区分于其他排序模型的特点:其一是相比于 Cross Encoder,ColBERT 仍采用双编码器策略,将查询和文档分别采用独立的编码器编码,因此查询的 Token 和文档的 Token 在编码时互不影响,这种分离使得文档编码可以离线处理,查询时仅针对 Query 编码,因此处理的速度大大高于 Cross Encoder;其二是相比于双编码器,ColBERT 输出的是多向量而非单向量,这是从 Transformer 的最后输出层直接获得的,而双编码器则通过一个 Pooling 层把多个向量转成一个向量输出,因此丢失了部分语义。在排序计算时,ColBERT 引入了延迟交互计算相似度函数,并将其命名为最大相似性(MaxSim),计算方法如下:对于每个查询 Token 的向量都要与所有文档 Token 对应的向量进行相似度计算,并跟踪每个查询 Token 的最大得分。查询和文档的总分就是这些最大余弦分数的总和。例如对于一个有 32 个 Token 向量的查询(最大查询长度为 32)和一个有 128 个 Token 的文档,需要执行 32*128 次相似性操作,如下图所示。因此相比之下, Cross Encoder 可以称作早期交互模型 (Early Interaction Model)。
下图从性能和排序质量上,分别对以上排序模型进行对比,并且也包含了全文搜索。图中的 Dense Encoder 既为双编码器,代表普通的向量搜索,它既可以做 Retriever,也可以做 Reranker。由于 ColBERT 的延迟交互机制,它既满足了对排序过程中查询和文档之间复杂交互的捕获,也能实现较快的排序性能,相同数据规模下, ColBERT 的效率可达 Cross Encoder 的 100 倍以上,兼顾了性能与效果,因此 ColBERT 是一种非常有前景的排序模型。
尽管如此,在使用上,ColBERT 仍然面临 2 个问题:
尽管采用了 MaxSim 延迟交互相似度函数,使得效率大大高于 Cross Encoder,但相比普通向量搜索,计算开销仍然很大:因为查询和文档之间的相似度,是多向量计算,因此 MaxSim 的开销是普通向量相似度计算的 M * N 倍 (M 为查询的 Token 数, N 为 文档的 Token 数)。除此之外,原始的 ColBERT 在排序质量上相比 Cross Encoder 略有差距,针对这些,ColBERT 作者在 2021 年推出了 ColBERT v2 (参考文献 [5]),通过 Cross Encoder 和模型蒸馏,改进了生成的 embedding 质量,并且采用压缩技术,对生成的文档向量进行量化,从而改善 MaxSim 的计算性能。基于 ColBERT v2 包装的项目 RAGatouille (参考文献 [6])成为高质量 RAG 问答的解决方案。然而,ColBERT v2 只是一个算法库,端到端的让它在企业级 RAG 系统使用,仍然是一件困难的事情。
由于 ColBERT 是预训练模型,而训练数据来自于搜索引擎的查询和返回结果,这些文本数据并不大,例如查询 Token 数 32 , 文档 Token 数 128 是典型的长度限制。因此将 ColBERT 用于真实数据时, 超过限制的长度会被截断,这对于长文档检索并不友好。
基于以上原因, Infinity 在 0.2 版本中提供了 Tensor 数据类型,并基于此原生地提供端到端的 ColBERT 方案。
首先,Tensor 作为一种数据类型,ColBERT 编码输出的多向量,可以直接用一个 Tensor 来存放,因此 Tensor 之间的相似度就可以直接得出 MaxSim 打分。针对 MaxSim 计算,Infinity 给出了 2 种方案, 一种是 binary 量化,它可以让原始 Tensor 的空间只需原始尺寸的 1/32 , 但并不改变 MaxSim 计算的相对排序结果。这种方案主要用于 Reranker,因为需要根据前一阶段排序的结果取出对应的 Tensor 。另一种是 Tensor 索引, Infinity 采用 EMVB 技术(参考文献 [7])实现了 Tensor Index。EMVB 可以看作是 ColBERT v2 的改进,它主要通过量化和预过滤技术,并在关键操作上引入 SIMD 指令来加速实现。Tensor 索引可以用来服务 Retriever 而非 Reranker,因此结合 Infinity 的多路召回能力,用户可以进行如下各种召回选择:例如可以选择直接用 Tensor 提供语义搜索,从而实现比向量搜索更高的排序质量,也可以组合 Tensor 和全文搜索,用来做高质量的 RAG 所必备的 2 路召回,甚至可以组合向量搜索和 Tensor ,前者用来在大规模数据上粗筛,然后用 ColBERT 来快速精排,等等。Infinity 提供了足够强大的能力可以满足对于各种搜索召回的需求。
其次,针对超过 Token 限制的长文本,Infinity 引入了 Tensor Array 类型:
一篇超过 ColBERT 限制的文档,会被切分成多个段落,分别编码生成 Tensor 后,都跟原始文档保存在一行。计算 MaxSim 的时候,查询跟这些段落分别计算,然后取最大值作为整个文档的打分。
从 0.2 release 开始, Infinity 提供了内置的 Tensor 数据类型,并解锁了端到端的 ColBERT 应用,这使得这种以延迟交互模型为代表的排序模型,可以在较大规模数据上直接提供高质量的排序结果,对于提升 RAG 的检索质量具有非常重要的意义。
有了这么多召回手段,我们需要在真实数据集上进行相应的评测,以验证这些手段的效果。下边是 Infinity 在 MLDR 数据集上进行的评测结果,这也是 MTEB 默认采用的数据集之一。MTEB 是评估 Embedding 模型质量最权威的 Benchmark,目前排行榜上排名前列的模型基本都是基于 Cross Encoder 的编码器。
从图中看到,混合 BM25 全文搜索,可以比单纯向量搜索有显著的提升。而采用全文搜索 + 向量搜索 + 稀疏向量,就是 Blended RAG[参考文献 3],确实可以比单路搜索,以及两路混合搜索,有更好的查询质量。在 3 路混合搜索的基础上,进一步添加 ColBERT 做 Reranker,可以有进一步大的提升。同采用外部的 Reranker (例如 MTEB 排名前列的那些编码器)相比,采用 Hybrid search + ColBERT Reranker,它可以在数据库内部完成重排序,有着更高的效率,因此混合搜索可以进一步扩大 Top K 的范围(例如扩大到 Top 1000)之后再重排序,从而既保证最终召回质量还不影响性能,因此是一种性价比很高的高召回混合搜索方案。下图是各种召回方式添加 ColBERT Reranker 之后的提升效果总揽。
需要说明的是,在不同数据集上,相同的召回手段可能得到不同的返回结果,但有一点是确定的,就是混合搜索的手段越多,返回质量越好。此外,上边的评测并没有涵盖用 Tensor Index 做 Ranker 组合,这是因为在具体实验中,我们发现用 Tensor 做 Reranker 的性价比要高很多,这会在我们后边的文章中详细阐述。因此推荐的最佳混合搜索方案是 Blended RAG + ColBERT Reranker。
Infinity 0.2 release,不仅提供了行业最全的混合搜索能力,还提供最快的混合搜索能力。下文来描述 Infinity 如何做到这一点。
Infinity 是一款在存储引擎和执行引擎层面都精细设计的数据库。如下是 Infinity 的执行引擎工作流程,可以看到,在完成针对 API 的查询绑定后,接下来执行计划会被编译成一个流水线执行计划。这种机制,常见于一些现代数据仓库,所不同的是,数据仓库的流水线执行,通常服务于并行执行,而 Infinity 的流水线,则同时服务查询的并行和并发执行,需要保证了高并发执行时查询算子的最佳调度策略和 CPU 亲和性,避免了无效上下文切换导致的开销。这种设计,使得查询的端到端开销非常小,完整的查询延迟并不会比运行单独的算法库增加多少。
上图的右边是一个多路召回的查询样例,图中包含 2 个数据 segment 上的向量搜索执行算子,2 个算子并行执行,以及一个向量搜索的 Top K 合并算子负责合并来自 2 个 segment 的向量搜索结果;还有一个全文搜索算子,两路召回最后是一个 Fusion 融合算子,这些算子在内存中形成一个 DAG 图,由查询执行器负责运行期调度。
存储引擎方面,Infinity 建立了完整的以列存为基础的索引体系,对于多路召回的每一路,都有相对应的索引负责高性能检索,这也使得 Infinity 添加新的类型支持变得非常方便。因此,Infinity 可以看做是一个以列存为基础的全索引数据库,这跟近期 OpenAI 收购的 Rockset 有着相似的特性,而在索引的类型上,Infinity 则提供更加丰富的选择。
下边来看 Infinity 的索引实现。
跟许多向量数据库一样,默认情况下 Infinity 也采用了 HNSW 作为向量索引的实现。但 Infinity 的 HNSW 进行了一系列深入优化,具体来讲,就是对每个需要建立索引的向量进行局部自适应量化——通过对每个向量进行缩放和量化操作来提升搜索性能,使得相似性计算速度极快,有效带宽降低,同时减少内存占用,却几乎不影响准确性,因此实质上是这一种压缩技术。
具体的,Infinity 对每个向量采用了两级量化:其中一级量化是针对每个向量和全部向量的均值之间的差值进行量化编码。一级量化主要在 HNSW 图遍历期间使用,通过将向量压缩到较少的规模,从而有效减少实际消耗的内存带宽,提高搜索性能。二级量化负责对前述差值后的残差进行量化编码,它主要用于最后的相似度比对,提高查询精度。局部量化技术因为只针对每个向量进行量化,并不改变任何向量之间的最近邻关系,因此具备随机内存访问模式,所以特别适合基于图索引的相似度搜索。在 Infinity 中,HNSW 索引只基于一级量化的结果来构建,所以查询性能和内存占用都大大优于传统的 HNSW 索引。除了局部量化技术之外,Infinity 采用大量 SIMD 针对距离做加速计算,得益于这些设计,Infinity 的向量搜索性能超出同类许多,下图是 Infinity 和其他向量数据库的 benchmark 对比:
针对全文索引,Infinity 也是采取了全新实现的方案,而没有引入一些流行的全文索引库如 Tantivy,Lucene 等等。这是因为全文索引只是 Infinity 的一个组件,它需要紧密地跟存储引擎和执行引擎高效率协同工作。这体现在几个方面:
全文索引需要支持实时数据插入。
全文索引包含倒排索引和前向索引,而前向索引的能力,跟数据库的功能是重叠的,因此简单地整合,必然导致不必要的冗余。
全文索引还需要跟其他索引返回的结果一起做融合排序,这需要搜索的逻辑跟执行引擎紧密配合。
除却以上因素,Infinity 是一款专用于 RAG 的数据库,对于 RAG 来说,它需要根据用户的提问搜索到答案,由于用户的提问可能会比较长,因此在默认情况下,查询的关键词之间不能提供“AND”语义而应提供“OR”语义,否则很容易导致零召回。然而,“OR”语义对性能是极大地损害,因为任何一个关键词命中的结果都会被打分,并送到最终的结果排序,所以全文索引需要采用动态查询剪枝技术,减少不必要的打分和排序。
例如近十年来学术界最佳的动态查询剪枝方案,是以 WAND,MaxScore 等为代表的系列技术。尽管全文搜索是一个相对成熟的领域,然而在当下,也只有 Lucene,Tantivy 等少数全文索引库具备生产级的算法实现。Infinity 实现了完整的 Block Max WAND 和 Block Max MaxScore 技术,两种查询动态剪枝策略适应的场景略有不同,在默认情况下,Infinity 选择采用 Block Max WAND (参考文献 [9])来作为首选剪枝策略。
WAND 是 Weak AND 的缩写,它针对全文搜索最常见的打分手段 BM25 进行查询时动态剪枝,通过计算每个关键词贡献的上限来估计最终 Top K 结果的上限,并以此为阀值来决定在倒排索引的上如何快速跳过不必要的文档 ID,从而得到提速的效果。每个关键词贡献的上限,根据该关键词的的 IDF(在多少文档中出现) 和最大 TF(在文档中出现的最大词频) 来确定。
上图是 Infinity 和 Elasticsearch 的全文搜索性能对比,测试方法如下:
索引数据集为 wikipedia 33M,数据集大小 32GB。
从数据集中根据词频生成词表,按照 IDF 词频分布百分比分别随机选取关键词生成查询。查询长度从 3 个 Term 到 19 个 Term,生成的查询文件在这里_(https://github.com/infiniflow/benchmark/tree/main/enwiki_queries)_。Infinity 和 Elasticsearch 均采用默认 Top-K Union 的语义(OR)进行查询。
Infinity 和 Elasticsearch 均给予一定预热时间,使得索引数据尽可能缓存在操作系统的 pagecache 中。
可以看到,不论是长查询,还是短查询, Infinity 相比 Elasticsearch 均具备压倒性优势,并且在测试过程中 Infinity 的内存消耗仅有 Elasticsearch 的 1/2。因此,提供 RAG 所必备的混合搜索能力(全文搜索 + 向量搜索),此前用户的唯一选择是 Elasticsearch(包括 Opensearch),而现在不仅仅多了 Infinity 这个选项,而且在性能上也远远超过了这些选择。
针对稀疏向量索引,Infinity 采用了跟全文搜索类似的设计,都采用倒排索引 + 查询动态剪枝的策略,所不同的是,稀疏向量首先按照区块组织成前向索引,倒排索引只用来存放跟固定区块有关的信息,查询时用来从一个区块跳转到另一个区块,而具体的相似度计算,则通过前向索引来进行。因此,稀疏向量索引并没有包含一个标准的倒排索引,而是基于 Block 的倒排索引跟前向索引的混合方案。该具体算法来源于 SIGIR 2024 的 Best Paper Runner Up 论文(参考文献 9)。
下图是 Infinity 跟知名向量数据库 Qdrant 在稀疏向量索引上的性能评测:
由此可见,在稠密向量、稀疏向量、全文搜索三种召回手段上, Infinity 的性能均达到了极致,再加上强大的多路召回能力,以及各种的 Reranker 尤其是基于张量的 Reranker,可以说 Infinity 不仅仅是目前最快的 RAG 专用数据库,也是最强大的 RAG 数据库选择。欢迎关注和 Infinity :https://github.com/infiniflow/infinity
参考文献
SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval, https://arxiv.org/abs/2109.10086 , 2021
Bge m3-embedding: Multi-lingual, multi-functionality, multi-granularity text embeddings through self-knowledge distillation, https://arxiv.org/abs/2402.03216
Blended RAG: Improving RAG (Retriever-Augmented Generation) Accuracy with Semantic Search and Hybrid Query-Based Retrievers, https://arxiv.org/abs/2404.07220 , 2024
Colbert: Efficient and effective passage search via contextualized late interaction over bert, SIGIR 2020.
Colbertv2: Effective and efficient retrieval via lightweight late interaction, arXiv:2112.01488, 2021.
RAGatouille https://github.com/bclavie/RAGatouille
Efficient Multi-vector Dense Retrieval with Bit Vectors, ECIR 2024.
Ding, Shuai and Suel, Torsten. Faster top-k document retrieval using block-max indexes. SIGIR 2011
Mallia, Antonio and Suel, Torsten and Tonellotto, Nicola, Faster learned sparse retrieval with block-max pruning. SIGIR 2024
今日好文推荐
剥离几百万行代码,复制核心算法去美国?TikTok 最新回应来了
GitHub 删除代码等于“任何人均可永久访问”!微软回应:我们有意为之
中科大保卫处要求硕士以上学历,校方回应:偏技术型;字节跳动“代码抄袭”案在美获受理;私人文档被“投喂”豆包?官方否认 | Q资讯
程序员三个月前就攻破并玩透的 SearchGPT,OpenAI 可算发布了
评论