如何在 Elasticsearch 中使用 BBQ 将 Jina v5 嵌入缩小 29 倍而不损失召回率

作者来自 Elastic Jeffrey Rengifo一项动手测试对 Elasticsearch 中的 BBQ 与 float32 向量索引进行了比较测量了五种语言下的内存占用、磁盘占用以及 recall10。立即通过 Search AI 的自助式动手学习亲自体验向量搜索。你可以立即开始免费的云试用或者在你的本地机器上体验 Elastic。BBQ 量化将Jina embeddings v5向量在Elasticsearch中的内存占用缩小了 29 倍。与全精度float32基线相比 Recall10 保持在 0.994。我们在一个包含五种语言的多语言新闻语料库上进行了这项测试使用jina-embeddings-v5-text-small构建了一个原始float32索引以及一个基于完全相同向量的bbq_hnsw索引。随后我们测量了两种索引的内存占用、磁盘占用以及检索质量。两种索引的磁盘占用几乎完全相同。真正决定你的集群是否能够容纳整个语料库的是内存占用而在本次测试中它从 12.71 MB 降低到了 0.44 MB。Jina v5 面向量化的训练正是召回率能够保持稳定的原因。前提条件Elasticsearch 9.x并且已提供jina-embeddings-v5-text-small推理端点。Python 3.10。Elasticsearch API 密钥。什么是量化嵌入是一组数字。默认情况下每个数字都是一个float32占用 4 字节。量化使用更少的位数来存储每个数字以空间换取精度。就像 JPEG 一样经过量化的向量是原始向量更小、保真度更低的副本但仍然能够完成相同的任务。JPEG 质量逐步降低图片由 Michael Gäbler 提供通过 Wikimedia Commons采用CC BY 3.0许可。名称每维字节数1024 维向量压缩率Float基线44096 B1xint811024 B4xint40.5512 B8xbbq~0.14142 B~29x什么是 BBQBetter Binary QuantizationBBQ是 Elasticsearch 针对稠密向量提供的 1 位量化模式。向量的每个维度都存储为一个比特再加上每个向量的少量校正字节。随后在查询时会执行重新评分rescore步骤。这使最终的检索质量能够接近全精度搜索。有关每一级量化背后的数学原理请参阅《标量量化入门 101》《优化标量量化》以及《https://elasticstack.blog.csdn.net/article/details/143701437》。BBQ 如何保持搜索准确性普通的 1 位量化本身会导致搜索质量下降过多。BBQ 通过以下三种机制保持较高的检索质量非对称精度存储的向量每个维度仅使用 1 位。校正因子每个向量使用少量浮点数记录舍入误差并在评分时对距离进行校正。过采样与重新评分BBQ 首先使用这些比特扫描候选结果然后使用更高精度对排名靠前的结果重新排序。获取前 10 个结果时通常会扫描约 30 个候选结果。最终得到的是体积约缩小 32 倍的向量同时检索质量仍然接近全精度。在本文下一部分中我们将在真实语料库上测量内存节省情况以及召回率。Jina embeddings v5 的工作原理Jina嵌入v5 是一个支持多语言嵌入的模型并采用了量化感知训练因此它天然适合与 Elasticsearch 中的 BBQ 搭配使用jina-embeddings-v5-text-small生成的 1024 维向量高于二进制量化仍能保持准确性的维度下限并且该模型经过训练使得 1 位量化几乎不会损失检索质量。它的主要特性包括一个模型支持多种任务v5 在单个基础模型之上使用小型低秩适配LoRA适配器每种任务对应一个适配器包括检索、文本匹配、聚类和分类。Elasticsearch 会在索引和查询时自动选择合适的适配器。Matryoshka 维度v5 经过训练允许你对向量进行截断1024、512 到 256 维同时尽可能减少搜索质量下降。这是另一种缩小向量的方法与量化相互独立。量化感知训练v5 专门针对 BBQ 进行了训练因此其 1 位向量几乎不会损失准确性。本文使用的是jina-embeddings-v5-text-small。该模型可通过Elastic Inference ServiceEIS使用输出 1024 维向量支持 32k令牌上下文窗口并支持 93 种语言。这使它高于384 维阈值低于该阈值时Elasticsearch 将不再默认使用bbq_hnsw。完整的模型介绍请参阅 Search Labs 上的 Jina v5 文章。设置 BBQ 与 float32 对比我们将创建两个索引两者共享映射配置变化的是index_options.type参数该参数告诉 Elasticsearch 如何存储稠密向量字段作为原始float32HNSW 或作为 1 位 BBQ索引 index_options 加载到内存vectors-float32hnsw未经过量化的原始float32基线vectors-bbqbbq_hnsw1 位 BBQ 量化 校正因子然后我们使用 Jina v5 对语料库进行一次嵌入处理将相同的向量分别索引到两个索引中并比较它们的磁盘占用、内存占用以及召回率。你可以通过完整的支持博客内容笔记本跟随操作。连接到 Elasticsearchfrom elasticsearch import Elasticsearch, helpers es_client Elasticsearch( ELASTICSEARCH_URL, api_keyELASTICSEARCH_API_KEY, request_timeout120 ) es_client.info()创建两个索引DIMS 1024 FLOAT_INDEX vectors-float32 BBQ_INDEX vectors-bbq def create_index(name, index_options): if es_client.indices.exists(indexname): es_client.indices.delete(indexname) es_client.indices.create( indexname, mappings{ properties: { text: {type: text}, lang: {type: keyword}, embedding: { type: dense_vector, dims: DIMS, index: True, similarity: cosine, index_options: index_options, }, } }, ) create_index(FLOAT_INDEX, {type: hnsw}) # raw float32 baseline create_index(BBQ_INDEX, {type: bbq_hnsw}) # 1-bit BBQ注意在生产环境中你可以使用semantic_text让 Elasticsearch 自动管理映射和推理端点。指向 Jina v5 推理端点我们直接调用模型jina-embeddings-v5-text-small无需创建推理端点来将文本转换为向量。INFERENCE_ID .jina-embeddings-v5-text-small def embed(texts, batch_size16): out [] for i in range(0, len(texts), batch_size): batch texts[i : i batch_size] try: resp es_client.inference.text_embedding( inference_idINFERENCE_ID, inputbatch ) except AttributeError: # older client versions resp es_client.inference.inference(inference_idINFERENCE_ID, inputbatch) out.extend(item[embedding] for item in resp[text_embedding]) return np.array(out, dtypenp.float32) embed([hello world]).shape # testing作为测试结果我们得到(1, 1024)加载多语言新闻数据集我们从hotchpotch/multilingual_cc_news流式读取真实新闻文章这是 CC-News 的 parquet 镜像。我们从五种语言中获取大约 1,000 篇文章总计约 3,000 个文档另外保留一小部分标题集合作为搜索查询。使用多种语言也可以让 Jina v5 展示其多语言能力。from datasets import load_dataset LANGS [en, de, ja, pt, ru] PER_LANG_DOCS 1000 PER_LANG_QUERIES 20 docs, queries [], [] for lang in LANGS: ds load_dataset( hotchpotch/multilingual_cc_news, lang, splittrain, streamingTrue ) rows [ r for r in ds.take(PER_LANG_DOCS PER_LANG_QUERIES) if r.get(maintext) and r.get(title) ] for row in rows[:PER_LANG_DOCS]: text (row[title] . row[maintext]).replace(\n, ).strip() docs.append({text: text[:1000], lang: lang}) for row in rows[PER_LANG_DOCS:]: queries.append({text: row[title], lang: lang}) # headlines as queries print(fCorpus: {len(docs)} docs | Queries: {len(queries)}) # RES: Corpus: 3102 docs | Queries: 18生成嵌入并批量索引我们只对语料库进行一次嵌入处理然后将这些完全相同的向量写入两个索引中。doc_vectors embed([d[text] for d in docs]) query_vectors embed([q[text] for q in queries])def index_docs(name): actions ( { _index: name, _id: i, _source: { text: d[text], lang: d[lang], embedding: doc_vectors[i].tolist(), }, } for i, d in enumerate(docs) ) helpers.bulk(es_client, actions, refreshTrue) for name in (FLOAT_INDEX, BBQ_INDEX): index_docs(name) es_client.indices.forcemerge(indexname, max_num_segments1) es_client.indices.refresh(indexname)我们执行强制合并到单个段以确保存储数据稳定且具有可比性。结果磁盘占用与内存占用磁盘使用 APIdisk usage API报告每个索引在向量knn_vectors上使用了多少字节。def vector_disk_bytes(name): du es_client.indices.disk_usage(indexname, run_expensive_tasksTrue) field du[name][fields][embedding] knn field.get(knn_vectors) if isinstance(knn, dict): return knn[size_in_bytes] return field[knn_vectors_in_bytes] float_disk vector_disk_bytes(FLOAT_INDEX) bbq_disk vector_disk_bytes(BBQ_INDEX) N len(docs) float_mem N * DIMS * 4 bbq_mem N * (DIMS // 8 14) print(fOn disk - float32: {float_disk/1e6:6.2f} MB | BBQ: {bbq_disk/1e6:6.2f} MB) print(fIn memory - float32: {float_mem/1e6:6.2f} MB | BBQ: {bbq_mem/1e6:6.2f} MB ({float_mem/bbq_mem:.0f}x smaller))结果On disk - float32: 12.80 MB | BBQ: 13.25 MB In memory - float32: 12.71 MB | BBQ: 0.44 MB (29x smaller)在磁盘上两个索引的大小基本相同。经过量化的索引仍然会保留原始float32向量用于重新评分以及合并期间的重新量化并在此基础上添加 1 位向量因此 BBQ 最终在磁盘上的占用会略大一些。真正的节省体现在内存中。HNSW 扫描只需要 RAM 中的 1 位向量而原始浮点向量会从磁盘中读取用于对顶部候选结果进行重新评分。我们使用文档中提供的kNN 内存公式来估算该占用float使用num_vectors × dims × 4而bbq使用num_vectors × (dims/8 14)。BBQ 在磁盘上的额外字节数应该与我们为内存计算出的 1 位负载相匹配。在这里计算结果为13.25 - 12.80 0.45 MB而计算得到的值为0.44 MB。两者基本一致。结果召回率为了检查量化索引返回的结果是否与 float 基线类似我们使用召回率recallk | BBQ top-k ∩ float32 top-k | / k并在所有查询上取平均值。我们调整过采样因子num_candidates / k该因子表示 BBQ 在使用 1 位向量扫描候选结果后在使用原始浮点向量对排名靠前的结果进行重新排序之前扫描的候选数量以找到仍然能够匹配float32的最低值。def search_ids(index, qvec, k10, num_candidates10): resp es_client.search( indexindex, sizek, _sourceFalse, knn{ field: embedding, query_vector: qvec.tolist(), k: k, num_candidates: num_candidates, }, ) return [h[_id] for h in resp[hits][hits]] K 10 # Ground truth: full-precision float32 with a wide candidate list (~exact) ground_truth [ set(search_ids(FLOAT_INDEX, qv, kK, num_candidates2000)) for qv in query_vectors ] oversamples [1, 2, 3, 5, 10] recalls [] for f in oversamples: num_candidates max(K * f, K) hits 0 for gt, qv in zip(ground_truth, query_vectors): got set(search_ids(BBQ_INDEX, qv, kK, num_candidatesnum_candidates)) hits len(got gt) recalls.append(hits / (len(query_vectors) * K)) print(foversample {f:2}x - recall{K} {recalls[-1]:.3f})结果如下BBQ 在 1x 过采样时的 recall10 从 0.994 开始在 3x 之前保持该水平然后在更高的因子下稳定在 0.989这意味着在所有过采样值下它返回的 top-10 文档中至少有 98.9% 与 float32 相同。有关量化情况下召回率如何在不同数据集之间变化的更多信息请参阅快速与准确测量量化向量搜索的召回率。BBQ 量化结果总结相同的向量两种存储格式以及一个实验磁盘大致相同12.80 MB对比13.25 MB。BBQ 为重新评分和合并保留了原始浮点向量。内存小 29 倍12.71 MB对比0.44 MB。这是决定你的集群是否能够容纳语料库的数字。Recall10在 1x 过采样时为0.994。量化感知训练发挥了作用。启用 BBQ 的条件如果你的维度数量高于384 维下限如果你的向量是主要的内存开销并且你能够接受额外的候选结果进行重新评分。那么可以启用 BBQ。对于 Jina v5 来说该模型专门针对这一场景进行了训练因此在大多数语料库上的召回率损失都很小。关于 BBQ 和向量量化的进一步阅读在支持博客内容仓库中运行本文完整的笔记本。有关 BBQ 背后数学原理的信息请参阅Lucene 和 Elasticsearch 中的 Better Binary Quantization。有关 Jina v5 架构的更多信息请参阅Search Labs 上的 Jina embeddings v5。有关采用 BBQ 的更广泛介绍请参阅如何将 BBQ 实现到你的使用场景中。有关 BBQ 原始研究请参阅论文RaBitQ具有理论误差界限的高维向量量化用于近似最近邻搜索。这篇内容对你有多大帮助原文BBQ vector quantization for Jina embeddings v5 - Elasticsearch Labs