NNDescent算法

算法的目标是为数据集中每个点构建一个近似的 k-近邻图（k-NN graph），即每个点维护 k 个最近邻的 id。

为什么需要近似：如果对 n 个点做精确 kNN（暴力）代价是 $O(n^2 d)$，数据很大时不可行；NNDescent 用局部传播，大幅减少距离计算量，但通常能得到高质量近邻。

NNDescent的主要思想是“邻居的邻居也是邻居”，想象一下，你要在一个陌生的城市里找到最好的餐馆。你不会盲目地随机寻找，而是会：

  1. 先问几个朋友他们喜欢的餐馆（初始随机邻居）
  2. 然后问这些餐馆的老板，他们推荐附近还有什么好餐馆（探索邻居的邻居）
  3. 不断重复这个过程，你的餐馆清单就会越来越优质

这就是NNDescent算法的核心思想——通过朋友的朋友来发现高质量的关系。

算法流程：

  1. 初始化：为每个点随机选一些邻居（不是目标k，而是初始化选多一些 l），这些作为候选集。
  2. 迭代传播：重复下面过程多次（iter)：
    • 对每个点 i，从它的候选邻居中选出“new”和“old”（new 是最近一次新加入的候选），并对这些邻居做成对比较（neighbor-of-neighbor）：
      • 比如比较新邻居之间（new-new）、新邻居与老邻居（new-old / old-new），以发现更近的点；
    • 将发现的更好邻居插入候选集中（并做去重、排序、截断到最多l个保存）。
  3. 收敛或到达迭代次数：最后取每个点候选集中距离最小的前 k 个作为最终kNN。

关键直觉：两点如果有共同近邻，它们本身更可能互为近邻；所以通过“邻居的邻居”探索空间比随机或全局搜索更高效。

那么该如何评判两个点之间的距离呢？下面来介绍欧式距离

欧式距离：

两个向量 a 和 b 之间的欧氏距离的定义：

$dist(a,b)=\sqrt{\sum_{i=1}^{n}(a_i-b_i）^2}$

直接计算需要做减法、平方、求和、开方，计算量较大。在近邻搜索中，我们通常更关心距离的相对大小而非绝对值，因此可以简化，使用平方欧氏距离来避免耗时的开方运算：
$dist(a,b)^2=\sum_{i=1}^{n}(a_i-b_i）^2$

进一步推导：
$$
dist(a,b)^2 =\sum_{i=1}^{n}(a_i-b_i）^2 =\sum_{i=1}^{n}(a_i^2-2a_ib_i+b_i^2)
\=\sum_{i=1}^{n}a_i^2+\sum_{i=1}^{n}b_i^2-\sum_{i=1}^{n}2a_ib_i
$$
现在，比较两个距离的大小，例如判断 a 离 b 更近还是离 c 更近：
$dist²(a, b) < dist²(a, c)$
代入上面的展开式：

$ \sum_{i=1}^{n}a_i^2+\sum_{i=1}^{n}b_i^2-\sum_{i=1}^{n}2a_ib_i<\sum_{i=1}^{n}a_i^2+\sum_{i=1}^{n}c_i^2-\sum_{i=1}^{n}2a_ic_i$

等式两边的 Σa_i² 可以抵消，不等式变为：

$ \sum_{i=1}^{n}b_i^2-\sum_{i=1}^{n}2a_ib_i<\sum_{i=1}^{n}c_i^2-\sum_{i=1}^{n}2a_ic_i$

为了更方便地与硬件计算单元（矩阵乘法器）匹配，我们调整一下形式。将不等式两边同时乘以 $-\frac{1}{2}$（注意乘以负数会翻转不等号方向）：

$ \sum_{i=1}^{n}a_ib_i-\frac{1}{2}\sum_{i=1}^{n}b_i^2>\sum_{i=1}^{n}a_ic_i-\frac{1}{2}\sum_{i=1}^{n}c_i^2$

令$score(a,b)=\sum_{i=1}^{n}a_ib_i-\frac{1}{2}\sum_{i=1}^{n}b_i^2$，那么寻找最小距离就变成了寻找最大score，计算 score(a, b)可以分解为两个部分：

  1. $\sum_{i=1}^{n}a_ib_i$：向量 a 和 b 的点积。这正好是一个矩阵乘法操作（$A * B^T$ 的一个元素），可以被NPU的Cube单元高效计算。
  2. $-\frac{1}{2}\sum_{i=1}^{n}b_i^2$：向量 b 的平方范数的一半，并取负。这是一个可以预先为每个向量计算好的常数。

在代码中可以预先计算好$-\frac{1}{2}\sum_{i=1}^{n}b_i^2$，这样计算最大值就可以通过一个矩阵乘法（计算点积 $A * B^T$）加上一个偏置来实现，这正是硬件所擅长的操作。在NPU上计算一个向量和所有其他向量之间的score矩阵就可以被快速地实现。

有了以上的思想，就可以构建我们的算法了

代码实现流程

算法首先需要定义一些基础的参数，例如我们要寻找Top-k，就是最终生成的图的度数，这里先定义k=128。然后还有向量数据库的一些参数比如说向量数据库的大小（base_cnt）和维度（d）。

定义邻居节点，包括id,dist和is_new是否是新邻居标签。然后用一个一维数组（graph）来存储这个图：

graph

L为迭代过程中每个节点保存的邻居数量，graph_cap为图为每个节点开辟的最大空间防止溢出。可通过graph[i * graph_cap + j]来寻找第i个节点的第j个邻居。

对于向量数据库中的每一个向量i，我们需要一些数组来存储它的相关数据，例如算法需要保存它的当前邻居数量（graph_size)、所有旧邻居的ID（old_id）、所有新邻居的ID（new_id），为了方便后面比较“距离”，还需要两个数组（old_sqr、new_sqr）来存储新旧邻居的平方和$-\frac{1}{2}\sum_{i=1}^{n}b_i^2$。

image-20250926205617104

graph_size[i]表示第i个点现在有多少邻居，new/old_cnt表示为新旧邻居开辟的最大容量

最重要的，需要存储每个点的新邻居之间的距离最小值，也就是说保存一个点的新邻居最近的新邻居（nn_min)，nn_min的长度是数据库的大小（base_cnt)$$新邻居的数量（new_cnt)$$2，这里新邻居的数量是算法最开始预定义的大小，与算法执行过程中每个点的实际邻居数量（new_id_cnt）不同。乘2是因为要存储id加距离。

image-20250929190556310

为每个节点开辟了new_cnt*2的空间存新邻居的邻居，可通过nn_min[i * new_cnt * 2 + j * 2 + 1]取出第i个邻居的第j个新邻居的id，+0表示取出距离

对应地on_min表示一个点的旧邻居最近的新邻居的最小距离和id。no_min和no_dist存储新邻居对应的最近的旧邻居的id和距离，这里分开来存储id和距离是为了方便后面的计算，因为nn和on都需要计算比较新邻居的距离。

初始化：计算数据库中每个节点的平方和存储在一个数组（base_sqr）方便后面取用，然后给每个节点随机分配初始邻居作为迭代的初始图，比如随机初始化st=160个随机邻居，这些随机邻居前new_cnt个邻居是新邻居，后面的都是旧邻居。

接下来开始**迭代**：

新旧邻居寻找

每次迭代首先要寻找每个节点的新旧邻居，便于计算各种距离，遍历节点i的所有邻居，如果这个邻居是new邻居（is_new=true）或者old邻居数量已满，就把这个邻居的id（graph[i * graph_cap + j].id）存入new_id[i * new_cnt + new_id_cnt[i]]。这里graph_cap是为每个节点预留的最大邻居数量，预定义好的一般很大，[i * graph_cap + j]就是节点i的第j个邻居。i * new_cnt + new_id_cnt[i]中new_cnt就是上文提到的预定义的新邻居数量，new_id_cnt[i]就是偏移，即i节点的第几个new邻居，同样地存起来这个new邻居的平方和，找到一个新邻居，new_id_cnt[i]数量加一供后续使用。

image-20250929190747389

取出id放入：

image-20250929194437449

同样的逻辑处理旧邻居。

距离计算

接下来计算new-new、old-new和new-old三类距离，这一块儿涉及NPU并行处理等下个文章详细说明。

插入图

遍历每个节点，然后遍历每个节点的new邻居，这里遍历每个new邻居用到寻找新旧邻居时使用的数组new_id_cnt[i]，也就是节点i有多少个新邻居，根据这个值，就能够从new_id[]中取出这个新邻居的id，这里设为jid，索引是new_id[i*new_cnt+j]，j遍历所有的new_id_cnt[i]就能找到所有的新邻居的offset。

找到所有新邻居的id:

image-20250929193914995

jid = new_id[i * new_cnt + j] //节点i的第j个new邻居，j的最大值是new_id_cnt[i]

然后使用nn_min[i * new_cnt * 2 + j * 2 + 1]取出与当前new邻居距离最小的new邻居在new_id里的下标，索引中+1是取id，+0是取距离。距离取-nn_min[i * new_cnt * 2 + j * 2]，这里取负数是因为我们上述将距离的比较转换成了一个score的比较，实际上nn_min中存的不是最小的距离而是最大的score（类似于相似度），为了后续排序的逻辑正确，所以取符号，也就是能把最小距离的邻居排在前面。

取出最近新邻居的id:

image-20250929191727150

有了id和距离就可以创建一个新的节点插入到图中，插入的位置是graph[jid * graph_cap + graph_size[jid]]，graph_size[jid]加一。对应地，要插入反向边，也就是在另一个新邻居的图中插入这个新邻居。

插入到图中：

image-20250929192333795

同样的逻辑，可以通过no_min来构建节点的新邻居与旧邻居之间的关系图，这两种距离关系都可以通过遍历新邻居时来完成。

接下来还剩下old-new邻居的关系，逻辑一样，只不过是遍历旧的邻居，然后处理对应的关系。

整理

经过上述步骤，每一个节点都或多或少新插入了邻居，最后再将新加入的邻居进行排序，然后与之前的邻居进行有序归并，最后截断至固定长度保留邻居。

迭代

排序归并后再次迭代，重新进行新旧邻居的整理和距离计算，然后插入图中再次整理，这样经过一定次数的迭代，图中保留的就几乎是所有节点的最近的邻居了。整个迭代过程大致如下：

while(iter–)

{

  1. 划分new/old邻居(CPU)：

    • 把邻居分为新旧两类，分别放进new_id和old_id，记录数量

    • is_new设为false,所有邻居下轮会变成old

    • 补零并保存当前邻居数量

  2. NPU计算距离(NPU)：

    • 把new/old邻居对发送到Ascend设备

    • 计算new-new、new-old、old-new三类距离

    • 返回最小距离以及邻居编号

  3. 插入新邻居(CPU)：

    • 遍历每个节点i的new邻居

    • 根据设备返回的最小距离结果生成候选边

    • 插入到邻居表graph，同时插入反向边

  4. 邻居表排序归并整理(CPU)：

    • 先对新邻居部分进行排序，然后和旧邻居进行有序归并

    • 去掉重复的邻居，并截断到前l个最优邻居

    • 拷贝回graph

}

Recall计算

$$Recall@1=\frac{1}{N}\sum_{q=1}^{N}1[GT最近邻居\in搜索Top-k]$$

意思就是在所有查询向量里面，有多少比例命中了真top1,也就是说有多少向量的真实最近的第一个邻居在构建的图中被找到了。

$$Recall@10=\frac{所有查询命中}{querycnt*10}$$

Recall@10就是gt的前十个只要有一个在128里面就算命中，然后分母乘10的意思是对于每个向量，gt前十大概有多少命中。

K 是最终每个节点在索引中保存的邻居数（也就是 index.nndescent.final_graph 每行的列数）。把 K 从 128 减到 10，会让图每个节点仅保留10条边，从而显著降低图搜索时可以“跳跃/扩展”的边数量 —— 导致搜索覆盖度变差、无法到达正确的近邻，召回率大降。时间不变化是因为只是改了输出度数，不影响NPU上实际的距离计算工作量（你的 nndescent 内部迭代参数 l、new_cnt、old_cnt 等可能并没有同步调整），K 越小，图越稀疏，搜索时就越容易漏掉真实近邻，Recall 会下降。

优化

暂未实现算法层面的优化，我将插入图过程中的锁给替换成了原子操作，时间加快了约2s，但是召回率略有下降，推测是插入顺序的原因。

image-20250929165728787

单独测了插入新邻居的部分时间，根据结果可以看到处理new-new和new-old边的时间和处理old-new的时间大致相当，原因是new-new和new-old在同一个循环下处理，即遍历节点的新邻居时处理，而old-new是遍历旧邻居时处理，而新邻居和旧邻居的数量整体上大致相当，所以这块的时间本质上和遍历规模有关，而不是算法逻辑有多复杂。

原子操作优化：

查看原始的代码在邻居插入时会有以下操作：

1
2
3
4
5
6
locks[jid].lock();//锁住new邻居jid，准备更新它的邻居列表
     if (graph_size[jid] < graph_cap) {
         graph[jid * graph_cap + graph_size[jid]] = x;//将x插入它的邻居列表
         graph_size[jid]++;//邻居数量加1
      }
locks[jid].unlock();

这里有两个共享变量graph和graph_size，由于程序运行是多线程的，多个线程可能同时向一个目标写入。如果两个线程几乎同时读取graph_size，就可能会得到相同的位置，于是写到同一个位置，导致一个邻居覆盖另一个，数据丢失，也可能会发生竞争条件，加锁的目的是保证一个线程完成“读size->写数据->更新size”这个流程之前，其他线程不能同时操作一个target，这样邻居不会丢失，也不会写错位置。如果不加锁就可能导致构建的图比预期稀疏，Recall降低，程序可能崩溃。但是加上锁之后很多线程会在某一个锁上面排队，造成时间消耗大

如果不使用lock的操作，可以使用原子操作来实现，这里不再使用锁而是先在每个线程内构建一个本地候选集合local_to_target，再用atomic fetch_add直接分配位置写入graph。这样避免了锁竞争提高性能。

local_upsert的作用是遍历当前本地缓存local_to_target，如果已有相同的target，只保留距离更小的，否则插入新候选；然后把local_to_target里的候选真正写入全局邻居表中。

但是local_upsert使用了线性查找去重，如果同一个邻居多次出现，只保留距离较小的一次，可能舍弃了距离相差很小、但后续迭代能“传递”更多信息的边。导致图的连通性降低，局部候选不足。而且写入顺序也是非确定的，保留下来的边可能不同，导致图质量变差。