也谈用机器学习索引替代B-Tree

昨天朋友圈里掀起了新的一阵膜拜Jeff Dean风，“The Case for Learned Index Structures”。大家都在讨论机器学习如何改变数据库中一些基本问题。

程序狗昨天也读了文章的前几章，想从算法和影响两个角度谈谈自己的认识，前3节快速浏览文章内容，后3节是实验和我的个人理解。

引言

每当我们需要快速访问大量数据时，最先想到一个优化方式就是建立索引。比方说，当我们想查找特定时间范围内的所有记录时，就会采用B-Trees；当我们想查找某一个键值的时候，就会采用Hash-Map；当我们只是查找这个简直是否存在的时候，就可以采用Bloom-filters。

在理论学习过程中，我们常常去分析一个数据结构均摊或最差的情况下时间复杂度是多少，工作的时候会忽略现实数据的形状。比方说，当我们知道要保存的数据均匀分布在1到100M之间，那么直接使用数字本身作为索引而不建立B-Tree可以将查询时间复杂度从O(log n)降到O(1)，将空间复杂度从O(n)也降到O(1)。换句话说，针对数据分布进行索引优化可以优化时间复杂度和空间复杂度。

然而，传统的数据库优化都是针对一般情况的优化，对于特定数据进行针对优化需要消耗不少的人力。所以Jeff等人就提出了机器学习的方法，这也是我个人认为本文最大的影响以及创新性所在。

B-Tree和机器模型的相通性

其实B-Tree本身也可以被看做一个机器学习模型，对于给定的键值预测他的位置，就如下图所示。

出于对索引效率的考虑，通常我们不会去索引每一个关键字，而是把大约每n个关键字放到同一个page，然后对page中第一个关键字进行索引。因此，我们也可以把B-Tree看做一个regression tree，将关键字映射到区间内，区间两端则是pos-min_err，和pos+max_err。

继续类比，这里的min_err和max_err就是机器学习模型在训练数据上的最小误差和最大误差。如果我们记录了每一个关键字最小误差和最差误差，那么对于任意一个关键字，就可以通过模型对数据位置进行预测，如果关键值存在，则可以在最小和最大误差之间搜索到。

机器学习Range Index

正如上文中提到的，我们可以设计一个模型，根据关键值预测出现位置。对于常见的范围查询，所有的数据都是排好序的，能想到一个简单的模型就是预测给定关键字的累计分布函数(cumulative distribution function):p=F(key)*N，这里的p就是预测位置，F(key)就是预测的累计分布函数。

作者文中提到了简单的预测分布函数所面对的挑战，其中既包括了Tensorflow本身的效率问题，又包括了预测函数的误差难以控制问题。为了克服上面的挑战，论文提出了几种学习框架，其中最为基本的是learning index framework (LIF)和recursive-model indexes (RMI)。

LIF主要是针对一个模型索引的优化，针对给定的模型规则，LIF生成索引配置，优化并且自动进行测试。它采用Tensorflow自带的神经网络模型去学习简单的函数，比如线性回归函数。在真正的数据库中应用时，LIF会生成基于C++的高效索引结构，将模型预测优化到30ns以内。

RMI则主要优化得到预测结果后last-mile查找准确性。举一个例子，要通过单一模型预测100m个关键字位置产生的最小和最大误差是很大的。但是如果我们采用两层模型，第一层模型将100m个关键字映射到10k个位置上，第二层将10k个关键字映射到1个位置上，则最小和最大误差会降低很多。

就像上图显示的模型结构那样，Model 1.1先对关键字进行预测，然后通过预测结果选择执行Model 2.1。另外要注意的是，这样的多层模型结构并不是一个查找树，多个stage 2模型可能会指向同一个stage 3模型。另外在最下层，为了保证最小和最大的偏差是可以控制的，我们也可以采用B-Tree来完成最后索引。

上面给出的Hybrid End-To-End Training就描述了如何逐层训练一个Hybrid的模型，其上层为两层ReLU神经网络，底层为B-Tree。

这样递归的模型有几个好处：1.它能够很好的学习出数据形状。2.通过分层，底层可以使用B-Tree或者decision tree得到非常高的查找效率。3.不同的stage可以使用偏移量调用模型，而不用搜索，这样更容易进行程序语言的底层优化。

实验效果之我见

Google团队出品的文章一大优点就是实验过程非常清晰，并且有说服力。他们测试了这么几种不同配置的算法，其中包括不同page size的B-Tree，不同大小2nd stage的机器学习算法 (需要特别指出的是这里的2nd stage并没有采用Hybrid的结构，而是简单的神经网络模型)，以及比较复杂的机器学习模型。

作者将模型时间分为Model时间，也就是在模型上检索花费时间，以及Search时间，在找到对应page后检索花费时间。

在这里其实有很多细节可以讨论。首先，不同Page size的B-Tree的模型大小和Page size成反比，128/16=8，所以Page size 16的Size Saving是-700%。尽管Page size 16最后的Search时间很短，但是由于模型比较深，除了CPU运行还要考虑很多内存切换时间，所以Model时间很长。

Learned Index对比Page size 128的Model时间平均提升了3倍以上。正如作者团队在之前估算的一样，两个模型的运算能够在非常短的时间内完成。于此同时，第二层的模型最小和最大误差在均摊情况下非常的好，所以Search时间一般没有比B-Tree差。

在结构方面，由于我们只要保存数10k~20k个模型，每个模型又非常简单，所以内存消耗比包含上百万节点的检索树要小一个数量级。

用AI来优化基本问题

我个人看完实验结果后第一个反应就是时间上的优化主要来自于Model时间，而非Search时间，并且应用的问题是Range Index，那么我们是不是可以把问题转换为一个更简单的算法问题，而非机器学习问题。

我们可以定义一个Hash函数，h=floor((key-a)/b)，我们需要找到a和b使得少于page size个key得到相同的h。这个Hash函数就是我们要做的模型。

一个Hash函数可能不够，因为得到的key范围可能很大，所以我们需要多个hash，分为多层去尝试解决这一个问题。

经过这两步，我也可以把B-Tree这一个复杂的检索模型通过几个简单的Hash模型替代。然而仔细想想，我提出的问题也不是一个容易解决的问题。Hash函数需要通过枚举得到最优的a和b，这样要O(n^2)时间复杂度，在现实中是行不通的。

这样来看，AI则是一个很好的工具，帮助我们在非常大的空间里搜索优化一个答案。从这个角度来看，索引就是一个自适应函数，有明确的目标，我们可以去训练。这突破了我们通过算法视角形成的思维方式，比如建立一颗二叉搜索树，对键值的插入是非常严格的，需要调整结构维护树的平衡，从而达到理论最优解，但是最后导致的结果就是模型非常复杂，如上面提到的B-Tree Model时间都要100ns以上。

再举两个例子，从前用模拟退火可以快速的求出旅行商问题的的一个近似解答。不断的找寻新的路径P(i+1)，然后通过模拟退火的概率选择接受路径P(i+1)，最后进行降温。反复上溯步骤，其实就是一个AI不断尝试更新向量，然后降低learning rate的过程。

现在人们炒的火热的AlphaGo Zero，使用神经网络来学习每一个下子的获胜概率，然后降低搜索空间。最新的版本通过自己和自己玩来进化算法，获得更准确的获胜概率预测，这一切本质上是将AI作为工具，帮助最优化算法减少搜索空间。

结语

这是一篇非常有意思的文章，最开始推敲后又觉得其实问题并没有Jeff团队说的那么复杂，但是和一些老师讨论后对论文的意义有了新的认识。

文章中间还有一个比较有意思的想法就是分治算法也可以重叠，并不一定真的需要分而治之。在同一个stage的不同模型，可能都对应到下一层的同一个子模型，这样的假设可以在降低模型复杂度的同时使算法容错率更高。告别传统树状模型，也算是点睛一笔。

总的来说，论文指出了大多数计算机科学家很容易忽略的一些非常实用的想法，这些想法在真实情况下非常有效。可是由于模型目前不可以被解释，复杂度不受控制的原因，暂时不能取代传统做法。

最后，引用一位老师的话，这个论文将一个fundamental building block用AI来优化，探究理论算法的实践极限，这估计是未来的一种趋势。还有哪些fundamental building block可以用AI优化？AI优化和理论算法在什么情况下会收敛？我很期待。

欢迎关注硅谷程序汪

原创文章，版权所有。