实战指南：在 CPU 上 200ms 内搜索 4000 万文档

实战指南在 CPU 上 200ms 内搜索 4000 万文档使用二进制嵌入和 Int8 重排序摘要这篇文章Search 40M documents in under 200ms介绍了一种在纯CPU环境下高效处理大规模语义搜索的技术方案。文章的技术亮点在于结合了二进制嵌入Binary Embeddings进行快速初筛和Int8嵌入进行重排序Rescoring1 引言4000 万文档。一个 CPU。没有 GPU。二进制搜索加上 int8 重排序如何悄悄改写大规模语义搜索的规则。当你处理嵌入足够长的时间时你会遇到一种非常具体的挫折感。你终于让检索工作正常了。它很准确。感觉很智能。你为它感到骄傲。然后你看了看账单。或者更糟糕的是RAM 使用率。或者比这更糟糕的是——当你意识到你的“简单”语义搜索需要180GB 的内存才能存在的那一刻。这通常是那种安静的假设潜入的地方“好吧……我想这只是正确进行向量搜索的成本。”但也许不是。因为这是我在滚动鼠标时停下来的原因你可以搜索4000 万个文本在 ~200ms 内在纯 CPU上使用8GB RAM和45GB 磁盘没有 GPU。没有奇异的硬件。没有魔法。只是一个非常深思熟虑的推理策略。一旦你看到它你就无法忽视它。2、关于 fp32 嵌入的不舒服的真相让我们大声说出来。我们要么默认使用float32 嵌入因为那是模型给我们的。我们将它们视为神圣的。不可触碰的。高精度或什么都没有。但想想我们在检索过程中实际上在做什么。我们不是在解物理方程。我们在对文档进行排名。而排名特别是在漏斗的顶部是令人惊讶地宽容的。这就是整个方法滑过的裂缝。3、技巧一旦你剥去神秘感设置几乎是无聊的二进制搜索用于召回。Int8 重排序用于精度。就是这样。你没有丢弃质量。你分阶段进行。你让廉价的表示先做繁重的工作然后只在重要的地方带回精度。这是完整流程的工作原理。4、推理策略一步步来让我们慢慢地过一遍就像我们在一起调试一样。5、嵌入查询我们从你期望的地方开始。一个密集的嵌入模型。一个标准的 fp32 向量。start_time time.time query_embedding model.encode_query (query) embed_time time.time() start_time还没有什么花哨的。这是你的参考信号。你稍后会信任的东西。6、将查询量化为二进制现在是思维方式的第一次转变。我们将那个 fp32 查询嵌入量化为二进制。同样的语义方向。32 倍小的表示。start_time time.time query_embedding_ubinary quantize_embeddings( query_embedding.reshape(1, 1), ubinary ) quantize_time time.time() start_time这一步很快。几乎快得可疑。它解锁了后面的一切。7、搜索二进制索引这就是规模突然不再可怕的地方。二进制索引很小。它们很快。它们快乐地生活在内存中。你可以在这里选择精确或近似搜索。index binary_ivf_index if use_approx else binary_index start_time time.time() _scores, binary_ids index.search query_embedding_ubinary, top_k rescore_multiplier ) binary_ids binary_ids[0] search_time time.time() start_time与其问“4000 万个文档中哪个最好”你在问“哪 ~40 个看起来有希望”这种重新框架很重要。8、为顶级候选者加载 int8 嵌入既然我们已经缩小了范围我们重新引入更多的信号。我们只为顶级的二进制命中加载int8 嵌入。start_time time.time() int8_embeddings np.array( title_text_int8_dataset [binary_ids][embedding], dtypenp.int8 ) load_int8_time time.time() start_time仍然很小。仍然很便宜。但比二进制更具表现力。9、使用原始 fp32 查询进行重排序这是精度回归的时刻。我们将原始 fp32 查询嵌入与int8 文档嵌入进行评分。start_time time.time() scores query_embedding int8_embeddings.T rescore_time time.time() start time这一步几乎不花费任何成本。它为你买回了你早期“丢失”的大部分质量。10、排序并挑选最佳结果start_time time.time() indices scores.argsort() [::-1][:top_k] top_k_indices binary_ids[indices] top_k_scores scores [indices] sort_time time.time() start_time此时你基本上完成了排名。11、加载实际文本只有现在你才接触标题、URL 和内容。start_time time.time() top_k_titles title_text_int8_dataset [top_k_indices] [title] top_k_urls title_text_int8_dataset [top_k_indices] [url] top_k_texts title_text_int8_dataset [top_k_indices] [text] top_k_titles [f [{title}] ({url}) for title, url in zip(top_k_titles, top_k_urls)] load_text_time time.time() start_time这正是你_应该_加载它们的时候。而不是更早。12、实际上重要的数字这是在4000 万个维基百科文本上的真实运行Embed Time: 0.0833 s Quantize Time: 0.0000 s Search Time: 0.0464 s Load int8 Time: 0.0263 s Rescore Time: 0.0006 s Sort Time: 0.0000 s Load Text Time: 0.0192 s Total Retrieval Time: 0.0925 s再读一遍。在 100ms 内端到端。在 CPU 上。你可以现在就自己尝试无需登录https://huggingface.co/spaces/sentence-transformers/quantized-retrieval13、为什么这行得通以及为什么起初感觉不对它感觉不对因为我们被教导要不惜一切代价保护精度。但检索不是分类。它不是回归。它甚至不是生成。它是分类。二进制嵌入在排除事物方面是现象级的。Int8 嵌入在排名竞争者方面很棒。Fp32 在你真正需要细微差别时是最好的——一旦你剩下 40 个文档这种情况很少见。因此与其到处支付 fp32 成本不如像聚光灯一样分阶段进行精度。先宽光束。后窄光束。14、存储和内存让我们将其与标准 fp32 设置进行比较。15、Float32 检索基线RAM: ~180GB磁盘: ~180GB速度: 慢成本: 痛苦16、二进制 int8 检索二进制索引:32 倍更小Int8 嵌入:4 倍更小RAM 使用: ~6GB磁盘使用: ~45GB这是关键所在。你只将二进制索引保留在内存中。那是真正的节省复合的地方。17、性能权衡让我们不要假装零损失。这是实验显示的| Representation | Index Size | Speed | Performance | | ---------------- | ----------- | ---------------- | ----------- | | float32 | 1× | 1× | 100% | | int8 / uint8 | 4× smaller | up to 4× faster | ~99.3% | | binary / ubinary | 32× smaller | up to 45× faster | ~96% |现在是重要的部分。通过二进制搜索检索更多候选者并使用 int8 对它们进行重排序你恢复了~99% 的 fp32 质量。那最后的 1% 几乎不花费你任何东西。18、这在实践中改变了什么老实说很多。这意味着你不需要 GPU 进行大规模语义搜索你可以在商品硬件上发布严肃的检索你不再到处支付 fp32 的“奢侈税”你可以扩展集合而无需重写所有内容也许更大的转变是心理上的。我们一直像对待易碎的玻璃一样对待嵌入。它们不是。它们是工具。当你停止对每颗钉子都使用一把锤子时工具效果最好。19、如果你要从中带走一件事分阶段思考。不要问“什么是最准确的表示”问“在管道的这个点上什么足够准确”二进制第一。Int8 第二。精度最后。一旦你这样看检索回到到处暴力 fp32 感觉……有点不负责任。也许这就是隐藏在这里的真正教训。不是量化很聪明。而是工程判断仍然胜过暴力即使在巨型模型的时代。参考文献实战指南在 CPU 上 200ms 内搜索 4000 万文档使用二进制嵌入和 Int8 重排序