Scaling HNSWs

Salvatore Sanfilippo

擴展 HNSW

我暫時停下 HNSW(階層式可導航小世界圖)的開發幾週(現在正在處理另一種資料結構,近期會有消息)。到這個階段,我為 Redis 新增的這個新類型已經夠穩定、夠完整,正是好好整理我在 HNSW 上學到的一切、並寫成一篇部落格文章的好時機。就是那種在 AI 時代之前很常見、現在或許變得稍微稀有的那種腦內傾倒式分享。經過將近一年對 HNSW 與向量相似度相關技術的思考與實作,是時候動筆了。不過這篇文章不會是 HNSW 的入門介紹:這類文章已經太多了。這篇要走的是「更進一步」的那一哩路。如果你已經了解 HNSW,我想與你分享一些更「進階」的發現,特別是在如何讓它夠快、足以提供「Redis 體驗」的脈絡下:你知道的,Redis 的設計目標是低延遲與高效能,而 HNSW 在某種程度上正好與此背道而馳,因此要將 HNSW 作為一種抽象資料結構來提供,確實面臨了不少挑戰。

這篇部落格文章會分成幾個小節。把它們想成是同一本書的不同頁面、同一次經驗的不同章節。喔對了,順帶一提,我其實已經寫過一次這篇文章,然後又把它搞丟了 :D[一段關於 MacOS 與壞習慣的漫長悲傷故事——自從 90 年代停電時期以來,我就沒再發生過這種事],所以這次最大的挑戰,多半是要回想幾天前寫過的內容,並趁機把當時不太滿意的地方改寫得更好。

關於 HNSW 現況的幾句話

在深入探討 HNSW 的內部機制與最佳化之前,我想先談談 HNSW 的現況。最初發表 HNSW 的論文是一篇出色的電腦科學文獻,HNSW 也是非常出色的資料結構,但是:我不認為它們是以貪婪式搜尋、依距離函式尋找鄰近向量的最終定論。這篇論文給人的感覺是缺少了一些「拼圖」,彷彿研究人員如果再多給半年時間,還有許多值得探索與闡述的內容。舉例來說,我自己就對論文做了修改與擴充,以支援實際移除項目——真正的刪除,而不只是用墓碑標記(tombstone deletion)將元素標示為已刪除、留待日後回收:刪除項目這件事在論文中是完全缺席的。同樣地,目前也有一些研究正在認真檢驗 HNSW 中那個「H」是否真的必要,以及是否改用只有單層的扁平資料結構,效能也差不多(希望未來我能對此有更多著墨:我的感覺是真相介於兩者之間,比較合理的做法是修改層級選擇函式,讓它只產生高於某個門檻的層級)。

說這些是想表達,如果你投入資料結構的研究,我認為一個很值得投入的方向是去想像 HNSW 的演進與改良,而不要陷入一種想法,認為所謂的演進就只是:來把它搬到磁碟上做一遍(可參考 Microsoft 的相關嘗試)之類的。好了,前言就到此為止,來談談真正的底層細節吧 :)

擴展記憶體

Redis 是一個記憶體內系統,而 HNSW 和向量都不幸地非常耗費空間。原因有三:1. HNSW 有大量的指標,像是 16、32 或更多指向鄰近節點的指標(這是 HNSW 的一個可調參數)。2. HNSW 有很多層,身為一種類似 skiplist(跳躍串列)的資料結構,這讓第一個問題更加惡化。3. HNSW 的衛星資料是一個浮點數向量,所以在原始情況下,每個分量佔 4 個位元組,而通常一個向量會有 300 到 3000 個分量,這是常見的範圍。

那麼,這裡學到的教訓是什麼?有些人會壓縮指標,因為在 64 位元系統中,許多指標(8 個位元組)的高位四個位元組很可能都相同。這很聰明,我還沒實作,因為在 Redis 中我需要追求速度,而這是在空間與時間之間的取捨:但或許值得,也或許不值得。我會再深入研究。

不過,如果你仔細計算,多層結構這件事其實沒有看起來那麼糟。平均而言,每個節點的多層結構只會讓情況惡化約 1.3 倍(如果層級選擇函式中層級遞增的機率是 0.25),因為許多節點其實只會在第 0 層。但 1.3 終究還是大於 1,而且如果 HNSW 中那個「H」其實沒那麼有用……[劇透一下,我發現如果所有東西都在第 0 層,搜尋時間會變長,貪婪搜尋的主迴圈會從較不理想的位置開始,最終雖然還是會抵達正確的群集,但會花費更多運算時間。不過這只是初期結果。]

所以這裡真正唾手可得的成果是:vector quantization(向量量化)。我發現如果使用 8 位元量化,你幾乎可以獲得 4 倍的速度提升、向量本身縮小 4 倍(但不是整個節點縮小 4 倍:指標還在,而且佔了很大空間),而在真實世界的使用情境中,召回率幾乎完全相同。這就是為什麼 Redis Vector Sets 預設使用 8 位元量化的原因。你可以透過 VADD 的選項指定想要全精度的向量或二進位量化向量——後者只取正負號——但我對於同時使用全尺寸向量與二進位量化向量抱持懷疑。在談論它們之前,先來看看我對 8 位元做了哪種量化。

我的做法是計算每個向量各分量的最大絕對值(所以量化是以每個向量為單位),然後使用有號 8 位元數值來表示從 -127 到 127 的量化值。這不如同時儲存最小值與最大值來得精確,但在計算 cosine similarity(餘弦相似度)時更快,因為我可以這樣做:

/* 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 命令的一個選項。

擴展速度:多執行緒與區域性

喔,你知道的,我得先說說我自己:當可以用單核心做很多事、再透過 shared-nothing(無共享)架構來運用多核心時,我並不是多執行緒系統的愛好者。但 HNSW 是不同的。它們*很慢*,而且在大多數使用情境中,幾乎總是以唯讀方式被存取。基於這個原因,我的 Vector Sets 實作是完全多執行緒的。不只是讀取,連寫入也部分採用多執行緒,你可能會好奇這怎麼可能不會搞得一團亂,特別是在像 Redis 這樣的系統中,鍵可能會被背景儲存程序、客戶端等以不同方式存取。

好吧,首先,讓我們聚焦在讀取上。情況是這樣的,只要沒有人在寫入這個資料結構,我們就可以產生執行緒來執行貪婪式蒐集鄰近向量,並將結果回傳給被阻塞的客戶端。然而,我的 HNSW 實作是從零開始寫的,我是說,從用 vim 開啟的空白 C 檔案開始,它與大多數其他系統所使用的兩種實作有 0% 的共用程式碼,所以有一些「新穎之處」。其中一個不同之處在於,為了避免重複造訪已經造訪過的節點,我在每個節點中儲存一個稱為「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 Sets 要直接將實際的資料結構暴露給使用者(我相信程式設計師很聰明,不需要過度保護,但原因*不只是*這樣)之前,我想先回頭再談談記憶體,因為關於這個特定面向,有個有趣的故事可以分享。

大多數 HNSW 實作在你從圖中刪除一個節點時,無法直接回收記憶體。我認為主要有兩個原因:

1. 人們以一種特定方式誤解了原始的 HNSW 論文:他們認為鄰居之間的連結可以不是相互的。而他們之所以這樣想,是有特定原因的。

2. 論文完全沒有提到節點刪除,以及在節點消失、連結的「網路」出現缺口後,該如何修復圖。

第一個問題我認為是論文不夠清晰,加上人們在實作 HNSW 時會遇到一個具體問題的綜合結果:當插入一個新節點、並在既有節點中搜尋合適的鄰居時,候選節點往往已經擁有最大數量的對外連結。這種情況下該怎麼辦?這個問題通常透過從我們正在插入的新節點單向連結到那些對外連結已「滿」的候選者來解決。然而,當你需要刪除一個節點時,你就無法再解析它所有的內向連結,因此也就無法真正回收記憶體。你只能用一個旗標將其標記為已刪除,之後有時會透過重建圖來「垃圾回收」過時的節點,有時則只是讓記憶體洩漏。

所以,首先,我在 Redis 中的實作採取了不同的做法,強制連結必須是雙向的。如果 A 連到 B,B 就連到 A。但是,考慮到 A 可能已經很忙,要怎麼做到這點呢?嗯,這就進入比較複雜的領域了,但做法是使用啟發式方法,從既有節點中移除與其他已良好連結的鄰居之間的連結,而且如果我們的節點對目標節點來說也是更好的候選者,就這麼做,如果情況並非如此,也有其他方法可以強制讓新節點至少擁有最少數量的連結,始終試圖滿足圖的小世界特性。

透過這種方式,當 Redis 從 Vector Set 中刪除一個節點時,它總是有辦法移除所有指向該節點的指標。然而,對於那些現在缺少一條連結的剩餘節點該怎麼辦呢?我的做法是在它們之間建立一個距離矩陣,試圖將舊節點的鄰居們彼此連結起來,並盡量讓平均距離最小化。基本上,對於矩陣中每一對 i、j 節點,我們會計算它們連結的好壞程度(它們的向量有多相似)以及將它們連結起來會對*剩餘*可能的配對造成多大的負面影響(因為如果我們將兩個特定節點連結起來,可能會讓某些元素找不到好的配對)。在建立這個分數矩陣後,我們接著進行貪婪配對步驟。

這效果非常好,你可以建立一個包含數百萬個元素的大型 HNSW,之後刪除其中 95% 的所有元素,而剩下的圖仍然保有良好的召回率,也不會出現孤立節點等問題。

這就是我說 HNSW 仍有空間讓新論文繼續投入研究時的意思。

將 HNSW 擴展到多個處理程序

當我開始著手 Redis Vector Sets 時,在 Redis 的生態系中已經有向量相似度的實作,具體來說是作為 RediSearch 的一種索引類型,而這也是大多數人對 HNSW 的看法:一種針對既有資料的索引形式。

然而,我想為 Redis 提供一種以完全不同方式暴露的新 HNSW 實作。猜猜是怎麼做的?當然是作為一種資料結構。這說明了這麼多年來我的腦袋有多麼「Redis 化」,或者或許從一開始就是 Redis 化的,而 Redis 才是依照我的腦袋塑形的,因為我立刻就構想出如何設計一個直接將 HNSW 暴露給使用者的 Redis 資料結構,並且對於 Redis 中處理向量的工作不是以這種方式進行感到困惑。

同時,當我把設計文件交給 Redis 的同事們時,我不能說他們立刻就「看懂」這是顯而易見的事。我的邏輯是:向量就像 Redis Sorted Sets 中的分數,只不過它們不是具有全序關係的純量分數。然而你可以 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 就是資料結構,那就輕而易舉了。你只要為每個項目建立一個 Vector Set 鍵,裡面只放少數幾個元素。當然,就像任何其他 Redis 鍵一樣,你可以為該鍵設定過期時間,讓它之後自動被移除。

這一切都可以濃縮成一個我認為在我們產業中應該更常見的原則:許多程式設計師都很聰明,如果你不是打造一個他們無法觸及的神奇系統,而是向他們展示資料結構、其中的取捨,他們就能打造出更多東西,並以特定的方式為其使用情境建模。而你的系統也會變得更簡單。

擴展載入時間

如果我不使用多執行緒,我的 HNSW 函式庫在使用單一執行緒的情況下,可以將 word2vec(每個向量 300 個分量)以每秒 5000 個元素的速度加入到 HNSW 中,並能以每秒 9 萬次查詢的速度查詢 resulting HNSW。如你所見,兩者之間存在很大的落差。

這意味著從 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 Sets 第一個可運作的實作感覺上完成的那一天。一切都如預期般運作,那是開始進行細部調整與額外功能的起點。

然而在過去幾週與幾個月裡,我在內部收到回饋,指出大多數使用情境都需要某種形式的混合搜尋:你想要與給定查詢向量相近的向量(例如與某部電影最相似的電影),但同時也需要某種篩選(僅限於 2000 年至 2010 年間發行的)。我的感覺是,你需要針對不同參數進行查詢的頻率,並沒有產品人員所想的那麼高,而且在大多數情況下,你可以透過更有效率的方式來達成——以這個特定案例來說,就是將每一年分別放到不同的 vector set 鍵中(這又是將 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 Sets 的特定情境下,想法是提供某種非常快速、易於使用的東西:記憶體內資料結構的彈性有助於實現這一點。所以問題歸結為:記憶體用量真的有那麼糟嗎?

使用預設的 int8 量化將 300 萬筆 Word2Vec 項目載入到 Redis 中需要 3GB 的 RAM,每個項目約 1KB。許多使用情境只有數千萬筆項目,或是少得多。而如果實作得當,並在記憶體中運作,你從 HNSW 獲得的是非常好的效能,這對於本質上就很慢的資料結構與工作負載來說至關重要。在我的 MacBook 上,我使用 redis-benchmark 對這個鍵(持有 word2vec 資料集)執行 VSIM,得到每秒 4.8 萬次操作。我的感覺是,記憶體內 HNSW 的記憶體用量對於許多使用情境來說是非常可接受的。即使在你希望將大部分向量放在磁碟上的使用情境中,儘管得付出較慢效能的代價,你的熱門資料集很可能還是應該由 RAM 來提供服務。

這也是我認為持續投入 HNSW 研究是個好主意的原因之一:我不認為它們在大多數使用情境中會很快被取代。更有可能的是,我們會繼續擁有分別適合 RAM 與磁碟、依使用情境與資料大小而定的不同資料結構。此外,我最近即使只是瀏覽 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 Sets,請隨意閱讀我自己撰寫的 README 檔案。也有官方的 Redis 文件,但我建議你從這裡開始:

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

感謝你閱讀這麼長的部落格文章!祝你有個愉快的一天。

參考資料。這是關於 HNSW 中「H」及其有用程度的論文 -> https://arxiv.org/abs/2412.01940

原文由 Salvatore Sanfilippo 發布

本文章由 muse-spark-1.2-contributor 進行翻譯