RAG中向量索引的五种方法:基于百万个维基百科嵌入的测试
在百万个维基百科嵌入数据上测试的RAG向量索引五种方法:扁平索引、IVFFlat、HNSW、BM25混合索引以及DiskANN,并附有延迟与召回率数据。
来源: **Scaling RAG Infrastructure: The Journey from Flat Search to DiskANN**
RAG系统的成败取决于检索性能,而随着语料库规模的扩大,检索难度也会随之增加。本篇将通过实际实验对比五种向量索引方法,因此方法的选择应基于数据结果而非习惯。
什么是向量索引以及它为何重要
向量索引是一种组织大量向量的方式,使得相似向量能够被快速找到,而无需将查询项与每一个向量逐一比较。关系型数据库中有类似的机制,如分桶和Z排序,这些方法可以减少查询需要扫描的数据量。向量数据库则运用专门的结构来以同样的方式缩小搜索范围。
精确且低延迟的检索是任何可扩展的RAG系统的基础。现有的大型产品都已依赖这一功能:亚马逊从包含数百万种商品的目录中检索信息,Netflix为全球用户提供数千部个性化内容,X则处理数十亿条帖子和账号数据。如今数据量每天以TB计,因此扩展性绝不能被搁置,而必须从最初的设计阶段就纳入架构之中。
下面列出的每种方法都会根据三项标准进行评估:
- 延迟:索引返回前K个结果的速度。
- 检索精度:结果与精确搜索所得结果的匹配程度。
测试环境设置
实验使用了CohereLabs/wikipedia-2023-11-embed-multilingual-v3数据集的前一百万行,这一规模足以体现不同方法之间的差异。查询通过mxbai-embed-large进行嵌入处理,该模型可生成与数据集嵌入向量一致的1,024维向量。索引则是使用PostgreSQL中的pgvector功能构建的。
方法1:扁平搜索
扁平搜索,也称为暴力搜索,是最简单的方案。它计算查询与每个向量之间的相似度,然后保留最匹配的结果。
- 延迟:在所有方法中最为糟糕,因为需要对比每一个向量。延迟随向量数量的增加而呈线性增长。
- 准确性:最高,因为搜索是精确的,能够返回真正的最近向量。
- 维护成本:无需维护,因为不存在索引结构。
测试中的平均延迟为6.25秒。无论是大型零售商还是小型初创企业都会立即拒绝这种方案。
方法2:IVFFlat
Inverted File Flat是一种近似最近邻算法。它利用K-Means方法将向量空间划分为多个簇。该算法不会扫描所有数据,而是找出与查询向量距离最近的簇中心,然后仅在这些簇内进行搜索,并使用精确的距离计算方法。
其处理流程如下:
Query → closest centroids → search selected clusters → Top-K
在pgvector中,有两个参数控制这一功能。首先需要确定集群的数量,即“列表”的数量,实验中选择了1,000个列表。在查询时,另一个参数“探针”决定要搜索多少个列表。当探针值设置为10时,就会扫描那些质心距离查询点最近的10个列表。
延迟方面表现尚可。在拥有100万个向量及1,000个集群的情况下,查询完成时间最多为一秒。
准确率较差,原因有二。首先是边界问题。相似度是根据质心而非实际向量来计算的,因此有些距离查询点非常近的向量,如果所属集群的质心距离查询点较远,就会被忽略。其次,探针机制将搜索范围限制在固定数量的集群内,导致未搜索集群中的相关向量根本无法被考虑。
将探测数量从10增加到500虽然会提升延迟,但在本次测试中增幅很小。不过无法提前判断10个、100个、500个还是900个探测才能获得最佳检索效果,因为真正的邻近向量可能位于被跳过的数据簇中。使用全部1,000个探测反而失去了意义,因为这会让索引回归到暴力搜索的模式。
维护方面也存在问题。为超过一百万个向量构建索引大约需要2.37分钟,这个时间还算可以接受。但问题出在更新时——每当有新数据批次到来,基于质心的索引可能就需要重新计算,或者这些向量需要重新索引。对于数十亿个向量的情况,这样的工程开销可能就不值得了。
对于这种应用场景,延迟与精度之间的权衡实在难以接受。接下来的问题是,是否还有其他方法能够在降低延迟的同时提升检索效果。
方法 3:HNSW
用邮差找房子来做类比很合适。给定一个地址后,他不会搜索每一条街道,而是先走通往目标区域的主干道,这些道路上的节点虽少但覆盖距离较远;接着再转走主干路,那里的节点密度更高;最后在本地街道上行走,节点密度最大,直到找到确切的门口。
HNSW 即分层可导航小世界,其工作原理类似。它从上层较稀疏的长距离链接开始,逐步过渡到下层密度逐渐增加的链接,最终找到最近的向量。如果道路层级结构合理,搜索就能直接指向目标,而良好的结构设计也能提升检索准确率。
延迟方面同样表现良好,处理一百万个向量时最多仅延迟一秒。
准确度表现良好,但存在因构建参数带来的限制。索引是通过如下格式的指令创建的:
CREATE INDEX wikipedia_embeddings_hnsw_idx
ON wikipedia_embeddings
USING hnsw (emb vector_cosine_ops)
WITH (
m = 16,
ef_construction = 64
);
该算法基于语料库构建索引。有三个参数至关重要:
m表示图中每个向量可以连接的邻居数量,它控制着图的连通性与密度。ef_construction是构建索引时考虑的候选邻居数量,它决定了构建过程中的搜索范围。ef_search则用于限制查询时的搜索范围,即决定在回答查询时需要考察的候选邻居数量。
那么图是如何选择其 m 个连接的呢?对于一百万个向量而言,每个向量都有999,999个潜在邻居。构建器并非随机选择,而是根据距离以及某种邻居选择策略对候选邻居进行加权,最终保留大约16个对遍历图有用的邻居。
为何这样能提升准确性?
- 每个向量对应的少量有用邻居构成了图中的“高速路”。
- 选择这些连接是为了保持图的连通性并便于导航。
- 在查询时,搜索会逐步沿着这些路径寻找与查询向量更为相似的向量。
- 由于该图能提供通向相关区域的高效路径,因此更有可能找到真正的最近邻居。
- 因此,通过精心构建和导航可以提高准确性,同时避免进行全面的比较。
简而言之,该图是经过刻意设计而非随机生成的。通过设置m和ef_construction可以决定图的构建方式,进而影响召回率。
维护问题是HNSW的短板。为超过一百万个向量构建索引竟然需要惊人的37分钟。在数据更新时,如果初始数据集规模较大且能很好地反映整体数据分布,那么就可以持续添加新数据而无需重新构建索引。但随着新数据量相对于原始数据集的增加,数据分布可能会发生变化,图结构也会变得不那么理想,最终还是需要重新构建索引。这种情况虽然比IVF少见,但确实会发生。
由于该算法基于合理的构建原理,尽管需要额外的工程工作,它仍比IVF更优。其延迟与召回率的平衡表现虽不算出色,但也处于可接受水平。
方法4:BM25结合向量搜索
有时候,最显而易见的答案就是正确的答案。当人们将注意力集中在越来越复杂的方法上时,有一种已有三十年历史的技术却一直被忽视。BM25是一种词汇排序算法,它通过查询词在文档中出现的频率、这些词在整个语料库中的稀有程度以及文档的长度来为文档打分。
其工作原理如下:
- 文档被拆分成各个词,然后通过倒排索引将每个词与包含它的文档对应起来。
- 词频用于衡量某个词在特定文档中出现的次数。
- 逆文档频率用于衡量该词在整个语料库中的稀有程度或出现频率。
- 在某个文档中出现频繁但在整体上较为罕见的词,会获得更高的权重。
- BM25将上述因素与文档长度结合,从而计算出文档的相关性得分。
虽然其公式看起来并不复杂,但原理很简单:找到与搜索词最相关的文档。
一旦获取到这些文档,搜索范围就会大幅缩小。随后通过对这一小部分文档进行简单的向量搜索,即可找到最佳结果。这就是混合检索方式。
需要注意的一点是:第一阶段是按精确关键词进行匹配的,因此用户的拼写错误可能会影响查询结果。在搜索开始前通过人工智能辅助的标准化处理来修正拼写错误和变体,就能轻松解决这个问题。
延迟。向量部分能在预期的毫秒级时间内给出响应,但在这种配置下,词汇部分的处理时间则在12到14秒左右。虽然缓存可以解决这个问题,但首次请求的响应速度依然较慢。这里需要特别提到缓存,因为它在BM25中表现尤为出色:可通过匹配关键词直接检索到精确的文档。而对于其他索引方法,缓存则需要具备语义理解能力,这需要设定相似度阈值,同时也可能略微影响检索的准确性。
维护。无需维护,因为并未构建向量索引。
在实际应用中,对于那些文档数量在5万篇左右的多数AI产品而言,结合BM25的混合检索方式一直是最佳选择,它能在该规模下实现延迟与准确率之间的良好平衡。
方法5:DiskANN与Vamana图
这种方法更为抽象,类似于嵌入向量本身。其处理流程分为八个步骤:
- 首先使用那些需要高效搜索的大量向量集合。
- 通过将每个向量与精心挑选的一组邻居相连来构建图结构,这就是Vamana图。
- 将图存储在磁盘上而非内存中,这样就能处理远大于RAM容量的数据集。
- 当查询到来时,从图中的某个入口点开始。
- 沿着向量之间的连接逐步向靠近查询结果的区域移动。
- 随着搜索的进行,从SSD中加载所需的节点。
- 在探索了图中的足够部分后,筛选出一小部分有潜力的候选项。
- 对这些候选项进行更精确的比较,最终返回距离最近的K个结果。
关键在于导航图存储在SSD上而非RAM中,因此容量限制不如内存成本那么重要。
先进行元数据过滤
有一种技术值得特别提及,它应优先于上述任何方法。元数据过滤的原理与关系数据库中的分桶类似:在执行向量搜索之前,通过类别、日期或所有者等属性缩小候选集范围。
如何选择
- 暴力搜索:适用于小型到中型数据集,文档数量约在20,000份以内。在可接受的延迟前提下具有最佳准确率。
- BM25结合向量搜索:适用于中型数据集,文档数量大约在20,000到100,000份之间。在词汇检索与语义检索之间实现了良好平衡。
- HNSW:适用于包含数百万个向量或更多的大型数据集。在构建复杂度较高的前提下,仍能保持良好的延迟与准确率。