scribble

Life of Xunzhang

About Story Talk Project Gallery Ideas Email

21 Dec 2017
Paper Notes: The Case for Learned Index Structure

前几周看了Jeff Dean在Nips17上演讲的Slides,其中提到一些有趣的工作,包括最近在网络媒体上出现次数比较多的这篇文章The Case for Learned Index Structure。周末仔细读了一遍,来分享一下思路和具体细节。

有没有想过用机器学习的方式解决系统设计中的问题?这是个蛮有趣的命题,它的难点在哪儿呢:系统里的算法或数据结构往往是确定性的,显式的(我们清楚地知道输入和输出间的映射关系),而机器学习中的模型是依赖于输入数据训练出来的(下文会看到事实上数据结构的建立也是依赖于数据的),并且存在一定的误差(而很多时候系统模块有严格限制条件)。虽然在系统设计时应该考虑所有(最坏)数据的情况,但是在做系统优化的时候,通常会通过一些特殊性来对系统做特定的设计。这里的特殊性包括数据特殊性和使用特殊性。数据特殊,比如对给定范围且没有重复的数组排序,可以直接用bitmap。使用特殊,例如由于业务场景,读写比例相差很多(每周只更新一次),实现某种数据结构的时候可以把读实现成O(1),写实现成O(n),而标准实现往往是在保证均摊最优情况下让各操作复杂度尽量接近,因为假设各操作使用频率相近。而大多数时候只可能知道使用特殊性,却很难知道数据的特殊性。这里很自然的思路就是通过数据训练出一个机器学习的模型(由于多层神经网络的泛化能力这里考虑神经网络模型)去encode这种特殊性(如果存在),然后对比看看用这个模型是不是能有些优势,如缩短计算映射的时间,减少存储等等。当然,如果这种映射关系能被人理解,也可以更好地去实现系统模块。

这篇paper就试图用这个想法去优化索引 (Index) 问题。首先,索引的过程本身就是模型:Range Index索引可以看做是从给定key到一个排序数组position的预测过程,Point Index索引可以看做是从给定key到一个未排序数组的position的预测过程,Existence Index索引可以看做是预测一个给定key是否存在(0或1)。这样一来,索引系统是可替换的,我们可以直接用ML模型去替换现有的索引模块实现。其次,数据特殊性对索引设计的影响明显。假设用于Range Index的key的范围是1-100M,那么可以用特殊的hash函数(直接把key本身当成是offset)去做索引而不必再用B-Tree。然而,用ML模型去做索引也可能存在一些问题:

   - Semantic Guarantees. 通过数据训练得到机器学习的模型是存在误差的,但是我们设计的索引经常需要满足严格的限制条件。举例来说,bloom-filter的优势之一是可能会有把不存在的错判成存在(False Positive)但绝不会把存在的错判成不存在(False Negative)。

   - Overfitting. 如果避免模型的overfitting?换句话说,如果当前训练的模型只能memorize当前数据形态而并不能generalize新的数据,该怎么办?

   - No Specialization. 如果数据并没有任何特殊性,如完全随机或者不能被我们的设计模型架构所归纳,该如何处理?另一种情况是,新的数据破坏了之前模型归纳出的特殊性,比如新的数据来源于一个新的或者是变化了的distribution。

   - Evaluate Efficiency. 模型的inference过程相比显式索引数据结构的计算更昂贵。

顺着上面的思路,带着这四个问题,具体来探索一下数据库系统中的三类索引:Range Index ,Point Index和Existence Index。

Range Index. 数据库系统通常使用B-Tree或B+-Tree来实现range index。查询时给定一个key(或一些定义range的keys),B-Tree会索引到包含该key的对应范围的叶子节点,在叶子节点内对key进行搜索。如果该key在索引中存在,就会得到其对应的position。这里position代表指向逻辑页的一个offset或pointer 。一般在一个逻辑页内的records会用一个key来index。如图所示,输入是一个key,输出是对应要查询record的position的区间[pos, pos + pagesize]。

index1

有趣的是,在机器学习的观点下,B-Tree索引实际上可以看成是一个regression tree模型,这个模型(B-Tree)的建立过程也是依赖数据的,只不过它不是通过梯度下降的方式得到的,而是通过预先定义的法则(insert函数)。如下图2,对于输入x,模型预测出一个pospred,若区间pospred - minerr, pospred + maxerr中包含这个record,那么预测准确,其中minerr的期望是0(这时pos就是record的位置),maxerr的期望是pagesize(这时record的位置在pos+pagesize)。那么minerr和maxerr在模型中具体如何计算呢?对于训练数据(现有的(key,pos)),计算出每一组pospred和和pos的正负差,然后取最大的正差就是maxerr,负差就是minerr。这里并不需要考虑不在索引中的key,因为B-Tree本身也不能保证在新的key insert后搜索出来的pos和insert前的pos的差距在 \( [0, pagesize] \)(insert前B-Tree根本找不到这个key!)。

index2

用图2的模型代替图1的B-Tree能否学到数据特殊性从而优化查询过程呢?假设有1M个key(取值从1M到2M),那么可以直接用一个等价关系(线性模型)代替B-Tree,也就能把log(n)的复杂度减少成常数复杂度。同样的,如果(key, pos)是一个抛物线,那么神经网络模型也可以通过常数复杂度进行预测。考虑到神经网络的泛化能力,如果数据具有特殊性,那么用模型代替B-Tree就能减少查询复杂度。具体B-Tree复杂度和神经网络模型复杂度的时间分析可以参考原文2.1。

更进一步,还可以从概率的角度来思考range index这个问题。Range index实际上是描述一个数据点到一个有序数组的函数关系,也就是说这个函数关系"模拟"出了数据的累计分布函数(CDF)。因此,可以这么来定义这个函数:\( p = CDF(key) * N\) ,其中p是position,CDF(x)是 \( Prob(X <= key)\),N是所有的key的总数。CDF函数可以简单理解成一种描述概率函数的wrapper函数,它描述随机变量X在小于一个key条件下的概率和。有了CDF,就能很容易地描述一个区间X的概率和:\( CDF(key2) - CDF(key1) = Prob(key1 <= X <= key2)\),它能够帮助进一步深入地研究概率分布理论。这么来理解此处的"模拟"就明白了:假设我们把pos和key的函数关系写成\(p = F(key) * N\),其中F是一个任意函数。那么F满足以下性质: $$设p1 = F(key1) * N和 p2 = F(key2) * N$$ $$p1 - p2 = (F(key1) - F(key2)) * N = F(key1 - key2) * N$$ 最后这个等式由range index满足,因为range index映射到的pos是一个在有序数列里。而巧合的是累计分布函数(CDF)也满足这个性质。虽然满足这个性质的函数不一定要是CDF,但是如果能把一个更强假设的函数(CDF)学习出来,那么一定也满足range index的功能!使用这个更强假设(CDF)的好处是:拟合数据的CDF函数是一个active的研究方向,可以受益于已有的研究成果。

基于这个想法,让我们先抛开拟合CDF,尝试用一个简单的两层全连接神经网络(每层32个neurons + ReLU acitivation)训练了一个200M的web-server log records,x是timestamp,label是对应B-Tree索引到的position。并测试了用训练得到的模型的预测时间,latency大约是80000 ns,throughput是每秒1250次预测。很遗憾,这比不用索引遍历还要慢。然而B-Tree的索引时间大概是300 ns。那么问题出在哪儿呢?首先,这里用了Tensorflow做训练和预测,Tensorflow是位较大规模的神经网络做的优化,但并不适用于这里的小网络,并且python层因为用swig做胶水层会造成一定的latency开销。其次,机器学习中的树模型(比如B-Tree模型)能够更好地记住数据特征,因为它能不断的split data space。而这里用的神经网络模型能够很快地学习出一个做地还不错的模型,但是很难做到局部数据上的准确(这里可以用overfitting来理解)。为了保证精度,我们要train更多的epoch,并且使用更深的网络来训练。这里还有一个原因导致需要用更深的网络来尽量达到树模型的精度:训练目标不仅要降低每个预测position和label间误差的平均,更重要的是要让训练出来模型的minerr和maxerr尽量小,也就是让预测出来的position在真正的page内(注意这里并没有把minerr和maxerr加到优化函数里,而只是假定预测精度越小,这两err越小)。同时,B-Tree是cache-efficient的:能把很少的上层节点放入cache。而对于全连接神经网络模型,要把模型都加载到内存,并且权重不是稀疏的!

为了尝试解决以上的问题,文中首先提出了Learning Index Framework(LIF)。LIF就是一个索引优化系统,给定一套索引数据,LIF自动生成多组超参数,并且自动地优化并测试,找出最优的一组配置。训练时,LIF任然用了TensorFlow,而在inference时LIF根据训练打包的模型,用code-generation(用特定的模型结构减少运行时开销,比如剪枝、去掉多余的if-else、direct memory access等,这里不展开。如果对数据库系统了解的同学应该比较熟悉)自动生成了相应C++的程序,这里实现具体的细节可以参考21。这样,可以把inference的开销减少到30 ns。

Recursive Model Indexes(RMI). 既然LIF可以在很短的时间里完成预测,我们就可以考虑用更复杂的模型来尝试替换B-Tree索引了。上文提到,单个全连接网络并不能达到很好地准确度。事实上,如下图所示,全连接神经网络模型虽然可以从整体上较好地拟合出整体数据分布,但由于每个数据个体本身的随机性,却很难同时在局部也准确拟合出个体数据,文中把这个问题称作the last-mile accuracy问题。换句话说,对于一个100M条记录的数据集,我们可以先用单个简单的模型把误差控制在10K并不困难,然后对于局部数据,再用另外的模型进行预测把误差从10K降低到100。因此这里的思路就是用多模型来解决the last-mile accuracy问题,这里的多模型的思路并不是boosting,而是来源于混合专家网络(下文用MoE,可以参考51)。

index3

MoE是90年代初提出的模型,下面先来介绍一下MoE模型的思想。真实情况下,数据往往来源于不同的概率分布,即使是单个概率分布也会随着时间变化。举个例子,比如股票价格,很难用一个连续的函数拟合分段函数(每天或每小时的变化曲线),也很难用一个分段函数去拟合一个连续函数(每月或每年整体的变化趋势曲线)。而如同上文所分析的,真实的数据分布函数往往是介于两者之间的:总体上满足一个变化趋势,但是局部却是分段函数(每天数据的分段函数可能都不同),因为这时随机因素更显著。MoE的想法是能否把数据partition,然后对于不同的数据,用不同的模型(称Expert)来进行拟合。这里需要特别指出的是,对数据本身的特征做partition并不一定能解决问题,而是应该根据不同的x和y之间的关系来对x做partition。

index4

比如在上图的例子中,8个数据点来自两个抛物线,如果只根据数据本身做partition(右侧黄线),并不能很好地把数据partition在两个抛物线上,而左边的蓝线才是一个更好的parition。所以MoE的想法就是把x和y之间的关系定义到目标优化函数中,然后训练出不同的Expert来拟合不同的数据。如下图里,Gating Network可以认为是赋予每个Expert的权重,比如可以从输入接一个神经网络得到Softmax概率或Attention当成权重。然后对于这一个特定的输入,只训练在这个Attention下(或者Softmax较大概率)的Expert模型。换句话说,当某个Expert1表现比另一个Expert2好,反向传播时(y_pred - y更接近)Expert1对应Attention的权重就会更高,这样下去对于这个输入MoE就更加依赖于Expert1了。就好比小学一年级老师给大家一份数学卷子,学生A比学生B做的好,以后只要是数学题,老师都会重点培养A,让他成为一个数学特长生。而B可能更擅长语文,这就是MoE的基本思想。这里需要注意两点,不同于boosting,在预测时MoE对于不同的输入分配给各个模型的权重是不同的(softmax的概率或者attention)。而boosting一旦模型训练完成,在预测时分配给每个模型的权重是完全一样的。另一点,MoE有个很好的性质:虽然在训练的时候由于模型很多,参数也很多。但是在预测时候模型的参数并不随着模型复杂度的增加而增加,因此计算开销也不回增加,因为我们总是能Gating Network找出几个相应的Experts来做预测。这种稀疏性对于能够成功替换B-Tree索引至关重要!

index5

回到论文,文中提出了一种特定的MoE模型(文中称作Recursive Regression Model):RMI包含从顶到下M层,每一层包含 \( M_{l} \) 个Experts和一个Gating策略 \( M_{l} * f_{l-1}(x) / N \) ,每一层的Loss function是 \( L_{l}=∑_{(x,y)}(f^{M_{l}f_{l-1}(x)/N}(x) - y)^2 \) ,也就是说对于特定的(x, y),RMI委托特定路径的M个experts来对特定范围的x进行训练和预测,很好地解决了the last-mile accuracy问题。可以看到,这个MoE的特殊性在于,除了最下层的experts,其他层的所有experts的训练目标都是为了做Gating。我们可以看成是学一个attention(只有一个非零)的Gating Network,而只有最下层的expert才是做预测真正的experts。这里能这么假设的原因在于对于最下层,相近position范围内的数据分布一致(上一段提到我们通过x和y的关系来对x来做partition,这里通过除了最下层的所有experts层重新组织x来做这种partition使得在一个连续范围内的y的x被partition在了一起),用单个模型就够了(这里可以做进一步优化:在最下层前学出一个稀疏向量或者真正的attention)。需要注意的是,这里的分层模型并不是树,因为每个父节点也可能指向其他父节点的孩子。另外,在顶层expert可以用ReLU neural这种泛化表达能力更强但拟合性稍差的网络,而底层expert可以用拟合性更强的线性回归模型。

indexsp

本文开始提到No Specialization的问题:如果数据本身是随机的,RMI模型并不能归纳出数据中的分布规律,应该怎么办?这里的思路就是rollback到B-Tree Index(下文会看到对于Point Index和Existence Index都会有类似处理来解决No Specialization问题)。由于RMI模型的稀疏性和LIF的高效率,即使对于这种情况,加入RMI的range index效率也并不会降低太多。整个训练过程见下图,文中是分stage训练的(也可以用end-to-end),训练时可以用grid search和一些best practice来tuning hyper-parameter。我们知道,在训练完RMI后,能得到每个experts的minerr和maxerr。我们可以用它们来和预先给定的threshold(比如设成pagesize)对比,当成rollback的条件。更重要的是,它们能够优化索引搜索过程。通常的B-Tree在找到输入key的offset后会用binary search(也有其他的方式如Quaternary Search但是效率提高不明显)。而RMI模型预测出来的position可以用作更准确的pivot来进一步提高搜索record的性能。论文中提出了Model Binary Search,Biased Search和Biased Quaternary Search,这部分比较简单,就不展开说了,可以参考论文3.4。

index6

论文中用了四个数据集来做了实验,分别是Weblogs(日志数据,用timestamp建索引,这个最符合真实复杂环境), Maps(地图数据用经度建索引,虽然这个也是真实数据,但是相对更规则), Web-documents(真实的document数据)和合成的Lognormal(采样生成的数据,为了模拟真实环境中的长尾分布)。其中Weblogs、Maps和Lognormal数据使用整数来当索引的,而Web-documents是用字符串来当索引的(因为不连续)。下面分别列出了四组数据的实验效果。在每组数据集的实验中,论文以pagesize=128作为baseline,分别给出了Learned Index与B-Tree Index索引时间(Model列表示搜索position的时间+Search列表示搜索record的时间)的对比、索引内存占用的对比以及索引误差。在最后一组Web-document数据的实验里,文中对比了不同的搜索策略,因为对于string的binary search会更耗时,所以上文提出的搜索策略能够更明显地提高效率。这里Learned Index只用了两层的RMI,Learned Indexed Complex表示第一层用了两层全连接的网络,而不带Complex的没有用hidden layer。而第二层都使用了linear model,论文对第二层里不同的experts数目做了不同的实验。

index7 index8 index9 index10

总的来说,在索引的内存占用上大大有了显著减少。在索引时间上,在特定的配置下也能有一些提升(最后一组实验比较有说服力,因为我们的初衷是找出人们发现不了数据特殊性)。可以看到,这里实验用的RMI只有两层,而且误差不稳定。猜想是如果更深的RMI,精度提出和时间效率提升以及内存占用提升相比并没有太好,当然未来可以进一步采样模型压缩来优化。

Point Index. Point Index顾名思义就是为了搜索出key对应的position,也就是Hashmap。Hashmap天然就是一个模型,输入是key输出是position。对于Hashmap来讲,为了减少冲突,往往需要使用比key本身数目大很多的slots数目(论文中提到在Google通常一个hashmap有78%额外的slots)。对于冲突,虽然有linked-list或者secondary probing的方式来处理,但是它们的开销都不算小。即使用sparse-hashmap的方式来优化内存占用,性能也要比hashmap慢3-7倍。因此能否通过学习出的模型一方面来减少额外内存开销的,另一方面减少一些冲突时的开销(Hashmap可以通过用更多的slots数目来避免冲突,如果冲突很少就很再难优化索引时间因为Hashmap本来就是常数复杂度的,所以说到底这里的目标是在保证索引时间的情况下减少Hashmap的内存开销)。如下图所示,Learned Hashmap可以在slots数目一样多的条件下既减少冲突,也就能更好地利用内存空间。

index11

文中通过实验验证了拟合累计分布函数CDF对学习Hashmap也起作用:\(h(Key) = F(Key) * M \),其中想训练处F去拟合数据的CDF函数。由于point index不需要保证records连续(不需要考虑maxerr和minerr),所以训练出的模型F拟合CDF更容易得到满足。如果这时仍有少量的冲突,可以用linked-list的方式来处理。下图给出了实验结果,仍然使用了B-Tree实验中的Map、Weblog和LogNormal数据集以及两层都是线性experts的RMI模型(最底层100K个experts)。结论是learned index能达到更少的(~80%)empty slots而搜索时间和原本Hashmap相近,这大大地提高了内存使用率。同样,若learned index不能很好地学出CDF(比如冲突很多),系统可以rollback到Hashmap。

index12

Existence Index. 对于Existence Index,DBMS采样Bloom-Filter来实现:对于一个key,算出k个hash函数得到的在bit array上的position并标记成1。Bloom-Filter能天然杜绝false negatives的存在,但有false positive的存在。为了降低false positive,比如对于一个100M记录的数据,Bloom-Filter通常需要比records本身大14倍的bits来达到0.1%的false positive准确度。

index13

这里learned index目标和之前讨论的B-Tree及Hashmap都不一样。之前我们想学习出现有的keys里的分布规律,而这里对于存在的keys我们却不需要区别对待。相反,我们想通过模型学出存在的keys和不存在的keys间的差异。换句话讲,如上图所示,我们期望存在的keys对应索引的positions尽量在一起(无所谓冲突不冲突),不存在的keys对应索引的positions尽量在一起,而存在key和不存在key索引出来的position尽量分开。同时,learned index应该满足尽量低的FPR和内存占用且保证FNR为0。

论文中提出了两种思路来建模,来满足上述要求。第一种方式可以把existence index看成是一个binary classification问题。考虑到existence index的使用场景通常是先判断record在cold storage存不存在,再通过disk或网络access数据。所以latency相对可以高一些,这样我们就可以用更复杂的RNN或CNN模型来当做分类器了。在预测的时候,模型会得到输入key是存在的可能性,需要给定一个threshold(τ)来做出实际判断。然而,这样训练出来的不同模型有其对应的FNR和FPR。我们可以通过threshold来降低FPR,但FPR降低FNR会增加,因为对于一个训练好的模型,调整τ整个False Rate是不变的(FR = FPR+FNR),这破坏了系统的Semantic Guarantees。论文解决的思路非常简单:对于模型要给出negative预测的keys,借助(rollback)传统的bloom filter来保证这个特性。由于这时的bloom filter只对< threshold的数据建,因此总体来说learned index可以很好地优化内存。即使FNR = 50%,理论上我们可以把内存优化到一半(类似Amdahl’s Law)。这样一来, threshold的选取就只需要考虑FPR了。

index14

第二种方式,还是把问题看成一个binary classification问题。训练得到了模型f以后,因为它满足不了FNR=0,我们可以把f来代替Bloom-Filter中的hash函数:\( d(f(key)) = m * f(key) \),其中m是bit array的大小。用这样的函数做成的bloom-filter可以大大减少bit array的大小。可以看到,两种方式都把上一段提到的目标当成一个binary classification问题,然后为了Semantic Guarantees,都是结合Bloom-filter本身来实现的,只是在不同level的hybird。

对于Existence Index,论文用了Google内部的1.7M个钓鱼网站URL当成是存在的key,把合法的URL当成是不存在的key,模型用了GRU。下图给出了在不同GRU配置下和不同FPR的threshold下的结果。实际上,相比B-Tree和Hashmap,这里实验给出的内存优化效果是比较明显的。并且模型预测越准确,内存降低地越明显,但这里论文中没有给出性能对比。

index15

回到本文最开始提到的四个问题。我们已经通过例子讨论了如何解决Semantic Guarantees和Evaluate Efficiency。那么learned index如何处理Overfitting和No Specialization问题呢?对于overfitting,论文的假设是要被新索引的key依然满足CDF。比如对于Learned Range Index的 insert操作,如果新的key也同样来源于训练出的CDF分布,那么我们并不需要重新训练,直接insert到预测出来的position就行了。而B-Tree则需要O(log N)甚至re-balance。但RMI解决the last-mile accuracy的方式决定了learned index并不能很好地generalize,这是可以进一步完善的地方。而对于No Specialization的数据索引,虽然learned index并不能generalize,但能在当前情况下提高查询效率,减少内存使用。对于OLAP为主的系统,更新索引的操作本来就是定期进行,可以这时再重新训练模型,这时训练的初值可以使用之前训练出来的权重。如果OLTP数据更新更频繁的系统,可以采用delta-index的一些技术,增量地更新learned index。文中还提到当CDF分布变化时的模型处理,这里不展开介绍了,详见3.7.1。另外对于面向disk索引的数据库系统来说,逻辑页并不是真正的物理页(从而不能满足前文分析的性质),论文中也给出了几个可能可行的方式优化,详见3.7.2。

说来巧合,其实类似的问题去年和同事讨论过。在Pivotal的时候,我和几个同事结合ML和DB做过一些尝试,比如用SQL写推荐系统,通过Transaction Log检测作弊用户和异常行为并集成到查询层。也经常讨论,如果用数据库系统去实现深度学习,可行吗?有哪些优势?我们从存储,计算,表达问题的泛化能力等几方面讨论过很多次,但每次都没有很好地去切入。现在想来,有点惭愧,因为我当时更关注的是如何用数据库里的一些技术做ML,而不是如何用ML来优化数据库设计。一次,明哥问我能不能用ML尝试优化plan,减少查询次数,能否对各类的join做些优化。我们讨论到索引优化,但遗憾的是当时并没有深入。不过,就像paper里所断言的,this is a fruitful direction for future research。


Hong Wu at 20:00

scribble