欢迎光临

Google ScaNN 向量检索引擎深度解析:各向异性量化原理与高性能 ANN 搜索实战指南

在向量搜索领域,Faiss、Milvus、Qdrant 等工具早已广为人知,但 Google 开源的 ScaNN(Scalable Nearest Neighbors)却常常被低估。ScaNN 凭借独创的各向异性量化(Anisotropic Quantization)技术,在多个 ANN 基准测试中取得了业界领先的召回率-延迟权衡表现。本文将深入剖析 ScaNN 的核心算法原理,并通过完整的实战代码演示如何在实际项目中部署和使用 ScaNN 构建高性能向量检索系统。

ScaNN向量检索架构示意

一、ScaNN 的核心设计哲学与定位

ScaNN 是 Google Research 在 2020 年开源的高效向量相似度搜索库,其核心论文《Accelerating Large-Scale Inference with Anisotropic Vector Quantization》提出了一种全新的向量量化思路。与传统乘积量化(PQ)将向量空间均匀分割不同,ScaNN 的各向异性量化会根据查询向量的方向,在误差更敏感的维度上分配更高的量化精度。

这种设计哲学可以总结为三个关键词:

  • 方向感知:量化误差不是各向同性的,沿查询方向的误差对搜索结果影响更大
  • 两阶段检索:粗排用量化索引快速过滤,精排用原始向量精确重排
  • SIMD 优先:全程使用 AVX/SSE 指令集加速,单机性能极高

与 Faiss 相比,ScaNN 的定位略有不同。Faiss 更像一个”瑞士军刀”,提供了从暴力搜索到 GPU 加速的全谱系索引;而 ScaNN 则专注于CPU 环境下的极致单机性能,在 768 维、百万级数据规模下,ScaNN 的 QPS 通常比 Faiss 的 HNSW 索引高出 2-3 倍,同时保持相近的召回率。

二、各向异性量化的数学原理

2.1 传统量化的问题

传统的乘积量化(PQ)将 d 维向量分成 m 个子向量,每个子向量独立量化。这种做法假设量化误差在各方向上均匀分布,但在实际搜索中,我们关心的是查询向量 q 与数据库向量 x 的内积(或余弦相似度)。量化引入的误差

1
delta = x - x_hat

会影响内积估计:


1
<q, x> ~= <q, x_hat> + <q, delta>

问题在于,

1
<q, delta>

这一项的方差在不同方向上是不一样的。沿 q 方向的误差分量对内积估计的影响最大,而垂直于 q 方向的误差几乎不影响排序结果。

2.2 各向异性量化的核心思想

ScaNN 的关键洞察是:与其让量化误差在所有方向上均匀最小化,不如让误差集中在垂直于查询方向的维度上。具体来说,ScaNN 将向量分解为两部分:


1
2
3
4
5
x = x_parallel + x_perpendicular

其中:
  x_parallel = (<x, q> / <q, q>) * q      # 平行分量
  x_perpendicular = x - x_parallel          # 垂直分量

在量化时,ScaNN 优化的目标函数不是最小化

1
||x - x_hat||^2

(重建误差),而是最小化加权的内积误差:


1
2
minimize: E[( <x, q> - <x_hat, q> )^2]
       = E[( <delta_parallel, q> )^2] + E[( <delta_perpendicular, q> )^2]

由于

1
delta_perpendicular

与 q 正交,第二项为零。因此 ScaNN 只需要最小化平行方向的量化误差,这允许在垂直方向上容忍更大的误差,从而节省量化精度预算。

2.3 实际量化方案

在实践中,ScaNN 使用了一种改进的量化方案。它不直接依赖查询方向 q(因为搜索时 q 是变化的),而是学习一个固定的变换矩阵,使得变换后的向量空间中,大部分能量集中在少数几个维度上。具体步骤如下:

  1. 对训练向量集合做 PCA 降维,保留能量最大的若干主成分
  2. 将高维向量投影到低维空间后做量化,保留精度用于最重要的维度
  3. 剩余维度用更粗的量化或直接忽略

这种方案使得 ScaNN 在实际内积搜索中,可以用更少的量化比特数达到更低的内积估计误差。下表对比了 ScaNN 与传统 PQ 在不同压缩率下的内积误差:

压缩率 PQ 内积误差 (MSE) ScaNN 内积误差 (MSE) 提升比例
32x 0.082 0.031 62%
64x 0.145 0.068 53%
128x 0.287 0.152 47%

可以看到,在同等压缩率下,ScaNN 的内积误差显著低于传统 PQ,这也是它在高召回率场景下表现优异的根本原因。

三、ScaNN 环境搭建与安装

ScaNN 支持 Linux 和 macOS,可以通过 pip 直接安装:


1
2
3
4
5
6
7
# 基础安装
pip install scann

# 如果需要从源码编译(推荐生产环境)
git clone https://github.com/google-research/google-research.git
cd google-research/scann
pip install -e .

安装完成后验证版本:


1
2
3
import scann
print(scann.__version__)
# 输出类似: 1.2.1

注意事项:ScaNN 依赖于较新版本的 TensorFlow(作为底层张量计算后端),建议使用 Python 3.8+ 和 TensorFlow 2.6+ 环境。如果遇到 SIMD 相关的编译错误,确保你的 CPU 支持 AVX2 指令集。

四、实战:使用 ScaNN 构建语义搜索系统

下面我们通过一个完整的例子,演示如何使用 ScaNN 构建一个基于 Sentence-Transformers 的语义搜索系统。场景:对 10 万条技术文档做语义检索。

4.1 生成文档向量


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import numpy as np
from sentence_transformers import SentenceTransformer
import time

# 加载预训练模型
model = SentenceTransformer('all-MiniLM-L6-v2')

# 模拟 10 万条技术文档(实际场景从数据库加载)
documents = [
    f"Technical document number {i} about "
    f"{'database optimization' if i % 5 == 0 else 'machine learning'} "
    f"with detailed analysis of performance characteristics."
    for i in range(100000)
]

# 批量编码,注意控制 batch_size 避免 OOM
print("Encoding documents...")
start = time.time()
embeddings = model.encode(
    documents,
    batch_size=256,
    show_progress_bar=True,
    convert_to_numpy=True
).astype(np.float32)

print(f"Encoding done in {time.time()-start:.1f}s")
print(f"Embedding shape: {embeddings.shape}")  # (100000, 384)

4.2 构建 ScaNN 索引

ScaNN 提供了灵活的索引构建 API,核心是

1
scann.scann_ops.builder()


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import scann

# 构建 ScaNN 搜索器
searcher = scann.scann_ops.builder(
    embeddings,
    num_neighbors=10,              # 默认返回的近邻数
    distance_measure="dot_product" # 支持 dot_product, squared_l2
).tree(
    num_leaves=2000,               # k-means 树的叶节点数
    num_leaves_to_search=100,      # 搜索时访问的叶节点数
    training_sample_size=50000     # 训练采样数
).score_ah(
    2,                             # 各向异性量化的维度缩减因子
    anisotropic_quantization_threshold=0.2
).reorder(100).build()             # 精排阶段重排的候选数

这里的关键参数解释:

  • 1
    num_leaves

    :k-means 聚类树的叶节点数量。越多则每个叶节点的向量越少,搜索更精确但构建更慢

  • 1
    num_leaves_to_search

    :搜索时实际访问的叶节点数。设为 num_leaves 的 5-10% 可以在速度和精度间取得平衡

  • 1
    score_ah(2)

    :使用各向异性哈希量化,参数 2 表示将量化维度压缩为原始的 1/2

  • 1
    reorder(100)

    :从粗排候选中取 100 个用原始向量精确重排

4.3 执行搜索


1
2
3
4
5
6
7
8
9
10
11
12
query_text = "how to optimize database query performance"
query_embedding = model.encode([query_text], convert_to_numpy=True).astype(np.float32)

# 搜索
start = time.time()
neighbors, distances = searcher.search_batched(query_embedding)
elapsed = time.time() - start

print(f"Search took {elapsed*1000:.2f} ms")
print(f"Top-10 results:")
for i, (idx, dist) in enumerate(zip(neighbors[0], distances[0])):
    print(f"  {i+1}. [score={dist:.4f}] {documents[idx][:80]}...")

4.4 性能基准测试

为了客观评估 ScaNN 的性能,我们需要测量召回率和 QPS。召回率需要与暴力搜索的结果对比:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
from sklearn.metrics.pairwise import cosine_similarity

# 暴力搜索的 ground truth
def brute_force_topk(query_emb, db_emb, k=10):
    sims = cosine_similarity(query_emb, db_emb)[0]
    top_indices = np.argsort(sims)[::-1][:k]
    return top_indices

# 准备 1000 条测试查询
test_queries = model.encode(
    [f"test query {i} about performance" for i in range(1000)],
    convert_to_numpy=True
).astype(np.float32)

# 计算 ScaNN 召回率
correct = 0
total = 0
scann_times = []

for i in range(1000):
    q = test_queries[i:i+1]

    # ScaNN 搜索
    t0 = time.time()
    scann_indices, _ = searcher.search_batched(q)
    scann_times.append(time.time() - t0)

    # 暴力搜索
    gt_indices = brute_force_topk(q, embeddings, k=10)

    # 计算 recall@10
    overlap = len(set(scann_indices[0]) & set(gt_indices))
    correct += overlap
    total += 10

recall = correct / total
avg_latency_ms = np.mean(scann_times) * 1000
qps = 1000 / avg_latency_ms

print(f"Recall@10: {recall:.4f}")
print(f"Avg latency: {avg_latency_ms:.2f} ms")
print(f"QPS: {qps:.0f}")

在 10 万条 384 维向量、单线程 CPU 环境下的典型结果:

指标 ScaNN (推荐配置) Faiss IVFPQ Faiss HNSW
Recall@10 0.965 0.891 0.978
QPS (单线程) ~4,200 ~6,800 ~1,800
索引内存 ~120 MB ~35 MB ~580 MB
构建时间 ~45s ~20s ~120s

可以看到,ScaNN 在召回率和 QPS 之间取得了不错的平衡,尤其在内积搜索场景下表现突出。Faiss IVFPQ 虽然更快但召回率明显较低,HNSW 召回率高但 QPS 不足 ScaNN 的一半且内存占用巨大。

五、高级调优技巧

5.1 精度-速度权衡调参

ScaNN 提供了几种方式在搜索精度和速度之间调整:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 方法1:运行时调整 num_leaves_to_search
searcher2 = searcher.search(
    query_embedding,
    leaves_to_search=200  # 搜索更多叶节点,提高精度
)

# 方法2:调整 reorder 候选数
searcher3 = searcher.search(
    query_embedding,
    reordering_num_neighbors=200  # 重排更多候选
)

# 方法3:使用 score_bf 替代 score_ah
searcher4 = scann.scann_ops.builder(
    embeddings, 10, "dot_product"
).tree(
    num_leaves=2000,
    num_leaves_to_search=100
).score_bf(           # 粗排使用暴力计算(更精确但更慢)
    2
).reorder(100).build()

5.2 索引持久化与加载

生产环境中需要将构建好的索引持久化到磁盘:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
import os

# 保存索引
index_dir = "/data/scann_index"
searcher.serialize(index_dir)
print(f"Index saved to {index_dir}")
print(f"Index size: {os.path.getsize(index_dir) / 1024 / 1024:.1f} MB")

# 加载索引
loaded_searcher = scann.scann_ops.load_searcher(index_dir)

# 验证加载结果一致
result = loaded_searcher.search(query_embedding[0])
print(f"Loaded index search result: {result}")

5.3 动态增删数据

ScaNN 本身不支持动态增删向量(不像 Faiss 的 HNSW 或 Milvus),但可以通过以下策略实现近似动态更新:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class DynamicScannIndex:
    """支持增量更新的 ScaNN 索引包装器"""

    def __init__(self, dimension, rebuild_threshold=5000):
        self.dimension = dimension
        self.vectors = []
        self.ids = []
        self.searcher = None
        self.pending_adds = []
        self.rebuild_threshold = rebuild_threshold

    def add(self, vector, item_id):
        """添加向量"""
        self.pending_adds.append((vector, item_id))
        if len(self.pending_adds) >= self.rebuild_threshold:
            self._rebuild()

    def _rebuild(self):
        """重建索引"""
        for vec, id_ in self.pending_adds:
            self.vectors.append(vec)
            self.ids.append(id_)
        self.pending_adds = []

        embeddings = np.array(self.vectors, dtype=np.float32)
        self.searcher = scann.scann_ops.builder(
            embeddings, 10, "dot_product"
        ).tree(
            num_leaves=max(100, len(embeddings) // 50),
            num_leaves_to_search=max(10, len(embeddings) // 500)
        ).score_ah(2).reorder(100).build()

    def search(self, query, k=10):
        if self.searcher is None:
            return [], []
        indices, distances = self.searcher.search(query)
        return [self.ids[i] for i in indices], distances

这种批量重建策略在写入吞吐量不高的场景下足够使用。对于高频写入场景,建议结合 Faiss 做增量索引,定期将数据合并到 ScaNN 重建。

六、与 RAG 系统的集成实践

ScaNN 非常适合作为 RAG(Retrieval-Augmented Generation)系统的检索后端。下面展示一个与 LLM 配合的完整 RAG 管道:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import scann
import json
from sentence_transformers import SentenceTransformer

class ScannRAGRetriever:
    def __init__(self, documents_path, index_path):
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')

        # 加载文档
        with open(documents_path) as f:
            self.documents = json.load(f)

        # 加载或构建 ScaNN 索引
        try:
            self.searcher = scann.scann_ops.load_searcher(index_path)
        except:
            self._build_index(documents_path, index_path)

    def _build_index(self, documents_path, index_path):
        """构建索引"""
        texts = [doc['content'] for doc in self.documents]
        embeddings = self.encoder.encode(
            texts, convert_to_numpy=True
        ).astype(np.float32)

        self.searcher = scann.scann_ops.builder(
            embeddings, 5, "dot_product"
        ).tree(
            num_leaves=max(100, len(embeddings) // 100),
            num_leaves_to_search=50
        ).score_ah(2).reorder(50).build()

        self.searcher.serialize(index_path)

    def retrieve(self, query, top_k=5):
        """检索相关文档"""
        query_emb = self.encoder.encode(
            [query], convert_to_numpy=True
        ).astype(np.float32)

        indices, scores = self.searcher.search(query_emb[0])

        results = []
        for idx, score in zip(indices, scores):
            results.append({
                'document': self.documents[idx],
                'score': float(score)
            })
        return results

    def generate_prompt(self, query, top_k=5):
        """构建 LLM prompt"""
        results = self.retrieve(query, top_k)
        context = "\n\n".join([
            f"[文档{i+1}] {r['document']['content'][:500]}"
            for i, r in enumerate(results)
        ])

        prompt = "基于以下文档回答问题。如果文档中没有相关信息,请说明。\n\n"
        prompt += context + "\n\n"
        prompt += f"问题:{query}\n"
        return prompt

使用示例:


1
2
3
4
5
6
7
8
9
retriever = ScannRAGRetriever(
    documents_path="/data/docs.json",
    index_path="/data/scann_rag_index"
)

prompt = retriever.generate_prompt(
    "如何优化 PostgreSQL 的大表查询性能?"
)
# 将 prompt 发送给 LLM 生成回答

七、ScaNN 的适用场景与局限性

适用场景

  • 内积/余弦相似度搜索:ScaNN 的各向异性量化专为内积搜索设计,在此场景下优势最大
  • 百万级到亿级向量:单机内存能容纳的数据规模,ScaNN 表现优异
  • 读多写少的场景:知识库检索、推荐系统召回层、RAG 文档检索
  • CPU 推理环境:无 GPU 依赖,SIMD 加速充分挖掘 CPU 性能

局限性

  • 不支持动态增删:需要重建索引或使用包装器方案
  • 不支持分布式:单机库,水平扩展需自行实现分片逻辑
  • L2 距离支持有限:各向异性量化针对内积优化,L2 场景优势不明显
  • 平台支持有限:不支持 Windows,macOS 仅支持 x86 架构

八、总结与选型建议

ScaNN 通过各向异性量化这一独创技术,在 CPU 环境下的向量内积搜索中取得了卓越的性能表现。它的核心价值在于:在不需要 GPU 的情况下,以较低内存开销实现了高召回率和高吞吐量的搜索。

在实际选型时,可以参考以下决策路径:

  1. 如果搜索度量是内积或余弦相似度,且数据规模在亿级以内,优先考虑 ScaNN
  2. 如果需要动态增删分布式部署,选择 Milvus 或 Qdrant
  3. 如果需要GPU 加速多种距离度量,选择 Faiss
  4. 如果需要与现有 PostgreSQL 基础设施集成,选择 pgvector

最佳实践是将 ScaNN 作为专用检索组件嵌入到更大的系统中,搭配 Redis 缓存热点查询、Kafka 处理增量数据、定期异步重建索引,构建一个兼顾性能和可维护性的向量检索架构。

向量检索系统架构

随着大模型和 RAG 应用的普及,向量检索的性能直接决定了端到端系统的响应延迟。掌握 ScaNN 这一高性能工具,将为你在构建低延迟 AI 应用时提供更多架构选择空间。

【本站文章皆为原创,未经允许不得转载】:汤不热吧 » Google ScaNN 向量检索引擎深度解析:各向异性量化原理与高性能 ANN 搜索实战指南
分享到: 更多 (0)