機器學習思路優化B Tree

最近一直在刷課, 刷課過程中,瞭解到機器學習作為一種工具,應用到改善傳統的的數據索引結構,覺得很有意思, 前兩年也應該在某個媒體的pr文章中看到類似的工作,但是一直沒有做了解,這次刷課,正好學習下

動機

為啥會提出使用機器學習模型,來改善甚至替換傳統的索引結構, 可以比較簡單的從以下幾點來說明:

1. CPU摩爾定律已死, GPU算力仍在成倍提升, 計算資源會越來越廉價,而存儲資源已到瓶頸;
2. 傳統數據結構對數據本身並不敏感,無法利用其中的數據分佈信息, 結構本身依賴的資源由數據量決定的,而數據本身的分佈暗含很多信息;
3. 機器學習模型能擬合數據分佈,其資源\\算力的消耗, 僅靠模型本身複雜性來決定而不靠數據量，而傳統數據結構，如B樹對資源的需求，靠數據量大小決定的;

B Tree

B-Tree是很常見的數據結構， 其規則如下：

1. 排序方式：所有節點關鍵字是按遞增次序排列，並遵循左小右大原則；
2. 子節點數：非葉節點的子節點數>1，且<=M ，且M>=2，空樹除外（注：M階代表一個樹節點最多有多少個查找路徑，M=M路,當M=2則是2叉樹,M=3則是3叉）；
3. 關鍵字數：枝節點的關鍵字數量大於等於ceil(m/2)-1個且小於等於M-1個（注：ceil()是個朝正無窮方向取整的函數 如ceil(1.1)結果為2);
4. 所有葉子節點均在同一層、葉子節點除了包含了關鍵字和關鍵字記錄的指針外也有指向其子節點的指針只不過其指針地址都為null對應下圖最後一層節點的空格子;

以下例子以字母為關鍵字，在B-Tree中排列

B-Tree To Learner Index

B樹用機器學習來描述其功能：給定Key，預測其值的位置，如下圖：

B樹為了提高效率，不是為每一條記錄，保存index，而是每n條給定一個index，這n條稱為page（注意這裡與內存管理中的頁不是一個概念）。舉個例子：有一個數據集，key大小為100M，page 為100， 第一層需要1M， 第二層需要10K，第三層需要100， 也就是最少要訪問四個節點， 每使用binary search需要50個時鐘週期，且很難並行，cpu每個時鐘週期做8-16個並行操作（具體可看sse、avx，基本原理是將多條數據放入寄存器中查找），另外，每次需要page在cache中，當數據量較大時， index數據無法完全存入cache中，需要從硬盤載入，則需要50-100額外的時鐘週期去查找，而機器學習至少從兩個方面打破這個問題：

1. 常數時間的模型的inference，不隨著數據量大小增加查詢時間；
2. 機器學習模型容易並行，在GPU等設備(60000個並行操作)下，計算數據極快，30個時間週期下，能進行百萬級的並行操作；

GPU現階段本身也存在一些問題，如在數據從內存讀入顯存，會存在耗時，這個不可忽略，但是目前無論是GPU還是TPU在本身計算上得到的收益還是很大。

Range Index ->CDF Models

B樹其實在有序的情況下去估計記錄的位置 ，因此可以將預測定義為CDF函數，

F(Key)是有序排列中查找所有小於Key的數據，實質是一個累積函數，N是指所有Key的條數。通過構造模型，我們的目的就是為了學習F(Key）這一函數分佈。

Naive Learned Index

文章使用200M的web-server日誌，按時間戳為key來，訓練了一個兩層的神經網絡，每層有32個神經元，激活函數為ReLU， 時間戳為輸入特徵，保存的位置為輸出label，訓練模型之後，選取隨機的key來查詢， 每秒能查詢約1250次，每次查詢需80ms，而作為對比，B tree查詢數據需要300ns，而二分查找需要900ns，差別很大，文章總結分析從以下幾個方面造成的性能差異：

1. TensorFlow設計是用來運行大模型，而我們實驗的模型極小，大模型下Python作為Graph的構造語言，性能差別影響不大，而當本身計算模型小時，使用Python作為前端造成的耗時佔比會大很多；
2. B樹，或者更通用來說決策樹，通過可以規則『分而治之』，將本身問題通過遞歸切分到更小空間去解決，而對於其他模型，其優勢的是從宏觀角度去擬合曲線，但是當我們把擬合曲線放大，我們會發現存在更多的極不規則的震盪數據出現，而對於神經網絡、或者回歸模型來說需要更多的算力去解決這樣的問題(文章中為last mile)；
3. B樹在設計時會採用cache和優化操作去保證經常被查詢的節點存在cache中，這會大大較少耗時，而神經網絡需要所有的模型參數去計算最終的label，相比較而已耗費大小時間來做運算。

RM-Index

基於上面3個問題，文章提出了相應的措施來解決相應的問題：

Learning Index Framework

LIF是一套整體的learned index的框架，內置簡單模型的計算邏輯，如linear regression，當需要使用神經網絡等複雜時，依賴TensorFlow來完成訓練，但是在做Inference時，LIF會從訓練的PB模型中拿到權重然後使用C++完成推理。LIF效率極高，能夠在簡單的模型上達到30ns的推理速度，當然LIF還有很多不足，比如為了保證公平對比僅使用編譯器的向量化邏輯，沒有使用特殊的SIMD並行邏輯。

Recursive Model Index

受樹模型『分而治之』的思路影響，為了解決『last mile』問題，文章轉變思路，不再解決從100M中精準找到百級別的，而是首先降低從100M中找到10K的錯誤，然後減低從10K中找到百級別的錯誤，這樣問題難度會變得簡單。

基於此思路，文章提出Recursive regression model，構建分層的模型，每一層模型輸入為key，而輸出為其下一步選擇的模型的ID，直到最後一層預測到最終的位置，數學建模為：

每一層中模型選擇進入某一個模型，最後一層模型輸出不同的位置，這樣設計的實質有點類似於B樹本身的邏輯，只是B樹本身通過樹去選擇保存下一步page的首地址，而Recursive Model Index是選擇過某條輸出，到最後一層預測最終的位置，而劃分的所有的執行邏輯是通過數據本身的分佈和上面構造的loss函數來定的，其實質在於替換B樹，以sorted的binary search的邏輯，而是通過擬合數據分佈，通過模型的inference去『算』出位置， 模型的中間輸出也類似，比如model2.1 其實就是把數據中的部分key『導流』在這一支路，B樹的邏輯是這部分支路選擇的一定是小於model2.2/2.3/...和model1.1本身的數據，而Recursive Model Index的邏輯是去構造整體模型去最小化loss。

Hybrid Indexs

Recursive Model Index的另一優勢在於能夠構造混合模型。例如：Recursive Model頂層可以使用神經網絡來擬合複雜的數據分佈，在底層使用數量眾多簡單的regression來擬合震盪的數據分佈，甚至於在數據分佈很難學習的時候使用B-Tree來構造底層的數據，Hybrid Iindex能夠保證learned index的最差性能也能pk掉B-Tree，甚至在極端情況下，直接講過所有模型替換為B-Tree。

Result

下列是在各種不同數據集上的對比實驗結果：

整型數據集：

字符串數據集

相關項目

有一個Learned Index的項目：https://github.com/yangjufo/Learned-Indexes, 重現了這部分工作，蠻有意思的，讀了一下代碼，實現的比較簡單， 我這邊跑了部分源碼結果如下：

看起來mean error 和search time都pk不了B-Tree，但是因為這部分只是很簡單的實現，其實真實去替換B-Tree是很複雜的工程問題，不僅在解決mean error的問題上，還是在解決工程Inference耗時上，都是有很多內容的，不過論文論證的機器學習作為一個工具，挑戰B-Tree這一套邏輯，十分值得肯定，況且Google已經在這部分做了很大的工作，也應用到很多query的場景

總結

機器學習確實從很多方向開始改變我們的計算方向，不像傳統的在算法業務上的改變，如推薦、計算機視覺、自然語言等等，在數據結構上的改變其優勢在於探索存儲數據本身的數據分佈並且加以利用，將性能目標數字化，構造合適的loss函數，通過機器學習模型去最小化性能目標。現在不妨我們來腦洞下相關場景：

在業界，線上做特徵拼接時，特徵數據在導入之後不會更改，除非下一次的更新，比如以redis為例，模型上線後會把諸如用戶特徵這部分的feature導入到redis中，然後做查詢，這部分工作是否能夠通過模型來預測呢？redis的性能會隨著key的數量增大而變差，且資源消耗會越來越大，而模型本身的計算時間為常數時間，且能擬合數據本身的分佈，學習好數據分身分佈後能夠更好地去分派數據到不同的內存地址，來做更高效的查詢。

閱讀更多 小石頭的碼瘋窩 的文章

關鍵字: 數據結構 人工智能 中央處理器