很高兴分享我们的近期工作TACHIOM已集成到PyLate生态系统中！https://arxiv.org/pdf/2604.28142 (@SilvioMartinico, @fmnardini, @rventurini_ )

基于令牌感知聚类与分层索引的高效多向量检索

来源：https://arxiv.org/html/2604.28142

摘要。

多向量检索模型通过细粒度的令牌级表示达到了最先进的效果，但其部署会带来巨大的计算和内存成本。当前的解决方案——基于著名的κ-均值聚类算法——将相似向量分组，从而实现有效的压缩和高效的检索。然而，标准的κ-均值在聚类数量和数据集规模增大时扩展性差，并且在训练过程中偏向高频令牌，而低频但具有判别性的令牌则代表性不足。在这项工作中，我们介绍了Tachiom，一种利用令牌级结构来显著加速聚类和检索的多向量检索系统。通过在质心分配时考虑令牌的分布，Tachiom能够轻松扩展到数百万个质心，使得仅使用质心就能实现高精度的文档评分，避免了昂贵的令牌级计算。Tachiom将基于图的质心索引与优化的乘积量化布局相结合，以实现高效的最终评分。在Ms Marco-v1和LoTTE上的实验表明，Tachiom的聚类速度比κ-均值快高达247倍，检索速度比最先进的系统快高达9.8倍，同时保持相当甚至更优的效果。

多向量检索，后期交互，聚类，效率。

1. 引言

基于Transformer的语言模型（Devlin等，2019; Vaswani等，2017; Karpukhin等，2020; Formal等，2021; Lassance等，2024）从根本上重塑了信息检索领域，将范式从词汇检索转向基于表示的检索。在此背景下，多向量模型，如ColBERT（Khattab和Zaharia，2020; Santhanam等，2022b），已成为效果的金标准。通过将文档编码为令牌级向量集合，并通过后期交互（即允许每个查询令牌与每个文档令牌交互）计算相关性，这些模型捕捉到了单向量稠密检索器经常忽略的细粒度语义细微差别。然而，多向量模型的实际部署仍受到其计算需求的严重限制。存储和搜索大量的令牌向量会产生巨大的内存占用和计算开销。为缓解这一问题，最先进的解决方案采用了“收集再精炼”策略，并结合令牌量化（Engels等，2023; Santhanam等，2022a; Nardini等，2024; Scheerer等，2025; Bian等，2025）。在收集阶段，对令牌向量进行轻量级搜索，以确定有希望的候选者；在精炼阶段，仅对选定的候选者计算完整的MaxSim分数。两个阶段都关键依赖于一个粗量化器——一组近似令牌嵌入空间的质心向量。在收集过程中，质心作为定位相关令牌的高效代理。对于压缩，每个令牌向量被分解为质心分配和残差向量，然后使用更精确的向量量化技术（例如乘积量化（Jégou等，2011））对残差进行压缩。这种两层方案显著改善了近似：质心分配以最低成本区分远距离向量，而残差压缩则捕捉局部细微差别。

ColBERTv2（Santhanam等，2022b）开创了这种基于质心的方法，压缩残差（向量与质心之间的元素差）而不是向量本身。Plaid（Santhanam等，2022a）通过质心交互和基于倒排文件的过滤（在聚集的令牌嵌入上）扩展了这一框架，以实现高效检索。后续工作进一步完善了这一框架：Emvb（Nardini等，2024）引入了位向量预过滤以加速候选选择，并优化了PQ变体（Ge等，2014; Fang等，2022）以改进残差压缩，相比Plaid实现了显著加速。Warp（Scheerer等，2025）将Plaid的质心交互与Xtr的轻量级架构（Lee等，2023）相结合；作者表明他们的方法也能推广到ColBERTv2嵌入。Igp（Bian等，2025）提出在质心集上构建图索引，以避免穷尽式质心评分。

在所有上述方法中，检索质量和压缩效果都直接与质心数量成正比：更精细的质心可以实现更精确的近似和更具选择性的收集。然而，标准的κ-均值聚类——质心选择的主要方法——造成了严重的瓶颈：它在数据集规模和聚类数量上扩展性差，而且关键的是，它忽略了多向量嵌入的令牌结构。在多向量集合上训练即使中等大小的质心集（例如262K）也需要在多核CPU或昂贵的GPU资源上花费数十小时，从根本上限制了可实现的粒度。在本文中，我们通过利用多向量嵌入的一个关键结构属性——令牌身份——来解决这个瓶颈。与通用的向量集不同，多向量集合在令牌类型和嵌入分布之间表现出强相关性。常见令牌（如停用词）在训练期间主导了κ-均值的损失，消耗了大部分质心，而低频的领域特定令牌——通常对检索更具判别性——却得到极少的表示。我们介绍了Tachiom，一个在聚类和检索阶段都利用这种结构的多向量检索系统。Tachiom的核心是采用了令牌感知聚类，它基于令牌频率和语义方差显式分配质心。通过将聚类工作负载划分到令牌组中，Tachiom相比标准κ-均值实现了显著更快的质心计算。这种效率使得能够扩展到数百万个质心，远超先前的工作。利用这些高分辨率的质心，Tachiom通过基于HNSW图索引（Malkov和Yashunin，2020）的质心级交互进行收集，从而完全绕过昂贵的令牌级计算。对于精炼，一种针对后期交互优化的缓存友好PQ布局实现了高效的最终评分。

总之，这项工作的新颖贡献如下：

• • 我们提出了Tac，一种利用令牌身份将全局聚类问题分解为独立的每令牌子问题的聚类算法，提高了可扩展性，同时产生了与检索目标更一致的质心。我们进一步推导了Tac相对于κ-均值计算加速的理论下界。
• • 我们介绍了Tachiom，一种检索架构，它利用Tac实现的聚类粒度，通过基于图的搜索进行仅质心收集，并使用专为后期交互定制的PQ进行高效全评分。
• • 我们在Ms Marco-v1和LoTTE上进行了全面的实验，证明了我们提出的架构提供了优越的效率-效果权衡。我们的结果显示，Tac的训练速度比标准κ-均值快高达247倍，而Tachiom的端到端检索速度比最先进的系统快高达9.8倍。为了促进可重复性，我们在GitHub上以Rust语言发布了Tac和Tachiom的实现：https://github.com/TusKANNy/tachiom。

2. 方法论

2.1. 令牌感知聚类

多向量模型如ColBERT（Khattab和Zaharia，2020）将每个文档表示为一组上下文相关的令牌嵌入。最先进的检索方法（Santhanam等，2022a; 2022b; Nardini等，2024; Scheerer等，2025; Bian等，2025; Fang等，2022）依赖于质心来近似这些大量的令牌向量集，通常使用κ-均值聚类来选择代表性质心。然而，标准的κ-均值将令牌嵌入视为通用向量集，丢弃了令牌身份提供的结构信息。令牌感知聚类（Tac）显式利用了多向量表示的令牌结构。Tac有两个目的：（i）通过将工作负载划分为独立的每令牌子问题，显著加速聚类，这对于ColBERT类模型产生的大规模向量集合至关重要；（ii）将质心分配向稀有令牌倾斜，这些令牌尽管对检索准确度重要，但在标准κ-均值中代表性不足。

κ-均值旨在找到一个聚类集C = {C₁, …, C_κ}，对数据集进行划分并最小化簇内平方和（WCSS）或惯性，即：

∑{i=1}^{κ} ∑{t ∈ C_i} ‖t - c_i‖²,

其中c_i表示聚类C_i的质心。这个目标对所有向量一视同仁，但多向量集合表现出严重的频率不平衡。常见令牌（如停用词）可能出现数百万次，而领域特定术语的频率则低得多（几百或几千次出现）。通过分析两个ColBERTv2多向量集合，我们观察到前100个最频繁的令牌占了所有向量的40%以上，具体来说，在Ms Marco-v1上为41%，在LoTTE-pooled上为45%。因此，WCSS损失由高频令牌主导，它们消耗了大部分质心，尽管对检索质量的贡献很小。与此同时，稀有令牌——通常携带最具判别性的语义信号（Sparck Jones, 1988）——却得到极少的表示，恰恰在最关键的地方降低了近似质量。

借鉴经典信息检索文献中长期以来使用逆文档频率（IDF）来偏袒稀有、判别性术语的做法（Robertson, 2004; Salton等, 1975），我们提出Tac来解决聚类中的这种不平衡。Tac通过根据令牌的频率和语义方差显式加权来将质心分配从统一的WCSS最小化解耦，以偏袒稀有令牌。Tac首先为每个令牌类型计算一个权重，表示其占总质心预算的份额，然后为每个令牌类型执行独立的κ-均值聚类。

Tac通过一个四阶段流程分配质心预算：（1）尾部处理隔离非常稀有的令牌，防止其完全边缘化；（2）阻尼评分计算剩余令牌的分配权重；（3）边界约束强制执行严格限制，以避免分配不足或过度；（4）预算调整将最终分配精确调整为匹配全局质心预算。

阶段1：尾部处理。我们根据频率n将令牌类型分为三类：微型令牌（n < μ）各获得1个质心，小型令牌（μ ≤ n < τ）各获得2个质心，活跃令牌（n ≥ τ）根据其频率和语义方差动态分配质心。这种固定分配防止了极稀有令牌的完全边缘化。

阶段2：阻尼评分。对于活跃令牌，我们计算一个衡量语义方差的散布度量s_j；为了计算效率，我们采用分量方差，即

s_j ≔ (1 / n_j) ∑{i=1}^{n_j} ‖t{j,i} - t̄_j‖²，

其中n_j是令牌j的频率，t_{j,i}是令牌j的第i次出现，t̄_j是其平均嵌入。然后我们计算一个阻尼权重，w_j = √n_j · s_j。平方根阻尼对高频令牌应用递减收益，防止它们垄断质心预算；散布系数s_j偏袒语义方差高的令牌。令牌j的质心数量κ_j与其权重成比例分配，

κ_j = ⌊ (w_j / ∑_{i=1}^{N_T} w_i) · B ⌋，

其中N_T是我们集合中不同令牌的数量，B是尾部处理后的剩余质心预算。

阶段3：边界约束。我们对所有活跃令牌强制执行一个硬下限ε，使得κ_j ≥ ε，以保证最小表示。此外，我们施加一个上限以防止过度分配：每个令牌必须满足 n_j / κ_j ≥ θ，确保平均每个质心至少有θ个向量。这种双重边界防止了低预算情况下稀有令牌的代表性不足，以及高预算情况下对已经近似良好的令牌的过度分配。

阶段4：预算调整。我们重新分配盈余或赤字质心，以精确匹配目标预算κ。一旦质心分配完成，我们为每个令牌j使用其分配的κ_j个质心执行独立的κ_j-均值聚类。

总之，Tac在两个方面与标准κ-均值有根本不同：1）它优先考虑检索质量而非纯重构误差：通过基于令牌频率和语义方差而不是统一的WCSS贡献来分配质心，Tac确保稀有、判别性令牌在分配的质心上获得充分的表示；2）它不是对包含N个向量和κ个质心的单个全局目标进行优化，而是为每个令牌解决独立的子问题，将令牌j的n_j个向量聚类为κ_j个质心。这种分解显著降低了每次迭代的计算成本，并实现了有效的并行化，因为每个令牌可以独立处理，而标准κ-均值需要对整个数据集进行共享访问。此外，阻尼质心分配减轻了瓶颈。