ryancodrai--turbovec
wehub-resource-sync
3ca67faf24
CI / Rust (macos-14) (push) Has been cancelled
CI / Rust (ubuntu-latest) (push) Has been cancelled
CI / Rust (windows-latest) (push) Has been cancelled
CI / Python (macos-14) (push) Has been cancelled
CI / Python (ubuntu-latest) (push) Has been cancelled
CI / Python (windows-latest) (push) Has been cancelled
245 行
14 KiB
Markdown
245 行
14 KiB
Markdown
<!-- WEHUB_ZH_README -->
|
||
> [!NOTE]
|
||
> 本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。
|
||
> [English](./README.en.md) · [原始项目](https://github.com/RyanCodrai/turbovec) · [上游 README](https://github.com/RyanCodrai/turbovec/blob/HEAD/README.md)
|
||
> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。
|
||
|
||
<p align="center">
|
||
<img src="docs/header.png" alt="turbovec — Google's TurboQuant for vector search" width="100%">
|
||
</p>
|
||
|
||
<p align="center">
|
||
<a href="https://github.com/RyanCodrai/turbovec/blob/main/LICENSE"><img src="https://img.shields.io/pypi/l/turbovec" alt="License"></a>
|
||
<a href="https://pypi.org/project/turbovec/"><img src="https://img.shields.io/pypi/v/turbovec?label=pypi&color=blue" alt="PyPI version"></a>
|
||
<a href="https://crates.io/crates/turbovec"><img src="https://img.shields.io/crates/v/turbovec?label=crates.io&color=blue" alt="crates.io version"></a>
|
||
<a href="https://arxiv.org/abs/2504.19874"><img src="https://img.shields.io/badge/paper-arXiv-b31b1b.svg" alt="TurboQuant paper"></a>
|
||
</p>
|
||
|
||
---
|
||
|
||
**1000 万文档语料库以 float32 存储需要 31 GB RAM。turbovec 将其压缩到 4 GB——且搜索速度比 FAISS 更快。**
|
||
|
||
turbovec 是一个带 Python 绑定的 Rust 向量索引，基于 Google Research 的 [**TurboQuant**](https://arxiv.org/abs/2504.19874)) 算法构建——一种数据无关（data-oblivious）的量化器，具有接近最优的失真且无需单独的训练阶段。
|
||
|
||
- **在线摄入（Online ingest）。** 添加向量即可建立索引——无需训练步骤、参数调优，语料库增长时也无需重建。
|
||
- **快速 SIMD 搜索。** 手写 NEON（ARM）和 AVX-512BW（x86）内核在 ARM 上比 FAISS IndexPQFastScan 快 10–19%；在 x86 上，4-bit 配置领先，2-bit 配置仅落后几个百分点。
|
||
- **搜索时过滤。** 向 `search()` 传入 id 白名单（或 slot 位掩码），内核会直接遵守。你始终从允许的集合中获得最多 `k` 条结果——无过度抓取，选择性过滤也不会影响召回。
|
||
- **纯本地。** 无托管服务，数据不会离开你的机器或 VPC。搭配任意开源嵌入模型，即可构建完全气隙隔离（air-gapped）的 RAG 栈。
|
||
|
||
在构建 RAG 时，隐私、内存或延迟很重要？**你来对地方了。**
|
||
|
||
## Python
|
||
|
||
```bash
|
||
pip install turbovec
|
||
```
|
||
|
||
```python
|
||
from turbovec import TurboQuantIndex
|
||
|
||
index = TurboQuantIndex(dim=1536, bit_width=4)
|
||
index.add(vectors)
|
||
index.add(more_vectors)
|
||
|
||
scores, indices = index.search(query, k=10)
|
||
|
||
index.write("my_index.tv")
|
||
loaded = TurboQuantIndex.load("my_index.tv")
|
||
```
|
||
|
||
需要能在删除后仍然保留的稳定 id？使用 `IdMapIndex`：
|
||
|
||
```python
|
||
import numpy as np
|
||
from turbovec import IdMapIndex
|
||
|
||
index = IdMapIndex(dim=1536, bit_width=4)
|
||
index.add_with_ids(vectors, np.array([1001, 1002, 1003], dtype=np.uint64))
|
||
|
||
scores, ids = index.search(query, k=10) # ids are your uint64 external ids
|
||
index.remove(1002) # O(1) by id
|
||
|
||
index.write("my_index.tvim")
|
||
loaded = IdMapIndex.load("my_index.tvim")
|
||
```
|
||
|
||
### 混合检索（过滤搜索）
|
||
|
||
将结果限制为另一系统（SQL、BM25、ACL、时间窗口等）产生的候选集：
|
||
|
||
```python
|
||
import numpy as np
|
||
from turbovec import IdMapIndex
|
||
|
||
idx = IdMapIndex(dim=1536, bit_width=4)
|
||
idx.add_with_ids(vectors, ids)
|
||
|
||
# Stage 1: external system narrows to candidate ids.
|
||
allowed = np.array(db.execute("SELECT id FROM docs WHERE tenant=?", (t,)).fetchall(),
|
||
dtype=np.uint64)
|
||
|
||
# Stage 2: dense rerank within the candidate set.
|
||
scores, ids = idx.search(query, k=10, allowlist=allowed)
|
||
```
|
||
|
||
过滤在 SIMD 内核内部以 32 向量块粒度进行：没有允许 slot 的块会在任何 LUT 查找或评分工作之前被短路跳过，已评分块内个别不允许的 slot 会在堆插入时被丢弃。因此，选择性白名单（仅允许索引的一小部分）能避免大部分 SIMD 开销，而不是先付出计算成本再丢弃结果。
|
||
|
||
输出长度为 `min(k, len(allowed))`——当白名单小于 `k` 时，你会得到恰好 `len(allowed)` 条结果，而非填充的回退值。
|
||
|
||
完整参考见 [`docs/api.md`](docs/api.md)。
|
||
|
||
### 框架集成
|
||
|
||
可作为各框架内置参考向量/文档存储的即插即用替代。相同的公开接口、相同的持久化语义、相同的检索器与流水线接线——只需更换 import，流水线其余部分保持不变。
|
||
|
||
- [LangChain](docs/integrations/langchain.md) — `pip install turbovec[langchain]` · 替代 `langchain_core.vectorstores.InMemoryVectorStore`
|
||
- [LlamaIndex](docs/integrations/llama_index.md) — `pip install turbovec[llama-index]` · 替代 `llama_index.core.vector_stores.SimpleVectorStore`
|
||
- [Haystack](docs/integrations/haystack.md) — `pip install turbovec[haystack]` · 替代 `haystack.document_stores.in_memory.InMemoryDocumentStore`
|
||
- [Agno](docs/integrations/agno.md) — `pip install turbovec[agno]` · 替代 `agno.vectordb.lancedb.LanceDb`
|
||
|
||
## Rust
|
||
|
||
```bash
|
||
cargo add turbovec
|
||
```
|
||
|
||
```rust
|
||
use turbovec::TurboQuantIndex;
|
||
|
||
let mut index = TurboQuantIndex::new(1536, 4).unwrap();
|
||
index.add(&vectors);
|
||
let results = index.search(&queries, 10);
|
||
index.write("index.tv").unwrap();
|
||
let loaded = TurboQuantIndex::load("index.tv").unwrap();
|
||
```
|
||
|
||
需要能在删除后仍然保留的稳定外部 id：
|
||
|
||
```rust
|
||
use turbovec::IdMapIndex;
|
||
|
||
let mut index = IdMapIndex::new(1536, 4).unwrap();
|
||
index.add_with_ids(&vectors, &[1001, 1002, 1003]).unwrap();
|
||
let (scores, ids) = index.search(&queries, 10);
|
||
index.remove(1002);
|
||
index.write("index.tvim").unwrap();
|
||
let loaded = IdMapIndex::load("index.tvim").unwrap();
|
||
```
|
||
|
||
## 召回率（Recall）
|
||
|
||
TurboQuant 对比 FAISS `IndexPQ`（LUT256，nbits=8）——论文第 4.4 节的基线。100K 向量，k=64。FAISS PQ 子量化器数量按 TurboQuant 的比特率匹配（2-bit 时 m=d/4，4-bit 时 m=d/2）。
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
在 OpenAI d=1536 和 d=3072 上，TurboQuant 在 2-bit 和 4-bit 的 R@1 上比 FAISS 高 0.2–1.9 个百分点，且两者在 k=8 时均达到 1.0（k=4 时已 ≥0.997）。GloVe d=200 是更难的场景——低维时渐近 Beta 假设更宽松。TurboQuant 在 4-bit 的 R@1 上比 FAISS 高 0.9 个百分点，2-bit 时基本持平（差距在 0.1 个百分点以内），到 k≈16 时两者都紧密跟踪 FAISS。
|
||
|
||
**关于基线的说明。** 我们与 FAISS `IndexPQ`（LUT256，nbits=8，float32 LUT）对比，因为它是大多数用户会选用的默认生产级 PQ。这比 [TurboQuant 论文](https://arxiv.org/abs/2504.19874)) 中的自定义 u8-LUT PQ 是更强的基线——FAISS 在评分时使用更高精度的 LUT，并用 k-means++ 训练码本。我们在 OpenAI d=1536 / d=3072 上复现了论文的 TurboQuant 数值，在低维嵌入上也达到了与其他社区参考实现相近的数值（见 [`turboquant-py`](https://pypi.org/project/turboquant-py/)) 在 d=384 的结果）。在 GloVe（d=200）上——渐近 Beta 假设最宽松的低维场景——TurboQuant 在 2-bit 与 FAISS 持平，4-bit 领先；TQ+ 校准弥补了基础算法在低维上留下的差距。
|
||
|
||
完整结果：[d=1536 2-bit](benchmarks/results/recall_d1536_2bit.json)、[d=1536 4-bit](benchmarks/results/recall_d1536_4bit.json)、[d=3072 2-bit](benchmarks/results/recall_d3072_2bit.json)、[d=3072 4-bit](benchmarks/results/recall_d3072_4bit.json)、[GloVe 2-bit](benchmarks/results/recall_glove_2bit.json)、[GloVe 4-bit](benchmarks/results/recall_glove_4bit.json)。
|
||
|
||
## 压缩
|
||
|
||
|
||
|
||
## 搜索速度
|
||
|
||
所有基准测试：100K 向量，1K 查询，k=64，5 次运行的中位数。
|
||
|
||
### ARM（Apple M3 Max）
|
||
|
||
|
||
|
||
|
||
|
||
在 ARM 上，TurboQuant 在所有配置下均比 FAISS FastScan 快 10–19%。
|
||
|
||
### x86（Intel Xeon Platinum 8481C / Sapphire Rapids，8 vCPUs）
|
||
|
||
|
||
|
||
|
||
|
||
在 x86 上，TurboQuant 在 4-bit 配置下领先最多约 5%（d=3072 多线程持平），在 2-bit 上略落后于 FAISS——最明显的是 d=1536 单线程（约 8%），其余配置差距在几个百分点以内——FAISS 的 AVX-512 VBMI 路径在短 2-bit 累加循环上更有优势。
|
||
|
||
## 工作原理
|
||
|
||
每个向量都是高维超球面上的一个方向。TurboQuant 用一条简单洞见压缩这些方向：施加随机旋转后，每个坐标都遵循已知分布——与输入数据无关。
|
||
|
||
**1. 归一化（Normalize）。** 从每个向量中剥离长度（norm）并将其存储为单个浮点数。现在每个向量都是超球面上的单位方向。
|
||
|
||
**2. 随机旋转（Random rotation）。** 将所有向量乘以同一个随机正交矩阵。旋转后，每个坐标独立服从 Beta 分布，在高维下收敛于高斯分布 N(0, 1/d)。这对任意输入数据都成立——旋转使坐标分布变得可预测。
|
||
|
||
**3. 逐坐标校准（TQ+）。** 步骤 2 中的 Beta 分布是渐近的——在有限维度下，单个坐标会偏离标准形状（尤其是低位宽和词向量风格的嵌入）。TQ+ 在首次 add 时为每个坐标拟合两个标量——一个平移和一个缩放——将每个坐标的经验 5/95% 分位数映射到标准 Beta 边缘分布上。随后 Lloyd-Max 码本针对其*设计目标*分布进行量化。校准在首次 add 后冻结，供后续 add 复用——无需重训练、无需重建、无需单独的训练阶段。召回增益：在漂移最严重的单元上 @1 最高可达 +1.4pp（例如 2-bit 的 GloVe）。
|
||
|
||
**4. Lloyd-Max 标量量化。** 由于分布已知，我们可以预先计算每个坐标的最优分桶方式。2-bit 为 4 个桶；4-bit 为 16 个桶。[Lloyd-Max 算法](https://en.wikipedia.org/wiki/Lloyd%27s_algorithm) 找出使均方误差最小的桶边界和质心。这些由数学一次性计算得出，而非从数据中学习。
|
||
|
||
**5. 位打包（Bit-pack）。** 每个坐标现在是一个小整数（2-bit 为 0-3，4-bit 为 0-15）。将这些整数紧密打包进字节。1536 维向量从 6,144 字节（FP32）压缩到 384 字节（2-bit），即 16 倍压缩。
|
||
|
||
**6. 长度重归一化评分。** 标量量化会系统性低估内积——重建的单位方向比原始向量略短。我们在编码时为每个向量计算一个标量——旋转后的单位向量与其自身质心重建的内积——并将 `||v|| / ⟨u, x̂⟩` 与每个压缩向量一并存储。搜索内核在堆插入前将该标量乘以每个候选分数，使内积估计器从向下偏置变为无偏，且搜索时零额外开销、零额外存储。召回增益在低位宽下最为明显，因为此时量化收缩最大。
|
||
|
||
编码开销：每个向量额外一次 `d` 维点积以计算 `⟨u, x̂⟩`。在 d=1536 的 100 万向量上，额外编码时间不到一秒——这是在摄入时一次性付出的成本，而非查询时。
|
||
|
||
**搜索。** 无需解压每个数据库向量，只需将查询一次旋转到同一域，并直接针对码本值评分。评分内核使用 SIMD 内建函数（ARM 上为 NEON，现代 x86 上为 AVX-512BW，并提供 AVX2 回退），配合半字节拆分查找表以实现最大吞吐。
|
||
|
||
Lloyd-Max 码本的失真在信息论下界（Shannon 失真-速率极限）的 2.7 倍以内；长度重归一化步骤消除了 Lloyd-Max 码本在内积估计器本身上引入的残余偏置。
|
||
|
||
## 构建
|
||
|
||
### Python（通过 maturin）
|
||
|
||
```bash
|
||
pip install maturin
|
||
cd turbovec-python
|
||
maturin build --release
|
||
pip install target/wheels/*.whl
|
||
```
|
||
|
||
### Rust
|
||
|
||
```bash
|
||
cargo build --release
|
||
```
|
||
|
||
所有 x86_64 构建通过 `.cargo/config.toml` 以 `x86-64-v3`（AVX2 基线，Haswell 2013+）为目标。任何能运行 AVX2 回退内核的 CPU 都能运行整个 crate——AVX-512 内核通过 `is_x86_feature_detected!` 在运行时门控，仅在支持该指令集的硬件上启用。
|
||
|
||
## 运行基准测试
|
||
|
||
下载数据集：
|
||
```bash
|
||
python3 benchmarks/download_data.py all # all datasets
|
||
python3 benchmarks/download_data.py glove # GloVe d=200
|
||
python3 benchmarks/download_data.py openai-1536 # OpenAI DBpedia d=1536
|
||
python3 benchmarks/download_data.py openai-3072 # OpenAI DBpedia d=3072
|
||
```
|
||
|
||
每个基准测试是 `benchmarks/suite/` 中的独立脚本。可单独运行任意一个：
|
||
```bash
|
||
python3 benchmarks/suite/speed_d1536_2bit_arm_mt.py
|
||
python3 benchmarks/suite/recall_d1536_2bit.py
|
||
python3 benchmarks/suite/compression.py
|
||
```
|
||
|
||
运行某一类别的全部基准测试：
|
||
```bash
|
||
for f in benchmarks/suite/speed_*arm*.py; do python3 "$f"; done # all ARM speed
|
||
for f in benchmarks/suite/speed_*x86*.py; do python3 "$f"; done # all x86 speed
|
||
for f in benchmarks/suite/recall_*.py; do python3 "$f"; done # all recall
|
||
python3 benchmarks/suite/compression.py # compression
|
||
```
|
||
|
||
结果以 JSON 保存到 `benchmarks/results/`。重新生成图表：
|
||
```bash
|
||
python3 benchmarks/create_diagrams.py
|
||
```
|
||
|
||
## 参考文献
|
||
|
||
- [TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate](https://arxiv.org/abs/2504.19874) (ICLR 2026) -- 本文实现的论文
|
||
- [RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search](https://arxiv.org/abs/2405.12497) (SIGMOD 2024) -- 步骤 5 中采用的逐向量长度重归一化修正的来源
|
||
- [FAISS Fast accumulation of PQ and AQ codes](https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)) -- turbovec 的 x86 SIMD 内核借鉴了 FastScan 的打包布局、半字节 LUT 评分和 u16 累加器策略
|