從刷臉打卡到各種應(yīng)用的 “猜你喜歡”，當(dāng)前機(jī)器學(xué)習(xí)（特別是深度學(xué)習(xí)技術(shù)）已經(jīng)廣泛應(yīng)用于我們?nèi)粘Ｉ畹姆椒矫婷妗Ｉ疃葘W(xué)習(xí)框架（如：TensorFlow，PyTorch等）和 AI專用芯片（如：TPU、NPU等）等軟硬件系統(tǒng)的設(shè)計極大地提升了機(jī)器學(xué)習(xí)的性能并擴(kuò)展了其應(yīng)用場景。與此同時，機(jī)器學(xué)習(xí)方法本身可否用于優(yōu)化計算機(jī)系統(tǒng)設(shè)計甚至是取代傳統(tǒng)設(shè)計模式呢？谷歌技術(shù)大神Jeff Dean領(lǐng)導(dǎo)的研究小組在2018年的SIGMOD學(xué)術(shù)會議上發(fā)表了一個有趣的工作：The Case for Learned Index   Structures [1]   。本文簡要介紹了Learned Index Structures的實(shí)現(xiàn)和優(yōu)缺點(diǎn)，希望可以給大家?guī)硪恍┫到y(tǒng)設(shè)計的啟發(fā)和思路。

一、什么是 Learned Index Structures？

Index（索引）是數(shù)據(jù)庫、文件系統(tǒng)等領(lǐng)域常見的數(shù)據(jù)結(jié)構(gòu)，最經(jīng)典的莫過于B-Tree。B-Tree是一種范圍索引（Range Index）數(shù)據(jù)結(jié)構(gòu)，查詢時給定一個key（或一些確定范圍的keys），B-Tree會索引到包含該key的對應(yīng)范圍的葉子節(jié)點(diǎn)，在葉子節(jié)點(diǎn)內(nèi)對key進(jìn)行搜索。如果該key在索引中存在，就會得到其對應(yīng)的位置。一般在一個邏輯頁內(nèi)的記錄會用一個key來索引。如圖1（a）所示，輸入是一個key，輸出是對應(yīng)要查詢記錄的位置區(qū)間。除了范圍索引，還有點(diǎn)索引（Point Index）如：哈希表（Hash Map），和存在索引（Existence Index）如：Bloom Filters，也都是常用的索引數(shù)據(jù)結(jié)構(gòu)。Learned Index Structures的主要思想是將這里B-Tree或是Bloom Filters等數(shù)據(jù)結(jié)構(gòu)替換為機(jī)器學(xué)習(xí)的模型——查找操作變成了根據(jù)key做索引數(shù)據(jù)位置的預(yù)測，如圖1（b）所示。我們以范圍索引為例，來詳細(xì)介紹下Learned Index Structures的設(shè)計和實(shí)現(xiàn)思路。

圖 1 B-Tree和Learned Index示意圖

圖 1就是一個傳統(tǒng)B-tree和Learned Index的對比?？梢钥吹?，Learned Index的輸入是Key，輸出是這個key對應(yīng)的檢索結(jié)果的位置（可能有誤差），誤差的上限是葉子節(jié)點(diǎn)一個數(shù)據(jù)頁中的結(jié)果條目（即所需檢索結(jié)果在數(shù)據(jù)頁的末尾）。

那么針對范圍索引，如何設(shè)計機(jī)器學(xué)習(xí)的模型呢？這里只考慮一維聚簇索引的情況（即數(shù)據(jù)是按照用于查找的 key來排序的，非聚簇索引可以通過聚簇索引加一層到真實(shí)數(shù)據(jù)排列的指針實(shí)現(xiàn)），Learned Index Structures很巧妙的給出了如圖2所示的洞察，索引位置實(shí)際上是隨著Key增長而增長的單調(diào)遞增函數(shù)。雖然具體的索引可能是離散的,整體上還是可以用一個函數(shù)來描述：

p = F(Key) * N

其中， p是估計得到的位置，N是索引Key的數(shù)量，F(xiàn)(key)是索引數(shù)據(jù)的累計分布函數(shù)（CDF）。F（key）的含義是小于等于key的索引數(shù)據(jù)條目總和（即key的位置估計）。本質(zhì)上體現(xiàn)了數(shù)據(jù)集的分布特征。相比B-Tree的通用設(shè)計，Learned Index Structures考慮了數(shù)據(jù)集的內(nèi)在分布特點(diǎn)并將其用于優(yōu)化索引的結(jié)構(gòu)。Learned Index誤差上限可控，只需要在誤差范圍內(nèi)根據(jù)預(yù)測的位置向左或向右二分查找即可準(zhǔn)確找到查找目標(biāo)。

圖 2 索引位置的累計概率分布（CDF）

那么如何學(xué)習(xí)得到 F(Key) 呢？ Learn Index Structures作者首先嘗試了用TensorFlow搭建一個每層32個神經(jīng)元，兩層全連接的神經(jīng)網(wǎng)絡(luò)，使用一個web server日志的數(shù)據(jù)集訓(xùn)練后發(fā)現(xiàn)效果遠(yuǎn)差于B-Tree。問題在于：

1）TensorFlow是用于大規(guī)模神經(jīng)網(wǎng)絡(luò)的訓(xùn)練的，小規(guī)模場景的調(diào)用開銷變得不可忽視。

2）欠擬合問題，如圖2所示機(jī)器學(xué)習(xí)模型可以很好的估計CDF的整體趨勢，但在單一數(shù)據(jù)項(xiàng)上很難得到精確的表示。而B-Tree可以簡單高效的使用if語句精確劃分范圍，為了優(yōu)化“最后一公里” 機(jī)器學(xué)習(xí)模型要付出較大的存儲空間和計算資源消耗。

3）B-Tree的CPU和cache行為是經(jīng)過高度優(yōu)化設(shè)計的，每次查找只需使用少量索引。機(jī)器學(xué)習(xí)模型則需要使用全部參數(shù)權(quán)重完成一次預(yù)測。

最終 Learn Index Structures的模型使用了如圖3所示的Staged Model實(shí)現(xiàn)。每一個Model都可以是任意一個機(jī)器學(xué)習(xí)模型，從最簡單的線性回歸（LR）到深度神經(jīng)網(wǎng)絡(luò)（DNN）都可以。實(shí)踐中，越簡單的模型越好（避免查找時在模型上花太多時間）。當(dāng)進(jìn)行查找時，最上層的模型（只有一個模型），將選擇一個第二層的模型來處理這個key。然后第二層的模型，會接著選擇一個下一層的模型來處理這個key，直到最底層的模型，才會給出這個key對應(yīng)的預(yù)測位置。但實(shí)際上上層每個模型輸出的都是預(yù)測位置，這個預(yù)測被用于選擇下層模型（模型id = (預(yù)測位置 / 記錄總數(shù)) * 該層模型數(shù)）。

整個 Staged Model分層訓(xùn)練，先訓(xùn)練最頂層，然后進(jìn)行數(shù)據(jù)分發(fā)。數(shù)據(jù)分發(fā)指的是，上層模型將key預(yù)測到哪個下層Model，該Model就擁有這條訓(xùn)練數(shù)據(jù)作為他的訓(xùn)練集。所以隨著層數(shù)的加深，以及每一層模型數(shù)量的提升，每個越底層模型擁有的訓(xùn)練數(shù)據(jù)是越少的。這樣的優(yōu)點(diǎn)是，底層模型可以非常容易的擬合這一部分?jǐn)?shù)據(jù)的分布（缺點(diǎn)是較少的數(shù)據(jù)量帶來了模型的選擇限制，復(fù)雜模型沒法收斂）。[1]中采用的結(jié)構(gòu)是：只在頂層使用神經(jīng)網(wǎng)絡(luò)模型，在其余層使用線性回歸模型。

圖 3 用于Learned Index Structures的Staged Model

點(diǎn)索引和存在索引的 Learned Index Structure 這里不再一一贅述，感興趣的話可以閱讀論文[1]，以及基于論文實(shí)現(xiàn)的RMI（Recursive Model Indexes）代碼[2]。

二、如何評價Learned Index Structures？

總的來說，Learned Index Structures向我們展示了機(jī)器學(xué)習(xí)在系統(tǒng)領(lǐng)域的巨大潛力，但還存在諸多待解決的問題（如：過擬合、索引的增刪改等問題）。對于過擬合，若新索引的key依然滿足CDF則并不需要重新訓(xùn)練，直接insert到預(yù)測出來的位置即可。若數(shù)據(jù)分布會發(fā)生變化，則需要嘗試在線學(xué)習(xí)（Online learning）的方法。對于數(shù)據(jù)更新頻繁的系統(tǒng)，可采用delta-index技術(shù)增量更新learned index。

實(shí)際上，以Learned Index Structures為代表的機(jī)器學(xué)習(xí)優(yōu)化方法并不是系統(tǒng)設(shè)計優(yōu)化的終結(jié)者。比如spline B-tree [6] 使用B-tree的每個葉子節(jié)點(diǎn)只存一個spline（即key和其位置），兩個spline之間的數(shù)據(jù)用兩點(diǎn)之間的直線來預(yù)測。這樣一個簡單的數(shù)據(jù)結(jié)構(gòu)，很多時候效果相當(dāng)于復(fù)雜的Learn Index，甚至更好。在Point Index 領(lǐng)域，Learned Index通過減少沖突實(shí)現(xiàn)的優(yōu)化可以被bucketized cuckoo hashing [7] 輕松打敗，該方法只是簡單的將每個key同時hash到兩個bucket而已。

但這并不能否定機(jī)器學(xué)習(xí)在系統(tǒng)設(shè)計上的 價值，通過機(jī)器學(xué)習(xí)可以啟發(fā)系統(tǒng)設(shè)計的優(yōu)化和思考，探索出之前未曾發(fā)現(xiàn)的系統(tǒng)設(shè)計思路。在優(yōu)化原理清晰、場景固定的情況下，顯然由人加以解釋和重新實(shí)現(xiàn)在效率和穩(wěn)定性上更勝機(jī)器學(xué)習(xí)方法一籌。在數(shù)據(jù)分布等特征動態(tài)變化的場景，機(jī)器學(xué)習(xí)方法 可以針對性優(yōu)化和適應(yīng)數(shù)據(jù)特征，理論上可以優(yōu)于通用的算法和數(shù)據(jù)結(jié)構(gòu)。

三、機(jī)器學(xué)習(xí)+系統(tǒng)設(shè)計 = ？

也許Learned Index Structures還存在很多不足，但無法忽略的是將機(jī)器學(xué)習(xí)應(yīng)用于計算機(jī)系統(tǒng)設(shè)計的趨勢已經(jīng)到來。如果說Learn Index Structures是機(jī)器學(xué)習(xí)打入計算機(jī)系統(tǒng)設(shè)計領(lǐng)域的一聲炮響，2019年發(fā)布的機(jī)器學(xué)習(xí)系統(tǒng)白皮書  [5]  就是正式確立了機(jī)器學(xué)習(xí)和計算機(jī)系統(tǒng)設(shè)計交叉研究方向的誕生。如Jeff Dean在SysML18會議上主旨演講所言：“使用啟發(fā)式技術(shù)的任何系統(tǒng)領(lǐng)域，都是可能應(yīng)用機(jī)器學(xué)習(xí)的好地方——編譯器、網(wǎng)絡(luò)、操作系統(tǒng)、芯片設(shè)計等”。要取得成功，關(guān)鍵點(diǎn)有兩個：

1) 找到一個能用數(shù)字精確表示的優(yōu)化指標(biāo)；

2) 有一個集成機(jī)器學(xué)習(xí)的清晰接口（模型的輸入輸出定義，訓(xùn)練、測試數(shù)據(jù)集的獲取等）。

對于計算機(jī)系統(tǒng)領(lǐng)域的優(yōu)化，這兩個要求似乎是比較容易實(shí)現(xiàn)的。