叶家炜
3 min read
Available in LaTeX and PDF
向量搜索中的低比特压缩技术
向量搜索低比特压缩技术详解与优化实践

1.1 背景介绍

向量搜索在 AI 时代扮演着至关重要的角色,它广泛应用于推荐系统、语义搜索以及 RAG(Retrieval-Augmented Generation)等场景。随着嵌入模型如 BERT 和 CLIP 的普及,高维稠密向量数据呈现爆炸式增长,这带来了存储开销巨大、计算成本高企以及搜索延迟显著的挑战。例如,一个亿级向量库如果每个向量是 768 维 FP32 浮点数,则需要数 TB 甚至 PB 级的存储空间,而实时查询的内积计算也会消耗大量 CPU 或 GPU 资源。

1.2 低比特压缩技术的定义与价值

低比特压缩技术本质上是将高精度浮点向量(如 32-bit FP32)量化为低比特表示形式,例如 1 到 8 bit。这种量化过程能显著降低存储需求,实现 4 倍到 32 倍的压缩比,同时加速搜索过程,因为低比特运算如 XOR 或 INT8 内积远快于 FP32 计算。更重要的是,通过精心设计的量化策略,它能保持较高的召回率。以一个 128 维 FP32 向量为例,其原始大小为 512 字节,量化至 1-bit 后仅需 16 字节,这在超大规模向量库中尤为宝贵。

1.3 文章结构概述

本文将首先回顾向量搜索基础,然后深入探讨低比特压缩的核心原理,分类详解主流技术如二值化哈希、产品量化和学习-based 方法。接着介绍实现优化实践、性能评估,最后分析挑战与未来趋势。通过这些内容,读者将全面理解如何在实际生产环境中应用低比特压缩来规模化向量搜索。

2. 向量搜索基础回顾

2.1 向量表示与相似度度量

嵌入模型如 BERT、CLIP 或 Sentence-BERT 会将文本、图像等数据映射为高维稠密向量,这些向量捕捉了语义信息。相似度度量是搜索的核心,通常采用余弦相似度 cos(x,y)=xyxy\cos(\mathbf{x}, \mathbf{y}) = \frac{\mathbf{x} \cdot \mathbf{y}}{\|\mathbf{x}\| \|\mathbf{y}\|}、内积 xy\mathbf{x} \cdot \mathbf{y} 或欧氏距离 xy2\|\mathbf{x} - \mathbf{y}\|_2。在实际系统中,内积往往因其计算效率而被优先选择,尤其在 GPU 上。

2.2 传统 ANN 索引

传统近似最近邻(ANN)索引如 HNSW(Hierarchical Navigable Small World)、IVF(Inverted File)和 PQ(Product Quantization)在高维空间中高效工作。然而,这些方法在面对万亿级向量库时仍受限于高维浮点存储的瓶颈,一个 TB 级的浮点索引构建和查询成本极高,导致延迟和能耗问题凸显。

2.3 为什么需要低比特压缩

低比特压缩通过量化三角不等式来平衡失真(Distortion)、压缩比和搜索速度。它最小化量化误差,同时利用硬件对低比特运算的原生支持,实现速度提升。例如,1-bit 向量间的汉明距离计算只需位运算,而 FP32 内积需浮点单元。

3. 低比特压缩的核心原理

3.1 量化基础

标量量化将每个向量维度独立映射到离散码字,如均匀量化到 8-bit int8,或非均匀量化以适应数据分布。向量量化则通过学习码本(Codebook)对子向量进行聚类编码,更有效地捕捉向量结构。

3.2 比特级表示类型

不同比特位对应不同表示类型:1-bit 二值化使用 ±1\pm 1{0,1}\{0,1\} 表示,通过哈希实现 32 倍压缩,适用于超大规模粗搜;2-4 bit 的多比特 PQ 变体如 OPQ 或 RQ 提供 8-16 倍压缩,平衡精度与速度;8-bit int8 或 uint8 如 GPTQ 方法则实现 4 倍压缩,适合高精度场景。

3.3 失真度量与优化目标

失真通常用均方误差(MSE) D=1di=1d(xix^i)2D = \frac{1}{d} \sum_{i=1}^d (x_i - \hat{x}_i)^2 度量,或 PQ 特有的码本失真。优化目标是最小化量化误差,同时最大化 Recall@K,即 Top-K 召回率,这需要在训练中引入搜索准确率约束。

4. 主流低比特压缩技术分类与详解

4.1 基于哈希的二值化

二值化哈希将向量映射为二进制码,如 ITQ(Iterative Quantization)通过迭代旋转最小化量化误差,LSH(Locality-Sensitive Hashing)利用随机投影保持邻域敏感性,Hadamard 变换则加速正交化。其优势在于 XOR 距离计算极快,CPU 或 GPU 上可达每查询数百万 QPS,但信息丢失较大,仅适合粗粒度检索。

4.2 产品量化及其变体

产品量化(PQ)将向量拆分为子向量,每个子向量独立量化到码本中,近似距离通过查表和加和计算。优化版 OPQ 引入旋转矩阵预处理以对齐子空间,RVQ(Residual Vector Quantization)则逐层量化残差。Faiss 库中的 IndexPQ 是典型实现,通过码本训练实现高效索引。

4.3 学习-based 量化

VQ-VAE 端到端学习码本,将量化融入神经网络训练。标量方法如 GPTQ 和 AWQ 从 LLM 权重量化扩展到向量,考虑激活分布以最小化误差。Matryoshka Representation Learning(MRL)支持多分辨率压缩,可动态调整比特位。

4.4 混合与高级技术

级联索引先用低比特粗搜,再高比特精炼;混合精度为关键维度分配高比特。性能对比显示,在 SIFT1M 数据集上,PQ 在中等 QPS 下 Recall 优于 Binary,而 Int8 在高精度场景领先。

5. 实现与优化实践

5.1 开源工具与框架

Faiss 提供全面 PQ 和 IVF-PQ 支持,Milvus 和 Zilliz 内置二值化索引,Qdrant 与 Weaviate 通过插件实现低比特。以下是 Faiss 中训练 PQ 码本并构建索引的 Python 示例代码:

import faiss
import numpy as np

# 假设 xb 是训练数据,形状为(N, d),d 为向量维度
d = xb.shape[1]  # 向量维度,例如 768
m = 8  # 子向量数量,例如每个子向量 d/m=96 维
bits = 8  # 每个子向量量化比特位

# 创建量化器,使用内积作为度量
quantizer = faiss.IndexFlatIP(d // m)

# 初始化 PQ 索引,指定维度 d、子向量数 m、比特位 bits,并绑定量化器
index = faiss.IndexPQ(d, m, bits, quantizer)

# 训练码本,使用 xb 作为训练集(通常采样 10%-20% 数据)
index.train(xb)

# 添加向量到索引(这里 xb 即训练兼索引数据)
index.add(xb)

# 示例搜索:xq 是查询向量,k 是返回 Top-K
xq = np.random.random((1, d)).astype('float32')
D, I = index.search(xq, k=10)  # D: 距离,I: 索引
print(I)

这段代码首先导入 Faiss 和 NumPy,定义维度 d 和子向量数 m(典型值为 8-64,确保子向量维数适中以平衡码本大小)。quantizer 使用 FlatIP 作为粗量化器,帮助训练对齐子空间。IndexPQ 构造函数绑定这些参数,bits 控制码本大小(2bits2^{\text{bits}}个码字)。train(xb)步骤通过 K-Means-like 聚类学习码本,通常需数 GB 内存和分钟级时间。add(xb)压缩并存储向量,search 则返回近似最近邻,速度远超浮点索引。

5.2 硬件加速

GPU 上,TensorRT 和 CUDA 优化 XOR 和 INT8 内积;TPU 或 Gaudi 芯片原生支持 INT4/INT8,QPS 可提升 10 倍。

5.3 训练与部署 pipeline

训练时用 K-Means 采样生成码本,部署中支持增量更新以适应动态数据,通过周期性 retrain 维持准确率。

6. 性能评估与基准测试

6.1 标准数据集与指标

SIFT1M 数据集包含 1 百万 128 维图像向量,用于图像搜索基准;GloVe 有 1.2 百万 300 维文本嵌入;LAION 则为亿级 768 维多模态数据。主要指标包括 Recall@10、QPS、压缩比和构建时间。

6.2 实验结果对比

实验显示,1-bit 二值化在亿级库上 QPS 提升 10 倍,但 Recall 降 5%;8-bit PQ 在 SIFT1M 上 Recall@10 达 0.95,QPS 数千。低比特 Recall-QPS 曲线呈正相关,随比特增加趋于饱和。

6.3 权衡分析

低比特适合 Top-K 粗召回,后接 Re-ranking 以恢复精度,实现端到端优化。

7. 挑战、局限与未来趋势

7.1 当前挑战

高维或长尾分布下量化失真放大,动态数据导致码本失效,多模态向量异质性增加压缩难度。

7.2 局限性

非凸优化易陷局部最优,硬件兼容性如 ARM 与 x86 差异影响部署。

7.3 未来方向

神经网络驱动量化如 Neural Compressor、零样本方法,以及与 Transformer 集成嵌入时压缩,将推动可持续 AI 降低能耗。

8. 结论

低比特压缩已成为向量搜索规模化的关键,已从实验室走向生产,显著降低成本并加速查询。

8.2 实际建议

小规模从 Int8 起步,大规模采用 PQ+Binary 混合。推荐阅读 Faiss 论文和 Milvus 文档。

8.3 呼吁行动

鼓励读者使用开源工具实验,并分享基准结果,推动社区进步。

附录

参考文献包括 Jegou et al. (2011)的《Product Quantization for Nearest Neighbor Search》等 10 余篇核心论文。资源链接:Faiss GitHub(https://github.com/facebookresearch/faiss)和 Milvus 低比特教程。术语表:ANN 指近似最近邻,PQ 为产品量化,Recall@K 为 Top-K 召回率。