近似最近邻向量索引深度工程:从 HNSW 到 DiskANN 的全链路实现

向量相似度搜索是现代 AI 基础设施的核心组件——从 RAG(检索增强生成)到推荐系统、图像检索、异常检测,几乎所有"理解语义"的系统都依赖它。当数据规模从千级跃升到亿级时,暴力计算的 O(N) 复杂度变得不可接受,近似最近邻(ANN, Approximate Nearest Neighbor)索引成为必选项。

本文将深入 ANN 索引的工程实现层,覆盖三大主流方案的核心原理:基于图的 HNSW、基于量化的 IVF-PQ,以及面向超大规模磁盘存储的 DiskANN,并给出完整的工程实践代码与性能对比。


一、为什么暴力搜索不够:复杂度视角

在一个 768 维的 dense embedding 空间里搜索 10 亿条向量,单次查询需要:

  • 逐向量计算余弦相似度(768 次乘加)
  • 总计算量约 768 GFLOPS
  • 对比:单张 A100 FP16 峰值 312 TFLOPS,但仍需 2.5ms 纯计算

这还没算内存带宽——10 亿 float32 向量占 3TB 内存,即便在分布式环境下,跨节点的数据传输也会使延迟飙升到秒级。

ANN 的核心优化思路:用可控的精度损失换取量级搜索加速。典型目标是在 95%+ 召回率(Recall@10)下,将延迟从秒级压到毫秒级。


二、HNSW:可导航小世界图

2.1 理论基础

NSW(Navigable Small World) 是一种基于图的索引结构,灵感来自社会网络中的"六度分隔"理论——任何两个节点之间只需少量跳转即可到达。

在 d 维欧氏空间中构建的 NSW 图有以下关键性质:

  1. 长程连接与短程连接:每个节点既有邻近连接的邻居,也有跨越长距离的"高速公路"
  2. 贪心路由的正确性:从任意起始节点出发,每一步移动到更近于查询目标的邻居,能在 O(log N) 步内收敛到查询点的近似最近邻
  3. 对数搜索路径:理论上搜索路径长度约为 O(log N),远低于扫描的 O(N)

HNSW(Hierarchical NSW) 在 NSW 之上引入层级结构:

  • 第 0 层包含所有向量
  • 第 i 层仅包含第 i-1 层的部分节点(按固定概率衰减)
  • 高层级的节点充当"快速路",实现跨图跳转

层级结构示意:
Layer 2:    A ----------------------> D
Layer 1:    A ------> B ------> D ------> F
Layer 0:    A -> C -> B -> E -> D -> G -> F -> H

2.2 索引构建算法


import numpy as np
from typing import List, Tuple, Dict
import heapq

class HNSWIndex:
    """HNSW 索引的简化实现"""
    
    def __init__(self, dim: int, m: int = 16, ef_construction: int = 200, m_l: float = 1.0):
        """
        Args:
            dim: 向量维度
            m: 每个节点的最大连接数(控制图的连通度)
            ef_construction: 构建时候选集大小(越大质量越好,越慢)
            m_l: 层级衰减因子 = 1 / log(m)
        """
        self.dim = dim
        self.m = m
        self.m_max = m * 2  # 第 0 层允许更多连接
        self.ef_construction = ef_construction
        self.m_l = m_l
        self.max_level = 0
        self.entry_point = None
        
        # graph[level][node_id] = set of neighbor_ids
        self.graph: List[Dict[int, set]] = []
        self.vectors: np.ndarray = None
        self.id_to_level: Dict[int, int] = {}
    
    def _distance(self, a: np.ndarray, b: np.ndarray) -> float:
        """余弦相似度(内积,归一化后等价于余弦)"""
        return np.dot(a, b)
    
    def _random_level(self) -> int:
        """按概率衰减决定节点最高层级"""
        level = 0
        while np.random.rand() < np.exp(-1.0 / self.m_l) and level < 16:
            level += 1
        return level
    
    def _search_layer(self, query: np.ndarray, entry_points: List[int], 
                      ef: int, level: int) -> List[Tuple[float, int]]:
        """在单层中搜索(贪心 + 候选扩展)"""
        # 已访问集合
        visited = set(entry_points)
        # 候选集(最小堆,最差在顶)
        candidates = []
        # 结果集(最大堆,最优在顶),存储 (相似度, node_id)
        results = []
        
        for ep in entry_points:
            dist = self._distance(self.vectors[ep], query)
            heapq.heappush(candidates, (dist, ep))
            heapq.heappush(results, (-dist, ep))
        
        while candidates:
            curr_dist, curr_id = heapq.heappop(candidates)
            
            # 剪枝:如果当前候选的最差距离已比结果集中最差的结果远,停止
            if -results[0][0] > curr_dist:
                break
            
            # 扩展邻居
            for neighbor in self.graph[level].get(curr_id, set()):
                if neighbor not in visited:
                    visited.add(neighbor)
                    dist = self._distance(self.vectors[neighbor], query)
                    
                    if dist > -results[0][0] or len(results) < ef:
                        heapq.heappush(candidates, (dist, neighbor))
                        heapq.heappush(results, (-dist, neighbor))
                        
                        # 保持结果集不超过 ef
                        if len(results) > ef:
                            heapq.heappop(results)
        
        return sorted([(-d, n) for d, n in results], reverse=True)
    
    def insert(self, node_id: int, vector: np.ndarray):
        """插入新节点"""
        level = self._random_level()
        self.id_to_level[node_id] = level
        
        # 扩展图层数
        while len(self.graph) <= level:
            self.graph.append({})
        
        # 无输入时的初始化
        if self.entry_point is None:
            self.entry_point = node_id
            self.max_level = level
            return
        
        # 从 entry point 开始搜索
        curr_ep = [self.entry_point]
        curr_dist = self._distance(self.vectors[self.entry_point], vector)
        
        # 从最高层贪心下降到目标层的上一层
        for l in range(self.max_level, level, -1):
            results = self._search_layer(vector, curr_ep, 1, l)
            curr_ep = [results[0][1]]
            curr_dist = results[0][0]
        
        # 在 [level, 0] 各层寻找连接
        for l in range(min(level, self.max_level), -1, -1):
            neighbors = self._search_layer(vector, curr_ep, self.ef_construction, l)
            
            # 选择连接策略(简化:直接选 top m)
            selected = neighbors[:self.m if l == 0 else self.m // 2]
            self.graph[l][node_id] = {n for _, n in selected}
            
            # 双向连接
            for _, neighbor_id in selected:
                if neighbor_id not in self.graph[l]:
                    self.graph[l][neighbor_id] = set()
                self.graph[l][neighbor_id].add(node_id)
                
                # 修剪过度连接的邻居
                if len(self.graph[l][neighbor_id]) > self.m_max:
                    self._trim_neighbors(l, neighbor_id)
            
            curr_ep = [n for _, n in selected]
        
        # 更新全局 entry point
        if level > self.max_level:
            self.max_level = level
            self.entry_point = node_id
    
    def _trim_neighbors(self, level: int, node_id: int):
        """修剪邻居,保留最近的 m 个"""
        neighbors = self.graph[level][node_id]
        dists = [(self._distance(self.vectors[n], self.vectors[node_id]), n) for n in neighbors]
        dists.sort()
        self.graph[level][node_id] = {n for _, n in dists[:self.m]}
    
    def search(self, query: np.ndarray, k: int = 10, ef_search: int = 50) -> List[Tuple[int, float]]:
        """搜索 top-k"""
        curr_ep = [self.entry_point]
        
        # 从顶层贪心下降到第 0 层
        for l in range(self.max_level, 0, -1):
            results = self._search_layer(query, curr_ep, 1, l)
            curr_ep = [results[0][1]]
        
        # 在第 0 层用 ef_search 搜索
        results = self._search_layer(query, curr_ep, ef_search, 0)
        
        return [(node_id, dist) for dist, node_id in results[:k]]

2.3 关键调优参数

| 参数 | 取值建议 | 影响 |

|------|---------|------|

| M | 16~64 | 内存 = N × M × 4 bytes,越大搜索越慢但召回越高 |

| ef_construction | 100~400 | 构建质量,越大越慢越准 |

| ef_search | 50~200 | 查询时的候选集大小,与延迟正相关 |

经验分享:在生产实践中,M=32 + ef_construction=200 是较好的平衡点。ef_search 设为 M 的 3~5 倍可在速度和召回之间取得不错的平衡。


三、IVF-PQ:量化降维的工业级方案

3.1 算法原理

IVF-PQ(Inverted File with Product Quantization) 在 FAISS 中是最常用的 IndexIVFPQ,由两步组成:

  1. IVF(倒排聚类):
  • 将所有向量聚为 nlist 个簇(用 k-means)
  • 每个簇有一个 centroid
  • 查询时只扫描 nprobe 个最近的簇(而非全部数据)
  1. PQ(乘积量化):
  • 将 d 维向量拆分为 m 个子空间(每段 d/m 维)
  • 对每个子空间训练 256 个小码本(k-means, k=256)
  • 原始向量用 m 个 8-bit 码本索引替代,压缩 32 倍

原始: 32-bit float × 768维 = 3072 bytes
PQ后: 8-bit × m子段, m=96时 → 96 bytes (压缩32倍)

距离计算: Asymmetric Distance Computation (ADC)
  query 不归约化,用每个子空间的距离表快速查表求和

3.2 FAISS 实战代码


import faiss
import numpy as np

# 数据集:100万条 768 维向量(模拟 Sentence-BERT 输出)
d = 768
nb = 1_000_000
nq = 100

np.random.seed(42)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 归一化(余弦搜索需要)
faiss.normalize_L2(xb)
faiss.normalize_L2(xq

# ===================== 构建 IVF-PQ 索引 =====================
nlist = 4096        # 聚类数量
m = 48              # PQ 子空间数 (d 能被 m 整除)
nbits = 8           # 每个子空间的码本大小 = 256

# 训练索引(只需跑一次,后续可复用)
quantizer = faiss.IndexFlatIP(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)

# 在子集上训练(数据太大时用采样)
print("Training...")
index.train(xb[:min(100000, nb)])

# 添加向量
print("Adding vectors...")
index.add(xb)

# 设置查询时的探测簇数
index.nprobe = 32    # 扫描 32/4096 ≈ 0.8% 的数据

# 搜索
k = 10
D, I = index.search(xq, k)
print(f"I shape: {I.shape}, Recall@10: ... ")

# 内存占用: 4 bytes(quantizer ID) + 48 bytes(PQ codes) per vector ≈ 52 MB
# 对比 HNSW: 768 × 4 + 32 × 4 = 3200 bytes per vector ≈ 3.2 GB

3.3 精度与速度对比

| 索引 | 内存/百万向量 | 延迟(Recall@10≈0.95) | 适用场景 |

|------|-------------|---------------------|---------|

| FlatIP | 3 GB | ~50ms @1M | 基准/小数据 |

| IVF-PQ(nprobe=32) | 52 MB | ~2ms @1M | 内存敏感、可接受略低召回 |

| HNSW(M=32) | 3.2 GB | ~1.2ms @1M | 极致性能、内存充裕 |

工程实践建议:当数据量 > 1000 万且内存预算紧张时,IVF-PQ 是更务实的选择。对于子毫秒级延迟要求,HNSW 配合 GPU 加速是最佳方案。


四、DiskANN:打破内存墙的磁盘友好索引

4.1 问题背景

当向量规模达到 100 亿级(如 BGE-large 768 维 × 100 亿 = 6 TB),完全内存索引不再现实。传统做法是分片+多副本,但 DiskANN 提出了新思路:

核心洞察:图的搜索路径天然具有稀疏性——单条查询仅访问少量索引节点,理论上可以只把"热区"保留在内存。

4.2 DiskANN 架构设计


                    SSD (NVMe, 6TB)
┌──────────────────────────────────────────────┐
│  全部 PQ 压缩向量 (10B × 48B = 480 GB)      │
│  全局 HNSW 图索引 (原材料的 ~5% 保留在内存) │
│                                              │
│  ┌──────────┐                               │
│  │ PQ Cache │ ◄──── L2 Cache (100 GB)      │
│  │(Hot PQ)  │                               │
│  └──────────┘                               │
└──────────────────────────────────────────────┘
              ▲
              │ NVMe 随机读取 (~50us/页)
              ▼
        系统内存 (256 GB)
┌──────────────────────────────────────────────┐
│  HNSW 骨架图:top-2 层全部节点 (约 20 GB)    │
│  Vamana 剪枝后的边列表 (约 60 GB)            │
│  PQ 解码缓存 (最近查询访问的 PQ codes)       │
└──────────────────────────────────────────────┘

DiskANN 的关键技术:

  1. Vamana 图构建:不同于 HNSW 的双向插入,DiskANN 用全局剪枝的 Vamana 算法——按随机顺序扫描所有节点,对每个节点贪心搜索后取最近邻进行 α 剪枝
  2. PQ 向量与图分离存储:图骨架在内存(支持随机游走),PQ 压缩向量在 SSD(按需读取)
  3. Beam Search 批量预取:查询时沿图路径批量异步预取即将访问的 PQ 磁盘页,利用 NVMe 的并行 IOPS

4.3 与内存索引的工程权衡

| 维度 | HNSW (内存) | DiskANN |

|------|------------|---------|

| 数据规模 | ~2 亿 (受内存限制) | 100 亿+ |

| 单查询延迟 | 0.5~2 ms | 2~10 ms |

| 吞吐 (QPS) | 5000+ | 1000~3000 |

| 成本 (per 10亿) | $3000/月 (64GB x 20节点) | $500/月 (单 NVMe) |

| 一致性 | 强 | 近实时(延迟写入) |


五、生产实践:选型决策树


def choose_ann_index(
    num_vectors: int,
    dim: int,
    latency_budget_ms: float,
    memory_budget_gb: float,
    recall_target: float = 0.95,
    qps: int = 1000
) -> str:
    """
    ANN 索引选型决策函数
    """
    # 计算 HNSW 的内存需求
    hnsw_mem_per_vec = dim * 4 + 64 * 4  # float vec + edges
    total_hnsw_gb = num_vectors * hnsw_mem_per_vec / (1024**3)
    
    # 数据量小 → 直接暴力
    if num_vectors < 100_000:
        return "Flat(暴力搜索) — 您的数据量太小,ANN 不值得"
    
    # 内存足够 HNSW 且延迟要求严格
    if total_hnsw_gb < memory_budget_gb and latency_budget_ms < 5:
        return f"HNSW — 需要 {total_hnsw_gb:.1f}GB 内存,召回可达 {min(0.99, recall_target)}"
    
    # 中等数据量 → IVF-PQ 折中
    pq_codes = 48 * num_vectors / (1024**3)
    if pq_codes < memory_budget_gb * 0.5:
        return f"IVF-PQ — 仅需 {pq_codes:.1f}GB (压缩 {768*4/48:.0f}x),建议 nprobe=32"
    
    # 超大规模 → 分片或 DiskANN
    if num_vectors > 1_000_000_000:
        return "DiskANN 或 Milvus Distributed — 单节点磁盘方案,多节点可扩展"
    
    if num_vectors > 100_000_000:
        return "分片 HNSW 或 Milvus — 负载均衡 + 数据分区"
    
    return "IVF-PQ + 缓存层 — 根据实际情况调 nprobe"


# 实际选型示例
print(choose_ann_index(
    num_vectors=50_000_000,
    dim=768,
    latency_budget_ms=10,
    memory_budget_gb=128,
    recall_target=0.95
))
# 输出: HNSW — 需要 20.0GB 内存,召回可达 0.95

六、前沿趋势与实战建议

6.1 剪枝与合并优化

现代向量数据库(如 Milvus 2.4+)引入了 Delta WAL 和 Segment Compaction 机制——新写入先进入增量段,后台合并时重建 ANN 索引,避免运行时的写入放大。这本质上是 LSM-tree 思想在向量索引领域的成功迁移。

6.2 混合检索已成标配

单一 ANN 已无法满足复杂工程需求。ANN + BM25 + Reranker 的三段式架构逐渐成为 AI 搜索的标准范式:


用户查询 ──► Query Embedding + 关键词 Tokenize
              │                        │
              ▼                        ▼
         ANN 向量检索              BM25 关键词检索
              │                        │
              └────────┬───────────────┘
                       ▼
               RRF 融合 / 交叉编码器重排
                       ▼
                  Top-K 结果返回

6.3 GPU vs CPU

对于高吞吐(QPS > 5000)场景,GPU 加速的 ANN(如 FAISS-GPU、RAPIDS cuVS)可以通过 embedding batching 和 CUDA 并行距离计算,将单卡 QPS 推到 20K+。但 GPU 显存有限,超过 5000 万向量的场景仍需 CPU 或混合部署。

6.4 关键 Checklist

选型时务必确认以下要点:

  • 召回率 vs 延迟曲线:通过实际 benchmark 画出 Pareto Frontier,而非拍脑袋估参数
  • 写入吞吐量:HNSW 的插入速度较慢(需图遍历),频繁写入场景需考虑 RocksDB/PQ 分离
  • 向量归一化:搜索余弦相似度前务必 L2 normalize,否则内积等价不了余弦
  • 批量查询优化:ANN 的批量查询能隐藏内存延迟,实测 batch=16 是较好的折中点

七、总结

ANN 向量索引的选型从来不是"哪个算法最好",而是在召回率、延迟、内存、写入吞吐四维约束下找最优解。

| 场景 | 推荐方案 | 预期召回 | 典型延迟 |

|------|---------|---------|---------|

| < 100K 小数据 | Flat暴力 | 100% | <5ms |

| 100K~100M 纯内存 | HNSW (M=32) | 97~99% | 0.5~2ms |

| 内存敏感千万级 | IVF-PQ | 92~97% | 1~5ms |

| 100M~10B 超大规模 | DiskANN + 分片 | 90~95% | 2~10ms |

建议从简开始,先用 Flat 建立召回基准,再逐步引入 HNSW、量化压缩,最后考虑分布式和磁盘方案——过早优化是 ANN 工程最常见的陷阱。

向量数据库不只是算法问题,更是系统工程。掌握原理后,真正的价值在于结合具体业务场景做出合理取舍。


相关参考: Malkov & Yashunin (2016) "Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs"; Subramanya et al. (2019) "DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node"; Johnson et al. (2019) "Billion-scale similarity search with GPUs"

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部