Si haces búsquedas sobre tus propios documentos con un modelo de estilo ColBERT, el índice probablemente sea más grande de lo que querrías. Estos modelos multivector guardan un vector por cada token en lugar de uno por documento, lo que los hace precisos y caros de almacenar. Una nueva entrada en el blog de Hugging Face muestra cómo Sentence Transformers puede ya entrenar un modelo para que siga siendo preciso después de comprimir esos vectores, y recuperar así la mayor parte de la calidad que suele perderse con la compresión.
El problema de almacenamiento es concreto. Con 4 bits por coordenada, 200 vectores de 128 dimensiones ocupan 12,8 KB por documento, o 12,8 GB para un millón de documentos. Reducir cada documento a 50 vectores lo deja en 3,2 GB. HierarchicalTokenPooling, de la biblioteca, fusiona vectores de tokens parecidos mediante agrupamiento de Ward, ya sea por un factor fijo o por un umbral de distancia que puede calibrarse para una media objetivo, y no toca los vectores de la consulta. El agrupamiento funciona con cualquier modelo multivector, pero uno que nunca se entrenó así puede perder distinciones importantes.
La novedad es un envoltorio de la función de pérdida, MultiVectorTokenPoolingLoss, que entrena con dos versiones de cada lote: los vectores completos del documento y una versión reducida a aproximadamente una cuarta parte. La idea recuerda a los embeddings Matryoshka, que entrenan prefijos más pequeños para que sigan siendo útiles. En el experimento del autor, dos modelos ModernBERT-base se entrenaron una época con las mismas 120.000 tripletas de MS MARCO y solo se diferenciaban en la pérdida. Con un factor de agrupamiento de 4, el modelo normal conservó el 94,6% de su puntuación NanoBEIR sin agrupar; el entrenado para ello, el 97,8%, con un nDCG@10 que subió de 49,03 a 51,23 y un 74,8% menos de vectores de documento. La calidad sin agrupar también mejoró algo. Un kernel de GPU para el agrupamiento sería 30-150x más rápido que SciPy.
Para los equipos que usan PyLate, la biblioteca de LightOn para recuperación de estilo ColBERT construida sobre Sentence Transformers, esto es directamente relevante: el trabajo parte del estudio de la propia LightOn sobre regularización del agrupamiento, y la receta de entrenamiento cabe en una sola A10G mediante Hugging Face Jobs en menos de cuatro horas. El consejo práctico es combinar cuantización, agrupamiento y truncado Matryoshka, y medir tamaño de índice y calidad con tu propio corpus.
Las salvedades son del propio autor. Es un solo experimento con un solo modelo base y sin búsqueda de hiperparámetros, evaluado en NanoBEIR, un conjunto de pruebas pequeño. La mejora es más clara con el factor de agrupamiento para el que se entrenó el modelo, aunque se traslada a otros ajustes. Los resultados con documentos largos o de un dominio concreto pueden ser distintos.