基于嵌入量化的大规模文档搜索

4000 万份文档。单 CPU。无需 GPU。

二元搜索加上 int8 重评分如何悄然改写大规模语义搜索的规则。

当你长时间使用嵌入技术时，你会遇到一种非常特殊的挫败感。

你终于实现了检索。它很准确。感觉很智能。你为此感到自豪。然后你看了看账单。

或者更糟，看看内存使用情况。

或者更糟——当你意识到你“简单”的语义搜索仅仅为了运行就需要 180GB 的内存时。

这时，一个悄悄的假设就会悄然出现：“好吧……我想这就是正确进行向量搜索的代价吧。”

但也许并非如此。

因为接下来这件事让我停下了滚动页面：

你只需 8GB 内存和 45GB 硬盘，就能在 CPU 上用大约 200 毫秒的时间搜索 4000 万条文本。

无需 GPU。无需特殊硬件。无需任何魔法。

只需一种精心设计的推理策略。

一旦你看到了这一点，就无法再视而不见。

1、关于 fp32 嵌入的令人不安的真相

让我们大声说出来。

我们大多数人默认使用 float32 嵌入，因为模型就是这么给出的。我们视它们为神圣不可侵犯。要么高精度，要么什么都不要。

但想想我们在检索过程中究竟在做什么。

我们不是在解物理方程，而是在对文档进行排序。

而且，排序，尤其是在漏斗顶端，容错率出奇地高。

这就是整个方法的漏洞所在。

2、诀窍

一旦揭开神秘面纱，你会发现这套设置几乎乏味：

二分查找用于召回率。Int8 重评分用于精确率。

就是这样。

你不会丢弃质量，而是对其进行优化。

你先让低成本的表示方法完成繁重的工作，然后在真正重要的地方再提升精确率。

以下是完整流程的工作原理。

3、推理策略，逐步详解

让我们慢慢来，就像一起调试一样。

步骤 1：嵌入查询

我们从你预期的地方开始。

一个密集嵌入模型。一个标准的 fp32 向量。

#1. Embed the query as float32
start_time = time.time
query_embedding = model.encode_query (query)
embed_time = time.time() start_time

目前还没有什么特别之处。这是你的参考信号，也是你之后会信赖的东西。

步骤 2：将查询量化为二进制

现在，思维方式发生了第一次转变。

我们将 float32 查询嵌入量化为二进制。语义方向相同，但表示形式缩小了 32 倍。

#2. Quantize the query to ubinary
start_time = time.time
query_embedding_ubinary = quantize_embeddings(
query_embedding.reshape(1, 1), "ubinary"
)
quantize_time = time.time() start_time

这一步非常快，快得几乎有点可疑。

而且它为后续所有步骤奠定了基础。

步骤 3：搜索二元索引

在这里，规模不再令人担忧。

二元索引体积小巧，速度快，而且内存占用低。

您可以选择精确搜索或近似搜索。

#3. Search the binary index (either exact or approximate)
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 个文档看起来有希望？”

这种重新定义问题至关重要。

步骤 4：加载最佳候选结果的 int8 词嵌入

现在我们已经缩小了范围，接下来需要引入更多信号。

我们仅加载排名靠前的二元匹配结果的 int8 词嵌入。

#4. Load the corresponding int8 embeddings
start_time = time.time()
int8_embeddings = np.array(
title_text_int8_dataset [binary_ids]["embedding"], dtype=np.int8
)
load_int8_time = time.time() start_time

数据量仍然很小，成本仍然很低，但比二元匹配更具表达力。

步骤 5：使用原始 fp32 查询重新评分

此时精度得以恢复。

我们使用原始 fp32 查询词嵌入，并将其与 int8 文档词嵌入进行比较评分。

#5. Rescore the top_k rescore_multiplier using the float32 query embedding and the int8 document embeddings
start_time = time.time()
scores = query_embedding @ int8_embeddings.T
rescore_time = time.time() start time

这一步几乎不耗费任何资源。而且它能帮你找回之前“丢失”的大部分质量。

步骤 6：排序并选择最佳结果

#6. Sort the scores and return the top k
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

至此，排名基本完成。

步骤 7：加载实际文本

现在才需要处理标题、URL 和内容。

#7. Load titles and texts for the top k results
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

这是加载文本的准确时间。

不能提前加载。

4、真正重要的数字

以下是对超过 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

再读一遍。

全程耗时不到 100 毫秒。基于 CPU。

你可以在这里立即亲自尝试，无需登录。

5、为什么这种方法有效

感觉不对劲是因为我们被教导要不惜一切代价维护精确性。

但检索并非分类，也非回归，甚至不是生成。

它是筛选。

二进制嵌入在排除某些内容方面表现出色。Int8 嵌入擅长对候选对象进行排序。FP32 最适合真正需要细微差别的情况——但当文档数量减少到 40 个时，这种情况就很少见了。

因此，与其处处消耗 FP32 的成本，不如像聚光灯一样，逐步提升精度。

先用宽光束，后用窄光束。

6、存储和内存

让我们将其与标准的 FP32 设置进行比较。

Float32 数据检索（基准）

• 内存：约 180GB
• 磁盘：约 180GB
• 速度：慢
• 成本：高

二进制 + int8 数据检索

• 二进制索引：缩小 32 倍
• int8 嵌入：缩小 4 倍
• 内存使用量：约 6GB
• 磁盘使用量：约 45GB

关键在于：您只需将二进制索引保存在内存中。

这才是真正节省成本的地方。

7、性能权衡

我们不能假装没有损失。

以下是实验结果：

现在来说说重点。

通过二分查找检索更多候选结果，并使用 int8 重新评分，您可以恢复约 99% 的 fp32 质量。

这最后的 1% 几乎不需要任何成本。

8、这在实际应用中会带来哪些改变？

说实话？改变很大。

这意味着：

• 您不再需要 GPU 来进行大规模语义搜索
• 您可以在普通硬件上实现强大的检索功能
• 您无需再为 fp32 支付“奢侈税”
• 您可以扩展数据集而无需重写所有代码

也许更大的转变在于思维方式。

我们一直把嵌入向量当作易碎的玻璃。

它们不是。

它们是工具。而工具只有在您不再用同一把锤子敲所有钉子时才能发挥最佳作用。

如果您只能记住一点：

分阶段思考。

不要问“最精确的表示是什么？”，而要问“在流程的这个阶段，什么精度才足够？”

二进制优先，Int8其次，精度最后。

一旦你理解了这种检索方式，再去到处用暴力破解的FP32格式就显得……有点不负责任了。

也许这才是真正值得我们深思的教训。

并非说量化有多么高明。

但即便在巨型模型时代，这种工程判断仍然胜过暴力破解。

原文链接：Search 40M documents in under 200ms on a CPU using binary embeddings and int8 rescoring.

汇智网翻译整理，转载请标明出处