感知池化训练：在保持检索质量的同时将嵌入存储需求降低 4 倍

像 ColBERT 这样的多向量嵌入模型在搜索方面功能强大，但由于文档中的每个标记（token）都会生成自己的向量，因此存在巨大的存储开销。Hugging Face 的一篇博客文章详细介绍了一种称为“感知池化训练”（token pooling-aware training）的新技术，通过在训练过程中压缩相似的标记向量，使开发人员能够在不牺牲检索准确性的情况下显著降低存储成本。

发生了什么

根据 Hugging Face 博客的介绍，多向量模型为每个查询和文档标记分配一个向量，从而通过 MaxSim 评分实现细粒度匹配。然而，这种方法非常消耗存储空间；例如，以 4 位精度存储每个文档 200 个 128 维的向量，对于一百万个文档的语料库，需要 12.8 GB 的空间。所提出的解决方案使用 HierarchicalTokenPooling，通过 Ward 分层聚类对相似的标记嵌入进行分组。通过合并这些组，模型减少了每个文档的向量数量。

关键创新在于训练编码器以感知这种池化过程。模型不再仅仅在完整表示上进行训练，而是使用包含未池化和池化检索目标的组合损失函数进行训练。这教会了编码器如何排列标记嵌入，以便有用的信息在压缩过程中得以保留。博客报告称，在 pool_factor 为 4 的情况下，普通训练的模型保留了 94.6% 的未池化 NanoBEIR nDCG@10 分数。通过感知池化训练，在相同的 4 倍存储缩减下，该保留率提高到了 97.8%。该文章提供了一个使用 Sentence Transformers 的可复现示例，演示了如何通过 MultiVectorTokenPoolingLoss 包装器实现这一点。

为什么这对构建搜索功能的开发人员很重要

对于构建内部搜索工具或 RAG 应用程序的小型团队来说，向量数据库的存储成本是一个日益严重的痛点。多向量模型比单向量模型提供更高的精度，但其成本往往使其在大规模语料库中难以应用。这项技术弥补了这一差距。通过将感知池化训练集成到您的流水线中，您可以以极低的基础设施成本部署高精度搜索。这对于使用 Sentence Transformers 等开源框架的团队尤为重要，因为该方法可以直接利用现有工具实现。它将权衡从“精度与成本”转变为更易于管理的“轻微精度下降与显著成本节约”，使得无法负担企业级向量存储的中型公司也能实现高级检索。

本周建议

1. 审计您当前的向量存储成本，并计算如果能将文档向量数量减少 50% 或 75% 可以节省的潜在成本。
2. 克隆与该博客文章相关的 Hugging Face 存储库，并在本地运行提供的训练脚本，以了解 MultiVectorTokenPoolingLoss 的实现细节。
3. 在现有语料库的一小部分上测试池化机制，以在进行全面的重新训练周期之前测量实际的质量下降（nDCG）。

来源

• Token Pooling-Aware Training for Multi-Vector Embedding Models — Hugging Face Blog