RyanCodrai/turbovec

**1000 万文档语料库以 float32 格式占用 31 GB 内存。turbovec 仅需 4 GB 即可容纳——并且搜索速度比 FAISS 更快。** turbovec 是一个带有 Python 绑定的 Rust 向量索引，基于 Google Research 的 [**TurboQuant**](https://arxiv.org/abs/2504.19874) 算法构建——这是一种具有接近最优失真率且无需单独训练阶段的数据无关量化器。 - **在线摄入。** 添加向量即被索引——无需训练步骤，无需调参，语料库增长时也无需重建。 - **快速 SIMD 搜索。** 手写的 NEON (ARM) 和 AVX-512BW (x86) 内核在 ARM 上比 FAISS IndexPQFastScan 快 10–19%；在 x86 上，它们在 4-bit 配置下胜出，而在 2-bit 配置下落后几个百分点。 - **搜索时过滤。** 向 `search()` 传递一个 id 白名单（或 slot 掩码），内核将直接遵照执行。您始终能从允许的集合中获得最多 `k` 个结果——无需超额抓取，对选择性过滤也没有召回率损耗。 - **纯本地。** 没有托管服务，数据不会离开您的机器或 VPC。与任何开源 embedding 模型搭配，即可构建完全气隙隔离的 RAG 技术栈。 正在构建对隐私、内存或延迟有严格要求的 RAG？**您找对地方了。** ## Python ``` pip install turbovec ``` ``` 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`： ``` 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、时间窗口等）生成的候选集中： ``` import numpy as np from turbovec import IdMapIndex idx = IdMapIndex(dim=1536, bit_width=4) idx.add_with_ids(vectors, ids) # 阶段 1：external system 缩小至 candidate ids。 allowed = np.array(db.execute("SELECT id FROM docs WHERE tenant=?", (t,)).fetchall(), dtype=np.uint64) # 阶段 2：在 candidate set 内进行 dense rerank。 scores, ids = idx.search(query, k=10, allowlist=allowed) ``` 过滤操作在 SIMD 内核内部以 32 个向量组成的块粒度进行：在执行任何 LUT 查找或评分工作之前，不包含任何允许 slot 的块会被短路处理，而在已评分的块内部，未被允许的单个 slot 则会在 heap 插入时被丢弃。因此，选择性白名单（仅允许索引中极小部分）可以避免大部分 SIMD 开销，而不是先付出开销再丢弃结果。 输出长度为 `min(k, len(allowed))`——当白名单小于 `k` 时，您将准确获得 `len(allowed)` 个结果，而不是带有填充的后备结果。 完整参考请见 [`docs/api.md`](docs/api.md)。 ### 框架集成 直接替换各框架中树内默认的向量/文档存储。保持相同的公共接口、相同的持久化语义、相同的检索器和 pipeline 连接方式——只需更改 import 即可保留您的 pipeline。 - [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 ``` cargo add turbovec ``` ``` 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： ``` 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(); ``` ## 召回率 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 下比 FAISS 高 0.9 个百分点，在 2-bit 下 R@1 基本打平（相差 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 数据，并在低维 embedding 上达到了与其他社区参考实现相近的数值（参见 d=384 的 [`turboquant-py`](https://pypi.org/project/turboquant-py/)）。在 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. 归一化。** 剥离每个向量的长度（范数）并将其存储为单个浮点数。现在每个向量都是超球面上的一个单位方向。 **2. 随机旋转。** 将所有向量乘以同一个随机正交矩阵。旋转后，每个坐标独立遵循 Beta 分布，该分布在高维下收敛于高斯分布 N(0, 1/d)。这对任何输入数据都成立——旋转使得坐标分布变得可预测。 **3. 逐坐标校准 (TQ+)。** 步骤 2 中的 Beta 分布是渐近的——在有限维度下，单个坐标会偏离标准形状（尤其是低比特和词向量类型的 embedding）。TQ+ 在首次添加时为每个坐标拟合两个标量——一个偏移量和一个缩放比例——将每个坐标的经验 5/95% 分位数映射到标准的 Beta 边缘分布上。随后 Lloyd-Max 码本针对其设计的*目标*分布进行量化。该校准在首次添加后被冻结，并被后续的添加操作复用——无需重新训练，无需重建，也无需单独的训练阶段。召回率提升：在偏移最严重的单元上（例如 2-bit 的 GloVe），@1 最高可提升 +1.4pp。 **4. Lloyd-Max 标量量化。** 由于分布已知，我们可以预计算每个坐标的最佳分桶方式。对于 2-bit，即 4 个桶；对于 4-bit，即 16 个桶。[Lloyd-Max 算法](https://en.wikipedia.org/wiki/Lloyd%27s_algorithm) 寻找能最小化均方误差的桶边界和质心。这些是直接通过数学计算得出的，而非基于数据。 **5. 比特打包。** 现在每个坐标都是一个小的整数（2-bit 为 0-3，4-bit 为 0-15）。将它们紧密打包成字节。一个 1536 维的向量从 6,144 字节 (FP32) 缩减至 384 字节 (2-bit)。这相当于 16 倍的压缩率。 **6. 长度重归一化评分。** 标量量化系统性地低估了内积——重构出的单位方向比原始方向略短。我们在编码时为每个向量计算一个标量——旋转后的单位向量与其自身质心重构的内积——并将 `||v|| / ⟨u, x̂⟩` 存储在每个压缩向量旁边。搜索内核在 heap 插入之前，将每个候选者的分数乘以此标量，从而将内积估计量从向下偏误转化为无偏估计，且不产生任何查询时开销和额外存储。这种召回率提升在低位宽下最为明显，因为此时量化收缩最为严重。 编码成本：每个向量额外进行一次 `d` 维点积以计算 `⟨u, x̂⟩`。对于 d=1536 的 100 万向量，这只是不到一秒的额外编码时间——这是在摄入时一次性付出的代价，而不是在查询时。 **搜索。** 我们无需解压每个数据库向量，而是将查询一次性旋转到相同的域中，并直接针对码本值进行评分。评分内核使用 SIMD 指令（ARM 上的 NEON，现代 x86 上的 AVX-512BW，以及 AVX2 回退方案），配合 nibble-split（半字节拆分）查找表以实现最大吞吐量。 Lloyd-Max 码本实现的失真率在信息论下限（Shannon 率失真极限）的 2.7 倍范围内；长度重归一化步骤消除了 Lloyd-Max 码本对内积估计量本身造成的残余偏误。 ## 构建 ### Python (通过 maturin) ``` pip install maturin cd turbovec-python maturin build --release pip install target/wheels/*.whl ``` ### Rust ``` cargo build --release ``` 所有 x86_64 构建均通过 `.cargo/config.toml` 针对 `x86-64-v3`（AVX2 基准，Haswell 2013+）进行设定。任何能够运行 AVX2 回退内核的 CPU 都可以运行整个 crate——AVX-512 内核在运行时通过 `is_x86_feature_detected!` 进行门控，仅在支持它的硬件上才会启动。 ## 运行基准测试 下载数据集： ``` 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/` 中的一个独立脚本。可单独运行任何一个： ``` python3 benchmarks/suite/speed_d1536_2bit_arm_mt.py python3 benchmarks/suite/recall_d1536_2bit.py python3 benchmarks/suite/compression.py ``` 运行某个类别的所有基准测试： ``` 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/`。重新生成图表： ``` 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 的打包布局、nibble-LUT 评分以及 u16 累加器策略

标签：Python, RAG, Rust, 可视化界面, 向量检索, 数据压缩, 数据库, 无后门, 本地部署, 网络流量审计, 逆向工具