RAG 实战教程(三):向量数据库检索算法,KNN、IVF、HNSW 与 Faiss 实战

0 阅读13分钟

把文档转成向量并存进向量数据库以后,接下来就是检索。数据只有几十、几百条时,直接让查询向量和所有文档向量计算距离并没有什么问题。但知识库里有十万、百万个 Chunk 以后,每次查询都全部计算一遍,速度就会越来越慢。

在这里插入图片描述

所以,向量数据库除了要决定怎样计算两个向量的距离,还要解决另一个问题:怎样少算一些,又尽量不要漏掉真正相关的结果。

这就是向量检索算法要处理的事情。常见的方案包括暴力检索、IVF 和 HNSW。下面先把它们的原理说明白,再使用 Faiss 对三种索引做一次实际对比。

一、相似度算法和检索算法不是一回事

在看具体算法以前,需要先分清两个概念。

余弦相似度、欧氏距离和点积,计算的是两个向量有多接近。它们属于距离或者相似度的计算方式。

Flat、IVF 和 HNSW 解决的是去哪里找候选向量,以及要比较多少个向量。它们属于检索和索引算法。

比如向量数据库中保存了 100 万个文档向量,余弦相似度可以告诉我们查询向量和其中某个向量是否接近,但是它不会自动减少计算次数。如果想把 100 万次比较缩小到几千次,就需要另外建立索引。

所以,一次完整的向量检索通常包含两部分:

  1. 使用索引找到一批可能相关的候选向量。
  2. 计算查询向量与候选向量之间的距离,返回 Top K 结果。

二、KNN 和 ANN 有什么区别

KNN 是 K-Nearest Neighbors,也就是查找距离查询向量最近的 K 个向量。

最直接的做法,是让查询向量和数据库中的每一个向量都计算一次距离,排好序以后取前 K 个。这种方式得到的是精确结果,也经常被叫作暴力检索或者穷举检索。

ANN 是 Approximate Nearest Neighbor,也就是近似最近邻。它不会扫描全部向量,而是通过提前建立的索引缩小搜索范围。这样有可能漏掉少量真实近邻,但是查询速度会快很多。在这里插入图片描述

这里的近似并不是随便返回几个差不多的结果。IVF 和 HNSW 都有对应的参数,可以在召回率和查询速度之间调整。参数开得越大,搜索范围通常越广,结果更接近暴力检索,不过需要的时间也会增加。

三、Flat:把所有向量都查一遍

Flat 是最容易理解的一种索引。查询进来以后,它会依次计算查询向量和全部文档向量之间的距离,然后返回距离最近的 K 个结果。

在这里插入图片描述

假设数据库中有 N 个向量,每个向量有 D 个维度,一次查询大致需要进行 N × D 规模的计算。数据量增加一倍,需要比较的向量数量也会跟着增加一倍。

Flat 的优点有很多

返回的是精确结果,可以作为其他索引的对照标准。 不需要训练,创建索引后就能添加向量。 代码简单,数据量较小时用起来很方便。

它的问题也很明显。数据多起来以后,每次都扫描全部向量,查询时间会越来越长。

因此,Flat 比较适合数据量较小、查询次数不多,或者必须拿到精确近邻的情况。在测试 IVF 和 HNSW 时,也可以先用 Flat 算出标准结果,再检查近似检索找回了多少。

四、IVF:先分组,再到附近的组里查

IVF 的全称是 Inverted File Index,中文一般叫倒排文件索引。它的思路并不复杂,可以理解为提前给向量分组,查询时只进入附近的几个组。

1. 构建索引

IVF 通常会使用 K-Means 对向量进行聚类,得到若干个中心点,也就是质心。每个向量会被分配到距离自己最近的质心下面,形成多个倒排列表。

这里有一个重要参数 nlist,表示一共划分多少个列表。

例如,数据库中有 10 万个向量,nlist=100 时,可以理解为把这些向量大致分进 100 个组。当然,每个组里的数量不一定完全相同。

IVF 在添加数据以前需要先执行 train()。训练的目的不是训练 Embedding 模型,而是根据样本向量找到这些质心。

2. 查询索引

收到查询向量以后,IVF 会先判断它靠近哪些质心,然后只扫描这些质心对应的倒排列表。

控制扫描列表数量的参数叫 nprobe

假设 nlist=100nprobe=5,一次查询只会进入最接近的 5 个列表,再在这几个列表中做精确距离计算。其余 95 个列表不会参与这次查询,所以计算量会少很多。 在这里插入图片描述

不过,这也带来了漏检的可能。真正的最近邻如果被分到了第 6 个列表,而本次查询只检查前 5 个列表,它就不会出现在结果中。

因此,nprobe 越大,召回率通常越高,查询也会更慢。当 nprobe 接近 nlist 时,IVF 扫描的范围越来越接近全部数据,速度优势也会逐渐变小。

3. IVF 适合什么情况

IVF 的索引结构比较清楚,搜索范围也容易通过 nlistnprobe 控制,适合数据量较大、可以接受离线训练索引的场景。

它不是每增加一条数据都必须重新训练。已有质心仍然可以接收新向量,但是如果后来加入的数据和原来的数据分布差别很大,原有聚类可能会越来越不均匀,这时就需要考虑重新训练和构建索引。

五、HNSW:沿着图中的近路逐层查找

HNSW 的全称是 Hierarchical Navigable Small World,通常翻译为分层可导航小世界图。Chroma 默认使用的向量索引就属于 HNSW 这一类。

HNSW 会把向量组织成一张多层图。最底层保存全部节点,越往上节点越少,连接也更稀疏。查询时从最高层的入口点开始,沿着距离查询向量更近的节点移动,找到当前层较近的位置以后,再下降到下一层继续搜索。

这个过程有点像查路线。要从一个城市的某个街道去另一个城市的某个小区,一般不会从家门口开始检查全国所有道路,而是先确定城市之间的路线,到达目标城市以后再找城区和街道。HNSW 上层负责快速跨过较远的区域,底层再做更细的搜索。

在这里插入图片描述

1. M

M 控制图中每个节点建立多少条连接。

M 较大时,一个节点可以通向更多邻居,通常有利于召回率,但是索引会占用更多内存,构建时间也会增加。

2. efConstruction

efConstruction 控制构建索引时搜索候选邻居的范围。

值越大,建立连接时考虑的候选节点越多,图的质量通常会更好,不过创建索引会更慢。这个参数主要影响索引构建阶段。

3. efSearch

efSearch 控制查询时保留和检查多少个候选节点。

值越大,搜索范围越广,召回率通常越高,同时查询时间也会增加。如果发现 HNSW 返回的结果和 Flat 相差较多,通常可以先提高 efSearch 再测试。

HNSW 一般有比较快的查询速度,也支持继续添加新向量,但是它需要保存额外的图连接,因此内存占用通常高于 Flat 和 IVFFlat。在 Faiss 的 HNSW 实现中,删除向量也不是一个直接支持的操作。

六、三种检索方式放在一起看

对比项FlatIVFHNSW
检索类型精确检索近似检索近似检索
基本思路扫描全部向量聚类分组后扫描部分列表在多层图中沿邻居搜索
是否需要训练不需要需要训练质心不需要单独训练,但构建图较慢
构建速度中等通常较慢
查询速度数据越多越慢较快大规模数据下通常较快
内存占用保存原始向量保存原始向量和倒排信息保存原始向量和图连接
常用参数基本没有nlistnprobeMefConstructionefSearch
适合场景小数据、精确结果、基准测试大批量数据、离线构建索引查询较多、内存充足、需要持续添加数据

这张表只能说明大致差别。实际项目中,数据分布、向量维度、机器配置和参数都会影响结果,不能只根据算法名称判断谁一定更快。

在这里插入图片描述

七、使用 Faiss 对比 Flat、IVF 和 HNSW

Faiss 是一个用于稠密向量相似度搜索和聚类的库。它提供了多种索引,也可以把索引写入文件再重新读取。不过它本身不是完整的向量数据库,不负责文档字段、Metadata 过滤、权限和服务接口这些功能。

这里使用 Faiss,是因为它能在同一套接口下切换 Flat、IVF 和 HNSW,比较起来比较方便。

在这里插入图片描述

图中的柱形只表示需要对三种索引进行比较,不代表固定的测试结果,实际数据以本机运行结果为准。

先安装 CPU 版本:

pip install faiss-cpu numpy

下面准备 10 万个 128 维向量,再生成 100 个查询向量。三种索引使用完全相同的数据,Flat 的结果作为标准答案。

from time import perf_counter

import faiss
import numpy as np

DIM = 128
DATA_SIZE = 100_000
QUERY_SIZE = 100
TOP_K = 10

rng = np.random.default_rng(42)

# Faiss 接收 float32 类型的二维数组
database_vectors = rng.random(
    (DATA_SIZE, DIM), dtype=np.float32
)
query_vectors = rng.random(
    (QUERY_SIZE, DIM), dtype=np.float32
)

def recall_at_k(exact_ids, result_ids):
    """计算近似结果相对于 Flat 结果的平均 Recall@K。"""
    total = 0.0

    for exact, result in zip(exact_ids, result_ids):
        total += len(set(exact) & set(result)) / len(exact)

    return total / len(exact_ids)

def search_time(index, queries, k, repeat=5):
    """预热一次,并返回多次查询的中位数耗时。"""
    index.search(queries, k)
    times = []

    for _ in range(repeat):
        start = perf_counter()
        distances, ids = index.search(queries, k)
        times.append(perf_counter() - start)

    return float(np.median(times)), distances, ids

# 1. Flat:精确检索
start = perf_counter()
flat_index = faiss.IndexFlatL2(DIM)
flat_index.add(database_vectors)
flat_build_time = perf_counter() - start

flat_search_time, _, flat_ids = search_time(
    flat_index, query_vectors, TOP_K
)

# 2. IVF:聚类后只搜索部分倒排列表
nlist = 256
nprobe = 16

start = perf_counter()
quantizer = faiss.IndexFlatL2(DIM)
ivf_index = faiss.IndexIVFFlat(
    quantizer,
    DIM,
    nlist,
    faiss.METRIC_L2,
)
ivf_index.train(database_vectors)
ivf_index.add(database_vectors)
ivf_index.nprobe = nprobe
ivf_build_time = perf_counter() - start

ivf_search_time, _, ivf_ids = search_time(
    ivf_index, query_vectors, TOP_K
)

# 3. HNSW:使用多层图检索
M = 32
ef_construction = 100
ef_search = 64

start = perf_counter()
hnsw_index = faiss.IndexHNSWFlat(DIM, M)
hnsw_index.hnsw.efConstruction = ef_construction
hnsw_index.add(database_vectors)
hnsw_index.hnsw.efSearch = ef_search
hnsw_build_time = perf_counter() - start

hnsw_search_time, _, hnsw_ids = search_time(
    hnsw_index, query_vectors, TOP_K
)

print(
    f"{'Index':<10} {'Build/s':>10} "
    f"{'Search/ms':>12} {'Recall@10':>12}"
)
print(
    f"{'Flat':<10} {flat_build_time:>10.4f} "
    f"{flat_search_time * 1000:>12.3f} {1.0:>12.3f}"
)
print(
    f"{'IVF':<10} {ivf_build_time:>10.4f} "
    f"{ivf_search_time * 1000:>12.3f} "
    f"{recall_at_k(flat_ids, ivf_ids):>12.3f}"
)
print(
    f"{'HNSW':<10} {hnsw_build_time:>10.4f} "
    f"{hnsw_search_time * 1000:>12.3f} "
    f"{recall_at_k(flat_ids, hnsw_ids):>12.3f}"
)

运行以后会输出三种索引的构建时间、查询时间和 Recall@10。不同电脑得到的时间不会一样,所以这里不直接写一个固定结果。我们需要观察的是它们之间的关系:

  • Flat 构建很快,Recall@10 固定为 1,但是数据变多以后查询时间上升明显。
  • IVF 需要先训练和分组,查询通常比 Flat 快,不过召回率受 nprobe 影响较大。
  • HNSW 构建图需要更多时间和内存,但是数据规模变大以后,查询速度一般比较稳定。

测试时不要只看一次查询的耗时。第一次运行可能受到缓存、线程初始化和系统负载影响,所以代码中先预热一次,再取五次查询的中位数。

八、调整 IVF 的 nprobe

IVF 最常调的参数是 nprobe。可以继续使用上面的索引,分别测试几个值:

for nprobe in [1, 4, 16, 64, 128]:
    ivf_index.nprobe = nprobe

    elapsed, _, result_ids = search_time(
        ivf_index,
        query_vectors,
        TOP_K,
    )

    recall = recall_at_k(flat_ids, result_ids)

    print(
        f"nprobe={nprobe:<3} "
        f"search={elapsed * 1000:>8.3f} ms "
        f"recall={recall:.3f}"
    )

如果 nprobe=1,IVF 只进入一个倒排列表,查询会比较快,但是更容易漏掉分布在相邻列表中的向量。

逐渐提高 nprobe 后,Recall@10 通常会上升。它并不是越大越好,因为扫描的列表过多以后,IVF 会越来越接近暴力检索。

在这里插入图片描述

nlist 也需要结合数据量调整。分组太少,每个列表会装入大量向量;分组太多,训练和定位质心的成本会增加,而且数据不足时还可能出现一些很小的列表。比较稳妥的办法,是先选几个 nlist 建立不同索引,再分别测试 nprobe,而不是只改一个参数。

九、调整 HNSW 的 efSearch

HNSW 查询时最常调的是 efSearch

for ef_search in [16, 32, 64, 128, 256]:
    hnsw_index.hnsw.efSearch = ef_search

    elapsed, _, result_ids = search_time(
        hnsw_index,
        query_vectors,
        TOP_K,
    )

    recall = recall_at_k(flat_ids, result_ids)

    print(
        f"efSearch={ef_search:<3} "
        f"search={elapsed * 1000:>8.3f} ms "
        f"recall={recall:.3f}"
    )

efSearch 较小时,HNSW 在图中检查的候选节点较少,查询更快。提高以后,搜索会覆盖更多路径,召回率一般也会提高。

在这里插入图片描述

MefConstruction 需要重新构建索引才能观察效果,因为它们会影响图是怎样建立的。实际测试时,可以先固定 M=32efConstruction=100,把 efSearch 调到一个能接受的范围。如果召回率仍然不够,再考虑提高构建阶段的参数。

十、索引也可以保存下来

Faiss 索引构建完成以后,可以保存到本地,不需要每次启动程序都重新创建:

faiss.write_index(hnsw_index, "rag_hnsw.index")

loaded_index = faiss.read_index("rag_hnsw.index")
loaded_index.hnsw.efSearch = 64

distances, ids = loaded_index.search(
    query_vectors,
    TOP_K,
)

保存的是向量索引。如果还需要根据返回的 ID 找到原始 Chunk、文件名和页码,这些内容仍然要单独保存,或者交给完整的向量数据库管理。

在这里插入图片描述

十一、实际项目中怎么选

如果知识库只有几千个 Chunk,而且更在意结果是否准确,可以先用 Flat。它没有训练和调参成本,也方便检查整个检索流程有没有问题。

数据量大起来以后,可以拿 Flat 的结果作为标准答案,再测试 IVF 和 HNSW。不要只比较查询时间,还要一起记录 Recall@K、构建时间和内存占用。

IVF 更适合数据相对稳定、可以集中训练和构建索引的情况。HNSW 的查询性能通常更好,也方便继续添加数据,不过会使用更多内存,构建索引也更慢。

在 RAG 中,检索算法没有一个固定的最佳参数。比较实用的做法,是先确定可以接受的召回率,再在这个范围内找查询时间较低的配置。如果只是把参数调得很快,却把真正相关的 Chunk 漏掉了,后面的重排和大模型也补不回来。