Proxima: 3D NAND中基于图的近似近邻搜索的近存储加速

向量检索实验室 2023-12-18

341

论文简要

Vector Search

本研究提出了一种基于 3D NAND 的近存储加速图形近似最近邻搜索解决方案，通过算法硬件协同设计，显著提高了图形近似最近邻搜索的吞吐量和能效。

论文链接：https://arxiv.org/pdf/2312.04257.pdf

背景信息

Vector Search

论文背景: 近似最近邻搜索在推荐系统、信息检索和语义搜索等应用中起着不可或缺的作用，而基于图形的近似最近邻搜索算法在大规模数据集上提供了卓越的准确性和可扩展性。

过去方案: 传统的最近邻搜索方法在处理大规模数据时效率低下，而现代的近似最近邻搜索算法通过减少搜索空间、距离计算和数据访问来提高效率，但仍存在精度下降的问题。

论文的 Motivation: 由于现有的近存储/内存图形加速器不适用于图形近似最近邻搜索，本研究旨在利用 3D NAND flash 和近数据处理来解决图形近似最近邻搜索面临的挑战。

方法

Vector Search

1.理论背景:

近似最近邻搜索（ANNS）在各种应用中的重要性以及基于图的 ANNS 算法面临的挑战。

使用 3D NAND 闪存进行数据密集型应用的优势。

Proxima 解决方案旨在通过在 3D NAND 闪存中进行算法硬件协同设计来加速基于图的 ANNS。其目标是降低图搜索的复杂性，提高效率和存储密度。

2.技术路线：

通过算法和硬件的协同优化来解决基于图的 ANNS 的挑战。

Proxima 图搜索方法改进了搜索效率。

基于 3D NAND 闪存和异构集成的 Proxima 近存储硬件架构被提出。

异构集成技术可以解决将处理元件集成到 3D NAND 芯片中所涉及的成本和面积问题。

3.具体实现：

a.Proxima 图搜索：

Proxima 图搜索是通用的，可以应用于各种图构建算法生成的图，如 HNSW、DiskANN 和 NSG。在查询搜索之前，Proxima 中的间隙编码首先对构建的图进行压缩，以减小索引大小。在查询搜索过程中，作者提出了三种策略，包括基于产品量化（PQ）的距离搜索、准确的距离排序和具有早期终止的动态列表，以进行准确和低复杂度的搜索。下图为 Proxima 搜索优化的总体流程。

b.基于距离的 PQ 搜索：

Proxima 使用 PQ 获得的近似距离遍历图形。PQ 将向量分割为 M 个子维度，并将每个向量 x 编码为一个 M • log 2 C 位的编码，其中 C 是每个子维度上的 k-means 聚类中心的数量。x 的每个子向量被量化成它的 k 均值质心，并用 log2c 位的质心索引进行编码。PQ 近似查询 q 和查询 x 之间的距离:

下图显示了 D=4，M=2，C=3 的情况下计算 ADT 和最终 PQ 距离的示例。值得注意的是，Proxima 只在搜索过程中使用 PQ。图形索引使用现有算法以保证正确结构，并使用全精度坐标构建。

c.精确的基于距离的重新排序：

作者通过在算法中引入一个新的参数 beta（PQ 误差比率）并将大小为 T 的原始搜索列表嵌套在大小为L的更大列表中的方式来改进重排序策略。beta 可以通过实证评估选择。

例如，通过对基础顶点进行采样并构建32字节的 PQ 码，作者观察到 SIFT 数据集中99%的 PQ 距离在其准确距离的1.06倍以内（beta）。然后，作者重新对所有 PQ 距离小于 d ist L[T] * beta 的顶点进行排序。最终可以在对 QPS 几乎没有影响的情况下提高召回率高达10％。如下图所示：

d.动态列表和提前终止：

提出了一种动态列表和提前终止的策略来利用图遍历过程中的信息。作者设计了一种新颖的提前终止技术，使用动态列表在图遍历过程中迭代调整 T。在每次迭代中，比较重新排序的前 k 个新候选项和旧候选项，如果它们连续 r 次迭代相同，则终止搜索；否则，将T增加T步长并继续。作者通过设置 r 和 T 步长来控制提前终止的结果。

存储计算得到的距离以分摊每次迭代中准确距离计算的开销。同时，保持 L 在相对较大的大小 L 上以保留在较小的T处的有用信息。这也使得该方法可以应用于优化的重新排序方案。作者将提前终止与优化的重新排序相结合，与表 I 中的6个数据集相比，对于相同的召回率，距离计算减少约10%。这表明提前终止可以很好地适应不同大小和分布的数据集。

e.顶点索引的间隙编码：

gap 编码包括两个步骤：首先，对每一行的NN索引按升序排序，然后将排序后的索引转换为除第一个索引外与前一个索引的差值。在这种情况下，位宽由最大差值的位数决定。因此，所需空间减少到 168-b。实验表明，1M 到 100M 数据集的图形需要 20-b 到 26-b 的间隙编码，导致至少19%到37%的图形索引数据压缩。压缩的图索引还有助于实现更快的图遍历，因为降低了索引获取的开销。

f.3D NAND 和异构集成来构建近存储加速器：

下图展示了 Proxima 加速器的架构，主要由 3D NAND Tiles 和 Search Engine 两部分组成。Proxima 是一个独立的近存储加速器，可以同时实现数据存储和高效的 ANNS 功能。这是通过利用先进的异构集成技术实现的，该技术将 3D NAND 晶圆与 CMOS 晶圆连接起来，以提高效率和存储容量。

3D NAND Tile：

3D NAND Tile 是实现在 3D NAND 晶圆上的，用于存储三种类型的图数据：1. 原始数据，2. 原始数据的 PQ 码，3. 压缩图 NN 索引。为了确保架构的可扩展性，采用了前期工作中的两级空间层次结构来构建 3D NAND Tile，总共有 N 个 Tile，每个 Tile 包含 N 个 Core。通过 Tile 或 Core H-tree 总线连接同一层次的 Tile 和 Core。在每个 Core 内部，有 N 个 3D NAND 亚阵列子块和外围电路。为了提升图搜索的数据访问速度，我们自定义了 3D NAND Core，通过在 BL 和页面缓冲区之间添加了 BL MUX，并选择了较小的阵列尺寸以降低读取延迟。

Search Engine ：

搜索引擎在 CMOS 晶圆上制造，以最大限度地提高逻辑密度。然后通过 H-tree 总线连接到 tile 控制器上，再通过 Cu-Cu 键合连接到存储器阵列上。该搜索引擎有效地实现了比邻图搜索算法，主要用于，调度和处理传入查询;从 3D NAND 内核获取图形数据。3D NAND 闪存核心是搜索引擎可以访问的最小单元，用于获取数据。

g.Proxima执行流程：

Proxima 加速器的执行主要由两个步骤组成：1. 图数据预加载和2. 查询图搜索。在图搜索之前，需要使用提出的数据映射方案将原始数据、原始数据的 PQ 码和 NN 索引预存到相应的物理地址中。

在所需的数据存储完毕后，图搜索引擎内部的数据流程如下图所示。整个图搜索包括四个步骤：

步骤1：是针对新查询的初始阶段，其中PQ模块根据查询数据计算出不对称距离表(ADT)。计算得到的ADT被传输到由调度器指定的目标队列中的ADT内存。

步骤2：队列中的候选列表弹出未评估的顶点，仲裁器生成相应的地址从3D NAND核心中获取NN索引。获取到的NN索引首先通过Bloom过滤器进行处理，以更新新访问的顶点并过滤掉之前计算过的顶点。然后，未访问顶点的PQ码经过距离计算模块进行获取和计算。

步骤3：需要进行排序以选择前L个候选项，在所有邻居都已访问完后。实现的Bitonic排序器是所有搜索队列共享的排序器，因为它能提供足够的吞吐量。

步骤4：在完成基于PQ距离的搜索后，候选列表内的顶点将使用其原始数据进行重新排序。候选列表还会检查是否满足提前终止条件。

h.3D NAND 核心设计：

基于 3D NAND 模拟器开发了一个模拟器，用于在下图中展示96层 3D NAND 的密度、面积和读取延迟之间的权衡，作为 Proxima 核心设计的指导。使用 N BL = 36768，N SSL = 4 和 N block = 64 来构建每个核心。大幅减小页面大小也会降低面积效率。同时，在页面缓冲区和 3D NAND 块之间实现了 BL MUX，以减小页面缓冲区的大小并获得更好的数据粒度。

i.搜索引擎设计：

图搜索的搜索引擎由以下组件组成:

（1）调度程序和仲裁程序：

调度器采用简单的 Round-Robin 策略，采用先到先服务的策略，将新查询分配到理想队列。调度器在 Nq-b 缓冲区中跟踪所有队列的状态。仲裁器用于将数据获取请求从队列分配到目标 3D NAND 内核。为此，仲裁器首先将队列中的顶点信息转换为相关的物理地址。如果目的核处于理想状态，则通过H-tree总线发送物理地址。否则，数据获取请求将被仲裁器暂时停止，并等待下一轮分配。

（2）PQ 模型：

PQ 模块包含一个码本存储器和总共32个 FP16 乘法和累积 (MAC) 单元来计算 C×M ADT。码本存储器存储通过离线 k-means 训练的 C × D PQ 码本。由于参数 C = 256, M = 32 对于不同的数据集是固定的，而向量维 D 的范围从96到128，我们使用 64kB 的 SRAM 作为码本存储器，16kB 的 SRAM 作为 ADT 存储器。复杂度为 0 (C·D) 的 ADT 计算会产生 8D (角距离)到 24D (欧氏距离)时钟周期的延迟。

（3）距离计算模块：

每个队列内部的距离计算模块由一个 ADT 内存、一个查询缓冲区和一个 MAC 单元组成。它支持两种类型的距离计算:近似 PQ 距离和精确距离。对于 PQ 距离，距离计算模块使用获取的 256-b PQ 代码在 ADT 存储器中查找 M 个子向量的部分距离。然后计算 M 个时钟周期内的累积。为了获得准确的距离，首先根据缓存在查询缓冲区中的查询数据获取和计算基本数据向量，总共需要 D 个时钟周期。

（4）队列：

Queue 独立工作以处理来自调度器或仲裁器的传入数据。来自 PQ 模块的数据包括 ADT 和查询向量，用于初始化 ADT 内存和查询缓冲区。队列将顶点信息发送给仲裁器以获取存储在 3D NAND 闪存块中的数据。读取数据 (PQ 代码、原始基向量或 NN 索引)通过仲裁器发送回相应的队列。高性能图搜索的一个关键设计是搜索队列。Proxima 加速器包含总共512个核心，提供高内部数据并行性。单队列无法充分利用这种并行性。为了提高核心利用率和搜索吞吐量，在图8的搜索队列模块中实现了 Nq 队列。

（5）布隆过滤器：

图搜索过程中，有些顶点可能会被重复访问。为了节省冗余计算，保存并检查访问顶点的索引。根据我们的模拟，访问顶点的数量随着候选列表大小 |L| 线性增加。对于 |L| = 250，最多需要保存8000个顶点。我们使用节省内存的 Bloom 过滤器来检测访问的顶点。布隆过滤器是一种假阳性为 (1−e kn/m) k 的概率数据结构，其中 k 为哈希函数个数，m 为大小位数组大小，n 为需要插入的元素个数。我们在 Bloom 滤波器模块中实现了一个 12kB 的 SRAM 和8个轻量级的 SeaHashes，以保证误报概率< 0.02%，提供可忽略不计的召回损失。

（6）双音分类器和候选列表：

在图遍历过程中，需要对每个搜索队列中的候选列表进行排序，以获得排序后的候选列表。在最坏的情况下，需要对超过200个距离进行排序。我们使用提供恒定延迟的并行 Bitonic 排序器来避免过多的排序延迟。Bitonic 排序器是分段流水线的，每个周期并行接受 Nsorter 输入。Nsorter 输入的排序延迟是恒定的 2 log2 Nsorter 时钟周期。Proxima 只实现一个全局 Nsorter = 256 点 Bitonic 排序器，满足使用列表大小 L 和队列的延迟和吞吐量需求。候选列表是一个 2kB 的缓冲区，用于存储算法1中的候选集 L，包括每个候选集的距离和相应的顶点索引。候选列表还存储上次迭代的先前 k-NN 结果，以支持提前终止。

j.数据映射策略：

数据映射策略需要向仲裁器提供低复杂度的地址转换逻辑；

数据映射策略应能使系统的查询搜索性能最大化；

数据映射策略应该是可扩展的，以支持不同的 ANNS 数据集规模。

实验

Vector Search

1.ANNS算法评价：

a.对图形基线的改进：

作者评估了 Proxima 图搜索，并将其与 DiskANN 和 HNSW 在 CPU 上进行了比较。下图显示了吞吐量 (QPS) 和召回率，其中 Proxima 在所有数据集上实现了相当的召回率和吞吐量，在 1M 数据集的相同吞吐量下，召回率比 DiskANN 和 HNSW 提高了10%。

b.与非图基线的比较：

作者评估了使用 PQ 压缩的非图算法 FAISS-IVF。尽管 FAISS-IVF 的内存占用更小，但有损耗的 PQ 压缩引入了显著的近似误差，并导致小数据集和大数据集的召回率分别达到90%和85%左右。基于图的方法始终比 FAISS-IVF 获得更好的召回率和吞吐量权衡。

2.硬件性能对比：

a.Proxima vs. CPU、GPU 和 ASIC：

我们将 3D NAND 中的 Proxima 近存储加速器与用于 ANNS 的最先进的 CPU、GPU 和 ASIC 设计进行比较。下图中，总结了两个小尺度数据集(SIFT 和 GLOVE)和两个大尺度 100M 数据集 (BIGANN-100M 和 DEEP-100M) 的吞吐量对比。GGNN 是针对基于图的人工神经网络优化的高性能 GPU 工具，而 ANNA 是基于 pq-ivf 的 ASIC 加速器。我们采用 HNSW 作为 CPU 基准。在所有基线中，Proxima 实现了最高的吞吐量，而基于 gpu 的 GGNN 是第二快的工具。与 ASIC 设计相比，ANNA， Proxima 在 1M 和 100M 数据集上的速度快6.6到13倍。与 CPU 基线相比，Proxima 的加速对于大规模数据集或高复杂性数据集更为显著，例如 GLOVE，需要更多的距离计算(6倍到8倍)才能达到相同的召回率。

同时，测量了比邻星的能源效率和下图中提到的基线。能源效率以 QPS/W 衡量。在所有数据集中，Proxima 的能源效率最高。由于 ANNA 需要大的片上 ram 和频繁的片外数据传输，其能效比 Proxima 低17倍。昂贵的数据移动增加了能源消耗。与 CPU 相比，Proxima 具有三个数量级的能效，因为 Proxima 实现在 3D NAND 附近，具有更短的数据访问延迟和更低的读取能量。

b.Proxima vs. Existing ANNS Accelerators：

在下表中比较了 Proxima 与最先进的 ANNS 加速度，包括基于 cpu 的 DiskANN-PQ、基于 gpu 的 GGNN、基于 asic 的 ANNA 和基于 nsp 的 VStore。Proxima 和 VStore 是仅有的两种可以在 3D NAND 闪存中保存图形数据的设计。与租户相比，Proxima 具有更紧凑的设计，更短的内存延迟，并且由于高速异构集成，峰值内存带宽提高了26倍。

c.硬件开销：

下表总结了 Proxima 近存储加速器的面积和能量分布。Proxima 的面积比 A40 GPU 的 628mm2 的芯片面积小2.4倍。

3.Proxima 优化的有效性：

a.间隙编码与提前终止：

在下图中分析了不同图 ANNS 搜索算法的内存流量。Proxima 中的间隙编码进一步减少了对神经网络索引的数据访问，而提前终止比 DiskANN-PQ 节省了10% - 25%的冗余操作。因此，Proxima 比 HNSW 实现了 1.9× 至 2.4× 的降噪比。

b.热节点重复：

下图显示了 100M 数据集上热节点百分比从0.0%到7.0%的运行时分解情况。当禁用热节点重复时，3D NAND 核和 H-tree 总线引起的数据访问延迟占总延迟的80%。增加1%的热节点(100M 数据集的 1M 个点)有助于将整体搜索延迟减少2.2倍。当热节点的比例进一步增加到3%时，与没有热节点的性能相比，性能提升约为3倍。然而，添加热节点> 3%的好处达到了一个平稳期。我们在 Proxima 中使用3%作为默认值。

c.对其他图 ANNS 算法的改进：

作者比较了 Proxima 与 HNSW 和 DiskANN-PQ 这两种最先进的产品的性能。下图显示了在 10M 和 100M 数据集上的吞吐量、能效、查询延迟。通过软件优化，具有间隙编码和提前终止的 Proxima 比 HNSW 和 DiskANN-PQ 实现了适度的吞吐量和效率提高。提前终止和间隙编码有助于减少计算冗余和搜索过程中的数据移动。HNSW 的性能最差，因为它需要许多精确的距离计算。采用热节点重复后，Proxima 的速度提高了2倍，能效提高了3倍，延迟降低了3倍。好处主要来自于图索引重排序和热排序带来的更好的数据局部性节点重复，减少图遍历期间的内存访问。