ColBERTSaR: Sparsified ColBERT Index via Product Quantization

Resumen Técnico: ColBERTSaR: Índice ColBERT esparcido mediante cuantización de producto

Planteamiento del problema
ColBERT es una arquitectura de recuperación densa de vectores múltiples altamente efectiva que modela documentos utilizando incrustaciones (embeddings) a nivel de token. Sin embargo, su dependencia del almacenamiento de un vector denso para cada token crea desafíos de escalabilidad significativos. Las implementaciones existentes, como PLAID, utilizan índices de archivos invertidos (IVF) con cuantización de producto para comprimir estos vectores. A pesar de la compresión, estos índices siguen siendo prohibitivamente grandes, requiriendo a menudo de cinco a diez veces el almacenamiento en disco del texto original bruto. Además, el proceso de recuperación en PLAID implica etapas de recolección (gathering) y descompresión ineficientes en el momento de la consulta. Si bien técnicas como el enrutamiento de tokens y la aproximación de puntuación han reducido la latencia, no resuelven el problema fundamental del crecimiento del índice, lo que limita el despliegue de modelos de tipo ColBERT en colecciones a gran escala.

Metodología
Los autores proponen ColBERTSaR, un método que transforma un índice ColBERT en un verdadero índice invertido esparcido al eliminar la dependencia de los vectores residuales. La metodología central implica los siguientes pasos:

1. Aproximación esparcida de MaxSim:
   La función de puntuación estándar de ColBERT, MaxSim, agrega el producto escalar máximo entre los tokens de la consulta y del documento. Los autores reescriben esta función utilizando cuantización de producto, donde los tokens de los documentos están representados por sus centroides de clúster más cercanos (anclas) y vectores residuales. Demuestran que si las normas de los residuales son pequeñas, los residuales pueden omitirse. Esto simplifica la función de puntuación para que dependa únicamente de los productos escalares entre los tokens de la consulta y los centroides de los tokens de los documentos.
   $$Score_S(q, d) = \sum_{i=1}^{|q|} \max_{k \in v_d} (q_i \cdot c_k)$$
   Aquí, $v_d$ representa el conjunto de anclas más cercanas a los tokens en el documento $d$. Esta formulación revela que ColBERT con cuantización es teóricamente equivalente a un modelo de recuperación esparcida aprendida (LSR), donde las anclas actúan como un vocabulario y los productos escalares actúan como pesos específicos de la consulta (similares a TF-IDF).

2. Optimización de Anclas:
   Para minimizar el error de aproximación introducido al eliminar los residuales, los autores proponen dos estrategias de optimización para ajustar la matriz de anclas $C$:

   • Optimización consciente de la consulta (Query-Aware): Minimizar el producto escalar entre los tokens de la consulta y los residuales de los tokens de los documentos que coinciden. Esto requiere entrenar consultas, las cuales no siempre están disponibles.
   • Optimización no supervisada: Para evitar la dependencia de registros de consulta externos, los autores utilizan tokens de documentos dentro del mismo lote (in-batch) como pseudo-consultas. Este enfoque minimiza la distancia entre los tokens de los documentos y sus anclas asignadas sin requerir datos de consulta adicionales, lo que lo hace robusto para las fases de indexación donde las distribuciones de las consultas son desconocidas.

3. Indexación y Recuperación:

   • Indexación: Los embeddings de los tokens de los documentos se mapean a sus IDs de anclas más cercanos, creando un índice invertido esparcido donde cada ancla apunta a un conjunto de IDs de documentos. También se construye un índice hacia adelante (forward index) para mapear los documentos de vuelta a sus anclas.
   • Recuperación: En el momento de la consulta, los tokens de la consulta se codifican y se comparan contra la matriz de anclas para identificar las $n$ anclas más cercanas (nprobe). Los documentos candidatos se recolectan de las listas de publicaciones de estas anclas. Una segunda etapa utiliza el índice hacia adelante para calcular la puntuación esparcida final utilizando los productos escalares precalculados entre la consulta y la ancla, evitando por completo la descompresión de los vectores residuales.

Contribuciones Clave

• Derivación Teórica: El artículo proporciona una prueba teórica de que ColBERT con cuantización de embeddings (sin residuales) es equivalente a un modelo de recuperación esparcida aprendida, diferenciándose únicamente en el mecanismo de puntuación.
• Compresión de Índice: El método elimina la necesidad de almacenar vectores residuales, que constituyen la mayor parte del almacenamiento en los índices PLAID.
• Reemplazo Directo (Drop-in Replacement): ColBERTSaR sirve como un reemplazo directo para PLAID, transformando ColBERT en un modelo de recuperación esparcida sin reentrenar la arquitectura de recuperación subyacente.
• Ajuste de Anclas No Supervisado: La introducción del uso de documentos in-batch como pseudo-consultas para la optimización de anclas permite una construcción de índice efectiva sin depender de distribuciones de consulta de entrenamiento específicas.

Resultados Experimentales
Los autores evaluaron ColBERTSaR en tareas de recuperación monolingües (BEIR), multilingües (NeuCLIRBench) y translingües (NeuCLIRTech).

• Efectividad: ColBERTSaR logra una efectividad competitiva en comparación con PLAID con compresión de residual de 1 bit. En el conjunto de datos BEIR, retiene aproximadamente el 92% de la efectividad de PLAID (usando ColBERT-Small-v1) y el 89% (usando PLAID-XLMR ET). En los benchmarks multilingües, rinde de manera comparable a los modelos LSR de vanguardia como MILCO.
• Eficiencia: El tamaño del índice se reduce entre un 50% y un 70% en comparación con un índice PLAID de 1 bit. Por ejemplo, en las colecciones de NeuCLIRBench, la reducción de tamaño oscila entre el 53% y el 77%, dependiendo del subconjunto de lenguaje.
• Limitaciones: La brecha de rendimiento entre PLAID y ColBERTSaR es ligeramente mayor en conjuntos de datos técnicos (NeuCLIRTech) y subconjuntos de preguntas y respuestas (por ejemplo, Quora, Hotpot QA). Los autores atribuyen esto a la dificultad de capturar significados de entidades específicas cuando los embeddings de los tokens se mapean a centroides que pueden mezclar múltiples conceptos semánticos.

Significancia y Reivindicaciones
El artículo afirma que ColBERTSaR ofrece una visión crítica sobre la relación entre la recuperación densa de vectores múltiples y la recuperación esparcida aprendida. Al demostrar que un modelo denso puede esparcirse mediante la cuantización de producto sin reentrenamiento, el trabajo cierra la brecha entre la expresividad de los modelos densos y la eficiencia de almacenamiento de los modelos esparcidos. Los autores posicionan a ColBERTSaR como una prueba de concepto que reduce sustancialmente el tamaño del índice manteniendo una efectividad de recuperación competitiva, sugiriendo que optimizaciones de ingeniería adicionales (como el empaquetado de bits y mejoras en el recorrido de índices) podrían hacer que este enfoque sea viable para sistemas de producción a gran escala. El trabajo reconoce modestamente que, si bien resuelve el problema del tamaño del índice, se requiere un análisis de latencia y más ingeniería para realizar plenamente su potencial en entornos de producción.