RaBitQ项目在MIP搜索场景下的优化实践

RaBitQ项目在MIP搜索场景下的优化实践

RaBitQ [SIGMOD 2024] RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search 项目地址: https://gitcode.com/gh_mirrors/ra/RaBitQ

背景介绍

RaBitQ作为一种高效的向量近似搜索算法，在欧式距离度量下表现优异。然而在实际应用中，我们发现当面对最大内积搜索(MIPS)场景时，特别是使用Cohere多语言嵌入模型生成的数据集时，算法的召回率出现了显著下降。本文将深入分析问题原因，并介绍基于RaBitQ核心思想进行MIPS优化的技术方案。

问题分析

在原始测试中，直接使用RaBitQ代码库处理Cohere的embed-multilingual-v2.0模型生成的数据集时，出现了召回率低于10%的情况。经过深入分析，我们发现核心问题在于：

1. 原生RaBitQ设计针对欧式距离优化，而Cohere模型需要支持最大内积(MIP)搜索
2. 向量幅值在MIP场景下具有重要影响，这与欧式距离的归一化处理存在本质差异
3. 直接使用欧式距离估计器处理内积运算会导致相关性显著降低

技术解决方案

基于RaBitQ的核心思想，我们推导出适用于MIP场景的数学变换公式。设原始数据向量为o_r，查询向量为q_r，中心向量为c，则内积运算可以分解为：

<o_r, q_r> = <o_r-c, q_r> + <c, q_r> = ||o_r-c||·||q_r||<o, q_r/||q_r||> + <c, q_r>

其中关键点在于：

1. ||o_r-c||可预先计算
2. ||q_r||和<c, q_r>可共享计算
3. 使用RaBitQ估计归一化后的内积项<o, q_r/||q_r||>

实现优化

在具体实现中，我们采用了以下优化措施：

1. 对查询向量进行中心化归一化处理
2. 保留原始向量幅值信息用于最终得分计算
3. 调整聚类参数以适应MIP场景的特性

经过优化后，在k=10、最大重排数=50的设置下，系统实现了超过97%的召回率，显著提升了在MIP场景下的搜索质量。

经验总结

通过本次实践，我们获得了以下重要经验：

1. 距离度量方式的转换需要严谨的数学推导
2. 向量幅值信息在MIP场景中不可忽视
3. RaBitQ的核心方法具有很好的扩展性，经过适当调整可支持多种相似度度量
4. 针对特定场景的参数调优对最终效果至关重要

该方案已在实际应用中验证了其有效性，为处理多语言嵌入等复杂场景提供了可靠的技术支持。

RaBitQ [SIGMOD 2024] RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search 项目地址: https://gitcode.com/gh_mirrors/ra/RaBitQ