一种基于混合索引的最近邻查找方法

彭永鑫，罗英

（1.商洛学院 数学与计算机应用学院/秦岭康养大数据陕西省高校工程研究中心,陕西商洛 726000；2.中国兵器工业信息中心，北京海淀 100089）

数据库系统中，一般使用索引技术来提升数据的存储性能。数据库系统通常把磁盘作为持久性的存储设备。然而使用磁盘会带来较高的访问延迟，磁盘输入和输出次数的降低成为提高数据库性能的重要指标。如何高效地在磁盘中进行数据检索成为数据库系统中被人们所关注的问题。传统方法通常从改进查找过程及更快命中数据两个角度来提高范围查找效率，但只能在一定程度上降低维度和数据量的增加对查找效率所带来的影响，所能取得的效果有限。作为树型索引结构的代表，KD[1]树由于具有准确度高、对异常数据不敏感和查找准确等优点广泛应用在最近邻查找之中。然而，在大数据环境下，传统的索引都面临诸如空间占用率高、查找过程需要多次间接搜索等问题，影响到了查询性能。尤其是树形结构，随着实际应用中数据维度的不断升高，在进行精确查找的过程中，会产生“维度灾难”[2]，使得搜索变为线型扫描。找到一种既能提高索引效率，同时又能满足一定索引精度的算法，成为新的研究方向。

随着人工智能的发展，研究人员开始使用机器学习技术来提高数据库的索引性能。关于可学习索引（learned index）[3]的工作为索引结构的研究开辟了新的方向。树型索引结构例如B树，可以被认为是一棵把键映射到索引条目排序数组中相应位置的回归树。这项工作开启了用基于数据分布构建的学习模型替换现有索引树可能性的研究，提出的RM-Index[3]在搜索性能和索引空间开销方面都优于传统的B树。

在可学习索引中，可以把索引本身当作一个机器学习模型。无论是B树还是可学习索引模型，输入值都为一个待查找的键。不同之处在于B树通过搜索查找，会得到这个键所在记录的位置。而可学习索引模型在待查找的键输入后，经过神经网络模型的一系列运算，会输出一个包含最大误差和最小误差的区域，并能在这个误差区域内使用二分法找到记录的确定位置。传统的索引结构往往没有考虑数据的分布，而可学习索引利用机器学习进行模型的构建，在训练的过程中能够通过不同层级的神经网络得到数据的分布信息。

Peng等[4]研究发现，相对于传统的KD树，使用基于可学习索引的可学习KD树模型可以有效地提高最近邻查找的查找效率。本文在可学习的KD树[4]的基础上，提出一种基于混合索引的最近邻查找方案。混合索引将传统KD树在查找精度和可学习KD树在查找效率上的优势相结合，能够根据对查找效率和查找精度的不同要求，更加灵活地解决空间中的最近邻查找问题，并表现出较好的试验结果。

1 KD树与混合索引

1.1 KD树及其改进

KD树的概念自从提出以来，就被广泛应用于最近邻查找之中，用来解决在k维数据空间中如何为数据集建立索引的问题。为了在已知的样本空间中高效地查找到输入点的最近邻点，通常采取建立索引的方式，“以空间换时间”。KD树采用分治法的思想，利用已有数据对k维空间进行切分。作为KD树的改进，Ball[5]树中划分特征空间后形成的超矩形体改为超球体，提高了查找效率。然而，树形结构需要对空间进行细分，对可能是精确的近邻点都要进行对比计算，会占用较高的计算资源。因此，Zhou等[6]利用了GPU的加速功能，让KD树在GPU上运行，目的是为了提高KD树的运算速率。其他一些方法也能在一定程度上解决由于维度升高造成的影响，但所能达到的效果有限。

1.2 可学习索引

RMI[3]采用了类似B树的分层机器学习模型来代替传统的索引结构，构建了一个层次化的模型，如图1所示。除去第一层外，每一层都由多个小模型构成，每一层模型都和上一层模型互相关联，上一层模型的预测结果将会决定接下来一层的数据如何进行划分。每一层模型都会缩小查找范围，直到最底层得到记录位置的最终预测。由于误差的存在，一般情况下还需要使用二分法定位记录的精确位置。相对于传统的B树，可学习索引在查找性能和减少内存占用上有了较大的提升。除了一些专注于可学习索引模型本身的研究，例如Xindex[7]和Colin[8]外，还有一些研究将注意力集中在多维范围的可学习索引上[9-10]。官嘉林等[11]研究了内存预分配策略，较好地降低了内存缺页中断率，提高了可学习索引的性能。

2 混合索引的构建

混合索引模型由可学习索引和传统KD树两部分构成。可学习索引部分通过构建机器学习模型，来代替传统索引结构中的查找和计算部分。可以将通过索引进行最近邻查找的过程看作是一个机器学习中的分类任务。在最近邻查找中，输入值和其所对应的最近邻点存在联系，这种联系以距离的形式体现，距离越接近的，联系越紧密，可以认为是进行聚类的过程。在构建模型的时候，采用监督学习的方式。模型训练的过程如下：对于处理好的数据集，首先将其构建成一棵KD树，随后，产生一系列随机数作为KD树的索引值，即为标签一，与构成KD树的每一个节点相对应。接着，通过KD树得到训练集和测试集的最近邻点，并将所对应的最近邻点转化成为索引值，即为标签二。标签二就是训练集和测试集的标签。准确率定义为正确查找到真实k近邻点的所有数据占所有待查询数据的比值。

模型训练阶段的算法如Algorithm 1所示：

影响可学习索引运行速度的一个重要因素是神经网络的层数。一般来说，隐藏层的数目越多，提取到的特征就会越准确，整个神经网络的误差就会越低，从而让精度得以上升。然而，当层数增加时，会使得总的参数量增加，导致神经网络更加复杂。同时也可能会使神经网络的训练出现过拟合现象，训练时间也会变长。为了使神经网络的性能得到一定程度的保证，同时保证模型的泛化能力，尽量避免过拟合的发生，在保证准确率的基础上，尽可能使用较少的神经网络层数。

使用KD树进行最近邻查找时，主要包含两个部分：首先是把最近邻点的叶子节点当作待查找点的近似邻最近点。其次是进行回溯操作，以待查找点和最近邻的近似点的距离沿树根部进行回溯和迭代。一般来说，KD树能够有效降低查找的次数。但是在某些场景中，待查询数据点所在的位置会需要在另外一棵子树上进行查找，造成查找次数增加，从而降低了查找的效果。随着数据维度的升高，使用KD树进行最近邻查找的性能会明显降低。通常情况下KD树在20维以内的数据集上较为有效。

在使用KD树进行最近邻查找的过程中，需要经过大量的回溯操作和距离计算，从而得到精确的结果。对于可学习索引来说需要做的是将这些回溯和计算用神经网络进行代替，同时辅以少量的计算和排序，在提高神经网络准确性的基础上，尽量得到相同或相似的结果。Peng等[4]详细介绍了基于可学习索引的可学习KD树的构建过程和在最近邻查找中相对于传统KD树的优势。

使用基于神经网络的可学习索引方法，利用神经网络并行计算的优势，可以有效地降低运算过程所花费的时间。同时，神经网络模型还可以通过设计更好的网络结构、降低参数数量等方法进行改进和优化，在确保模型准确率的情况下提高模型的运算速度。然而，使用这种方法，往往难以达到100% 的准确率。虽然使用传统的树型结构查找更为精确，但在查找速度上容易受到较多因素的制约。因此，将二者结合，通过构建混合索引，发挥两种方法各自的优势，有理由相信会取得更好的效果。

混合索引的工作流程如图2所示。对于原始的KD树，首先利用可学习索引，使用神经网络进行训练。随后，把构成KD树的数据点一分为二，并使用任意一半的数据点重新构建一棵KD树。对于待查找的输入数据，同时输入神经网络模型和新构建的KD树之中，使之并行计算。两者都输出k个近邻点。接着，将二者输出的结果进行混合，得到2k个近邻点，这2k个近邻点中包含了正确的k个近邻点。最后，对这2k个点相对于输入点的距离进行排序，根据k值得到最终的k近邻点。

图2 混合索引的工作流程

对于一条输入数据，假设使用KD树进行k近邻查找所需时间为t1，使用一半节点构成的KD树进行查找所需的时间为t2，使用神经网络查找所需要的时间为t3，后续计算所需的时间为tn，准确率为acc1。使用混合索引时后续产生的计算量的用时为tmix，总的准确率为acc2。一般来说，使用KD树进行k近邻查找需要花费的时间最长。使用混合索引时，要对2k个数据进行计算和排序，而使用可学习的KD树，需要计算和排序的数据个数为k，因此使用混合索引产生的计算量比可学习KD树所产生的计算量要大，花费的时间也会更多。即tn

传统KD树是精确查找，一般情况下，使用一棵构建好的节点数为初始节点数一半的新KD树，会使得整个查找结果准确率提高（因为对这一部分数据，得到的结果是准确的）。因此，使用混合索引，相对于使用可学习KD树进行k近邻查找，在查找时间增加的情况下，一般都会获得准确率的提升：

由于使用可学习KD树模型所花费的时间一般比传统KD树要低，而使用部分节点构成的KD树比使用所有节点构成的KD树进行k近邻查找所花费的时间也要低，所以相对于单一的使用可学习KD树模型，会在原本运算速度较快但准确率较低的情况下，牺牲一些运算时间以达到精度上的提高。在合适的情况下，混合索引的方案应该是可行的。

3 结果与分析

将通过试验来评估混合索引的性能，并与传统KD树和可学习KD树在最近邻查找的性能上进行对比，从而验证使用混合索引进行最近邻查找的可行性。试验使用的框架为Keras。Keras是一个使用python实现的高度模块化的神经网络组件库，提供了许多API模块封装了来自TensorFlow的诸多小组件，作为用户，只需要将API构建的模块排在一起就可以设计出各种类型的神经网络。试验中所使用的Keras版本为2.2.4。试验在一块NVIDIA GTX1080Ti 11G GPU上进行训练，采用的编程语言及版本为Python 3.6.5。

试验数据集采用不同数据量和不同维度组合的生成数据集。该数据集满足正态分布，数据量n为10万条数据，数据的维度d分别为10维、20维、50维。从查找准确率和查找时间两个角度来评价模型的性能，可以认为传统KD树的查找准确率为100% 。可学习索引的神经网络模型采用全连接网络，神经网络的层数分别为5层、6层和7层，所使用的激活函数为Adam。首先验证数据集在可学习KD树模型上所能取得的效果，随后对准确率较低的试验组使用混合索引模型进行改进。采用10次试验结果的平均值作为最终的试验结果，如图3所示。实线部分为可学习KD树的试验结果，虚线部分为经过混合索引模型改进后的结果。

图3 可学习索引模型和混合索引模型的准确率

从图3实线部分可以看出，使用可学习索引模型进行最近邻查找，查找的准确率随着维度的升高而降低。当维度较低，为10维时，无论神经网络的层数如何，模型的查找精度均表现得较好，准确率均在97% 以上。当维度为20维时，准确率开始下降，层数为6层时准确率低至90.3% 。当维度为50维时，准确率出现了明显的下降，最低至90% 以下。因此，在维度为20维，层数为6层和维度为50维的情况下，使用混合索引模型进行改进，结果如图3虚线部分所示。可以看到，经过混合索引模型的改进后，查找的准确率有了一定的提升，分别由90.3% ，88.6% ，90.1% ，92.3% 提高到了97.9% ，94.2% ，96.5% ，95.1% ，准确率分别提升了7.6% ，5.6% ，6.4% ，2.8% 。

图4展示了在20维和50维的情况下查找时间的变化。可以看到，无论是可学习索引模型还是混合索引模型，其查找时间相对于KD树都有较为明显的优势，查找时间随着神经网络层数的增加而增加。使用混合索引模型改进后，相对于可学习索引模型，查找时间略有增加，分别由8.9，17.5，18.1，22.3 s增加到 12.5，20.3，22.6，25.9 s。

图4 KD树可学习索引模型和混合索引模型运行时间的对比

总的来说，相对于单一使用KD树或者可学习KD树模型，混合索引能够综合两种方法的优势，实现了查找效率和查找时间的平衡。即相比于单独使用可学习索引模型，使用混合索引时，付出了一定的查找时间增加的代价（这个代价相对于使用传统KD树进行最近邻查找所花费的时间来说是比较小的），可以获得准确率的进一步提升。极端的状况是，正确的k近邻点全都没有分布在新构成KD树的数据中，则混合索引从准确率上退化至使用神经网络，时间上相比神经网络反而花费更长。在实际运用的过程中，可以根据对准确率和查找时间的不同需要，选择不同的方式进行最近邻查找。

4 结语

大数据时代，数据检索的快慢是衡量存储系统性能的一个关键因素。传统方法的查找结果通常较为精确但检索效率较低。为了解决这个问题，Kraska等[3]提出了可学习索引模型，在提高查询效率和降低内存占用上取得了良好的效果。但在查找效率获得较大提升的情况下，需要进一步提高查找的准确率。不足之处在于新构建KD树的节点数目、神经网络层数和节点数目的选择都是基于经验等。未来，将会从理论的角度出发，研究在更高维度和实际数据下该算法所能获得的效果。