近似最近邻向量索引深度工程:从 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 图有以下关键性质:
- 长程连接与短程连接:每个节点既有邻近连接的邻居,也有跨越长距离的"高速公路"
- 贪心路由的正确性:从任意起始节点出发,每一步移动到更近于查询目标的邻居,能在 O(log N) 步内收敛到查询点的近似最近邻
- 对数搜索路径:理论上搜索路径长度约为 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,由两步组成:
- IVF(倒排聚类):
- 将所有向量聚为
nlist个簇(用 k-means) - 每个簇有一个 centroid
- 查询时只扫描
nprobe个最近的簇(而非全部数据)
- 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 的关键技术:
- Vamana 图构建:不同于 HNSW 的双向插入,DiskANN 用全局剪枝的 Vamana 算法——按随机顺序扫描所有节点,对每个节点贪心搜索后取最近邻进行 α 剪枝
- PQ 向量与图分离存储:图骨架在内存(支持随机游走),PQ 压缩向量在 SSD(按需读取)
- 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"

发表评论 取消回复