【AI理论】Jeff Dean出品：用机器学习索引替代B-Trees，3倍性能提升，10-100倍空间缩小

译者 | 马卓奇

编辑 | Emily

数据库最开始是统一的，一刀切的“黑箱”问题。随着时间的推移，这一观点细化到了“标准尺寸”的 OLAP数据库和 OLTP数据库。

数据库使用索引来快速访问数据。B-tree和哈希映射是常用的实现索引的技术。但从“黑箱”的观点来看，数据库处理将数据视为不透明的，盲目地使用这些索引方法，而不对数据做任何假设。然而，很明显的一点是，不清楚数据分布会使算法不能发挥它最大的作用。考虑一下这样的实验：如果关键字从 0到 500M，那么直接用关键字进行索引要比用哈希要快。如果我们知道数据的累积分布函数（cumulative distributed function，CDF），这种观察可以扩展到其他的数据分布。我们可以概括成“ * CDF 关键字 * 记录大小”可以得到关键字指向的记录的大致位置。

那么，通过了解数据分布，我们可以获得性能提升。但是当我们把这个问题视作全“白盒”问题时，会失去可重用性。我们不能全“白盒”，因为每次都从头检查数据、设计索引的代价太大了。

这篇文章研究表明，利用神经网络来学习数据的分布，我们可以有一个“灰盒”的方法来进行索引的设计，并且通过将索引设计为数据感知的，可以获得性能优势。

目前较为主流的索引结构可以根据各种访问模式的不同需求分为如下三类：

• B-tree，用以处理范围查询

• 哈希映射，用于点查找查询

• 布隆过滤器，用于包含集合的查找

我们主要关注用可学习结构替代 B-tree结构的部分。对于哈希映射，学习到的结构是基于数据累积分布函数的直接函数。

B-tree提供了一种层级结构的有效索引方法。

为什么可以用神经网络模型来代替 B-tree呢？从概念上讲，B-tree本身就是一个模型，类似机器学习中的回归树模型：它可以将查找键映射到页面中的位置，并且最小误差为 0，最大误差为页面尺寸，而且如果该键存在，则保证它一定能被定位。因此，我们可以用其他机器学习模型，包括深度学习模型来替代这种索引，只要它们能提供同样的最小和最大误差的保证。

下一步，如何用其它机器学习模型提供同样的误差保证？我们唯一需要做的就是为每个键执行模型，并记住一个位置的最坏预测。给定一个键，该模型对数据的位置进行预测；如果该键存在，则保证它位于由最小和最大误差定义的预测范围内。我们用已有的数据来训练模型。因此，我们能够用任何其他类型的回归模型取代 B-tree，包括线性回归和神经网络。

我们用学习到的模型来替代 B-tree能获得如下好处：

• 更小的索引：更少的内存或 L1缓存

  *更快的查找：因为索引较小，提升了查找速度

• 更多的并行性（TPU），如果是 B-tree则是层次性。

这里的关键思想是以提高计算为代价减少内存，寄希望于计算成本越来越低的市场趋势（如果你能在 TPU或 GPU上使用它，你会得到更多的好处）

为了克服上述挑战，并探索可学习模型作为索引结构的替代或改进的潜力，作者提出了可学习索引框架（Learning Index Framework，LIF）、递归模型索引（Recursive-Model Indexes，RMI），以及基于标准误差的搜索策略。作者主要关注简单、全连接的神经网络，仅仅是因为它们的简单性，但许多其他类型的模型也是可以考虑的。

LIF可以看作是一个索引合成系统，给定一个索引规范，LIF可以生成不同的索引配置，优化它们，并自动测试它们。给定一个训练好的 Tensorflow模型，LIF从模型中自动提取所有的权重，并且基于模型规范在 C++中生成有效的索引结构。实验证明，运行简单模型的时间只需 30纳秒。

用单一模型来降低最小最大误差是十分困难的，但是仅用简单模型来代替 B-tree的前两层，来得到同样的预测精度提升却是十分容易做到的。类似的，如果模型可以只关注子数据集，那么降低误差也就更简单了。

基于以上观察，作者提出了递归回归模型（如下图）。

作者建立了模型的层次结构。在每一层，模型将查询键作为输入，并基于它选择下一阶模型，直到最后一阶模型预测出位置。每个模型对键的位置进行了一定的预测，并利用预测结果来选择下一个模型，该模型负责键空间的某个区域，并且以较低的误差进行更好的预测。

该模型有以下几个优点：

（1）它利用了数据分布的整体形状具有很容易学习的特点。

（2）该结构有效地将空间划分成更小的子范围，类似 B-tree或决策树，能够在更少的操作次数内达到所要求的“最后一步”的准确度。

（3）在各阶段之间不需要搜索过程。

该递归模型索引的另外一个优点是可以构造混合模型。例如，在模型顶层，一个简单的 ReLU神经元或许是最好的选择，因为它通常能够学习各种复杂的数据分布，而底层的模型可以是大量的简单线性回归模型，因为它们在空间和执行时间上消耗都很低。

注意，混合索引可以将可学习索引结构在最坏情况下的性能保留在 B-tree水平。也就是说，当数据分布很难学习得到时，所有的模型都会自动替换成 B-tree，让它几乎成为一个完整的 B-tree。

论文提出了几个搜索策略：

（1）模型二进制搜索：默认搜索策略，传统的二进制搜索不同之处仅在于第一个中间点被设置为模型所预测的值。

（2）偏置搜索：这种搜索策略是默认搜索策略的改进，通过在每次迭代中不均匀地从中间到左右位置进行范围分割。

（3）偏置四进制搜索：最后，我们开发了一种新的搜索策略，在每次迭代中并不是选择一个新的中间点进行测试，而是三个新的数据点，即四进制搜索。

论文主要关注索引实值键，但是许多数据库依靠字符串索引，但幸运的是，大量机器学习研究都在关注字符串建模。

将字符串转化为可输入模型的特征向量需要进行词条化操作。由于大多数机器学习模型在输入尺寸相同的条件下表现更好，作者设置了最大输入长度 N。作者使用了一个相对较小的前馈神经网络的层次结构。与之前的模型不同之处在于输入为一个向量，而不是实值。

根据字符串的累计分布函数（CDF）设计出通用的机器学习模型是十分困难的，未来的大量研究工作可以围绕优化字符串索引的可学习模型展开。

为了对比可学习的索引和 B-tree索引，作者创建了 4个二级索引，3个真实世界数据库，（1） Weblogs，（2）Maps，（3）web-documents，和一个合成的数据库，（4）Lognormal。

Weblogs数据库，可学习索引结构与 B-tree索引对比结果：

Maps数据库，可学习索引结构与 B-tree索引对比结果：

Lognormal合成数据库，可学习索引结构与 B-tree索引对比结果：

作者主要对比了总查询时间和局部搜索时间，以及索引结构大小、节省空间、模型误差以及误差方差。从实验结果中可以看出，几乎所有的配置情况下，可学习索引结构的表现都领先于 B-tree，速度上快了 3倍，尺寸上小了一个数量级。

web-documents，基于字符串的索引数据库，可学习索引结构与 B-tree索引对比结果：

从实验结果中可以看出，对于字符串索引，可学习索引结构的优势并不是很突出。作者分析认为运行模型的计算量太大，这个问题可以用 GPU或 TPU解决。

学习得到的字符串索引在不同搜索策略下的对比结果：

（1）模型泛化性和“最后一步”准确度的折衷。模型的“最后一步”预测越准确，模型就会越容易过拟合，并且很难泛化到新的数据条目上。

（2）如果数据分布发生变化怎么办？可学习索引是否可以检测到数据变化？是否可以像 B-tree一样可以一直保证搜索和插入操作的复杂度为 O(logn)？

作者认为处理插入的另一种更简单的方法是构建 delta索引：所有的插入都放在缓冲区中，并且不时地合并到模型的再训练中。这种方法已经被大量使用，并且可以得到 GPU或 TPU的硬件加速支持。除此之外，也可以通过使用以前的解决方案作为出发点，热启动每个模型的训练。特别是依赖于梯度下降优化的模型可以从这种优化方法中获益。

对于存储在磁盘上的数据索引而言，将数据分割成存储在磁盘中不同区域的较大页面是很常见的。这种情况下，学习数据的 CDF分布不再满足建模要求。

作者认为有以下几种替代方法：

（1）利用 RMI结构，最小化模型覆盖区域的重叠部分。

（2）使用翻译表。

（3）用更复杂的模型学习页面的实际指针。

更多算法细节，请参考原论文：

https://arxiv.org/abs/1712.01208

查看英文原文：

http://muratbuffalo.blogspot.com/2017/12/paper-summary-case-for-learned-index.html

转自：https://mp.weixin.qq.com/s/MpuUdZi8CWcu0b-ij-bHjA