Los modelos de embeddings multivectoriales, como ColBERT, ofrecen una gran precisión en la búsqueda, pero el número de vectores almacenados por documento hace que el costo del índice aumente rápidamente. Una técnica de agrupamiento, el token pooling, reduce este tamaño, y un entrenamiento adaptado permite, según la fuente, conservar mejor la calidad de los resultados con un mismo presupuesto de almacenamiento. Con pool_factor=4, un modelo entrenado normalmente conserva el 94,6 % de su nDCG@10 no comprimido en NanoBEIR, frente al 97,8 % para el modelo entrenado teniendo en cuenta el pooling.
Por qué los modelos multivectoriales almacenan tantos vectores
Un modelo de embeddings clásico resume un documento completo en un solo vector. Los modelos multivectoriales, por el contrario, conservan un conjunto de vectores contextuales, uno por token (una unidad de texto, a menudo una palabra o parte de una palabra), lo que permite que diferentes partes de una consulta coincidan con diferentes partes de un documento. La fuente toma el ejemplo de una consulta sobre un ordenador portátil de larga duración, que puede encontrar un pasaje sobre el dispositivo en una parte del documento y un pasaje sobre la autonomía de la batería en otra.
Este funcionamiento tiene un costo. Según el cálculo presentado, con una cuantificación de 4 bits por coordenada residual, 200 vectores de 128 dimensiones representan 12,8 KB por documento, es decir, 12,8 GB para un millón de documentos. Reducir cada documento a 50 vectores lleva este total a 3,2 GB, con dimensión y precisión idénticas.
El método se basa en una observación: muchos vectores de un mismo documento pueden estar próximos, ya que varios tokens describen a menudo nociones similares. HierarchicalTokenPooling, la biblioteca Sentence Transformers, agrupa estos vectores mediante clustering de Ward, que fusiona en cada etapa los dos grupos que provocan el menor aumento de la dispersión interna. Cada grupo se reemplaza luego por su media, y se normaliza. Con pool_factor=4, un documento conserva aproximadamente un cuarto de sus vectores iniciales, de modo que un documento más largo guarda proporcionalmente más vectores.
Entrenar el modelo para la compresión
Cuanto mayor es la compresión, más fácil es perder información útil. El entrenamiento que tiene en cuenta el pooling tiene precisamente como objetivo conservar la información importante cuando los vectores se fusionan. Cada lote de entrenamiento combina dos pérdidas: la búsqueda con los embeddings de documentos originales, y la búsqueda después de la compresión a aproximadamente un cuarto de los vectores. Estas dos pérdidas se promedian, y ambos objetivos actualizan el mismo codificador.
Dado que el clustering toma decisiones discretas, la fuente indica que el entrenamiento utiliza un estimador straight-through, que transmite los gradientes de los vectores agrupados hacia los tokens que los componen, sin diferenciar la asignación a los clusters. La clase MultiVectorTokenPoolingLoss encapsula una pérdida de búsqueda existente, aquí CachedMultiVectorMultipleNegativesRankingLoss, con pesos de 0,5 para cada variante. El modelo de partida es answerdotai/ModernBERT-base, entrenado en 120 000 tripletes de MS MARCO durante una época.
Resultados en NanoBEIR
Con pool_factor=4, el nDCG@10 pasa de 49,03 a 51,23 para el modelo entrenado con pooling, al tiempo que reduce en un 74,8 % el número de vectores de documentos en comparación con una búsqueda no comprimida. El modelo no comprimido también mejora, de 51,84 a 52,39. Para un promedio de 29,58 vectores por documento, el pooling adaptativo conserva el 97,3 % de la puntuación no comprimida del modelo entrenado, frente al 93,3 % para el modelo entrenado normalmente.
Para reducir el tamaño de un índice, la fuente sugiere primero la cuantificación de los embeddings, luego el token pooling, posiblemente con un ajuste fino que tenga en cuenta el pooling, y finalmente los embeddings Matryoshka para reducir la dimensión de cada vector. Estos enfoques pueden combinarse. Los trabajos en los que se basa el artículo se describen en el estudio Learn to Pool de Stefan Josef.
