比 pgvector 快 20 倍的 Postgres 向量运算插件：pg_embedding

本文翻译自 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%。

file

虽然使用 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 中需要更多探测来实现更大的召回。然而，探测数量越大，搜索就越快地收敛到顺序扫描。

file

结果和分析

指标

我们根据以下指标比较了两个扩展的性能：

执行时间：以执行 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 索引将数据集分成多个簇，并为每个簇维护一个倒排列表。在搜索过程中，只有选定数量的簇被检查，与平坦索引相比，这大大加快了搜索过程。

file

默认情况下，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))。在搜索过程中，它通过这些图进行导航以快速找到最近邻居。

file

HNSW 算法中有三个主要参数：

m: 该参数表示在构建图形过程中，每个新元素创建的双向链接的最大数量。
efConstruction: 该参数用于索引构建阶段。较高的 efConstruction 值会创造更高质量的图和更准确的搜索结果。然而，这也意味着索引构建需要更长时间。下图表显示了使用 GIST-960 Euclidean 数据集时根据 ef_construction 值的构建时间。

file

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);

你该选哪个索引？

我们根据以下五个指标比较了这两个索引：

搜索速度
准确性
内存使用量
索引构建速度
距离度量

file

在 pg_embedding 和 pgvector 之间的选择取决于你的具体用例和需求：

内存限制 (pgvector)：如果你在严格的内存限制下工作，可以选择 IVFFlat，因为它通常比 HNSW 消耗内存更少，代价是搜索速度和准确性。
搜索速度 (pg_embedding)：如果你主要关注快速检索最近邻的速度，特别是在高维空间中，由于其基于图形的方法，pg_embedding 很可能是更好的选择。
准确性和召回 (pg_embedding)：如果对你的应用程序来说实现高准确性和召回至关重要，则 pg_embedding 可能是更好的选择。与 IVFFlat 相比，HNSW 基于图形的方法通常产生更高水平的召回。
距离度量 (pgvector)：无论是 pgvector 还是 pg_embedding 都支持 L2 距离度量。此外，pgvector 还支持内积和余弦距离。