ANN 与 NNDescent 浅析

ANN 介绍

在模式识别领域中,K - 近邻算法(K-Nearest Neighbor, KNN)是一种用于分类和回归的非参数统计方法。K - 近邻算法非常简单而有效,它的模型表示就是整个训练数据集。就原理而言,对新数据点的预测结果是通过在整个训练集上搜索与该数据点最相似的 K 个实例(近邻)并且总结这 K 个实例的输出变量而得出的。KNN 可能需要大量的内存或空间来存储所有数据,并且使用距离或接近程度的度量方法可能会在维度非常高的情况下(有许多输入变量)崩溃,这可能会对算法在你的问题上的性能产生负面影响。这就是所谓的维数灾难。

近似最近邻算法(Approximate Nearest Neighbor, ANN)则是一种通过牺牲精度来换取时间和空间的方式从大量样本中获取最近邻的方法,并以其存储空间少、查找效率高等优点引起了人们的广泛关注。

严格地讲,ANN 是一种在 NN 搜索过程中允许少量误差的算法。但在实际的应用中,真实的邻居数量比被搜索的 K 近邻数量要多。与暴力 KNN 相比,人工神经网络可以在短时间内获得卓越的准确性。

比较出名的有 Annoy、ScaNN、Faiss、HNSW 等。

在图中搜索

让我们来看看如何在图上利用近邻搜索去搜索。

假设我们得到了一个最近近邻的图。这种图是用边连接的一组节点。这种图或者是这种网络是由存在于某种空间中的数据推导出来的,对于这些数据,我们有一个关于相似性或不相似性或距离的定量概念。具体而言,该图是通过将每个数据点作为一个节点,并通过边将该数据点与最相似、最不相似或最接近的其它点连接起来而得出的。我们想用这个图作为搜索的索引结构,来找到一个新的查询点(可能不是图中已有的一个点)在该图中相似的邻居们。我们怎样才能做到这一点呢?简单来说我们可以通过广度优先搜寻算法来进行查找并保留搜寻到的 \(k\)  个最佳的点。

广度优先搜索算法(英语:Breadth-First-Search,缩写为BFS),又译作宽度优先搜索,或横向优先搜索,是一种图形搜索算法。简单的说,BFS是从根节点开始,沿着树的宽度遍历树的节点。如果所有节点均被访问,则算法中止。广度优先搜索的实现一般采用 open-closed表。
  • 选择一个图中的点进行开始(可能是随机的)并作为候选节点。
  • 寻找图中没有被寻找过的且与这个点所有与它有边连接的其它节点作为候选节点。
  • 将这些所有节点放入候选节点池子中。
  • 对候选池子中的点与需要查询的点进行相似度排序。
  • 将池子截断成 \(k\)  大小。
  • 然后从第二步继续开始,直到尝试了整个池子中的候选人。

我们来看一个简单的例子,我们有个非常小的图并且我们希望  \(k=2\)  。橙色的 x 点就是我们希望查询的 query 位置。我们来看看可视化后是如何寻找的呢?

0:00
/

我们可以在上面的视频中发现的确是在找寻过程中慢慢逼近了查询的点的位置并且找到了邻居。我们来看看在更复杂和大的图下的效果。

0:00
/

尽管是从一个不太好的位置触发的点,但它仍然在沿着查询点的位置不断靠近。当然这个上面算法只是一个演示,有很多细节要优化,比如需要一个数据结构来记录你已经计算过相似度的节点等等。以及通过视频可以发现较小的 epsilon 可以提供更快的搜索,但精度较低,而大的 epsilon 可以提供良好的精度,但代价是搜索较慢。

ANN 的本质其实就是要建立这么一张图,然后使用这么一张图来进行近邻搜索。

本文将介绍 NNDescent 建立图和搜寻的方法来举例。

NNDescent

建图

就像我们前面描述的一样,我们需要一个好用的图来作为索引,对数据集中的每个点进行搜索。假设我们有个不太精确带但是能用的 k 近邻图,我们利用它进行搜寻,虽然一开始效果不好,但由于搜索的性质,我们最终可以为一个节点找到新的邻居,这些邻居比我们最开始的图中的邻居更接近。我们可以利用这些信息来更新我们的图。在更新了图之后,我们可以对数据集中的每一个点再次进行搜索。由于图更好了,我们将可以更好的找到近似的邻居甚至有可能找到一些新的邻居,那么在再次找到新邻居后再次更新这个图。图越好搜寻速度越快、越准确,因此理论上每一次迭代都会比上一次快,我们最终会得到一个真正好用的 k 近邻图。

一般来说为搜索时的点找到一个比较好的初始节点是比较难的一个问题,就像前面的视频里的一样其实初始的点的位置都比较糟糕。所以我们可以使用搜寻节点本身作为初始的点,并且尽快的进行搜索和更新图。我们可以先搜索这个点本身的邻居,然后搜索这些邻居的邻居,在这样的两轮搜索后可能可以找比当前图更近的邻居,那么我们就可以更新图。

  • 我们生成一个随机的图,每个节点连接着 k 个随机的节点。
  • 对每个节点进行下面操作:
  • 测量它与附近邻居的位置,如果有更近的就会更新图只保留最近的 k 个点。
  • 对图进行了任意的更新后就会回到第二步。

我们可以用下面的可视化整个过程:

0:00
/

我们可以在视频中看到最开始是杂乱无章的连接在一起的,后面逐渐的变得更加清晰并且连接的边变得更加的少。每次迭代较上次迭代的内容其实会越来越少,最后剩下的图会是一个比较好的图。

搜索应该是在一个无向图中进行的,这意味着我们寻找一个节点的邻居时不仅要考虑节点有 k 个邻居,还需要考虑它可能也会作为其它 k 个邻居的邻居,也就是所谓的 “reverse nearest neighbors“。虽然跟踪每个节点的 k 个邻居很容易,但跟踪每个节点的反向近邻就很难了。所以在迭代开始的时候就需要计算每个节点的邻居和反向邻居,并在每次迭代都用上它。

第二个是我们在迭代过程中需要查询每个节点的邻居和邻居的邻居,如果从绿色的点出发大概长成下面这样。

但我们从蓝色的点角度来看,其实绿色和红色都是它的邻居。

所以我们可以将绿色和红色添加到彼此的近邻列表中,让它不需要从绿穿越到蓝再穿越到红。我们可以为每个节点找到像上面这样的绿色和红色节点,让每个可能的配对都连接到一个共同的节点。

基于上述两个考虑,我们可以得到一种方法,首先为每个节点生成一组邻居(包括反向邻居),然后对每组都进行距离计算来更新图,所有的这些节点的计算都是可以独立运行的并且保存每个邻居集的更新,然后重新进行排序和更新图。这样做很利于并行处理,性能会更好。同时因为每个都是独立的,所以我们可以只对新的集合进行距离运算或者限制元素数量等等。

我们现在已经有了一个简单就能生成的图了。但是图里的不能太大,太多的话可能导致性能下降。

建树

随机投影的理论依据是 J-L Lemma,公式的核心思想总结一句话就是:在高维欧氏空间里的点集映射到低维空间里相对距离得到某误差范围内的保持

而随机投影树(Random projection trees)是一种对空间数据(即具有向量空间表示的数据)进行近似近邻搜索的有效方法。这个想法很简单:首先选择一个随机超平面将数据一分为二(可以通过选择两个随机点并将超平面作为它们之间的正交超平面来安排分割--在欧几里得或角度意义上,取决于距离度量);对于每一半数据选择一个不同的随机超平面来分割该数据子集;不断重复这个过程,直到每个子集最多是某个选定的大小(通常称为树的 "叶大小")。这种简单的方法能够递归地将数据分割成大小合适的桶。随机超平面有两个目的:首先,通过随机的方向,产生的树能很好地适应高维数据(不像其他树的方法,如 kd 树);其次,通过随机化,我们可以很容易地建立许多这样的树,它们都是不同的。用随机投影树进行近邻搜索的问题是,对于一个给定的查询,近邻可能在这些随机分割的错误一侧,如果我们只是搜索查询所处的叶子中的点桶,我们会错过这些近邻。解决办法是有许多树。由于每棵树都是随机的,所以每棵树都会在随机分裂方面有一定误差。总的来说,通过搜索一整片随机投影树的森林,找到真正的近邻的几率变得相当大。跟下面的图的分割过程比较相似。

我们这里以 Annoy 举例,看看是如何进行的:

通过多次递归迭代划分的话,最终原始数据会形成类似下面这样一个二叉树结构。二叉树底层是叶子节点记录原始数据节点,其他中间节点记录的是分割超平面的信息。Annoy建立这样的二叉树结构是希望满足这样的一个假设:  相似的数据节点应该在二叉树上位置更接近,一个分割超平面不应该把相似的数据节点分割二叉树的不同分支上。

那么,为什么不直接用这个来做近似的近邻搜索呢?为了获得足够的准确性,我们需要使用一个非常大的叶子大小(并通过一个叶子中的所有点进行强行搜索),或者使用大量的树。在实践中,这往往比上面描述的基于图的搜索慢得多。不过,树的建造成本很低,而且对于一棵树来说,搜索速度也非常快,所以我们能不能以某种方式使用它们?回顾一下,我们的图是以随机初始化开始的,但是我们的图越好,每次迭代的速度就越快,而且迭代的过程中越能找到好邻居。我们能不能用随机投影树来为图形找到一个好的起点?不一定要非常好,只要比随机好就行了。比随机好也不是那么难的。因此,我们可以建立一个非常小的随机投影树森林(以你想要合理的准确性查询而使用的树的数量为标准的小),并用它来初始化我们的图。如果我们不在新的查询点上使用这些树,我们如何使用这些树呢?树的每个叶子节点都是一个桶,我们可以对其进行全对距离计算,图中的每个节点都可以根据它所在的桶的结果获得一个初始的 k 边。对于少量的树来说,这一切都可以非常迅速地完成,而且完全是并行的。

细化图

现在我们有办法有效地建立一个 k 近邻图,有什么办法可以微调这个图以使我们的近邻搜索运行得更快?导致搜索速度变慢的其中一件事就是向候选库中添加大量的新节点,所有这些节点都需要计算与查询点的距离。如果节点有很多连接的边,那么我们最终会有非常大的候选池,搜索速度会自然。乍一看,这似乎是一个简单的案例,因为我们建立了一个 k 邻接图,每个节点都有 k 条边,不同的是,搜索是在一个无向图上进行的,这意味着一个节点不仅有与所有它认为是其邻居的节点的边,而且还有与所有认为它是其邻居的节点的边。如下图,每个节点都有一条通往其最近邻居的边。

我们转为无向图就会看到下面这样的效果:

我们看到中心节点有四条边,尽管这是一个 1-近邻图。这可能是出乎意料的普遍现象,特别是对于高维数据来说,由此产生的无方向的 k-近邻图的节点可能比你预期的 k 条边多出几个数量级。修剪图中的边,以减少任何一个节点可能连接的边的数量,将是加快搜索速度的最好方法。当然,问题是删除边会使搜索的准确性降低。那么,我们完善图形的目标是在尽可能多地删除边的同时,尽可能不降低准确性。

一种直接的可能性是简单地使用有向图,并扔掉所有在无向情况下诱发的反向近邻边。虽然这可能是有效的,但它确实降低了搜索的准确性;事实上,许多反向近邻边是非常有用的。相反,我们可以采取有向图加上一定数量的反向邻接边的方式。在一个节点很少有这些边的情况下,我们可以保留所有的边,但是对于节点的情况,比如上面例子中的中心节点,有相对较多的反向邻居边,我们可以简单地取其中最上面的几个。

在我们进行这种修剪之前,我们可以看看是否有一些其它的对搜索质量影响不大的边也可以放弃,答案是肯定的——特别是我们可以放弃三角形的长边。假设我们有一个由这样的边组成的三角形:

短边是用来寻找它们所连接的每个节点的近邻的,如果我们按照我们的搜索算法,我们将不需要长边,因为搜索将遍历两条短边并找到相同的结果,只是多了一个步骤。

通过以这种方式对图进行预处理,首先去除三角形的长边,然后修剪回每个节点的最大固定边数,这样的预处理是一个非微不足道的计算工作几乎不怎么消耗算力,但我们可以大大改善图的搜索性能。

总结

总的来说这是一个简单的、高效的建图的索引算法。但这种算法可能没考虑巨大无比的索引(如 tb、pb 级别)的性能问题,所以最近我参考了 NNDescent 进行改进,基于 ANN 搜索的特性设计了一种新的分片算法,能够将一个完整的数据集拆分为 n 个分片,并且尽量只查一个分片。

最多只查一个分片(准确率为 93.84%):

最多只查两个分片(准确率为 97.10%):

参考资料

How to use PyNNDescent — pynndescent 0.5.0 documentation
KNN (K-Nearest Neighbors) is Dead!
Long live ANNs for their whopping 380X speedup over sklearn’s KNN while delivering 99.3% similar results.
Johnson–Lindenstrauss lemma - Wikipedia
Locality-sensitive hashing - Wikipedia
GitHub - spotify/annoy: Approximate Nearest Neighbors in C++/Python optimized for memory usage and loading/saving to disk
Approximate Nearest Neighbors in C++/Python optimized for memory usage and loading/saving to disk - GitHub - spotify/annoy: Approximate Nearest Neighbors in C++/Python optimized for memory usage an...