Scaling HNSWs

Salvatore Sanfilippo

扩展 HNSW

原文由 Salvatore Sanfilippo 发布,订阅该博客

我暂停了几周的 HNSW 开发(现在正在做另一种数据结构,很快会有消息)。到目前为止,我为 Redis 新增的类型已经足够稳定和完整,正是时候回顾一下我对 HNSW 的所学所悟,并把它写成一篇博客。算是那种在 AI 时代之前很常见的“脑内倾倒”,如今或许已经变得少见了。经过近一年对 HNSW 和向量相似度相关工作的思考与实现,也是时候动笔写点东西了。不过这篇不会是 HNSW 的入门介绍——这类文章已经太多了。这篇要讲的是“更进一步”的内容。如果你已经了解 HNSW,我想分享一些更“进阶”的发现,尤其是在如何让它足够快、从而带来“Redis 式”体验的背景下:大家都知道,Redis 追求低延迟和高性能,而 HNSW 多少有点“抗拒”这一点,所以要把 HNSW 作为一种抽象数据结构暴露出来,遇到了不少挑战。

这篇博客会分成几个小节。可以把它们看作同一本书的不同页、同一段经历的不同章节。顺便说一句,我其实已经写过一遍这篇文章,后来又弄丢了 :D [关于 macOS 和坏习惯的一段漫长而悲伤的故事——自从 90 年代停电丢稿之后,我就再没丢过这种东西],所以这次主要的难题是回忆几天前写过的内容,顺便把之前不太满意的地方重新组织一下。

关于 HNSW 现状的几句话

在深入 HNSW 的内部原理和优化之前,我想先聊聊 HNSW 的现状。提出 HNSW 的原始论文是一篇非常出色的计算机科学文献,HNSW 本身也是极其精妙的数据结构,但是:我并不认为它是基于距离函数以贪心方式搜索近邻向量的最终答案。这篇论文给人的感觉是少了几块“拼图”,好像研究者如果再多给半年时间,还有很多可以探索和阐述的东西。举个例子,我自己就对论文做了扩展,以支持真正的条目删除,而不是仅仅打上墓碑标记、稍后回收的那种“假删除”——论文里完全没有涉及删除。同样,目前也有人在认真验证 HNSW 中那个“H”(分层)是否真的必要,单层的扁平结构是否也能达到差不多的性能(希望以后能多聊聊这个:我的直觉是答案在中间,比较合理的做法是修改层级选择函数,只保留高于某个阈值的层)。

说这些是想表达,如果你热衷于数据结构研究,HNSW 的演进和改进是一个很有前景的方向,不必局限于“把它搬到磁盘上”这类思路(可以参考微软的相关工作)之类的演进。好了,前言就到这里,我们进入正题,聊聊底层的细节吧 :)

扩展内存

Redis 是内存系统,而 HNSW 和向量都有一个不幸的特点:非常占空间。原因有三:1. HNSW 有大量指针,比如 16、32 甚至更多指向邻居节点的指针(这是 HNSW 的一个可调参数)。2. HNSW 有很多层,它本质上是一种类似跳表的数据结构。这让第一个问题更加严重。3. HNSW 的卫星数据是一个浮点数向量,所以在原生情况下,每个分量占 4 字节,而通常一个向量会有 300 到 3000 个分量,这是常见范围。

那么,这里有什么经验教训?有人会压缩指针,因为在 64 位系统里,很多指针(8 字节)的高 4 字节很可能是相同的。这很聪明,我还没实现,因为在 Redis 里我需要追求速度,这是在空间和时间之间的权衡:也许值得,也许不值得。我会再深入研究一下。

不过,如果你算一下,多层结构其实并没有看起来那么糟糕。平均来说,每节点多层的额外开销大约只有 1.3 倍(如果层级选择函数中层级递增的概率是 0.25),因为很多节点其实只在第 0 层。但 1.3 终究比 1 大,如果 HNSW 里的那个“H”真的没那么有用……[剧透一下,我发现如果所有东西都在第 0 层,搜索时间会更长,贪心搜索的主循环会从更不理想的起点开始,虽然最终还是能到达正确的簇,但会花费更多计算时间。不过这只是初步结果。]

所以在这里,*真正*唾手可得的优化是:向量量化。我的发现是,如果使用 8 位量化,你能获得近 4 倍的速度提升,向量本身缩小 4 倍(但整个节点不会缩小 4 倍:指针还在那里,占很大空间),而在真实场景中的召回率几乎不变。这就是为什么 Redis Vector Set 默认使用 8 位量化。你可以通过 VADD 的选项指定使用全精度向量或二进制量化向量(只取符号位),但我对同时使用全精度和二进制量化都持怀疑态度。在谈论它们之前,先看看我对 8 位量化的具体做法。

我的做法是,对每个向量计算其分量绝对值的最大值(所以量化是按向量进行的),然后用有符号的 8 位数来表示从 -127 到 127 的量化值。这不如同时存储最小值和最大值那么精确,但在计算余弦相似度时会更快,因为我可以这样做:

/* Each vector is quantized from [-max_abs, +max_abs] to [-127, 127]
 * where range = 2*max_abs. */
const float scale_product = (range_a/127) * (range_b/127);

然后在整数域内做乘累加(实际上代码中的主循环是展开的,并使用了多个累加器,以让现代 CPU 更忙碌)

for (; i < dim; i++) dot0 += ((int32_t)x[i]) * ((int32_t)y[i]);

最后再回到浮点距离:

float dotf = dot0 * scale_product;

详情可以查看 vectors_distance_q8(),但我想你已经明白思路了:从整数量化域回到未量化的点积,只需要非常简单的操作。

所以,8 位量化非常划算,而全精度则是一个*必要*的功能,因为总会有人处理的向量生成方式使得每一个微小的精度差异都很重要(不过,对于学习得到的向量来说,并非如此……)但是,为什么要有二进制量化?因为我希望用户在*原始*信息本身就是二进制时,有一种简单的方式来节省空间。想象你有一组用户,他们拥有是/否这样的属性,你想找到相似的用户、物品之类。好吧:这就是应该使用二进制量化的场景,它同样只是 VADD 命令的一个选项。

扩展速度:多线程与局部性

哦,得跟你说说我自己:我其实不太喜欢多线程系统,只要单核能做很多事,再通过无共享架构去利用多核,我更倾向于后者。但 HNSW 不同。它们*很慢*,而且在大多数用例中几乎总是以只读方式被访问。正因如此,我的 Vector Set 实现是完全多线程的。不只是读,连写也是部分多线程的,你可能会好奇,在 Redis 这样的系统中,键可能被后台持久化进程、客户端等不同方式访问,这怎么可能不乱套。

好吧,先从读说起。只要没有人在写入数据结构,我们就可以派生线程去贪心地收集近邻向量,并把结果返回给被阻塞的客户端。不过,我的 HNSW 实现是从零开始写的,我是说,从用 vim 打开的空 C 文件开始,与其他系统常用的两种实现没有任何共享代码,所以有一些“新花样”。其中一个不同之处在于,为了避免重复访问已经访问过的节点,我在每个节点里存了一个叫“epoch”(世代)的整数,而不是用另一种数据结构(比如哈希表)来标记已访问节点。那种做法我觉得相当慢。而 epoch 则是节点本地的,全局数据结构会在每次搜索时递增 epoch。所以在每次搜索的上下文中,我们可以确保找到的 epoch 都 <= 当前 epoch,而当前 epoch 就可用来标记已访问节点。

但有了线程,就会有多个搜索同时进行!没错,我需要的是一组 epoch 数组:

typedef struct hnswNode {
    uint32_t level;         /* Node's maximum level */
    … many other stuff …
    uint64_t visited_epoch[HNSW_MAX_THREADS];
}

这就是你在 hnsw.h 里能看到的内容。这同样是空间换时间的权衡,而这一次,时间再次战胜了空间。

那么,多线程写入是怎么实现的?诀窍在于,HNSW 插入时,大量时间都花在寻找候选邻居上。所以写入被拆成了“读取阶段”和“提交阶段”,只有后者需要写锁,并且有一些技巧来确保如果在此期间 HNSW 发生了变化,第一阶段累积的候选会被丢弃,一些节点可能已经不再有效。还有另一个问题。如果用户在后台线程还在处理某个值的时候删除了对应的键怎么办?针对这种情况,我们有一个函数,会等待后台操作完成后再真正回收对象。有了这些技巧,在真实的向量负载下很容易达到每秒 5 万次操作,而且这些数字还是来自 redis-benchmark 本身,包含了所有开销。扁平的 HNSW 库本身的原始性能要高得多。

扩展内存:正确地回收

在讨论如何通过多实例来扩展 HNSW 到大规模场景,以及为什么 Redis Vector Set 要把底层数据结构直接暴露给用户之前(我相信程序员很聪明,不需要过度保护,但也不*仅仅*是出于这个原因),我想先回到内存这个话题,因为关于这一点有一个很有意思的故事。

大多数 HNSW 实现在从图中删除节点时无法直接回收内存。我认为主要有两个原因:

1. 人们对原始 HNSW 论文在某个特定点上存在误解:他们认为邻居之间的链接可以不是相互的。而他们之所以这样想,是有特定原因的。

2. 论文完全没有提及节点的删除以及在节点消失、连接网络出现缺口后如何修复图。

第一个问题我认为是论文表述不够清晰,加上人们在实现 HNSW 时会遇到一个具体问题:插入新节点并在现有节点中搜索合适的邻居时,候选节点往往已经拥有最大数量的出边。这时该怎么办?常见的做法是,从新插入的节点到那些已经“满”的候选节点建立单向链接。然而,当你需要删除节点时,就无法找到所有指向它的入边,也就无法真正回收内存。你只能打上删除标记,稍后再进行某种图重建来“垃圾回收”陈旧节点,有时甚至直接泄漏内存。

所以首先,我在 Redis 中的实现采用了不同的做法,强制链接是双向的。如果 A 链接到 B,B 就链接到 A。但是,考虑到 A 可能已经很忙,该如何做到这一点?这就进入了比较复杂的领域,但大致的做法是,使用启发式策略从现有节点中丢弃一些与其他邻居连接良好的边,如果我们的节点即使对目标节点来说也是更好的候选,就替换掉,而如果不是这样,也有其他办法确保新节点至少拥有最少数量的链接,同时始终尽量满足图的小世界特性。

这样一来,当 Redis 从 Vector Set 中删除一个节点时,总能找到并移除所有指向它的指针。然而,对于那些因此少了一条链接的剩余节点该怎么办?我的做法是在它们之间构建一个距离矩阵,尝试将旧节点的邻居们彼此连接起来,尽量最小化平均距离。基本上,对于矩阵中的每一对 i、j 节点,我们会计算它们连接的好坏程度(向量有多相似)以及连接它们对*剩余*可能配对的影响(因为如果连接了某两个特定节点,可能会剩下一些元素找不到好的配对)。构建完这个评分矩阵后,再进行贪心的配对步骤。

这套方法效果非常好,你可以构建一个包含数百万元素的大型 HNSW,之后删除其中 95% 的元素,剩下的图依然保持良好的召回率,不会出现孤立节点等问题。

这就是我所说的,HNSW 还有空间留给新的论文来继续这项工作。

将 HNSW 扩展到多进程

当我开始做 Redis Vector Set 时,在 Redis 生态中已经有了一种向量相似度实现,具体是作为 RediSearch 的一种索引类型,而这也是大多数人对 HNSW 的看法:一种对现有数据的索引形式。

然而我想为 Redis 提供一种完全不同的 HNSW 实现。猜猜是怎么暴露的?当然是作为一种数据结构。这说明了我的脑袋有多“Redis 化”——这么多年下来,或者也许一开始就是 Redis 化的,而 Redis 本身也是按我的脑袋塑造的,因为我立刻就构想出了如何设计一个直接向用户暴露 HNSW 的 Redis 数据结构,并且很困惑为什么 Redis 中对向量的处理不是这样做的。

同时,当我把设计文档交给 Redis 的同事时,我不能说他们立刻就觉得这是显而易见的事。我的逻辑是:向量就像 Redis 有序集合(Sorted Set)中的分数,只不过它们不是具有全序的标量分数。然而你可以 VADD、VREM 元素,然后调用 VSIM 而不是 ZRANGE 来获取*相似*的元素。这不仅在 API 层面说得通,我还认为 HNSW 是高度可组合的,并不绑定于某个特定用例(不特定于文本嵌入、图像嵌入,甚至不一定是*学习得到的*嵌入)。你只需执行:

VADD my_vector_set VALUES [… components …] my_element_string

所以无论你的分量里是什么,Redis 都不关心,当你调用 VSIM 时,它会返回相似的元素。

但这也意味着,如果你把同一用例的不同向量分散在不同的实例/键中,你可以用同一个查询向量对所有实例调用 VSIM,加上 WITHSCORES 选项(会返回余弦距离),然后在客户端合并结果,你就神奇地把数亿向量扩展到了多个实例上,把数据集拆成了 N 份[关于这种用例,一个有意思的点是,如果你的客户端库足够智能,你可以通过多路复用来并行查询这 N 个实例]。

把 HNSW 以这种原始方式暴露出来的另一个非常显著的好处是,你可以非常轻松地扩展写入。只需将你的元素哈希后对 N 取模,然后定位到对应的 Redis 键/实例。多个实例可以同时吸收(虽慢,但在 HNSW 标准下仍算快的)写入,并行化这个原本非常缓慢的过程。

这种暴露 HNSW 的方式在“缩小规模”上也同样意义重大:有时你想为每个用户/物品/产品/你正在处理的任何对象都维护一个 HNSW。如果你的 HNSW 是构建在某种东西之上的索引,这很难建模,但如果你的 HNSW 本身就是数据结构,那就很简单了。你可以为每个对象都创建一个 Vector Set 键,里面只放少量元素。当然,就像任何其他 Redis 键一样,你可以为键设置过期时间,让它稍后自动被删除。

这一切可以浓缩成一条我认为在行业里应该更受重视的规则:许多程序员都很聪明,如果你不是构建一个他们无法触及的黑盒魔法系统,而是把数据结构、权衡直接展示给他们,他们能构建出更多东西,以特定方式建模自己的用例。而你的系统也会更简单。

扩展加载速度

如果不使用多线程,我的 HNSW 库以单线程处理 word2vec(每个向量 300 个分量)时,能以每秒 5000 个元素的速度向 HNSW 中添加数据,而查询构建好的 HNSW 时能达到每秒 9 万次查询。如你所见,两者差距巨大。

这意味着,从 Redis 的转储文件将包含数百万元素的 HNSW 重新加载到内存中会花费很长时间。而且这个时间也会影响复制。很不理想。但,只有当我们以最朴素的方式把元素从磁盘添加到内存——即在磁盘上存储“元素,向量”然后在内存中重建 HNSW——才会如此。这里有另一个值得吸取的教训。使用 HNSW 时,你需要按原样序列化节点和邻居关系,这样你就可以在内存中只需分配内存并把邻居 ID 转成指针来重建一切。这带来了 100 倍的速度提升。

但你真的以为故事到这里就结束了吗?嘿嘿。最近 Redis 拥有更强的安全特性,即使 RDB 文件被攻击者破坏,也要避免做危险的事情。所以我需要确保 HNSW 在加载后无论序列化数据结构中存在何种错误和损坏都是有效的。这涉及许多技巧,但我想在这里任性地直接贴出一段我写的注释,因为我觉得其中关于相互性校验的部分特别巧妙:

/* Second pass: fix pointers of all the neighbors links.
 * As we scan and fix the links, we also compute the accumulator
 * register "reciprocal", that is used in order to guarantee that all
 * the links are reciprocal.
 *
 * This is how it works, we hash (using a strong hash function) the
 * following key for each link that we see from A to B (or vice versa):
 *
 *      hash(salt || A || B || link-level)
 *
 * We always sort A and B, so the same link from A to B and from B to A
 * will hash the same. Then we xor the result into the 128 bit accumulator.
 * If each link has its own backlink, the accumulator is guaranteed to
 * be zero at the end.
 *
 * Collisions are extremely unlikely to happen, and an external attacker
 * can't easily control the hash function output, since the salt is
 * unknown, and also there would be to control the pointers.
 *
 * This algorithm is O(1) for each node so it is basically free for
 * us, as we scan the list of nodes, and runs on constant and very
 * small memory. */

扩展用例:JSON 过滤

我还记得第一个可用的 Vector Set 实现完成的那一天。一切都按预期工作,这是一个开始进行细化和添加额外功能的起点。

然而在过去的几周和几个月里,我在内部收到反馈说大多数用例都需要某种混合搜索:你想查询与给定查询向量相近的向量(比如与某部电影最相似的电影),但同时也要进行某种过滤(比如只限于 2000 到 2010 年间上映的)。我的感觉是,需要按不同参数查询的频率其实没有产品人员以为的那么高,而且在很多情况下,通过为每一年创建一个不同的向量集键可以更高效地实现(这又是将 HNSW 作为数据结构而非某种索引来暴露其可组合性的另一个体现)。

不过我开始思考 HNSW 贪心搜索的主循环,它大概是这样的:

// Simplified HNSW greedy search algorithm. Don’t trust it too much.
while(candidates.len() > 0) {
    c = candidates.pop_nearest(query);
    worst_distance = results.get_worst_dist(query);
    if (distance(query,c) > worst_distance) break;
    foreach (neighbor from c) {
        if (neighbor.already_visited()) continue;
        neighbor.mark_as_visited();
        if (results.has_space() OR neighbor.distance(query) < worst_distance) {
            candidates.add(neighbor);
            results.add(neighbor);
        }
    }
}
return results;

于是我开始琢磨为每个节点添加一组 JSON 元数据的想法。如果每个节点都有像 {“year”: 1999} 这样的信息,是否就足以在执行贪心搜索时进行过滤?当然,搜索需要是有界的,但这里有一个关键洞察:我首先想要的是*靠近*查询向量的元素,所以如果满足 JSON 属性条件的节点不多,我其实不需要遍历整张图。我会让用户来指定搜索的“投入”,而且无论如何,那些虽然匹配过滤条件但距离很远的结果本来就没什么用。

所以这也是我的 HNSW 与众不同的又一之处:它支持通过类似于在编程语言的“if”语句中编写的表达式来进行过滤。而 Vector Set 中的元素可以关联 JSON 块来描述其属性。然后你可以做这样的事情:

VSIM movies VALUES … your vector components here… FILTER '.year >= 1980 and .year < 1990'

关于内存占用的几句话

HNSW 理论上的致命问题是——它们通常需要在内存中提供服务。实际上,你也可以在磁盘上实现 HNSW,尽管从磁盘访问延迟的角度来看,有更好的数据结构。然而,在 Redis 和 Vector Set 的特定场景下,目标是提供一种非常快速、易于使用的东西:内存数据结构的灵活性对此很有帮助。所以问题归结为:内存占用真的那么糟糕吗?

将 300 万条 Word2Vec 数据以默认的 int8 量化方式加载到 Redis 中需要 3GB 内存,平均每条约 1KB。许多用例只有几千万条甚至更少。而你从实现良好的内存 HNSW 中获得的是非常好的性能,这对于本身就比较慢的数据结构和工作负载来说至关重要。在我的 MacBook 上,使用 redis-benchmark 对这个容纳了 word2vec 数据集的键执行 VSIM,能达到每秒 4.8 万次操作。我的感觉是,对于许多用例来说,内存 HNSW 的内存占用是完全可以接受的。即使在你希望将大部分向量放在磁盘上的用例中,即便要付出性能变慢的代价,你的热数据集很可能仍然应该由内存来提供。

这也是我认为投身于 HNSW 研究是个好主意的原因之一:我不认为它们会在短期内被大多数用例的其他方案所取代。更可能的是,我们会根据用例和数据规模,继续拥有分别适用于内存和磁盘的不同理想数据结构。此外,就在最近,哪怕只是浏览一下 Hacker News 的首页,也能看到有人用几百万条数据却在与比实际需要更慢或更复杂的系统作斗争。HNSW 并以正确的方式谨慎地暴露它们,可以避免所有这些麻烦。

结论

我喜欢 HNSW,实现它们的过程真的很愉快。我认为向量非常适合 Redis,即使在没有 AI 的世界里也是如此(比如几个月前,我就曾用它来为 Hacker News 用户做指纹识别,复现了过去在 HN 上发表的一项工作)。HNSW 对于众多用例来说实在是太酷、太强大了,而有了 AI 和学习型嵌入,这一切更是扩展到了无数的潜在用例。然而,和 Redis 中的大多数功能一样,我预计在人们意识到它们的有用和强大、以及如何使用它们之前,还需要很长时间(不,这不仅仅是 RAG 的问题)。Streams 也是如此:经过这么多年,终于迎来了大规模采用。

如果你对 HNSW 及其实现更感兴趣,我认为代码是相当易读的,并且有大量注释:

https://github.com/redis/redis/blob/unstable/modules/vector-sets/hnsw.c

如果你想了解更多关于 Redis Vector Set 的信息,请随意阅读我自己写的 README 文件。也有官方的 Redis 文档,但我建议你从这里开始:

https://github.com/redis/redis/tree/unstable/modules/vector-sets

感谢你读完这么长的博客文章!祝你有愉快的一天。

参考文献。本文是关于 HNSW 中“H”及其作用的论文 -> https://arxiv.org/abs/2412.01940

本文章由 muse-spark-1.2-contributor 进行翻译

评论