MUVERA:让多向量检索像单向量检索一样快速
在当今数字化信息爆炸的时代,信息检索技术已经成为我们日常生活中不可或缺的一部分。从搜索引擎到推荐系统,从自然语言处理到数据挖掘,信息检索技术的高效与否直接关系到我们获取信息的速度和准确性。随着深度学习技术的发展,神经嵌入模型逐渐成为现代信息检索系统的核心组件。这些模型能够将数据点转换为单个向量嵌入,使得语义相似的数据点在向量空间中也具有相似的数学表示,从而可以通过高效的内积最大值搜索(MIPS)算法进行检索。
神经嵌入模型与多向量检索的挑战
神经嵌入模型的出现极大地推动了信息检索技术的发展。它们通过将数据点映射到单个向量嵌入中,使得语义相似的数据点在向量空间中也具有相似的数学表示,从而实现了高效的信息检索。然而,近年来,多向量模型(如 ColBERT)的出现为信息检索任务带来了显著的性能提升。
多向量模型与单向量模型不同,它为每个数据点生成一组向量嵌入,而不是单一的向量嵌入。例如,ColBERT 模型为每个查询或文档中的每个标记生成一个嵌入向量,然后使用 Chamfer 相似性(也称为 MaxSim 操作)来计算查询和文档之间的相似性。Chamfer 相似性通过计算查询向量集合和文档向量集合之间的最大内积之和来衡量它们的相似性。
尽管多向量模型在准确性上具有优势,但它们也带来了显著的计算挑战。首先,每个标记生成一个嵌入向量,这使得数据集中嵌入的数量呈数量级增加。其次,Chamfer 相似性的非线性特性使得多向量检索的计算成本远高于单向量检索。此外,目前缺乏针对多向量检索的高效优化系统,这使得多向量检索在实际应用中面临巨大的挑战。
MUVERA:一种高效的多向量检索算法
为了应对多向量检索的挑战,我们介绍了 MUVERA(Multi-Vector Retrieval via Fixed Dimensional Encodings),一种新颖的多向量检索算法。MUVERA 通过将多向量检索问题简化为单向量最大内积搜索(MIPS)问题,从而显著提高了多向量检索的效率。
MUVERA 的核心思想是将查询和文档的多向量集合转换为固定维度的编码(FDEs)。这些 FDEs 是单个向量,它们的内积可以近似多向量集合之间的相似性。通过这种方式,MUVERA 能够利用现有的单向量 MIPS 算法进行高效检索,然后在重新排序阶段使用原始的多向量相似性进行优化,从而实现高效且准确的多向量检索。
FDEs 的生成过程
FDEs 的生成过程涉及以下几个关键步骤:
-
空间划分:使用 Locality Sensitive Hashing(LSH)技术(如 SimHash)将向量空间划分为多个区域。每个区域中的向量在空间上更接近彼此。 -
向量聚合:对于查询和文档中的每个向量集合,MUVERA 分别在每个空间区域内进行向量求和(对于查询)或平均(对于文档)。 -
重复执行:为了提高准确性,MUVERA 会多次重复上述过程,每次使用不同的随机种子进行空间划分和向量聚合。 -
最终投影:为了进一步降低维度,MUVERA 可以对生成的 FDEs 应用最终投影,例如使用 AMS Sketch 技术。
MUVERA 的理论保证
MUVERA 提供了强大的理论保证,证明了 FDEs 能够以高概率给出多向量相似性的 ε-近似解。具体来说,MUVERA 的 FDEs 在内积上提供了一种单边误差估计器,这意味着它们不会高估真实的 Chamfer 相似性。通过合理设置参数(如重复次数、空间划分的粒度等),MUVERA 能够在理论和实践中都提供高质量的近似。
实验结果
我们在 BEIR 基准测试的多个信息检索数据集上对 MUVERA 进行了评估。实验结果表明,MUVERA 在召回率和延迟方面均优于现有的最先进的多向量检索方法 PLAID。
-
✦ 召回率提升:MUVERA 在平均召回率上比 PLAID 高出 10%,并且在多个数据集上表现更为稳定。 -
✦ 延迟降低:MUVERA 的平均延迟比 PLAID 低 90%,在某些数据集上甚至降低了 5.7 倍。 -
✦ 压缩效果:通过使用乘积量化(PQ)技术,MUVERA 的 FDEs 可以被压缩 32 倍,同时对检索质量的影响微乎其微。
MUVERA 的实际应用与优势
MUVERA 的优势不仅体现在理论和实验结果上,还在于其在实际应用中的广泛适用性和高效性。以下是一些具体的应用场景和优势:
搜索表现在不同数据集上的一致性
MUVERA 在多个 BEIR 数据集上都展现出了稳定且优越的性能。例如,在 MS MARCO 数据集上,MUVERA 与 PLAID 的召回率几乎相当(相差不到 0.4%),同时延迟显著降低。在其他数据集如 HotpotQA 上,MUVERA 的召回率比 PLAID 高出 1.56 倍,同时延迟更低。这种一致性使得 MUVERA 成为一个可靠的多向量检索解决方案,适用于各种不同类型和规模的数据集。
高效的内存使用与压缩技术
MUVERA 通过使用乘积量化(PQ)技术,能够显著降低 FDEs 的内存占用。例如,使用 PQ-256-8 方案,MUVERA 可以将存储每个 FDE 的内存从 10240 维的浮点数压缩为 1280 字节,压缩比达到 32 倍。这种高效的压缩技术不仅减少了存储需求,还提高了检索速度,因为更小的内存占用意味着更快的数据访问速度。
简化的调优过程
与复杂的多阶段检索管道相比,MUVERA 的设计更为简洁,调优过程也更为简单。在我们的实验中,MUVERA 使用相同的 10240 维 FDEs 和 PQ-256-8 压缩方案,在所有测试数据集上都取得了优异的性能。这表明 MUVERA 对参数的选择并不敏感,可以在不同的数据集上快速部署,而无需进行大量的参数调整。
支持流式数据处理
由于 MUVERA 的 FDE 生成过程是数据无关的(data-oblivious),它对数据分布的变化具有很强的鲁棒性,并且非常适合流式数据处理场景。这意味着即使数据的统计特性随时间发生变化,MUVERA 仍然能够保持稳定的性能,无需重新训练或调整模型。
Python 实现指南
为了使 MUVERA 的 FDE 算法更加易于访问,我们提供了一个 Python 实现,该实现与原始的 C++ 实现保持完全一致。以下是 Python 实现的一些关键组件和使用示例。
配置类
Python 实现中定义了几个配置类,用于指定 FDE 生成的各种参数:
class EncodingType(Enum):
DEFAULT_SUM = 0 # 对于查询:在每个分区中求和向量
AVERAGE = 1 # 对于文档:在每个分区中平均向量
class ProjectionType(Enum):
DEFAULT_IDENTITY = 0 # 不进行降维
AMS_SKETCH = 1 # 使用 AMS Sketch 进行降维
@dataclass
class FixedDimensionalEncodingConfig:
dimension: int = 128 # 原始向量维度
num_repetitions: int = 10 # 独立运行次数
num_simhash_projections: int = 6 # 控制分区粒度
seed: int = 42 # 随机种子
encoding_type: EncodingType = DEFAULT_SUM
projection_type: ProjectionType = DEFAULT_IDENTITY
projection_dimension: Optional[int] = None
fill_empty_partitions: bool = False
final_projection_dimension: Optional[int] = None
核心算法
核心算法 _generate_fde_internal()
实现了 FDE 生成的主要逻辑:
def _generate_fde_internal(point_cloud: np.ndarray, config: FixedDimensionalEncodingConfig):
# 步骤 1:验证输入(与 C++ 参数验证匹配)
# 步骤 2:计算维度
# 步骤 3:对于每个重复:
# - 应用 SimHash 进行空间划分
# - 应用可选的降维
# - 按分区聚合向量
# - 对于文档 FDE 应用平均操作
# 步骤 4:可选的最终投影
使用示例
以下是一个基本的使用示例,展示了如何生成查询和文档的 FDE 并计算它们的相似性:
import numpy as np
from fde_generator import FixedDimensionalEncodingConfig, generate_query_fde, generate_document_fde
# 1. 创建配置
config = FixedDimensionalEncodingConfig(
dimension=128, # 向量维度
num_repetitions=10, # 独立划分次数
num_simhash_projections=6, # 创建 2^6 = 64 个分区
seed=42
)
# 2. 准备数据
# 查询:32 个 128 维向量
query_vectors = np.random.randn(32, 128).astype(np.float32)
# 文档:80 个 128 维向量
doc_vectors = np.random.randn(80, 128).astype(np.float32)
# 3. 生成 FDEs
query_fde = generate_query_fde(query_vectors, config)
doc_fde = generate_document_fde(doc_vectors, config)
# 4. 计算相似性(近似 Chamfer 相似性)
similarity_score = np.dot(query_fde, doc_fde)
print(f"相似性分数:{similarity_score}")
高级用法
您还可以通过指定投影类型和最终投影维度来应用降维:
from fde_generator import ProjectionType, replace
# 使用 AMS Sketch 进行内部投影
config_with_projection = replace(
config,
projection_type=ProjectionType.AMS_SKETCH,
projection_dimension=16 # 从 128 维降维到 16 维
)
# 使用 Count Sketch 进行最终投影
config_with_final_projection = replace(
config,
final_projection_dimension=1024 # 最终 FDE 将是 1024 维
)
结论
MUVERA 作为一种新颖的多向量检索算法,通过将多向量检索问题简化为单向量 MIPS 问题,有效地解决了多向量检索的效率挑战。MUVERA 的 FDEs 不仅具有强大的理论保证,能够在内积上提供高质量的多向量相似性近似,而且在实际实验中也展现出了卓越的性能。与现有的最先进的方法相比,MUVERA 在召回率和延迟方面均有显著提升,同时通过乘积量化技术实现了高效的内存使用。
MUVERA 的 Python 实现进一步降低了其使用门槛,使得研究人员和开发者能够更轻松地将这一先进的多向量检索技术应用于实际项目中。无论是在搜索引擎、推荐系统还是自然语言处理领域,MUVERA 都有潜力显著提升信息检索的效率和准确性,为用户提供购买有价值的技术解决方案。