Elasticsearch测试：BBQ量化让Jina v5内存占用降29倍，召回率达0.994！

突发Elasticsearch向量索引对比测试一项动手测试对Elasticsearch中的BBQ与float32向量索引进行了比较测量了五种语言下的内存占用、磁盘占用以及recall10。用户可通过Search AI的自助式动手学习亲自体验向量搜索能立即开始免费的云试用或者在本地机器上体验Elastic。BBQ量化效果显著BBQ量化将Jina embeddings v5向量在Elasticsearch中的内存占用缩小了29倍。与全精度float32基线相比Recall[10]保持在0.994。测试在包含五种语言的多语言新闻语料库上进行构建了原始float32索引和bbq_hnsw索引测量发现两种索引磁盘占用几乎相同而内存占用从12.71 MB降低到了0.44 MB。Jina v5面向量化的训练使召回率保持稳定。前提条件进行测试需满足以下条件Elasticsearch 9.x并且已提供 jina-embeddings-v5-text-small推理端点Python 3.10Elasticsearch API密钥。什么是量化嵌入是一组数字默认每个数字是float32占用4字节。量化使用更少的位数来存储每个数字以空间换取精度。经过量化的向量如同JPEG是原始向量更小、保真度更低的副本但仍能完成相同任务。什么是BBQBetter Binary QuantizationBBQ是Elasticsearch针对稠密向量提供的1位量化模式。向量的每个维度存储为一个比特加上每个向量的少量校正字节查询时会执行重新评分步骤使最终检索质量接近全精度搜索。有关每一级量化背后的数学原理可参阅相关文章。BBQ如何保持搜索准确性普通的1位量化会使搜索质量下降过多BBQ通过三种机制保持较高检索质量非对称精度存储的向量每个维度仅使用1位校正因子每个向量用少量浮点数记录舍入误差并在评分时校正距离过采样与重新评分先使用比特扫描候选结果再用更高精度对排名靠前的结果重新排序获取前10个结果时通常扫描约30个候选结果。最终得到的向量体积约缩小32倍检索质量接近全精度。Jina embeddings v5的工作原理Jina嵌入v5是支持多语言嵌入的模型采用量化感知训练适合与Elasticsearch中的BBQ搭配使用。其生成的1024维向量高于二进制量化仍能保持准确性的维度下限且模型训练使1位量化几乎不损失检索质量。它的主要特性包括一个模型支持多种任务在单个基础模型上使用小型低秩适配适配器每种任务对应一个适配器Elasticsearch会自动选择合适的适配器Matryoshka维度允许对向量进行截断减少搜索质量下降量化感知训练专门针对BBQ训练1位向量几乎不损失准确性。本文使用的jina-embeddings-v5-text-small模型可通过Elastic Inference Service使用输出1024维向量支持32k令牌上下文窗口和93种语言高于384维阈值低于该阈值时Elasticsearch将不再默认使用bbq_hnsw。设置BBQ与float32对比创建两个索引共享映射配置通过index_options.type参数告诉Elasticsearch如何存储稠密向量字段。使用Jina v5对语料库进行一次嵌入处理将相同向量分别索引到两个索引中比较磁盘占用、内存占用和召回率。可通过支持博客内容笔记本跟随操作。连接到Elasticsearch使用相关代码连接到Elasticsearchfrom elasticsearch import Elasticsearch, helperses_client Elasticsearch(ELASTICSEARCH_URL, api_keyELASTICSEARCH_API_KEY, request_timeout120)es_client.info()创建两个索引定义相关参数和函数创建索引DIMS 1024FLOAT_INDEX vectors-float32BBQ_INDEX vectors-bbqdef 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 baselinecreate_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-smalldef 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 versionsresp 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流式读取真实新闻文章获取五种语言约1000篇文章总计约3000个文档保留一小部分标题集合作为搜索查询。代码如下from datasets import load_datasetLANGS [en, de, ja, pt, ru]PER_LANG_DOCS 1000PER_LANG_QUERIES 20docs, queries [], []for lang in LANGS:ds load_dataset(hotchpotch/multilingual_cc_news, lang, splittrain, streamingTrue)rows [rfor 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 queriesprint(fCorpus: {len(docs)} docs | Queries: {len(queries)})结果为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)执行强制合并到单个段确保存储数据稳定且具有可比性。结果磁盘占用与内存占用使用磁盘使用API报告每个索引在向量上使用的字节数代码如下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 * 4bbq_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 MBIn memory - float32: 12.71 MB | BBQ: 0.44 MB (29x smaller)。磁盘上两个索引大小基本相同BBQ因保留原始float32向量并添加1位向量占用略大。真正的节省体现在内存中HNSW扫描只需RAM中的1位向量原始浮点向量从磁盘读取用于重新评分。使用kNN内存公式估算占用BBQ在磁盘上的额外字节数与计算出的1位负载基本一致。结果召回率使用召回率检查量化索引返回结果与float基线的相似度公式为recall[k] | BBQ top-k ∩ float32 top-k | / k并在所有查询上取平均值。调整过采样因子找到匹配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 10ground_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 0for 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过采样时的recall[10]从0.994开始3x之前保持该水平更高因子下稳定在0.989意味着所有过采样值下其返回的top-10文档中至少有98.9%与float32相同。BBQ量化结果总结相同向量两种存储格式的实验结果磁盘上大致相同12.80 MB对比13.25 MBBBQ为重新评分和合并保留原始浮点向量内存小29倍12.71 MB对比0.44 MB这决定集群是否能容纳语料库Recall10在1x过采样时为0.994量化感知训练发挥作用。启用BBQ的条件为维度数量高于384维下限向量是主要内存开销且能接受额外候选结果重新评分Jina v5针对此场景训练多数语料库上召回率损失小。关于BBQ和向量量化的进一步阅读可在支持博客内容仓库中运行本文完整的笔记本了解BBQ背后数学原理的信息获取Jina v5架构的更多信息查看采用BBQ的更广泛介绍阅读BBQ原始研究论文。