LEMUR与均值中心化:txtai中后期交互检索的技术突破解析

1 阅读

在现代信息检索系统中,后期交互模型保持了一个有用的细节层次:它们不是立即将查询或文档折叠成一个嵌入,而是为每个token保留一个向量。MaxSim分数比较每个查询token与文档token,保留每个查询token的最强匹配,并将这些匹配相加。权衡在于操作方面。多向量表示不能像传统密集检索管道使用的固定向量索引那样自然地适应。

txtai中的LEMUR(Learned Multi-Vector Retrieval)关闭了这一差距。它学习多向量表示的固定维度编码,因此后期交互模型可以使用与标准嵌入模型相同类型的向量索引。第二个变化是可配置的均值中心化,解决了在测试LEMUR与现代LateOn模型时发现的不同故障模式:token向量可能如此各向异性,以至于固定维度编码器几乎没有有用的变异来保留。

LEMUR代表学习多向量检索。在txtai中,其工件包含特征编码器、输出归一化统计和token向量样本。在推理时,查询成为其学习token特征的总和。文档成为存储样本上的普通最小二乘权重集。它们的固定向量内积近似原始后期交互分数。

训练管道从语料库文本开始。它使用模型的数据编码器对每个目标文档进行编码,构建标准化的MaxSim目标,学习特征映射,并将推理工件保存为config.json加上model.safetensors。该工件是语料库特定的,必须在txtai嵌入索引加载之前进行训练。

一个训练细节比预期更重要:用于学习特征映射的分布。后期交互模型可以以不同的方式编码查询和文档。在nfcorpus上,从数据编码器token向量学习产生了一个训练MLP,甚至低于未经训练的ELM回退。从查询编码器token向量学习恢复了大部分改进,基于验证的选择增加了较小的收益。

在nfcorpus配置中,从数据编码器token学习的MLP得分为0.15870,未经训练的ELM得分为0.19187,从查询编码器token学习的MLP得分为0.24868,带有验证选择的查询token MLP得分为0.25534。这个四行消融研究显示,学习分布占从数据token运行到最终运行测量增益的93%,而验证选择占剩余的7%。

这就是为什么LemurTrainer默认将learncategory设置为"query"的原因,同时仍允许调用者选择"data"或提供单独的学习可迭代对象。训练的MLP是面向质量的路径。txtai还要求训练选择明确:epochs=100选择记录的MLP设置,而epochs=0选择确定性随机ELM特征作为低成本回退。

对齐矩阵使用colbert-ir/colbertv2.0构建,配备NVIDIA GeForce RTX 4080 SUPER,torch 2.13.0+cu130和精确Faiss IDMap,Flat搜索。该表比较了默认的10,240维MUVERA编码、减少到相同2,048维预算的MUVERA和训练的2,048维LEMUR编码。

在nfcorpus数据集上,MUVERA 10,240维得分为0.23544,MUVERA 2,048维得分为0.16299,LEMUR MLP 2,048维得分为0.25524。在scifact数据集上,相应得分分别为0.50021、0.36757和0.54910。在arguana数据集上,得分分别为0.34614、0.26280和0.42556。

在相等的向量大小下,LEMUR相对于匹配大小的MUVERA在nfcorpus上提高了56.6%,在scifact上提高了49.4%,在arguana上提高了61.9%。相对于MUVERA的完整默认向量,增益分别为8.4%、9.8%和22.9%。测量的LEMUR索引使用了默认MUVERA索引五分之一的存储空间,因为向量宽度和记录的索引大小都从10,240缩放到2,048维。

这些数字证明了特定配置的证据,而不是一般排名声明。五个数据集中的三个在请求集中进行了测量;fiqa和scidocs未完成。矩阵涵盖一个模型和一台机器。它还强制执行精确搜索,这对部署很重要。

txtai的Faiss后端使用精确搜索直到5,000行,超过此阈值则切换到IVF索引。在测量的scifact运行中,默认IVF相对于精确搜索将LEMUR NDCG@10降低了43%,而MUVERA降低了25%。对于更大的LEMUR语料库,在精确搜索实用时将faiss.components固定到IDMap,Flat,或者为语料库调整IVF而不是假设默认索引将保留精确搜索结果。

第一个LEMUR基准测试使用ColBERTv2。在转向lightonai/LateOn时,原始token向量显示了不同的几何形状。在5,000对样本中,平均成对token余弦为0.9508,MaxSim传播为0.0559。文章Regularizing ColBERT models to fix efficient ANN methods描述了各向异性和这如何成为MUVERA评分不佳的罪魁祸首。经过集合平均减法和重新归一化后,这些值变为0.0033和0.4772。这本身不是检索指标,但它表明中心化暴露了更多方向变异供固定维度编码器使用。

检索测量证实了该信号的有用部分,但它们也排除了通用开关。正如David Mezzetti所说,中心化"似乎伤害了一些模型,同时帮助了其他模型。"他要求一个可配置的center参数,并建议当加载的模型有多于一个线性层时默认启用它。

基础LateOn模型加载了五个线性层,而ColBERTv2加载了一个。中心化强烈帮助了LateOn矩阵。比较了三个范围:存储的集合平均、每个文档的平均和当前模型批次的平均。这些CUDA运行使用torch 2.13.0+cu130和每次基准编码批处理32个文本。批量范围在下面所有四个基础LateOn单元格中都是最强的简单范围。

在nfcorpus数据集上,LEMUR 2,048维的关闭状态为0.00000,集合状态为0.31473,批量状态为0.33309。MUVERA 10,240维的相应值分别为0.03639、0.13369和0.18382。在scifact数据集上,LEMUR 2,048维的关闭状态为0.04985,集合状态为0.68624,批量状态为0.69016。MUVERA 10,240维的相应值分别为0.00269、0.29293和0.37252。

每个文档中心化是混合的:它在两个nfcorpus单元格中落后于集合平均,在scifact LEMUR中基本平局,但在scifact MUVERA中有所改进。批量范围避免了语料库传递和存储平均,同时在此矩阵中领先集合结果。David选择了这种权衡:当他明确启用中心化时,他会"选择批量作为默认值"。

一层检查就是为什么自动规则被描述为保守而不是通用的原因。在ColBERTv2上,集合中心化对LEMUR略有帮助或持平,但在两个测量数据集上使MUVERA倒退。仅层数无法解释每个结果;编码器选择也很重要。因此,发布的默认设置在模型加载多个torch.nn.Linear层时启用批量中心化,而零层和一层模型保留其先前输出。调用者可以在任一方向覆盖该默认设置。

中心化发生在token向量归一化之后和MUVERA或LEMUR之前,然后是另一个归一化。center: true选择批量范围。字典可以选择document、batch或collection;集合范围接受内联平均或包含center.mean的Safetensors文件。center: false禁用操作。

最新的已发布txtai版本仍然是v9.12.0,而两个合并的更改针对v9.13.0。对于预发布源代码,请在隔离环境中使用pip install "git+https://github.com/neuml/txtai.git@master"。

这是一个轻量级ELM示例。重要的是一致性:如果索引将中心化token向量,则将相同的vectors设置传递给训练器,以便LEMUR工件适合该表示。

from txtai.pipeline import LemurTrainer

corpus = [ "Late interaction compares token embeddings.", "Dense indexes search fixed dimensional vectors.", "Mean centering changes token-vector geometry.", ]

Description of image

vectors = {"center": True}

LemurTrainer()( "neuml/colbert-bert-tiny", corpus, "lemur-model", gpu=False, vectors=vectors, epochs=0, )

通过嵌入向量配置加载工件。此示例固定精确Faiss搜索,使其索引行为明确。

from txtai import Embeddings

embeddings = Embeddings( { "path": "neuml/colbert-bert-tiny", "content": True, "vectors": { "center": True, "lemur": {"path": "lemur-model"}, }, "faiss": {"components": "IDMap,Flat"}, } )

embeddings.index(corpus) print(embeddings.search("dense indexes"))

对于面向质量的工件,用记录的MLP设置替换ELM选择并使用保留的验证拆分。工件仍然是一个单独的训练优先步骤;之后,txtai可以从本地目录或Hugging Face Hub路径加载它。

后续PR将训练整合和进度反馈移入LemurTrainer。它不会改变此处显示的加载时间API。

txtai实现基于Elias Jääsaari、Ville Hyvönen和Teemu Roos的LEMUR(ICML'26论文)。David通过#1164和#1168的审查线程塑造了集成,包括配置边界、批量默认和存储平均的Safetensors路径。他还标记了Elias Jääsaari、Raphael Sourty和Antoine Chaffin,因为这些更改建立在他们的工作基础上。

实际结果是一对控件而不是一个承诺,即一种设置在所有地方获胜。LEMUR提供后期交互的紧凑学习固定向量。中心化提供了一种显式方法来修复模型需要时的各向异性token几何。精确搜索结果令人鼓舞,但更广泛的模型和数据集覆盖,以及调整的近似索引测量,仍然是下一个要收集的证据。

LEMUR技术的核心优势在于它能够将复杂的多向量表示压缩为固定维度的编码,同时保持检索质量。这种方法解决了传统密集检索系统在处理后期交互模型时面临的主要挑战。通过学习token特征的映射关系,LEMUR能够在保持高精度的同时显著减少存储需求和计算复杂度。

均值中心化技术的引入进一步优化了系统的性能表现。当token向量表现出强烈的各向异性特征时,中心化处理能够有效地调整向量的空间分布,使得固定维度编码器能够更好地捕捉和利用有用的信息变化。这种几何变换对于提高检索系统的鲁棒性具有重要意义。

在实际部署中,开发者需要根据具体的应用场景和性能要求来选择合适的配置参数。批量中心化通常是最优选择,因为它避免了额外的语料库遍历和存储开销,同时能够提供良好的性能表现。对于不同的模型架构,需要仔细评估是否启用中心化功能,以确保获得最佳的检索效果。

未来的发展方向包括扩展对更多模型类型和数据集的支持,优化近似索引算法的性能,以及探索更高效的训练策略。随着技术的不断演进,LEMUR和均值中心化方法有望在更大规模的信息检索应用中发挥重要作用,为用户提供更加精准和高效的服务体验。