Efficient Multivector Retrieval with Token-Aware Clustering and Hierarchical Indexing

1. Planteamiento del Problema

Los modelos de recuperación de multivectores (por ejemplo, ColBERT) han logrado una efectividad de vanguardia en Recuperación de Información (RI) al codificar documentos como conjuntos de vectores a nivel de token y utilizar "interacción tardía" (MaxSim) para calcular la relevancia. Sin embargo, su despliegue se ve obstaculizado por:

• Altos Costos Computacionales y de Memoria: Almacenar y buscar miles de millones de vectores de tokens requiere recursos masivos.
• Agrupamiento Ineficiente: Las soluciones actuales dependen del agrupamiento $k$-means estándar para generar vectores centroides para la compresión y la recolección de candidatos.
  • Cuello de Botella de Escalabilidad: $k$-means escala mal con el tamaño del conjunto de datos y el número de clústeres ($k$), requiriendo a menudo decenas de horas o GPUs costosas para entrenar incluso conjuntos moderadamente grandes (por ejemplo, 262k centroides).
  • Desequilibrio de Tokens: $k$-means estándar trata todos los vectores por igual. En consecuencia, los tokens frecuentes (palabras vacías) dominan la función de pérdida y consumen la mayoría de los centroides, mientras que los tokens raros y altamente discriminativos están subrepresentados, degradando la calidad de la recuperación.

2. Metodología

Los autores proponen Tachiom, un sistema de recuperación que comprende dos innovaciones centrales: Agrupamiento Consciente de Tokens (Tac) y un Índice Jerárquico para una recuperación optimizada.

A. Agrupamiento Consciente de Tokens (Tac)

Tac aborda las limitaciones de $k$-means estándar explotando la propiedad estructural de la identidad del token en las incrustaciones de multivectores. En lugar de agrupar todos los vectores globalmente, Tac descompone el problema en subproblemas independientes por token.

• Estrategia de Asignación de Centroides: Tac asigna el presupuesto de centroides basándose en la frecuencia del token y la varianza semántica, en lugar de una distribución uniforme de vectores.
  • Manejo de la Cola: Los tokens extremadamente raros ($n < \mu$) tienen garantizado al menos 1 centroide; los tokens pequeños ($\mu \le n < \tau$) obtienen 2.
  • Puntuación Atenuada: Para los tokens activos, se calcula un peso $w_j$ como $\sqrt{n_j} \cdot s_j$, donde $n_j$ es la frecuencia y $s_j$ es la varianza semántica (varianza componente a componente). La raíz cuadrada atenúa la influencia de los tokens de alta frecuencia, evitando que monopolicen el presupuesto.
  • Acotación y Reconciliación: Los pisos y techos duros aseguran una representación mínima para los tokens raros y evitan la sobreasignación. El presupuesto final se reconcilia para coincidir con el objetivo $k$.
• Paralelización: Al resolver problemas independientes de $k_j$-means para cada tipo de token, Tac permite un paralelismo masivo, reduciendo drásticamente la complejidad computacional en comparación con $k$-means global.
• Compresión de Residuos: Después del agrupamiento, los residuos (vector - centroide) se normalizan y comprimen utilizando Cuantización de Producto (PQ).

B. Índice Jerárquico para Recuperación (Tachiom)

Tachiom aprovecha los centroides de alta resolución generados por Tac para optimizar el pipeline de recuperación en dos fases:

1. Fase de Recolección (Puntuación Solo con Centroides):
   • Se construye un grafo Hnsw sobre los centroides.
   • Cada centroide mantiene una lista invertida de identificadores de documentos.
   • Para una consulta, el sistema recorre el grafo Hnsw para encontrar los centroides más cercanos, recupera los documentos asociados y calcula una puntuación MaxSim aproximada utilizando solo interacciones de centroides.
   • Esto elude por completo los cálculos costosos a nivel de token durante la selección inicial de candidatos.
2. Fase de Refinamiento (MaxSim Optimizado con PQ):
   • Los mejores candidatos se refinan utilizando tanto centroides como residuos comprimidos con PQ.
   • Diseño Optimizado para Caché: Los autores introducen un diseño de memoria especializado para las tablas de distancia PQ. En lugar de accesos a memoria dispersos para cada token de consulta, las distancias se organizan en Macro-bloques (subespacios), Bloques de centroides y Micro-bloques de tokens de consulta. Esto permite recuperar las distancias para todos los tokens de consulta contra un centroide específico como un único bloque de memoria contiguo, mejorando significativamente la localidad de la caché y la utilización de SIMD.

3. Contribuciones Clave

1. Agrupamiento Consciente de Tokens (Tac): Un algoritmo novedoso que descompone el agrupamiento global en subproblemas por token. Garantiza teóricamente una aceleración con un límite inferior determinado por la relación entre el peso total del vocabulario y el peso máximo de un solo token.
2. Arquitectura Tachiom: Un sistema de recuperación que combina la recolección de centroides basada en Hnsw con un diseño de PQ optimizado para caché, adaptado para la interacción tardía, permitiendo una puntuación completa eficiente.
3. Validación Teórica y Empírica: El artículo proporciona un límite inferior teórico para la aceleración de Tac y demuestra que el método escala a millones de centroides mientras mantiene o mejora la efectividad de la recuperación.

4. Resultados Experimentales

Los experimentos se realizaron en Ms Marco-v1 (8.8M pasajes) y LoTTE (2.4M pasajes) utilizando incrustaciones de ColBERTv2.

• Eficiencia de Agrupamiento:
  • Tac logró hasta 247$\times$ más rápido en entrenamiento que $k$-means estándar (Faiss) y 84$\times$ más rápido que Faiss con aceleración MKL.
  • Tac agrupó con éxito 600M vectores en 262k clústeres en 8 minutos, mientras que los competidores alcanzaron un tiempo de espera de 24 horas antes de llegar a 131k–524k clústeres.
  • Tac escala sin problemas a más de 4 millones de centroides en ~100 minutos.
• Efectividad de Recuperación:
  • Tac logró un MRR@10 igual o superior al de $k$-means estándar con presupuestos de centroides fijos, demostrando que priorizar los tokens raros mejora la calidad de la aproximación.
• Velocidad de Recuperación:
  • Tachiom logró hasta 9.8$\times$ más rápido en recuperación de extremo a extremo en comparación con sistemas de vanguardia (Emvb, Warp, Igp).
  • Igualó la efectividad máxima de Emvb (que utiliza variantes de PQ supervisadas costosas) mientras usaba PQ estándar y logró una aceleración de 5.5$\times$ sobre Emvb.

5. Significado

Este trabajo cambia fundamentalmente el paradigma para la eficiencia de la recuperación de multivectores:

• Escalabilidad: Elimina el cuello de botella computacional del agrupamiento, permitiendo que los sistemas utilicen millones de centroides (granularidad más fina) que anteriormente eran inviables de entrenar.
• Compromiso Calidad-Eficiencia: Al alinear la asignación de centroides con los objetivos de recuperación (favoreciendo tokens raros discriminativos) en lugar del error de reconstrucción puro, logra una mejor efectividad con menos recursos.
• Despliegue Práctico: La combinación de entrenamiento rápido e inferencia ultra-rápida (mediante recolección Hnsw y PQ optimizada) hace viable la recuperación de multivectores de alto rendimiento para entornos de producción a gran escala.

Los autores han liberado su implementación en Rust de Tac y Tachiom en GitHub para garantizar la reproducibilidad.