土法炼钢兴趣小组的算法知识备份

ScaNN 与 DiskANN:两大厂的向量检索之路

目录

一、引言:十亿向量检索的两条路

当向量数据库需要服务十亿级别的向量时,工程师面临一个根本性的选择:把所有数据压进内存,还是让磁盘参与计算?Google 和 Microsoft 分别给出了截然不同的答案。

Google 的 ScaNN(Scalable Nearest Neighbors)走的是量化压缩路线。其核心洞察是:传统的向量量化方法在优化目标上存在根本性缺陷——它们优化的是量化误差的均方差,而非检索质量本身。ScaNN 提出了各向异性向量量化(Anisotropic Vector Quantization),让量化过程直接感知最终排序分数,从而在同等压缩比下获得显著更高的召回率。

Microsoft 的 DiskANN 走的是另一条路。它的核心思想是:NVMe SSD 的随机读取延迟已经降到 10 微秒级别,如果能巧妙地设计图索引的磁盘布局,并在内存中只保留极少量的辅助数据,就能以极低的内存开销服务十亿级别的向量检索。DiskANN 设计了 Vamana 图索引算法,配合 SSD 上的扇区对齐存储和内存中的 PQ 压缩向量,实现了亚毫秒级别的查询延迟。

这两个系统代表了向量检索领域的两个重要方向。ScaNN 追求的是「在有限内存预算下,最大化检索质量」;DiskANN 追求的是「以磁盘换内存,大幅降低硬件成本」。理解这两个系统的设计思路,对于理解现代向量检索的工程实践至关重要。

本文将从数学原理出发,深入分析 ScaNN 和 DiskANN 的核心算法,并提供 Vamana 图构建的完整 C 实现。最后我们会讨论 SPANN、基准测试对比以及生产环境中的实际考量。

二、背景:从乘积量化到图索引

在进入 ScaNN 和 DiskANN 的细节之前,我们需要回顾两个基础技术——乘积量化(Product Quantization, PQ)和图索引(Graph-based Index)。

乘积量化回顾

乘积量化将高维向量分割为 M 个子空间,每个子空间独立做 k-means 聚类,生成 k 个码本。原始向量被编码为 M 个码本索引,每个索引用 8 bit 存储,因此一个 d 维浮点向量从 4d 字节压缩到 M 字节。

查询时使用非对称距离计算(ADC):

d_ADC(q, x) = Σ_{m=1}^{M} || q_m - c_m(x_m) ||^2

其中 q_m 是查询向量的第 m 个子空间分量,c_m(x_m) 是数据向量 x 在第 m 个子空间的量化中心。可以预计算查询向量到所有码本中心的距离表,每次距离计算只需 M 次查表和加法。

图索引回顾

图索引的基本思想是将数据集建模为一个导航图。每个向量是一个节点,通过贪心搜索在图上导航到查询向量的近邻。HNSW 是当前最流行的图索引方案,它构建多层跳表结构,上层稀疏用于快速定位,下层稠密用于精确搜索。

HNSW 的主要问题是内存占用:对于 10 亿个 128 维浮点向量,仅原始数据就需要约 512 GB,加上图的邻接表,总内存需求可能超过 600 GB。这在绝大多数生产环境中都是不可接受的。

两条优化路径

面对这个挑战,ScaNN 和 DiskANN 走了不同的路:

维度 ScaNN DiskANN
核心思路 更好的量化,减少内存中的数据量 将数据放到 SSD,只在内存保留索引
数据存储 内存(量化后) SSD(全精度) + 内存(PQ 压缩)
索引结构 IVF + 量化 Vamana 图
距离计算 各向异性量化距离 PQ 估算 + SSD 精确重排
10 亿向量内存 ~60 GB ~32 GB
硬件依赖 大内存 NVMe SSD

这张表给出了一个宏观的对比。接下来我们分别深入两个系统的技术细节。

三、ScaNN:各向异性向量量化

ScaNN 的核心创新是各向异性向量量化(Anisotropic Vector Quantization, AVQ)。要理解 AVQ 的价值,我们必须先理解传统量化方法的缺陷。

为什么非对称距离计算还不够

传统 PQ 使用非对称距离计算(ADC),其量化误差为:

ε = d(q, x) - d_ADC(q, x)
    = ||q - x||^2 - ||q - Q(x)||^2
    = -||x - Q(x)||^2 + 2⟨q - Q(x), x - Q(x)⟩

其中 Q(x) 是 x 的量化近似。传统量化方法(k-means、OPQ 等)优化的是量化误差 ||x - Q(x)||^2 的期望。但请注意上面展开式中的第二项——它包含了查询向量 q 与量化残差的交叉项。

这意味着,即使量化误差 ||x - Q(x)||^2 很小,交叉项也可能导致排序反转。具体来说,如果量化残差 x - Q(x) 的方向恰好与 q - Q(x) 对齐,即使残差的模长很小,也会产生显著的距离估计偏差。

更关键的是,对于最终检索结果影响最大的是那些「边界附近」的向量——它们的真实距离非常接近,微小的量化误差就可能导致排序反转。而传统量化方法对所有向量一视同仁地最小化量化误差,并没有将优化资源集中在这些关键向量上。

分数感知的量化损失

ScaNN 的洞察是:量化的目标不应该是最小化量化误差本身,而应该是最小化量化对内积估计的影响。对于最大内积搜索(MIPS)问题,我们关心的是:

⟨q, x⟩ ≈ ⟨q, Q(x)⟩

量化误差对内积估计的影响为:

⟨q, x⟩ - ⟨q, Q(x)⟩ = ⟨q, x - Q(x)⟩

这个误差取决于量化残差 x - Q(x) 在查询方向上的投影。ScaNN 将量化残差分解为平行于 x 的分量和垂直于 x 的分量:

x - Q(x) = η_∥ · x̂ + η_⊥

其中 是 x 的单位方向向量,η_∥ 是平行分量的标量,η_⊥ 是垂直分量。

对于真正的近邻(与查询向量内积大的向量),查询方向 q 与 x 的方向大致对齐。因此:

⟨q, x - Q(x)⟩ ≈ η_∥ · ⟨q, x̂⟩ + ⟨q, η_⊥⟩

由于 q 与 x 方向接近,第一项 η_∥ · ⟨q, x̂⟩ 的贡献远大于第二项。这就是「各向异性」的来源——平行于数据向量方向的量化误差比垂直方向的误差重要得多。

AVQ 的数学形式化

ScaNN 定义了各向异性量化损失函数:

L_AVQ(Q) = E_x [ ||x - Q(x)||_W^2 ]

其中 ||·||_W 是加权范数,权重矩阵为:

W(x) = (1 - t) · x̂ x̂^T + t · I

参数 t ∈ [0, 1] 控制各向异性的程度: - t = 1 退化为标准量化(各向同性) - t = 0 只惩罚平行方向的误差(完全各向异性) - 实践中 t ∈ [0.1, 0.3] 效果最好

这个加权范数的效果是:如果量化残差主要在垂直于 x 的方向上,那么惩罚较小;如果残差在平行于 x 的方向上,惩罚会很大。

实现细节

AVQ 的训练过程如下:

import numpy as np

def train_avq_codebook(data, M, K, t, n_iter=20):
    """
    训练各向异性向量量化码本。
    data: (N, d) 数据矩阵
    M: 子空间数量
    K: 每个子空间的码本大小
    t: 各向异性参数
    """
    N, d = data.shape
    sub_dim = d // M
    codebooks = []

    for m in range(M):
        sub_data = data[:, m * sub_dim : (m + 1) * sub_dim]
        centers = sub_data[np.random.choice(N, K, replace=False)]

        for iteration in range(n_iter):
            # 计算加权距离分配
            assignments = np.zeros(N, dtype=np.int32)
            for i in range(N):
                x = data[i]  # 需要完整向量计算方向
                x_norm = np.linalg.norm(x)
                if x_norm < 1e-8:
                    x_hat = np.zeros(d)
                else:
                    x_hat = x / x_norm

                # 子空间方向分量
                x_hat_sub = x_hat[m * sub_dim : (m + 1) * sub_dim]
                sub_x = sub_data[i]

                best_dist = float('inf')
                best_k = 0
                for k in range(K):
                    residual = sub_x - centers[k]
                    # 各向异性加权距离
                    parallel = np.dot(residual, x_hat_sub) * x_hat_sub
                    perp = residual - parallel
                    w_dist = (1 - t) * np.dot(parallel, parallel) + t * np.dot(perp, perp)
                    if w_dist < best_dist:
                        best_dist = w_dist
                        best_k = k
                assignments[i] = best_k

            # 更新码本中心(加权质心)
            for k in range(K):
                mask = assignments == k
                if np.sum(mask) > 0:
                    centers[k] = np.mean(sub_data[mask], axis=0)

        codebooks.append(centers)

    return codebooks

ScaNN 的完整检索流程

ScaNN 使用两阶段检索:

  1. 粗筛阶段(IVF):使用倒排索引(IVF)将查询路由到最近的若干聚类中心,缩小候选集到总数据的 1%~5%。

  2. 精排阶段(AVQ):在候选集上使用各向异性量化距离计算,选出 top-k 结果。

def scann_search(query, ivf_index, codebooks, top_k, n_probe):
    """ScaNN 两阶段检索"""
    # 阶段一:IVF 粗筛
    candidates = ivf_index.probe(query, n_probe)

    # 阶段二:AVQ 精排
    # 预计算查询到码本中心的距离表
    M = len(codebooks)
    dist_tables = []
    for m, cb in enumerate(codebooks):
        sub_q = query[m * sub_dim : (m + 1) * sub_dim]
        dist_tables.append(np.dot(cb, sub_q))  # 内积

    # 对候选集计算 AVQ 距离
    scores = []
    for idx in candidates:
        score = sum(dist_tables[m][codes[idx][m]] for m in range(M))
        scores.append((idx, score))

    # 返回 top-k
    scores.sort(key=lambda x: -x[1])
    return scores[:top_k]

ScaNN 论文报告的结果显示,在 glove-100 数据集上,AVQ 在相同的召回率下比 OPQ 快 2 倍以上;在 10 亿规模数据集上,ScaNN 的 达到 0.95 时,QPS 比标准 PQ 高 3~4 倍。

四、DiskANN 与 Vamana 图索引

DiskANN 的核心是 Vamana——一种专为磁盘访问模式优化的图索引算法。与 HNSW 的多层结构不同,Vamana 只构建单层图,但通过精心设计的剪枝策略,确保图的直径足够小,从而将搜索路径长度控制在较少的跳数内。

Vamana 图的设计目标

Vamana 图需要满足以下约束:

  1. 低直径:从任意起点到任意目标的最短路径不超过 O(log n) 跳。这保证了搜索效率。
  2. 有界出度:每个节点的出度不超过 R,控制存储开销和每跳的计算量。
  3. 磁盘友好:每次贪心搜索步只需读取一个节点的邻居列表和对应的全精度向量,可以合并为单次 4KB 页面读取。

为实现这些目标,Vamana 使用了鲁棒剪枝(Robust Pruning)策略。该策略通过参数 α ≥ 1 控制长边和短边的平衡。

鲁棒剪枝的直觉

考虑一个场景:节点 p 需要从候选邻居集合中选择最多 R 个邻居。标准方法是选择距离 p 最近的 R 个节点。但这种方法有一个问题——所有选出的邻居可能都聚集在 p 附近的一个小区域,导致图的覆盖范围不足。

鲁棒剪枝的核心思想是:如果候选邻居 v 离 p 很近,但已经被 p 的某个已选邻居 u 很好地「覆盖」了(即 d(u, v) 远小于 d(p, v)),那么 v 就是冗余的,应该被跳过。参数 α 控制这个「覆盖」的阈值:

如果存在已选邻居 u,使得 α · d(u, v) ≤ d(p, v),则 v 被剪枝。

当 α > 1 时,剪枝条件更容易满足,图倾向于保留更多的长边(long-range edges),有助于降低图的直径。实践中 α = 1.2 是一个好的默认值。

Vamana 图构建算法

Vamana 图构建过程如下:

  1. 初始化一个随机 R-正则图。
  2. 计算数据集的 medoid(几何中心最近的点)作为导航节点 s。
  3. 对数据点进行随机排列,依次处理每个点。
  4. 对于每个点 p,执行贪心搜索从 s 出发找到 p 的近邻候选集。
  5. 将搜索路径上的所有访问节点与 p 的现有邻居合并为候选集。
  6. 对候选集执行鲁棒剪枝,选出 p 的新邻居列表。
  7. 为保持图的无向性,将 p 也添加到其新邻居的邻居列表中(必要时对这些邻居也执行剪枝)。

算法需要执行两遍:第一遍使用 α = 1,构建基础图;第二遍使用 α > 1(如 1.2),在第一遍的基础上添加长边。

贪心搜索

贪心搜索是 Vamana 图上的基本操作,用于图构建和查询:

GreedySearch(s, q, k, L):
  输入: 起始节点 s, 查询 q, 返回数 k, 搜索列表大小 L ≥ k
  V ← {s}              // 已访问集合
  C ← {s}              // 候选列表(按距离排序)
  while C 中有未展开的节点:
    p* ← C 中距离 q 最近的未展开节点
    标记 p* 为已展开
    for each neighbor n of p*:
      if n ∉ V:
        V ← V ∪ {n}
        if d(n, q) < d(C[L], q) or |C| < L:
          C ← C ∪ {n}
          if |C| > L:
            C ← C 中距离最近的 L 个节点
  return C 中距离最近的 k 个节点, V

搜索列表大小 L 是一个重要参数,它控制了搜索的广度。L 越大,召回率越高,但搜索开销也越大。在 DiskANN 中,L 对应磁盘 I/O 次数的上限。

磁盘布局

DiskANN 将每个节点的数据存储在 SSD 上,布局按 4KB 页面对齐:

┌──────────────────────────────────────────┐
│  Node i (4KB aligned page)               │
├──────────────────────────────────────────┤
│  full vector (128 × 4B = 512B)           │
│  neighbor count (4B)                     │
│  neighbor IDs (R × 4B, e.g. 64×4=256B)  │
│  padding to 4KB                          │
└──────────────────────────────────────────┘

这种布局确保读取一个节点的全部信息只需一次 4KB 随机读取。搜索时对一个节点的处理包括:读取其全精度向量用于精确距离计算,读取其邻居列表用于导航。

DiskANN Architecture

五、Vamana 图构建的完整 C 实现

以下是 Vamana 图构建算法的完整 C 实现。该实现包含数据结构定义、鲁棒剪枝、贪心搜索和完整的图构建流程。

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <float.h>
#include <time.h>

/* ===== 常量与配置 ===== */
#define MAX_DEGREE    64   /* R: 最大出度 */
#define SEARCH_LIST  128   /* L: 搜索列表大小 */
#define DIM          128   /* 向量维度 */

/* ===== 数据结构 ===== */
typedef struct {
    int   id;
    float vec[DIM];
} Point;

typedef struct {
    int neighbors[MAX_DEGREE];
    int degree;
} AdjList;

typedef struct {
    int   id;
    float dist;
} Candidate;

typedef struct {
    Point   *points;
    AdjList *adj;
    int      n;
    int      medoid;
} VamanaGraph;

/* ===== 距离计算 ===== */
static float l2_distance(const float *a, const float *b)
{
    float sum = 0.0f;
    for (int i = 0; i < DIM; i++) {
        float d = a[i] - b[i];
        sum += d * d;
    }
    return sum;  /* 返回平方距离,省去 sqrt */
}

/* ===== 候选列表操作 ===== */
static int cand_cmp(const void *a, const void *b)
{
    float da = ((const Candidate *)a)->dist;
    float db = ((const Candidate *)b)->dist;
    return (da > db) - (da < db);
}

/* 在有序候选列表中插入新候选,保持大小 ≤ max_size */
static int cand_insert(Candidate *list, int *size, int max_size,
                       int id, float dist)
{
    if (*size >= max_size && dist >= list[*size - 1].dist)
        return 0;  /* 距离太大,不插入 */

    int pos = *size;
    while (pos > 0 && list[pos - 1].dist > dist)
        pos--;

    if (*size < max_size) {
        memmove(&list[pos + 1], &list[pos],
                (*size - pos) * sizeof(Candidate));
        (*size)++;
    } else {
        memmove(&list[pos + 1], &list[pos],
                (max_size - 1 - pos) * sizeof(Candidate));
    }
    list[pos].id   = id;
    list[pos].dist = dist;
    return 1;
}

/* ===== 贪心搜索 ===== */
/* 返回: result 中存放距离最近的 k 个节点,visited 记录所有访问节点 */
static int greedy_search(const VamanaGraph *g, const float *query,
                         int start, int k, int L,
                         Candidate *result, int *result_size,
                         Candidate *visited, int *visited_size)
{
    /* visited_set 用位图或哈希表;此处简化为数组标记 */
    char *seen = calloc(g->n, sizeof(char));
    if (!seen) return -1;

    Candidate *cand_list = malloc(sizeof(Candidate) * (L + 1));
    int cand_size = 0;
    int expand_idx = 0;  /* 下一个待展开的位置 */

    /* 初始化 */
    float d = l2_distance(query, g->points[start].vec);
    cand_insert(cand_list, &cand_size, L, start, d);
    seen[start] = 1;

    *visited_size = 0;
    visited[(*visited_size)].id   = start;
    visited[(*visited_size)].dist = d;
    (*visited_size)++;

    while (expand_idx < cand_size) {
        int p = cand_list[expand_idx].id;
        expand_idx++;

        const AdjList *adj = &g->adj[p];
        for (int i = 0; i < adj->degree; i++) {
            int nb = adj->neighbors[i];
            if (seen[nb]) continue;
            seen[nb] = 1;

            float nb_dist = l2_distance(query, g->points[nb].vec);

            /* 记录到 visited */
            visited[*visited_size].id   = nb;
            visited[*visited_size].dist = nb_dist;
            (*visited_size)++;

            cand_insert(cand_list, &cand_size, L, nb, nb_dist);
        }
    }

    /* 输出 top-k */
    *result_size = (cand_size < k) ? cand_size : k;
    memcpy(result, cand_list, *result_size * sizeof(Candidate));

    free(cand_list);
    free(seen);
    return 0;
}

/* ===== 鲁棒剪枝 ===== */
static void robust_prune(VamanaGraph *g, int p_id, Candidate *cands,
                         int cand_size, float alpha, int R)
{
    /* 按距离排序候选集 */
    qsort(cands, cand_size, sizeof(Candidate), cand_cmp);

    int new_neighbors[MAX_DEGREE];
    int new_degree = 0;
    char *pruned = calloc(cand_size, sizeof(char));

    for (int i = 0; i < cand_size && new_degree < R; i++) {
        if (pruned[i]) continue;

        int v = cands[i].id;
        if (v == p_id) continue;

        new_neighbors[new_degree++] = v;

        /* 剪枝: 如果后续候选被 v 覆盖,则标记为 pruned */
        for (int j = i + 1; j < cand_size; j++) {
            if (pruned[j]) continue;
            int u = cands[j].id;
            float d_vu = l2_distance(g->points[v].vec,
                                     g->points[u].vec);
            float d_pu = cands[j].dist;
            if (alpha * d_vu <= d_pu) {
                pruned[j] = 1;
            }
        }
    }

    /* 更新邻居列表 */
    g->adj[p_id].degree = new_degree;
    memcpy(g->adj[p_id].neighbors, new_neighbors,
           new_degree * sizeof(int));

    free(pruned);
}

/* ===== 计算 medoid ===== */
static int compute_medoid(const Point *points, int n)
{
    /* 计算质心 */
    float centroid[DIM] = {0};
    for (int i = 0; i < n; i++)
        for (int d = 0; d < DIM; d++)
            centroid[d] += points[i].vec[d];
    for (int d = 0; d < DIM; d++)
        centroid[d] /= n;

    /* 找最近质心的点 */
    float best_dist = FLT_MAX;
    int   best_id   = 0;
    for (int i = 0; i < n; i++) {
        float d = l2_distance(centroid, points[i].vec);
        if (d < best_dist) {
            best_dist = d;
            best_id   = i;
        }
    }
    return best_id;
}

/* ===== 初始化随机图 ===== */
static void init_random_graph(VamanaGraph *g, int R)
{
    for (int i = 0; i < g->n; i++) {
        int deg = (R < g->n - 1) ? R : (g->n - 1);
        g->adj[i].degree = 0;
        char *used = calloc(g->n, sizeof(char));
        used[i] = 1;
        while (g->adj[i].degree < deg) {
            int r = rand() % g->n;
            if (!used[r]) {
                used[r] = 1;
                g->adj[i].neighbors[g->adj[i].degree++] = r;
            }
        }
        free(used);
    }
}

/* ===== Fisher-Yates 洗牌 ===== */
static void shuffle(int *arr, int n)
{
    for (int i = n - 1; i > 0; i--) {
        int j = rand() % (i + 1);
        int tmp = arr[i];
        arr[i]  = arr[j];
        arr[j]  = tmp;
    }
}

/* ===== 反向边插入辅助函数 ===== */
static void add_reverse_edge(VamanaGraph *g, int from, int to,
                             float alpha, int R)
{
    AdjList *adj = &g->adj[from];

    /* 检查是否已存在 */
    for (int i = 0; i < adj->degree; i++)
        if (adj->neighbors[i] == to) return;

    if (adj->degree < R) {
        adj->neighbors[adj->degree++] = to;
    } else {
        /* 邻居已满,收集候选后重新剪枝 */
        int cand_size = adj->degree + 1;
        Candidate *cands = malloc(sizeof(Candidate) * cand_size);
        for (int i = 0; i < adj->degree; i++) {
            cands[i].id   = adj->neighbors[i];
            cands[i].dist = l2_distance(g->points[from].vec,
                                        g->points[adj->neighbors[i]].vec);
        }
        cands[adj->degree].id   = to;
        cands[adj->degree].dist = l2_distance(g->points[from].vec,
                                              g->points[to].vec);
        robust_prune(g, from, cands, cand_size, alpha, R);
        free(cands);
    }
}

/* ===== Vamana 图构建主函数 ===== */
void vamana_build(VamanaGraph *g, float alpha, int R, int L)
{
    /* 步骤 1: 初始化随机图 */
    init_random_graph(g, R);

    /* 步骤 2: 计算 medoid */
    g->medoid = compute_medoid(g->points, g->n);

    /* 步骤 3: 随机排列 */
    int *perm = malloc(sizeof(int) * g->n);
    for (int i = 0; i < g->n; i++) perm[i] = i;
    shuffle(perm, g->n);

    /* 步骤 4: 逐点插入 */
    int max_visited = g->n;
    Candidate *result  = malloc(sizeof(Candidate) * L);
    Candidate *visited = malloc(sizeof(Candidate) * max_visited);

    for (int iter = 0; iter < g->n; iter++) {
        int p = perm[iter];
        int result_size = 0, visited_size = 0;

        /* 从 medoid 出发做贪心搜索 */
        greedy_search(g, g->points[p].vec, g->medoid,
                      1, L, result, &result_size,
                      visited, &visited_size);

        /* 合并现有邻居到候选集 */
        int cand_cap = visited_size + g->adj[p].degree;
        Candidate *cands = malloc(sizeof(Candidate) * cand_cap);
        int cand_size = 0;

        for (int i = 0; i < visited_size; i++) {
            if (visited[i].id != p) {
                cands[cand_size].id   = visited[i].id;
                cands[cand_size].dist = visited[i].dist;
                cand_size++;
            }
        }
        for (int i = 0; i < g->adj[p].degree; i++) {
            int nb = g->adj[p].neighbors[i];
            float d = l2_distance(g->points[p].vec,
                                  g->points[nb].vec);
            cands[cand_size].id   = nb;
            cands[cand_size].dist = d;
            cand_size++;
        }

        /* 鲁棒剪枝 */
        robust_prune(g, p, cands, cand_size, alpha, R);

        /* 添加反向边 */
        for (int i = 0; i < g->adj[p].degree; i++) {
            add_reverse_edge(g, g->adj[p].neighbors[i], p, alpha, R);
        }

        free(cands);
    }

    free(perm);
    free(result);
    free(visited);
}

/* ===== 使用示例 ===== */
int main(void)
{
    srand((unsigned)time(NULL));

    int N = 10000;
    VamanaGraph g;
    g.n      = N;
    g.points = malloc(sizeof(Point) * N);
    g.adj    = calloc(N, sizeof(AdjList));

    /* 生成随机数据 */
    for (int i = 0; i < N; i++) {
        g.points[i].id = i;
        for (int d = 0; d < DIM; d++)
            g.points[i].vec[d] = (float)rand() / RAND_MAX;
    }

    printf("Building Vamana graph (N=%d, R=%d, L=%d)...\n",
           N, MAX_DEGREE, SEARCH_LIST);

    /* 两遍构建: pass 1 用 alpha=1.0, pass 2 用 alpha=1.2 */
    vamana_build(&g, 1.0f, MAX_DEGREE, SEARCH_LIST);
    printf("Pass 1 complete (alpha=1.0)\n");

    vamana_build(&g, 1.2f, MAX_DEGREE, SEARCH_LIST);
    printf("Pass 2 complete (alpha=1.2)\n");

    /* 查询测试 */
    float query[DIM];
    for (int d = 0; d < DIM; d++)
        query[d] = (float)rand() / RAND_MAX;

    Candidate result[10];
    Candidate visited[10000];
    int result_size = 0, visited_size = 0;

    greedy_search(&g, query, g.medoid, 10, SEARCH_LIST,
                  result, &result_size, visited, &visited_size);

    printf("Top-10 results (visited %d nodes):\n", visited_size);
    for (int i = 0; i < result_size; i++)
        printf("  rank=%d  id=%d  dist=%.4f\n",
               i + 1, result[i].id, result[i].dist);

    free(g.points);
    free(g.adj);
    return 0;
}

代码的要点说明:

  1. greedy_search 返回搜索路径上所有访问过的节点,这些节点将作为鲁棒剪枝的候选集。
  2. robust_prune 实现了 α-剪枝:对候选按距离排序后,贪心选择邻居,跳过被已选邻居「覆盖」的候选。
  3. add_reverse_edge 在添加反向边时,如果目标节点度数已满,会触发剪枝。
  4. vamana_build 需要执行两遍,第一遍 α=1.0 构建局部连接,第二遍 α=1.2 增加长距离连接。

在理解了 Vamana 图结构之后,我们来看 DiskANN 如何在 SSD 上高效执行查询。这是 DiskANN 工程上最精彩的部分。

内存中的 PQ 压缩向量

DiskANN 在内存中维护所有向量的 PQ 压缩版本。对于 128 维浮点向量(512 字节),使用 M=32 的乘积量化可以压缩到 32 字节,压缩比 16 倍。10 亿个向量的 PQ 数据约占 32 GB 内存。

PQ 压缩向量的作用是:在图搜索过程中,当需要展开某个节点的邻居时,先用 PQ 距离估算候选邻居与查询的距离,只将最有希望的候选提交给 SSD 读取。这大幅减少了实际的 SSD I/O 次数。

Beam Search with SSD

DiskANN 的查询流程结合了贪心搜索和批量 SSD 读取:

DiskANN_Search(query, k, W, B):
  输入:
    query: 查询向量
    k: 返回结果数
    W: beam width(搜索宽度)
    B: 每轮最大 SSD 读取次数

  初始化:
    pq_dists ← 预计算 query 到所有 PQ 码本中心的距离表
    frontier ← {medoid}  // 初始前沿
    result   ← {}        // 结果集
    visited  ← {}        // 已访问集合

  while frontier 非空:
    // 阶段 1: PQ 距离估算
    for each node p in frontier:
      读取 p 的邻居列表(来自 SSD 缓存或本轮 I/O)
      for each neighbor n not in visited:
        visited ← visited ∪ {n}
        计算 PQ 距离 d_pq(query, n) 使用距离表
        加入全局候选集

    // 阶段 2: 选择 SSD 读取目标
    从全局候选集中选择 PQ 距离最小的 B 个未读取节点
    批量发起 B 次异步 SSD 读取

    // 阶段 3: 精确距离计算
    等待 SSD 读取完成
    对每个读取到的节点:
      计算精确 L2 距离
      更新 result(保持 top-k)
      将该节点加入下一轮的 frontier

    // 收敛判断
    if frontier 中最近节点的距离 > result[k] 的距离:
      break

  return result

异步 I/O 的关键作用

NVMe SSD 的单次随机读取延迟约 10 微秒,但吞吐量可达 500K+ IOPS。DiskANN 通过批量异步 I/O 充分利用这一特性:

/* 使用 Linux io_uring 的伪代码 */
struct io_uring ring;
io_uring_queue_init(QUEUE_DEPTH, &ring, 0);

/* 批量提交读取请求 */
for (int i = 0; i < batch_size; i++) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&ring);
    io_uring_prep_read(sqe, fd,
                       buffers[i],        /* 4KB 对齐缓冲区 */
                       PAGE_SIZE,          /* 读取 4KB */
                       node_offset[i]);    /* 节点在文件中的偏移 */
    io_uring_sqe_set_data(sqe, &contexts[i]);
}
io_uring_submit(&ring);

/* 等待所有完成 */
for (int i = 0; i < batch_size; i++) {
    struct io_uring_cqe *cqe;
    io_uring_wait_cqe(&ring, &cqe);
    /* 处理完成的读取 */
    process_node(io_uring_cqe_get_data(cqe));
    io_uring_cqe_seen(&ring, cqe);
}

每次查询的 SSD 读取次数通常在 W × C 次左右,其中 W 是 beam width(通常 4~8),C 是平均搜索深度(通常 3~5 跳)。因此每次查询约 12~40 次 SSD 读取,延迟约 0.1~0.4 毫秒。

与 HNSW 的对比

在 10 亿 128 维向量上,DiskANN 与纯内存 HNSW 的对比如下:

指标 HNSW DiskANN
内存占用 ~600 GB ~32 GB
查询延迟 (P50) ~0.5 ms ~1 ms
查询延迟 (P99) ~2 ms ~5 ms
0.99 0.95+
硬件成本(月) ~$3000 ~$400

DiskANN 在延迟上有约 2 倍的劣势,但在成本上有近 10 倍的优势。对于大多数应用场景,这是一个非常划算的交易。

七、FreshDiskANN:流式更新

原始 DiskANN 的一个重要限制是它假设数据集是静态的——需要一次性构建完整索引。FreshDiskANN(也称 Fresh-DiskANN)解决了这个问题,支持在不停止服务的情况下进行流式插入和删除。

架构设计

FreshDiskANN 使用双缓冲架构:

  1. 主索引(SSD):存储大部分数据的 Vamana 图,只读。
  2. 临时索引(内存):一个小型的内存 Vamana 图,存储最近插入的数据。
  3. 合并线程:后台将临时索引中的数据合并到主索引中。
          ┌─────────────┐
  insert → │  Temp Index  │ ← 内存中,支持实时插入
          │  (小型 Vamana)│
          └──────┬───────┘
                 │ 定期合并
          ┌──────▼───────┐
          │  Main Index   │ ← SSD 上,大型 Vamana 图
          │  (Frozen)     │
          └──────────────┘

插入流程

新向量到达时:

  1. 将向量添加到临时索引中。
  2. 在临时索引中执行贪心搜索,找到近邻候选。
  3. 对新节点执行鲁棒剪枝,连接到临时索引的图中。
  4. 更新 PQ 压缩向量(内存中)。

查询时同时搜索主索引和临时索引,合并结果返回 top-k。

删除与合并

删除操作使用惰性删除——标记要删除的节点,在查询结果中过滤掉它们。当临时索引的大小达到阈值时,触发合并:

  1. 将临时索引中的所有节点批量写入 SSD。
  2. 更新主索引的图结构(添加新节点的边,修剪删除节点的边)。
  3. 清空临时索引。

这个过程可以在后台执行,不阻塞查询服务。代价是合并期间的内存使用量会暂时增加。

流式更新的性能影响

FreshDiskANN 论文报告的数据显示:

这使得 DiskANN 适用于数据持续增长的场景,如推荐系统中不断新增的用户和物品向量。

八、SPANN:混合倒排索引与图索引

SPANN(Scalable and Parameter-free Approximate Nearest Neighbor)是 Microsoft Research 提出的另一种面向磁盘的向量索引方案。它结合了倒排索引和图索引的思想,在某些场景下比 DiskANN 更具优势。

核心思想

SPANN 的设计基于一个观察:对于十亿级数据集,聚类中心的数量通常在几万到几十万的量级,可以完全放入内存。数据点则按聚类分组存储在 SSD 上。

查询时:

  1. 内存中:用图索引(如 HNSW)在聚类中心之间做导航,找到最相关的 K 个聚类。
  2. SSD 上:读取这 K 个聚类的倒排列表,在列表中做精确距离计算。
查询 q
  │
  ▼
内存: HNSW 导航 (聚类中心)
  │ 找到 top-K 聚类
  ▼
SSD: 读取 K 个倒排列表
  │ 精确距离计算
  ▼
返回 top-k 结果

与 DiskANN 的对比

维度 DiskANN SPANN
索引结构 单层 Vamana 图 聚类 + 倒排列表
SSD 访问模式 多次 4KB 随机读 少量大块顺序读
适合的数据特征 通用 数据有明显聚类结构时更优
更新难度 较复杂 较简单(追加倒排列表)
查询延迟 ~1 ms ~2 ms(取决于倒排列表大小)

SPANN 的优势在于 SSD 访问模式:它执行少量大块顺序读(每个倒排列表通常几十 KB 到几百 KB),而 DiskANN 执行多次小块随机读。在某些 SSD 型号上(尤其是 SATA SSD),顺序读的吞吐量优势更明显。

倒排列表的重叠分配

SPANN 的一个关键设计是允许数据点被分配到多个聚类的倒排列表中。对于位于聚类边界附近的点,它们会被同时分配到 2~3 个最近的聚类。这种冗余存储增加了约 30% 的磁盘空间,但显著提高了召回率,因为边界点不会因为聚类分配错误而被遗漏。

def assign_to_clusters(point, centroids, max_assign=3, threshold=1.2):
    """将点分配到多个聚类"""
    dists = [(i, l2_dist(point, c)) for i, c in enumerate(centroids)]
    dists.sort(key=lambda x: x[1])

    assignments = [dists[0]]  # 最近的聚类一定分配
    min_dist = dists[0][1]

    for i in range(1, len(dists)):
        if len(assignments) >= max_assign:
            break
        if dists[i][1] <= threshold * min_dist:
            assignments.append(dists[i])
        else:
            break

    return [a[0] for a in assignments]

九、基准测试:ScaNN vs DiskANN vs Faiss

为了公平比较这三个系统,我们使用标准的 ANN 基准测试数据集和评估框架。以下数据综合了多篇论文和 ann-benchmarks.com 的结果。

数据集

数据集 维度 数据量 类型
SIFT-1M 128 1M SIFT 特征
DEEP-1B 96 1B 深度学习嵌入
SPACEV-1B 100 1B Bing 搜索嵌入
BigANN-1B 128 1B SIFT 特征

SIFT-1M (百万级) 结果

在百万级数据上,三个系统的差异不大,因为数据完全可以放入内存:

系统 =0.95 时 QPS =0.99 时 QPS 内存 (MB)
ScaNN 42,000 18,000 180
Faiss-IVF-PQ 35,000 12,000 200
Faiss-HNSW 50,000 28,000 520
DiskANN (in-mem) 45,000 22,000 580

DEEP-1B (十亿级) 结果

十亿级数据是真正的分水岭:

系统 QPS 内存 (GB) 磁盘 (GB)
ScaNN (IVF-AVQ) 0.92 2,500 55 0
Faiss-IVF-PQ (OPQ) 0.88 1,800 60 0
DiskANN 0.95 1,200 30 480
DiskANN (高精度) 0.98 500 30 480
SPANN 0.93 800 25 520

关键观察:

  1. ScaNN 的 QPS 最高:得益于 AVQ 的优越量化质量,ScaNN 在相同召回率下需要扫描更少的候选,因此 QPS 更高。但它需要约 55 GB 内存。

  2. DiskANN 的召回率最高:因为 DiskANN 在最终阶段使用全精度向量做 rerank,召回率的上限更高。在需要高精度的场景下(如 recall > 0.95),DiskANN 的优势更加明显。

  3. Faiss 是最灵活的基线:Faiss 提供了最丰富的配置选项,可以在不同的精度/速度/内存权衡点上工作。但在极端场景下,它不如专门优化的 ScaNN 和 DiskANN。

延迟分布

对于延迟敏感的在线服务,P99 延迟比平均延迟更重要:

系统 P50 (ms) P95 (ms) P99 (ms)
ScaNN 0.4 0.8 1.5
Faiss-IVF-PQ 0.5 1.2 2.5
DiskANN 1.0 2.5 5.0
SPANN 1.5 3.0 8.0

DiskANN 的 P99 延迟较高,主要来自 SSD I/O 的尾延迟。在极端情况下,SSD 可能因为垃圾回收或热管理导致单次读取延迟飙升到毫秒级别。

十、成本分析:内存与 SSD 的经济学

在选择向量检索方案时,硬件成本是一个不容忽视的因素。以下分析基于 2025 年的云服务器定价。

内存方案(ScaNN/Faiss-HNSW)

以 10 亿 128 维 float32 向量为例:

原始数据:      128 × 4B × 1B = 512 GB
ScaNN (AVQ):   M=32 子空间 × 1B = 32 GB PQ
               + IVF 倒排索引 ≈ 20 GB
               总计 ≈ 55 GB

HNSW:          512 GB 数据 + ~120 GB 图 = ~640 GB

云服务器内存定价(参考 AWS r6g 系列):

方案 实例类型 内存 月费用 (USD)
ScaNN r6g.4xlarge × 2 128 GB × 2 ~$1,600
HNSW r6g.metal 768 GB ~$4,800

SSD 方案(DiskANN)

内存:  PQ 压缩向量 32 GB + 元数据 ≈ 35 GB
SSD:   512 GB 数据 + 图邻接表 ≈ 560 GB
方案 实例类型 内存 / SSD 月费用 (USD)
DiskANN i3en.xlarge × 2 32 GB + 2.5 TB NVMe × 2 ~$700

总拥有成本对比

方案 月费用 QPS/$ 适用场景
HNSW $4,800 ~10 极低延迟需求
ScaNN $1,600 ~30 高 QPS,可接受适度内存
DiskANN $700 ~20 成本敏感,可接受稍高延迟

混合方案

在实际生产中,许多团队采用混合方案:

这种分层架构可以在成本和性能之间取得最佳平衡。

十一、生产实践:从论文到工程

实际应用场景

ScaNN 和 DiskANN 都已在大规模生产环境中得到验证:

Google Search:ScaNN 是 Google 内部向量检索的核心组件。Google 搜索在多个环节使用向量检索——查询理解、文档匹配、广告排序等。ScaNN 的各向异性量化使得 Google 可以在有限的内存预算下服务数百亿级别的向量索引。Google 还将 ScaNN 开源为 TensorFlow 生态的一部分(google-research/scann)。

Bing Search:DiskANN 最初就是为 Bing 搜索开发的。Bing 的网页索引包含数百亿文档,每个文档有多个向量表示。DiskANN 使得 Bing 可以用 SSD 服务这些向量的近邻搜索,大幅降低了硬件成本。

Azure Cognitive Search / Azure AI Search:DiskANN 是 Azure 向量搜索服务的底层引擎。用户通过 REST API 上传向量,Azure 在后端使用 DiskANN 构建和查询索引。这是 DiskANN 最直接的商业化应用。

工程踩坑指南

以下是在生产环境中部署向量检索系统时常见的陷阱和对策:

问题 表现 原因 对策
SSD 尾延迟飙升 P99 延迟从 5ms 突增到 50ms SSD 垃圾回收(GC)阻塞读取 预留 20% SSD 空间给 GC;使用 over-provisioned SSD
内存 OOM 索引构建到 80% 时被 OOM Killer 杀死 图构建过程中邻居列表动态扩展 预分配固定大小邻接表;分批构建
召回率随数据增长下降 新插入数据后 recall 从 0.95 降到 0.88 图结构退化,长边不足 定期重建索引;使用 FreshDiskANN 的合并策略
PQ 码本失效 数据分布漂移后检索质量下降 新数据的分布与训练码本时不一致 监控量化误差;定期重训练码本
多线程竞争 QPS 随线程数增加反而下降 共享数据结构的缓存行争用 每线程独立的搜索上下文;NUMA-aware 内存分配
查询热点 少数向量被频繁命中,SSD 读放大 数据分布不均匀,热点向量在图中度数高 将热点向量缓存到内存;LRU 缓存 SSD 页面
磁盘对齐问题 读取速度比预期慢 2~3 倍 节点数据跨越 4KB 页面边界 确保节点数据 4KB 对齐;使用 O_DIRECT 读取
构建时间过长 10 亿向量构建需要 3 天 两遍图构建 + 随机 I/O 使用分片并行构建;第二遍可以只处理高度数节点
向量维度过高 256 维以上时性能严重下降 PQ 子空间过多,距离估计精度不足 先做降维(PCA/随机投影到 128 维以内)再量化
距离函数不匹配 使用余弦相似度但索引用 L2 构建 忘记将向量归一化 对于余弦相似度,预处理时将所有向量 L2 归一化

监控指标

在生产环境中,以下指标是必须监控的:

# 核心性能指标
vector_search_latency_p50_ms
vector_search_latency_p99_ms
vector_search_qps
vector_search_recall_at_10  # 需要定期用 ground truth 评估

# 资源指标
ssd_read_iops
ssd_read_latency_p99_us
memory_usage_bytes
pq_quantization_error_mean

# 数据质量指标
index_freshness_lag_seconds  # 最新插入到可查询的延迟
index_size_vectors           # 索引中的向量总数
deleted_ratio                # 已删除但未清理的向量比例

十二、个人观点与总结

我的看法

在向量检索这个领域摸爬滚打了几年之后,我形成了一些不太「政治正确」的观点。

ScaNN 的真正价值不在各向异性,而在 Google 的工程能力。 AVQ 的数学很漂亮,但实际效果提升在很多数据集上只有 10%~20%。真正让 ScaNN 在 Google 内部大放异彩的,是 Google 对 SIMD 指令的极致优化、对内存分配的精细控制、以及对 TCMalloc 的深度定制。开源版本的 ScaNN 虽然也不错,但和 Google 内部版本的性能差距可能超过 2 倍。

DiskANN 是一个被低估的系统。 很多人看到 DiskANN 的延迟比 HNSW 高就不感兴趣了。但在十亿级别的真实场景中,DiskANN 的成本优势是压倒性的。一台配备 NVMe SSD 的中等服务器就可以服务 HNSW 需要整个机架才能搞定的数据量。对于大多数互联网公司来说,成本才是第一约束。

Vamana 的理论优雅性常被忽视。 与 HNSW 的多层结构相比,Vamana 的单层图加鲁棒剪枝在数学上更简洁。α 参数精确地控制了局部连接和全局连接的平衡,而 HNSW 的层数分布依赖于概率跳表,理论分析上不如 Vamana 严谨。在学术研究中,Vamana 应该获得更多关注。

混合方案才是未来。 纯内存方案太贵,纯磁盘方案太慢。未来的趋势是智能分层——根据数据的访问频率和重要性,自动将热数据放在内存、温数据放在 SSD、冷数据放在云存储。CXL(Compute Express Link)等新硬件技术的出现,可能会进一步模糊内存和磁盘的界限。

不要忽视简单方案的力量。 在很多实际场景中,数据量只有几百万,一个优化得当的 Faiss Flat Index 加上 GPU 加速,就可以达到亚毫秒延迟和完美召回率。花大力气部署 ScaNN 或 DiskANN 只有在数据量确实达到十亿级别时才有必要。过度工程是向量检索领域的常见病。

技术选型建议

数据量 < 100 万:
  → Faiss Flat Index (暴力搜索 + GPU)
  → 不需要近似,精确结果

数据量 100 万 ~ 1 亿:
  → HNSW (Faiss 或 HNSWlib)
  → 内存充足时优先选择
  → ScaNN 如果需要更高 QPS

数据量 1 亿 ~ 10 亿:
  → ScaNN: 内存预算 50~100 GB
  → DiskANN: 内存预算 < 50 GB,有 NVMe SSD
  → SPANN: 数据有明显聚类结构

数据量 > 10 亿:
  → DiskANN + 分片
  → SPANN + 分片
  → 内存方案的成本通常不可接受

展望

向量检索技术仍在快速发展。几个值得关注的方向:

  1. 硬件协同设计:CXL 内存扩展、计算型 SSD(Computational Storage)、专用向量检索加速器等硬件创新,可能根本性地改变内存与磁盘的权衡。

  2. 学习索引:用机器学习模型替代传统的索引结构(如 learned IVF、learned graph),根据数据分布自适应地优化索引参数。

  3. 多模态检索:随着多模态大模型的发展,向量检索需要同时处理文本、图像、音频等不同模态的向量,对索引结构提出了新的挑战。

  4. 流式场景的深度优化:实时插入、删除、更新向量的场景越来越普遍,FreshDiskANN 只是开始,未来需要更成熟的流式索引方案。

ScaNN 和 DiskANN 分别代表了「更好的压缩」和「更好的存储层次利用」两条技术路线。在可预见的未来,这两条路线将继续并行发展,并在不同的应用场景中各有胜负。理解它们的核心思想和工程权衡,对于做出正确的技术选型至关重要。


参考文献

  1. Guo, R., et al. “Accelerating Large-Scale Inference with Anisotropic Vector Quantization.” ICML 2020.
  2. Subramanya, S. J., et al. “DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node.” NeurIPS 2019.
  3. Singh, A., et al. “FreshDiskANN: A Fast and Accurate Graph-Based ANN Index for Streaming Similarity Search.” arXiv 2021.
  4. Chen, Q., et al. “SPANN: Highly-efficient Billion-scale Approximate Nearest Neighbor Search.” NeurIPS 2021.
  5. Jayaram Subramanya, S., et al. “Rand-NSG: Fast Accurate Billion Point Nearest Neighbor Search on a Single Node.” 2018.
  6. Jegou, H., Douze, M., Schmid, C. “Product Quantization for Nearest Neighbor Search.” IEEE TPAMI 2011.
  7. Malkov, Y. A., Yashunin, D. A. “Efficient and Robust Approximate Nearest Neighbor Using Hierarchical Navigable Small World Graphs.” IEEE TPAMI 2020.

上一篇: 乘积量化与 IVF-PQ 下一篇: 从零实现向量搜索引擎 相关阅读: - HNSW - B+tree 与 LSM-tree


By .