Scaling HNSWs

Salvatore Sanfilippo

擴展 HNSW

原文由 Salvatore Sanfilippo 發布,訂閱此部落格

我暫停 HNSW 的開發幾個星期(現在在做另一種資料結構,很快就會有消息)。到這個階段,我為 Redis 新增的這個型別已經夠穩定、夠完整,正是好好整理一下我對 HNSW 學到的東西、寫成一篇部落格文章的好時機。那種在 AI 時代以前很常見的腦內傾倒文,現在或許已經變得有點稀少了。總之,經過將近一年對 HNSW 與向量相似度相關技術的思考與實作,是時候寫點東西了。不過這不會是一篇 HNSW 入門文:那種文章已經太多了。這篇要走的是「更進一步」的路。如果你已經了解 HNSW,我想跟你分享一些更「進階」的發現,尤其是在如何讓它夠快、足以提供「Redis 體驗」的脈絡下:你知道的,Redis 的設計目標是低延遲與高效能,而 HNSW 多少有點跟這個目標背道而馳,所以要把它當成一個抽象資料結構暴露出來,確實遇到了不少挑戰。

這篇文章會分成幾個章節。把它們想成同一本書的不同頁、同一段經歷的不同章就好。喔對了,順帶一提,我其實已經寫過這篇文章,然後又把它弄丟了 :D[一個關於 MacOS 與壞習慣的漫長悲傷故事——我上次搞丟這種東西,已經是 90 年代停電的時候了],所以這次大部分的挑戰,是要回想起幾天前寫了什麼,並趁機把當時不太滿意的地方重寫得更好。

關於 HNSW 現況的幾句話

在深入 HNSW 的內部機制與最佳化之前,我想先談談 HNSW 的現況。原始那篇介紹 HNSW 的論文是一篇很棒的電腦科學文獻,HNSW 也是非常厲害的資料結構,但是:我不認為它是根據距離函式、以貪婪方式搜尋相近向量的最終答案。這篇論文給人的感覺像是少了幾塊「拼圖」,好像研究人員如果再多六個月,就還有很多可以探索與發表的地方。舉例來說,我就自己修改並擴充了論文的做法,以支援真正刪除資料,而不只是墓碑式刪除——也就是把元素標記為已刪除、之後再回收:論文裡完全沒有提到如何刪除項目。類似地,現在也有一些研究在認真檢驗 HNSW 裡的「H」是否真的必要,改用只有一層的扁平資料結構,效能是否也差不多(希望未來能再多談談這個:我的感覺是真相在中間,比較合理的做法是修改層級選擇函式,只保留大於某個門檻的層級)。

說這些是想表達,如果你對資料結構研究有興趣,我認為一個很棒的方向是去想像 HNSW 的演進與改良,而不要被侷限在「演進就等於:來把它搬到磁碟上吧(可參考微軟的嘗試)」這類想法裡。好,前言就到這裡,來談點真正底層的東西吧 :)

擴展記憶體

Redis 是記憶體內系統,而 HNSW 和向量都有個不幸的特性:非常占空間。原因有三:1. HNSW 有很多指標,像是 16、32 甚至更多指向鄰近節點的指標(這是 HNSW 的一個可調整參數)。2. HNSW 有很多層,本身就是一種類似 skip list 的資料結構。這讓第一個問題更加嚴重。3. HNSW 的附屬資料是浮點數向量,所以在最原始的情況下,每個分量要 4 個位元組,而通常一個向量會有 300 到 3000 個分量,這是常見的範圍。

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

不過,如果你實際算一下,多層結構其實沒有看起來那麼糟糕。平均來說,每個節點的多層結構只會讓情況變糟約 1.3 倍(如果層級選擇函式中層級遞增的機率是 0.25),因為很多節點其實只在第 0 層。但 1.3 終究還是大於 1,而且如果 HNSW 裡的「H」真的沒那麼有用……[劇透一下,我發現如果全部都放在第 0 層,搜尋時間會變長,貪婪搜尋的主迴圈會從比較不理想的位置開始,最終還是會找到正確的群集,但需要更多運算時間。不過這只是初步結果。]

所以這裡真正唾手可得的成果是:向量量化。我的發現是,如果使用 8 位元量化,你幾乎可以獲得 4 倍的速度提升,向量大小縮小 4 倍(但不是整個節點縮小 4 倍:指標還在,而且占了很多空間),而在真實世界的使用情境中,召回率幾乎不變。這就是 Redis Vector Sets 預設使用 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 指令的一個選項。

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

喔,你知道的,我得先跟你說說我自己的想法:當單核心就能做很多事、然後可以用 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 的同事們時,我不能說他們立刻就覺得這是理所當然的事。我的 reasoning 是這樣的:向量就像 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 函式庫在單執行緒下每秒可以把 5000 個 word2vec(每個向量 300 個分量)加入到 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 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

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

留言