從刷臉打卡到各種應用的 “猜你喜歡”，當前機器學習（特別是深度學習技術）已經廣泛應用于我們日常生活的方方面面。深度學習框架（如：TensorFlow，PyTorch等）和 AI專用芯片（如：TPU、NPU等）等軟硬件系統的設計極大地提升了機器學習的性能并擴展了其應用場景。與此同時，機器學習方法本身可否用于優化計算機系統設計甚至是取代傳統設計模式呢？谷歌技術大神Jeff Dean領導的研究小組在2018年的SIGMOD學術會議上發表了一個有趣的工作：The Case for Learned Index   Structures [1]   。本文簡要介紹了Learned Index Structures的實現和優缺點，希望可以給大家帶來一些系統設計的啟發和思路。

一、什么是 Learned Index Structures？

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

圖 1 B-Tree和Learned Index示意圖

圖 1就是一個傳統B-tree和Learned Index的對比。可以看到，Learned Index的輸入是Key，輸出是這個key對應的檢索結果的位置（可能有誤差），誤差的上限是葉子節點一個數據頁中的結果條目（即所需檢索結果在數據頁的末尾）。

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

p = F(Key) * N

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

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

那么如何學習得到 F(Key) 呢？ Learn Index Structures作者首先嘗試了用TensorFlow搭建一個每層32個神經元，兩層全連接的神經網絡，使用一個web server日志的數據集訓練后發現效果遠差于B-Tree。問題在于：

1）TensorFlow是用于大規模神經網絡的訓練的，小規模場景的調用開銷變得不可忽視。

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

3）B-Tree的CPU和cache行為是經過高度優化設計的，每次查找只需使用少量索引。機器學習模型則需要使用全部參數權重完成一次預測。

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

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

圖 3 用于Learned Index Structures的Staged Model

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

二、如何評價Learned Index Structures？

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

實際上，以Learned Index Structures為代表的機器學習優化方法并不是系統設計優化的終結者。比如spline B-tree [6] 使用B-tree的每個葉子節點只存一個spline（即key和其位置），兩個spline之間的數據用兩點之間的直線來預測。這樣一個簡單的數據結構，很多時候效果相當于復雜的Learn Index，甚至更好。在Point Index 領域，Learned Index通過減少沖突實現的優化可以被bucketized cuckoo hashing [7] 輕松打敗，該方法只是簡單的將每個key同時hash到兩個bucket而已。

但這并不能否定機器學習在系統設計上的 價值，通過機器學習可以啟發系統設計的優化和思考，探索出之前未曾發現的系統設計思路。在優化原理清晰、場景固定的情況下，顯然由人加以解釋和重新實現在效率和穩定性上更勝機器學習方法一籌。在數據分布等特征動態變化的場景，機器學習方法 可以針對性優化和適應數據特征，理論上可以優于通用的算法和數據結構。

三、機器學習+系統設計 = ？

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

1) 找到一個能用數字精確表示的優化指標；

2) 有一個集成機器學習的清晰接口（模型的輸入輸出定義，訓練、測試數據集的獲取等）。

對于計算機系統領域的優化，這兩個要求似乎是比較容易實現的。