基于神经网络的可学习Kd树

彭永鑫

（云南大学 软件学院，云南昆明 650000）

近几十年来，高维数据在数据仓库、信息检索、数据挖掘等方面的应用越来越广泛。Kd 树作为一种用于查询高维键值的流行算法，由于其准确性高、可扩展性强、查询速度较快，通常被用来进行多维空间关键数据的检索。使用Kd树进行k 近邻搜索，就是对于给定的一个查询点q，需要从一个构成Kd 树的数据集D 里找到距离q 最近的k 个数据。当k=1 时，就是最近邻搜索。

使用Kd 树在高维向量空间中进行搜索时，由于高维向量的距离计算需要花费相当大的代价，使得查询在一定程度上变为线性搜索，极大影响了查找效率。为了减少距离的计算，提高执行效率，目前已有很多文献对此提出了不同的解决方法。这些方法能够有效的回答低维和中维空间中的最近邻查找问题[1-4]，但是由于“维度灾难”[5]，依然不能很好地应用于中高维空间的搜索。

随着大数据和gpu 技术的繁荣，人们能够探索更高的计算能力和更灵活的数据结构。然而传统上的数据结构更多是为了适应cpu 而设计的，其特点是按照固定的模式来组织数据，并不考虑数据的分布。如何设计高效灵活的数据结构以降低原有模型的复杂度成为一个急需解决的问题。神经网络经常被用于实现各种复杂的功能，使用Kd 树进行最近邻查找可以看作是一个分类问题，这与神经网络完成的工作没有本质上的区别。因此，使用神经网络来代替Kd 树应该是一种具有可行性的工作。传统的索引结构是按固定的方式构建的，而机器学习模型是在训练数据的基础上建立的，但这两者本质上都是对空间位置的定位和寻找，潜在来说神经网络和索引是具有一定联系的。文献[6]提出了一种基于机器学习的可学习索引的方法，该方法具有学习数据分布的能力。它探索了使用神经网络在一定程度上代替传统索引结构的可行性，并进一步讨论了可学习哈希映射索引与传统哈希映射索引之间的区别。在之前的工作[7]中，探索利用神经网络建立倒排索引，表明了无监督的可学习索引相比传统索引有着明显的优势。在此基础上，本文提出了基于神经网络的可学习Kd 树模型，用于解决中高维空间中的最近邻搜索问题。在本文的方案中，首先构建一棵Kd 树，并对构成Kd 树的数据构建索引。随后通过Kd 树找到输入数据的最近邻点，将最近邻点所在的分类作为标签进行训练。

这种基于神经网络的可学习的Kd 树方案新颖而且可扩展，为最近邻查找提供了一种全新的思路。该方案使用神经网络，能够并行的解决搜索问题，能在更短的时间内进行较为精确的查找，运行时间更具优势。

1 传统Kd树

传统上研究人员使用Kd 树进行最近邻搜索。Kd 树是每个节点为一个k 维向量的二叉树。每个非叶子节点可以看作一个超平面, 而这个超平面将多维空间分割为两个子平面。在这个超平面左侧的点被分为左子树, 右侧的点则被分为右子树。决定这个超平面方向的方式如下：每个 Kd 树中的点都与k 维向量中的一个特定维度相关联，而这个维度正是垂直于分割空间的超平面的轴的维度。举个例子，如果这个轴是X 轴，那么x 的值小于决定这个超平面的点的剩余点被分到左子树，所有x 的值大于这个点的被分到右子树。

使用Kd 树进行最近邻搜索方法如下：假设有一棵已经构建好的维度为d 的Kd 树，对这个数据集进行最近邻搜索，首先，给定查询点q。需要做的是从树上找到一个节点o，使得点o 到查询点q 之间的欧式距离，小于这棵树上除点o 之外任意一点oi到q 的距离。具体做法是比较查询点q 和分裂节点的分裂维的值，小于等于就进入左子树分支，等于就进入右子树分支直到叶子结点。顺着二叉树的搜索路径进行搜索。然后再回溯搜索路径，并判断搜索路径上的结点的其他子结点空间中是否可能有距离查询点更近的数据点，如果有可能，则需要跳到其他子结点空间中去搜索（将其他子结点加入到搜索路径）。重复这个过程直到搜索路径为空。其中，空间中两个点 x1(x11,x12,…,x1k)与 x2(x21,x22,…,x2k)间的欧氏距离 d（x1,x2）定义为：

使用传统Kd 树进行最近邻查找存在诸多问题，例如占用存储空间大，需要进行大量的回溯操作等。尤其是当数据的维度很高时，Kd树的查找效率变得特别低，性能甚至不如线性扫描，产生了“维度灾难”。其他用于k 近邻查询的树形结构例如R 树[8]，采用了空间分割的理念，其核心思想是聚合距离相近的节点并在树结构的上一层将其表示为这些节点的最小外接矩形，这个最小外接矩形就成为上一层的一个节点。因为所有节点都在它们的最小外接矩形中，所以跟某个矩形不相交的查询就一定跟这个矩形中的所有节点都不相交。在此基础上，R* 树[1]作为 R 树的一种变体，提升了 R树的性能。然而，随着维度的升高，这些索引技术的效率会变得较低，线性扫描成为唯一可行的解决方案。在一些实际应用中，由于并不需要严格查找到完全准确的k 近邻点，因此可以牺牲一定的搜索精度来实现实际的运行时间[9-11]。在这种情况下，NN（Nearest neighbor）和 kNN 问题分别被转化为 ANN（Approximate nearest neighbor）和kANN 问题。作为一种流行的 kNN 查询方法，局部敏感哈希（LSH）[12]被广泛使用。

2 LK模型

2.1 LK模型架构及执行过程

可学习的Kd 树，本质上是利用训练好的深度学习模型替换掉传统Kd 树中的搜索部分，并且，训练好的深度学习模型可以仍保持原有数据的局部空间不变性。所以，如何训练这个深度学习模型，将是模型中最重要的研究部分（如图1 所示）。

图1 LK 基本架构

首先，生成一棵维度为d 的Kd 树，这棵Kd树由n 条数据构成。随后，随机为构成这棵Kd树的数据从0 开始打上标签，则可以得到对应的从0 开始直到n-1 一共n 个类别。训练的时候，将训练数据和测试数据输入Kd 树，得到最近邻点，根据最近邻点得到相应的类别。最后，将这n 个类中最近邻点的索引值的位置为1，其余置为0，作为训练集和测试集的标签。由于神经网络能够对空间进行定位和寻找，因此使用神经网络来构建模型，用来模拟Kd 树的结构。神经网络的输出值是预测到的索引值所在的位置。

模型分为五个部分：输入阶段、映射阶段、标签阶段、索引阶段、计算阶段。如图2 所示。

图2 LK 框架和执行过程

输入阶段：输入的数据要求具有相同的维度。一般来说，输入的数据不应和构成Kd 树的任何一条数据相同。这些数据可以直接输入神经网络而不需要进行额外的处理。如果是图片或者音频，则需要进行数据的转化才能输入神经网络。一条输入数据为 I=（i1,…,id），其中 d 为输入数据的维度。则整个输入数据为I*=I1,...,Ix，其中x 为输入数据的尺寸，I1，…，Ix中 d 值都相同。

映射阶段：这个阶段实现了传统Kd 树的查找功能。将输入数据（假设为m 条）输入神经网络之后，由神经网络进行运算处理。这是整个模型最为重要的一个阶段, 由多层神经网络构成。神经网络的层数和每一层神经网络的节点数会根据实验的结果不断的调整，使得神经网络能够输出预期的结果，每一层神经网络的节点数可能不相同。每条数据都会输出一个1*n 的向量。因此，整个神经网络会输出一个m*n 的矩阵，矩阵中每个值都为一个0 到1 之间的数。

标签阶段：在这个阶段，需要对神经网络的输出进行处理，并不是直接输出最近邻点的位置，而是输出可能是最近邻点的位置，从而能够从这些潜在的最近邻中找到真正的最近邻。不同的输入数据，可能是最近邻点的数量也不相同。本文规定，如果某个位置输出值为1，则这个位置在Kd 树中所对应的点，就是潜在的最近邻点；如果这个位置输出值为0，这个位置就不是最近邻点，在接下来的阶段不需要对这个位置进行处理。由于使用了sigmoid 激活函数，因此需要一个阈值，将分布在0 到1 之间的数转换为0 或者1。这个阶段的输出也是一个m*n 的矩阵，矩阵中的每个值，要么为0，要么为1。

索引阶段：在构建Kd 树的标签时，构成Kd树的每一个节点，都被划分为 0，1，…，n-1 中的某个类别。因此，为了直观展示和便于计算，可以将上个阶段所输出的值为1 的位置，转化成与之对应的实际的类别，即根据值为1 的位置，得到潜在最近邻点的索引值。对于每一条数据I，值为1 的个数可能不同，因此得到索引值的数量也不同。这表示神经网络对于不同的输入数据，会做出误差不同的预测。输出值为1 的个数越多，表示需要进行更多次的计算，才能从潜在的最近邻点中找到真实的最近邻点。

计算阶段：实际上，神经网络输出的不会总是完美的得到满足条件的最近邻点，大多数情况下，需要在输出的所有可能是最近邻点的几个点中，通过计算，找到真实的最近邻点。因此，需要根据输出的索引值，通过Kd 树，得到其原本表示的数据。最后，将输入数据I 和这几个点所代表的存在于Kd 树中的每一条数据，依次计算欧氏距离并排序，便能得到真实的最近邻点。

本文中，准确率是这样定义的：假设待查询的数据有M 条，其中有N（N≤M）条查找到了正确的最近邻点，则准确率为N/M×100%。

2.2 调整参数

本文中使用sigmoid 函数作为神经网络最后一层的激活函数。由于需要将所有的输出值映射成0 或者1，而sigmoid 函数的输出值都在0 和1 之间，所以需要确定一个阈值，将输出划分为0 或者1。实验结果表明，阈值选取的越小，准确率会越高，但需要进行查询比较的时间就会越长。极端情况下，当阈值为0 时，几乎所有的输出都被置为1，这意味着基本需要将输入数据和所有构成Kd 树的点进行一一对比，虽然能够达到接近百分之百的准确率，但查找所耗费的时间远超Kd 树和线性查找没有区别。

阈值决定了输出中值为1 的个数，也就决定了需要进行计算比较的次数，而计算比较的次数又决定了查找时间的长短。在第3 部分中进行了相关的实验，展示了阈值的选取和其他指标的关系。

2.3 时间复杂度

Kd 树在进行最近邻搜索时，首先由根结点从上到下找到对应包含查询点的叶子结点，其中要计算该区域内的点到查询点的最小距离（若是k 近邻搜索，需要计算最小的k 个距离）。对于一个具有n 个节点的d 维Kd 树，最近邻查找的时间复杂度为tworst=O(d*n1-1/d)。

使用神经网络进行最近邻搜索是一个三次索引加上一次计算的过程，第一次是数据输入神经网络后得到的输出值，即Kd 树上的点所对应的索引值的位置；第二次是根据这个位置得到Kd 树中点的索引值；第三次是根据这些索引值得到Kd 树中的点。索引完成后，将输入的点和查找到的点进行一一比较，得到最近邻。因此，时间复杂度为O(1)+O(2)+O(3)+O(4)，分别对应了三次索引过程和计算过程。在进行最近邻搜索时，相对于Kd 树需要进行多次的计算和回溯操作，模型只需要进行少量的计算。使用Kd 树进行最近邻查询会在计算和回溯的过程中消耗大量的时间，而使用神经网络在一定程度上代替了这个过程。当构成Kd 树的数量和维度数变大时，此模型应该会表现出更明显的优势。

3 结果与分析

所有的实验均使用python 编写的代码在Intel（R）Core（TM）i7-4770 上完成，该机器具有3.4GHz CPU，32GB RAM 和运行 Linux Ubuntu 12.04 的2TB 外部磁盘空间。

本文从准确率和运行时间两个方面来评估模型。一方面将神经网络查找到的结果和传统Kd 树找到的结果进行对比；另一方面比较两者查找到最近邻所需的时间。本文所使用的最近邻点，都是通过传统Kd 树得到的，可以认为使用Kd 树进行查找的准确率为100%。

在进行设计和分析的过程中，发现阈值和维度的选取会对实验结果产生较大的影响。在实验中，分别使用了 3 个阈值：0.001，0.005，0.01，在1 000 条数据和10 000 条数据的情况下来进行实验观察（都为符合正态分布的随机数）。

图3 阈值和维度对准确率的影响

首先在神经网络层数相同，维度分别为10，15，20，25，30 维的情况下进行了实验。最终结果为10 次实验结果的平均值（实验结果如图3）。结果表明，当n 为1 000 时，在相同的维度下，阈值越低，准确率越高。随着维度的升高，不管阈值如何，准确率都在持续下降；而当n 为10 000时，除了在维度为10 的情况下出现了异常，相同维度下，阈值越低，准确率越高。同时，随着维度的增加，准确率出现了一定程度的波动。当维度为10，阈值为0.001 时，两者的准确率甚至达到了99%以上；当维度为30 维时，可以看到，极端情况下，当阈值为0.01 时，准确率甚至会低于70%。由于只需要查找最近邻点（k=1），k 的值远远小于构成Kd 树本身的数据量n，因此最终的预测值，值为1 的个数远远少于值为0 的个数。在神经网络的输出中，只要输出值不为0，便倾向于这个位置是可能的最近邻点。阈值设置的越小，就会把更多潜在的最近邻点包含其中，使得准确率升高。一般来说，随着维度的升高，准确率会降低；相同维度下，阈值越小，准确率越高。

图4 展示了搜索时间和维度、阈值的关系。其中图4(a)，图4(b)，图4(c)分别是神经网络层数为 5，n 为 10 000，维度分别为 10，20，30 维的结果。横坐标为不同的阈值，纵坐标为所用的时间。观察相同维度下运行时间的变化，可以发现，运行时间随着阈值的变大而减少。阈值设置的越小，就会把更多潜在的最近邻点包含其中，会导致需要进行更多次的计算。

图4 阈值和维度对搜索时间的影响

本文将使用神经网络进行最近邻搜索的时间和使用传统Kd 树进行最近邻搜索所花费的时间，在n 为10 000，阈值为0.005 的情况下进行了对比，此时，由图3b 可知，模型的准确率均在90%以上。两者的运行时间如表1。从表1 可以看出，相比于传统的Kd 树，LK 模型在运行时间上更具优势。

表1 运行时间对比

4 结论

本文解决了在实际计算时间内的最邻搜索的问题，同时表现出了较高的质量。总的来说，在应用中，根据不同的需求，需要在准确率和运行时间上做出取舍。可以得到：模型的准确率随着阈值的增加而降低，但与此同时，运行时间也在减少。在搜索中，如果需要更高的准确率，那么就需要花费更多的时间，但在一定的情况下，所用的时间也会比传统的Kd 树花费的时间更少；如果对结果要求不是非常精确，那么此模型就会具有更大的时间优势。未来，将继续研究如何在更高的维度、更多的数据量上取得更好的实验效果。