本文翻译自 Neon 博客 20x faster than pgvector: introducing pg_embedding extension for vector search in Postgres and LangChain,作者 Raouf Chebri。
Neon 前不久发布了适用于 Postgres 和 LangChain 的 pg_embedding 扩展,相比于 pgvector,它将 Postgres 中基于图形的近似最近邻搜索速度提升了 20 倍,准确度达到了 99%。
虽然使用 IVFFlat 索引的 pgvector 扩展一直是一个受欢迎的选择,Neon 新推出的 pg_embedding 扩展使用 Hierarchical Navigable Small Worlds (HNSW) 索引,在高维相似性搜索中提供了更高的效率。
你可以轻松在应用程序中集成 pg_embedding 实现向量搜索,对向量索引有先验知识不是一定的。运行以下命令开始使用:
CREATE EXTENSION IF NOT EXISTS embedding;
CREATE INDEX ON items USING hnsw (embedding) WITH (maxelements = 1000000, dims=1536, m=32);
你还可以用 LangChain 中的扩展 PGEmbedding 向量存储:
from langchain.vectorstores import PGEmbedding
db = PGEmbedding.from_documents(
embedding=embeddings,
documents=docs,
collection_name="state_of_the_union",
connection_string=CONNECTION_STRING,
)
db.create_hnsw_index(max_elements= 10000, dims = 1536, m = 8, ef_construction = 16, ef_search = 16)
如果你对 IVFFlat 和 HNSW 在基于 Postgres 应用中的内部工作原理和差异感兴趣,我们在 Neon Postgres 上进行了基准测试,以比较他们的性能,请继续阅读!
基准测试结果
基准测试使用 GIST-960 Euclidean 数据集,比较了使用 pg_embedding 的 HNSW 以及使用 pgvector 的 IVFFlat 索引的性能。该数据集提供了一个包含 100 万个 960 维向量的训练集,以及一个包含 1000 个向量的测试集。每次搜索返回 k=100 个向量。
为了提高效率,HNSW 索引存储在内存中。为了进行可比较的测试,我们使用 pg_prewarm 扩展将 IVFFlat 索引也存储在内存中:
SELECT pg_prewarm('3448836', mode => 'buffer');
下图采用对数刻度。pg_embedding 在相同的召回内执行速度快 5-30 倍。如果想要达到更高的准确性,IVFFlat 的执行时间会更长。这是因为在 pgvector 和 IVFFlat 中需要更多探测来实现更大的召回。然而,探测数量越大,搜索就越快地收敛到顺序扫描。
结果和分析
指标
我们根据以下指标比较了两个扩展的性能:
- 执行时间:以执行 100个 最近邻搜索查询所需的平均时间来衡量。
- 准确率/召回率:以每个查询返回的真实最近邻的比例来衡量。
设置
- 数据集:我们使用 GIST-960 Euclidean 数据集进行基准测试。该数据集被广泛用于基准测试相似性搜索算法。
- 环境:基准测试在一个具有 4 个虚拟处理器 和16GB RAM 的 Neon 实例上进行。
配置:
- 对于 pgvector,我们在 [1, 2, 10, 50] 范围内变化 probes 参数,并在 [1000, 2000] 范围内变化 lists 参数。
- 对于 HNSW,我们用了 m 为 [32, 64, 128],efConstruction 为 [64, 128, 256],efSearch 为 [32, 64, 128 ,256 ,512]。
为什么 HNSW 比 IVFFlat 快?
我们先了解一下 IVFFlat 和 HNSW 是什么以及它们的原理。
使用 pgvector 的 IVFFlat
用 pgvector 可以在数据库中直接进行向量搜索。其中一种索引技术叫做 IVFFlat。IVFFlat 索引将数据集分成多个簇,并为每个簇维护一个倒排列表。在搜索过程中,只有选定数量的簇被检查,与平坦索引相比,这大大加快了搜索过程。
默认情况下,IVFFlat 探测 = 1,这意味着索引将在最近的簇质心(或列表)中进行搜索。然而,由于查询向量更接近簇的边缘,这可能导致不准确性。
提高准确性的一种方法是增加探测次数。最佳列表和探测次数分别为 sqrt(rows) 和 sqrt(lists),从而得到时间复杂度为 O(sqrt(rows))。下图显示随着探测次数增加,召回呈指数增长。
HNSW 和 pg_embedding 组合
HNSW(Hierarchical Navigable Small World)最初由 Yu A Malkov 和 Dmitry A Yashunin 在他们的论文中引入,题为:Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs。
HNSW 是一种基于图的高维数据索引方法。它构建了一个图层次结构,其中每个层次都是前一个层次的子集,从而导致时间复杂度为 O(log(rows))。在搜索过程中,它通过这些图进行导航以快速找到最近邻居。
HNSW 算法中有三个主要参数:
- m: 该参数表示在构建图形过程中,每个新元素创建的双向链接的最大数量。
- efConstruction: 该参数用于索引构建阶段。较高的 efConstruction 值会创造更高质量的图和更准确的搜索结果。然而,这也意味着索引构建需要更长时间。下图表显示了使用 GIST-960 Euclidean 数据集时根据 ef_construction 值的构建时间。
- efSearch: 此参数用于搜索阶段。与 efConstruction 类似,较大的 efSearch 值会在增加搜索时间的代价下提供更准确的搜索结果。该值应等于或大于 k(你想要返回的最近邻居数量)。
- 默认情况下,efConstruction 和 efSearch 均为 32,但可以在创建 HNSW 索引时修改它们的值:
CREATE INDEX ON vectors_hnsw USING hnsw (vec) WITH (maxelements = 1000000, dims=1536, m=32, efconstruction=32, efsearch=32);
你该选哪个索引?
我们根据以下五个指标比较了这两个索引:
- 搜索速度
- 准确性
- 内存使用量
- 索引构建速度
- 距离度量
在 pg_embedding 和 pgvector 之间的选择取决于你的具体用例和需求:
- 内存限制 (pgvector):如果你在严格的内存限制下工作,可以选择 IVFFlat,因为它通常比 HNSW 消耗内存更少,代价是搜索速度和准确性。
- 搜索速度 (pg_embedding):如果你主要关注快速检索最近邻的速度,特别是在高维空间中,由于其基于图形的方法,pg_embedding 很可能是更好的选择。
- 准确性和召回 (pg_embedding):如果对你的应用程序来说实现高准确性和召回至关重要,则 pg_embedding 可能是更好的选择。与 IVFFlat 相比,HNSW 基于图形的方法通常产生更高水平的召回。
- 距离度量 (pgvector):无论是 pgvector 还是 pg_embedding 都支持 L2 距离度量。此外,pgvector 还支持内积和余弦距离。
结论
借助 Postgres 的 pg_embedding 扩展,你现在可以在 Postgres 中高效处理高维向量相似性搜索。HNSW 算法基于图形结构,相比 IVFFlat 索引,在搜索速度、准确性和设置方面具有多个优势。
使用 pgvector 的 IVFFlat 索引在内存限制严格的应用中仍然是一个可行的选择,但会牺牲召回率。
当然,选择 pg_embedding 还是 pgvector + IVFFlat 组合取决于你的具体需求。可以都尝试一下,找到最适合你的。Neon 也期待看到你使用 pg_embedding 开发出的创新应用和反馈!
你可以访问官网:https://www.bytebase.com/,免费注册云账号,立即体验 Bytebase。