¿Qué es ColBERT y la Interacción Tardía y Por Qué Son Importantes en la Búsqueda?

Jina ColBERT v2: Recuperador de Interacción Tardía Multilingüe para Embedding y Reranking

Jina ColBERT v2 soporta 89 idiomas con rendimiento superior de recuperación, dimensiones de salida controladas por el usuario y longitud de token de 8192.

El viernes pasado, el lanzamiento del modelo ColBERT por Jina AI en Hugging Face generó un entusiasmo significativo en la comunidad de IA, particularmente en Twitter/X. Si bien muchos están familiarizados con el revolucionario modelo BERT, el revuelo alrededor de ColBERT ha dejado a algunos preguntándose: ¿Qué hace que ColBERT destaque en el saturado campo de las tecnologías de recuperación de información? ¿Por qué la comunidad de IA está entusiasmada con ColBERT de longitud 8192? Este artículo profundiza en las complejidades de ColBERT y ColBERTv2, destacando su diseño, mejoras y la sorprendente efectividad de la interacción tardía de ColBERT.

Reranker API

Maximiza la relevancia de búsqueda y la precisión RAG con facilidad

tag¿Qué es ColBERT?

El nombre "ColBERT" significa Contextualized Late Interaction over BERT, un modelo proveniente de la Universidad de Stanford, que aprovecha la comprensión profunda del lenguaje de BERT mientras introduce un nuevo mecanismo de interacción. Este mecanismo, conocido como interacción tardía, permite una recuperación eficiente y precisa al procesar consultas y documentos por separado hasta las etapas finales del proceso de recuperación. Específicamente, hay dos versiones del modelo:

• ColBERT: El modelo inicial fue creación de Omar Khattab y Matei Zaharia, presentando un enfoque novedoso para la recuperación de información a través del artículo "ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT". Su trabajo fue publicado en SIGIR 2020.

ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT

Recent progress in Natural Language Understanding (NLU) is driving fast-paced advances in Information Retrieval (IR), largely owed to fine-tuning deep language models (LMs) for document ranking. While remarkably effective, the ranking models based on these LMs increase computational cost by orders of magnitude over prior approaches, particularly as they must feed each query-document pair through a massive neural network to compute a single relevance score. To tackle this, we present ColBERT, a novel ranking model that adapts deep LMs (in particular, BERT) for efficient retrieval. ColBERT introduces a late interaction architecture that independently encodes the query and the document using BERT and then employs a cheap yet powerful interaction step that models their fine-grained similarity. By delaying and yet retaining this fine-granular interaction, ColBERT can leverage the expressiveness of deep LMs while simultaneously gaining the ability to pre-compute document representations offline, considerably speeding up query processing. Beyond reducing the cost of re-ranking the documents retrieved by a traditional model, ColBERT's pruning-friendly interaction mechanism enables leveraging vector-similarity indexes for end-to-end retrieval directly from a large document collection. We extensively evaluate ColBERT using two recent passage search datasets. Results show that ColBERT's effectiveness is competitive with existing BERT-based models (and outperforms every non-BERT baseline), while executing two orders-of-magnitude faster and requiring four orders-of-magnitude fewer FLOPs per query.

arXiv.orgOmar Khattab

• ColBERTv2: Basándose en el trabajo fundamental, Omar Khattab continuó su investigación, colaborando con Barlas Oguz, Matei Zaharia y Michael S. Bernstein para introducir "ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction", presentado en SIGIR 2021. Esta siguiente iteración de ColBERT abordó limitaciones anteriores e introdujo mejoras clave, como la supervisión sin ruido y la compresión residual, mejorando tanto la efectividad de recuperación del modelo como su eficiencia de almacenamiento.

ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction

Neural information retrieval (IR) has greatly advanced search and other knowledge-intensive language tasks. While many neural IR methods encode queries and documents into single-vector representations, late interaction models produce multi-vector representations at the granularity of each token and decompose relevance modeling into scalable token-level computations. This decomposition has been shown to make late interaction more effective, but it inflates the space footprint of these models by an order of magnitude. In this work, we introduce ColBERTv2, a retriever that couples an aggressive residual compression mechanism with a denoised supervision strategy to simultaneously improve the quality and space footprint of late interaction. We evaluate ColBERTv2 across a wide range of benchmarks, establishing state-of-the-art quality within and outside the training domain while reducing the space footprint of late interaction models by 6--10$\times$.

arXiv.orgKeshav Santhanam

tagEntender el Diseño de ColBERT

Dado que la arquitectura de ColBERTv2 permanece muy similar a la del ColBERT original, con sus innovaciones clave girando en torno a técnicas de entrenamiento y mecanismos de compresión, primero profundizaremos en los aspectos fundamentales del ColBERT original.

tag¿Qué es la interacción tardía en ColBERT?

"Interacción" se refiere al proceso de evaluar la relevancia entre una consulta y un documento comparando sus representaciones.

La "interacción tardía" es la esencia de ColBERT. El término se deriva de la arquitectura y estrategia de procesamiento del modelo, donde la interacción entre las representaciones de la consulta y el documento ocurre tarde en el proceso, después de que ambos han sido codificados independientemente. Esto contrasta con los modelos de "interacción temprana", donde los embeddings de consulta y documento interactúan en etapas anteriores, típicamente antes o durante su codificación por el modelo.

La interacción temprana puede aumentar la complejidad computacional ya que requiere considerar todos los pares posibles de consulta-documento, haciéndola menos eficiente para aplicaciones a gran escala.

Los modelos de interacción tardía como ColBERT optimizan la eficiencia y escalabilidad al permitir el precálculo de representaciones de documentos y emplear un paso de interacción más ligero al final, que se centra en las representaciones ya codificadas. Esta elección de diseño permite tiempos de recuperación más rápidos y demandas computacionales reducidas, haciéndolo más adecuado para procesar grandes colecciones de documentos.

tagSin interacción: similitud del coseno entre embeddings de documento y consulta

Muchas bases de datos vectoriales prácticas y soluciones de búsqueda neural dependen de la coincidencia rápida de similitud del coseno entre embeddings de documentos y consultas. Si bien es atractivo por su sencillez y eficiencia computacional, se ha encontrado que este método, a menudo denominado "sin interacción" o "no basado en interacción" tiene un rendimiento inferior en comparación con modelos que incorporan alguna forma de interacción entre consultas y documentos.

La limitación principal del enfoque "sin interacción" radica en su incapacidad para capturar los matices complejos y las relaciones entre los términos de consulta y documento. La recuperación de información, en su esencia, consiste en comprender y hacer coincidir la intención detrás de una consulta con el contenido dentro de un documento. Este proceso a menudo requiere una comprensión profunda y contextual de los términos involucrados, algo que los embeddings únicos y agregados para documentos y consultas tienen dificultades para proporcionar.

tagCodificadores de consulta y documento en ColBERT

La estrategia de codificación de ColBERT se basa en el modelo BERT, conocido por su profunda comprensión contextual del lenguaje. El modelo genera representaciones vectoriales densas para cada token en una consulta o documento, creando un conjunto de embeddings contextualizados para una consulta y un conjunto para un documento, respectivamente. Esto facilita una comparación matizada de sus embeddings durante la fase de interacción tardía.

tagCodificador de consultas de ColBERT

Para una consulta $Q$ con tokens ${q_1, q_2, ..., q_l}$, el proceso comienza tokenizando $Q$ en tokens WordPiece basados en BERT y anteponiendo un token especial [Q]. Este token [Q], posicionado justo después del token [CLS] de BERT, señala el inicio de una consulta.

Si la consulta es más corta que un número predefinido de tokens $N_q$, se rellena con tokens [mask] hasta $N_q$; de lo contrario, se trunca a los primeros $N_q$ tokens. La secuencia rellenada luego pasa por BERT, seguida de una CNN (Red Neuronal Convolucional) y normalización. La salida es un conjunto de vectores de embedding denominados $\mathbf{E}_q$ a continuación:
$$\mathbf{E}_q := \mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[Q]},q_0,q_1,\ldots,q_l\mathtt{[mask]},\mathtt{[mask]},\ldots,\mathtt{[mask]}\right)\right)$$

tagCodificador de documentos de ColBERT

De manera similar, para un documento $D$ con tokens ${d_1, d_2, ..., d_n}$, se antepone un token [D] para indicar el inicio del documento. Esta secuencia, sin necesidad de relleno, pasa por el mismo proceso, resultando en un conjunto de vectores de embedding denominados $\mathbf{E}_d$ a continuación:
$$\mathbf{E}_d := \mathrm{Filter}\left(\mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[D]},d_0,d_1,...,d_n\right)\right)\right)$$

El uso de tokens [mask] para rellenar consultas (acuñado como aumento de consulta en el documento) asegura una longitud uniforme en todas las consultas, facilitando el procesamiento por lotes. Los tokens [Q] y [D] marcan explícitamente el inicio de consultas y documentos, respectivamente, ayudando al modelo a distinguir entre los dos tipos de entradas.

tagComparación de ColBERT con codificadores cruzados

Los codificadores cruzados procesan pares de consultas y documentos juntos, haciéndolos altamente precisos pero menos eficientes para tareas a gran escala debido al costo computacional de evaluar cada posible par. Sobresalen en escenarios específicos donde es necesaria la puntuación precisa de pares de oraciones, como en tareas de similitud semántica o comparación detallada de contenido. Sin embargo, este diseño limita su aplicabilidad en situaciones que requieren recuperación rápida de grandes conjuntos de datos, donde los embeddings precalculados y los cálculos eficientes de similitud son primordiales.

En contraste, el modelo de interacción tardía de ColBERT permite el precálculo de embeddings de documentos, acelerando significativamente el proceso de recuperación sin comprometer la profundidad del análisis semántico. Este método, aunque aparentemente contraintuitivo en comparación con el enfoque directo de los codificadores cruzados, ofrece una solución escalable para tareas de recuperación de información en tiempo real y a gran escala. Representa un compromiso estratégico entre la eficiencia computacional y la calidad del modelado de interacción.

tagEncontrando los top-K documentos usando ColBERT

Una vez que tenemos embeddings para la consulta y los documentos, encontrar los K documentos más relevantes se vuelve directo (pero no tan directo como calcular el coseno de dos vectores).

Las operaciones clave incluyen un producto punto por lotes para calcular similitudes término a término, max-pooling a través de términos de documento para encontrar la similitud más alta por término de consulta, y suma a través de términos de consulta para derivar la puntuación total del documento, seguido de ordenar los documentos basados en estas puntuaciones. El pseudo código en PyTorch se describe a continuación:

import torch

def compute_relevance_scores(query_embeddings, document_embeddings, k):
    """
    Compute relevance scores for top-k documents given a query.
    
    :param query_embeddings: Tensor representing the query embeddings, shape: [num_query_terms, embedding_dim]
    :param document_embeddings: Tensor representing embeddings for k documents, shape: [k, max_doc_length, embedding_dim]
    :param k: Number of top documents to re-rank
    :return: Sorted document indices based on their relevance scores
    """
    
    # Ensure document_embeddings is a 3D tensor: [k, max_doc_length, embedding_dim]
    # Pad the k documents to their maximum length for batch operations
    # Note: Assuming document_embeddings is already padded and moved to GPU
    
    # Compute batch dot-product of Eq (query embeddings) and D (document embeddings)
    # Resulting shape: [k, num_query_terms, max_doc_length]
    scores = torch.matmul(query_embeddings.unsqueeze(0), document_embeddings.transpose(1, 2))
    
    # Apply max-pooling across document terms (dim=2) to find the max similarity per query term
    # Shape after max-pool: [k, num_query_terms]
    max_scores_per_query_term = scores.max(dim=2).values
    
    # Sum the scores across query terms to get the total score for each document
    # Shape after sum: [k]
    total_scores = max_scores_per_query_term.sum(dim=1)
    
    # Sort the documents based on their total scores
    sorted_indices = total_scores.argsort(descending=True)
    
    return sorted_indices
Copiar

Tenga en cuenta que este procedimiento se utiliza tanto en el entrenamiento como en la reclasificación durante la inferencia. El modelo ColBERT se entrena utilizando una pérdida de clasificación por pares, donde los datos de entrenamiento consisten en triples $(q, d^+, d^-)$, donde $q$ representa una consulta, $d^+$ es un documento relevante (positivo) para la consulta, y $d^-$ es un documento no relevante (negativo). El modelo busca aprender representaciones tales que la puntuación de similitud entre $q$ y $d^+$ sea mayor que la puntuación entre q y $d^-$.

El objetivo de entrenamiento puede representarse matemáticamente como minimizar la siguiente función de pérdida: $$\mathrm{Loss} = \max(0, 1 - S(q, d^+) + S(q, d^-))$$

, donde $S(q, d)$ denota la puntuación de similitud calculada por ColBERT entre una consulta $q$ y un documento $d$. Esta puntuación se obtiene agregando las puntuaciones de máxima similitud de los embeddings mejor emparejados entre la consulta y el documento, siguiendo el patrón de interacción tardía descrito en la arquitectura del modelo. Este enfoque asegura que el modelo esté entrenado para distinguir entre documentos relevantes e irrelevantes para una consulta dada, fomentando un margen mayor en las puntuaciones de similitud para pares de documentos positivos y negativos.

tagSupervisión sin ruido en ColBERTv2

La supervisión sin ruido en ColBERTv2 refina el proceso de entrenamiento original seleccionando negativos desafiantes y aprovechando un codificador cruzado para la destilación. Este método sofisticado de aumentar la calidad de los datos de entrenamiento involucra varios pasos:

1. Entrenamiento Inicial: Utilización de los triples oficiales del conjunto de datos MS MARCO, que comprende una consulta, un documento relevante y un documento no relevante.
2. Indexación y Recuperación: Empleo de la compresión de ColBERTv2 para indexar pasajes de entrenamiento, seguido de la recuperación de los top-k pasajes para cada consulta.
3. Reclasificación con Codificador Cruzado: Mejora de la selección de pasajes mediante reclasificación por un codificador cruzado MiniLM, destilando sus puntuaciones en ColBERTv2.
4. Formación de Tuplas de Entrenamiento: Generación de tuplas de w-vías para entrenamiento, incorporando pasajes de alta y baja clasificación para crear ejemplos desafiantes.
5. Refinamiento Iterativo: Repetición del proceso para mejorar continuamente la selección de negativos difíciles, mejorando así el rendimiento del modelo.

Nótese que este proceso representa una mejora sofisticada del régimen de entrenamiento de ColBERT en lugar de un cambio fundamental en su arquitectura.

tagHiperparámetros de ColBERT

Los hiperparámetros de ColBERT se resumen a continuación:

tagLa estrategia de indexación de ColBERT

A diferencia de los enfoques basados en representación que codifican cada documento en un vector de embedding, ColBERT codifica documentos (y consultas) en conjuntos de embeddings, donde cada token en un documento tiene su propio embedding. Este enfoque inherentemente significa que para documentos más largos, se almacenarán más embeddings, lo cual es un punto débil del ColBERT original, y que posteriormente fue abordado por ColBERTv2.

La clave para gestionar esto eficientemente radica en el uso que hace ColBERT de bases de datos vectoriales (por ejemplo, FAISS) para indexación y recuperación, y su detallado proceso de indexación que está diseñado para manejar grandes volúmenes de datos eficientemente. El paper original de ColBERT menciona varias estrategias para mejorar la eficiencia de la indexación y recuperación, incluyendo:

• Indexación Offline: Las representaciones de documentos se calculan offline, permitiendo el pre-cálculo y almacenamiento de embeddings de documentos. Este proceso aprovecha el procesamiento por lotes y la aceleración GPU para manejar grandes colecciones de documentos eficientemente.
• Almacenamiento de Embeddings: Los embeddings de documentos pueden almacenarse usando valores de 32 bits o 16 bits para cada dimensión, ofreciendo un equilibrio entre precisión y requerimientos de almacenamiento. Esta flexibilidad permite a ColBERT mantener un balance entre efectividad (en términos de rendimiento de recuperación) y eficiencia (en términos de costos de almacenamiento y computación).

La introducción de la compresión residual en ColBERTv2, que es un enfoque novedoso no presente en el ColBERT original, juega un papel clave en reducir la huella espacial del modelo en 6-10× mientras preserva la calidad. Esta técnica comprime aún más los embeddings capturando y almacenando efectivamente solo las diferencias desde un conjunto de centroides de referencia fijos.

tagEfectividad y Eficiencia de ColBERT

Uno podría inicialmente asumir que incorporar la comprensión contextual profunda de BERT en la búsqueda inherentemente requeriría recursos computacionales significativos, haciendo tal enfoque menos factible para aplicaciones en tiempo real debido a la alta latencia y costos computacionales. Sin embargo, ColBERT desafía y revierte esta suposición a través de su uso innovador del mecanismo de interacción tardía. Aquí hay algunos puntos destacables:

1. Ganancias Significativas en Eficiencia: ColBERT logra una reducción de órdenes de magnitud en costos computacionales (FLOPs) y latencia comparado con modelos de ranking basados en BERT tradicionales. Específicamente, para un tamaño de modelo dado (por ejemplo, codificador transformer "base" de 12 capas), ColBERT no solo iguala sino que en algunos casos supera la efectividad de los modelos basados en BERT con demandas computacionales dramáticamente menores. Por ejemplo, a una profundidad de re-ranking de k=10, BERT requiere casi 180× más FLOPs que ColBERT; esta brecha se amplía a medida que k aumenta, alcanzando 13900× en k=1000 y hasta 23000× en k=2000.
2. Mejora en Recall y MRR@10 en Recuperación End-to-End: Contrario a la intuición inicial de que sería necesaria una interacción más profunda entre las representaciones de consulta y documento (como se ve en los modelos de interacción temprana) para un alto rendimiento de recuperación, la configuración de recuperación end-to-end de ColBERT demuestra una efectividad superior. Por ejemplo, su Recall@50 supera el Recall@1000 del BM25 oficial y casi todos los Recall@200 de otros modelos, subrayando la notable capacidad del modelo para recuperar documentos relevantes de una vasta colección sin comparación directa de cada par consulta-documento.
3. Practicidad para Aplicaciones del Mundo Real: Los resultados experimentales subrayan la aplicabilidad práctica de ColBERT para escenarios del mundo real. Su rendimiento de indexación y eficiencia de memoria lo hacen adecuado para indexar grandes colecciones de documentos como MS MARCO en pocas horas, manteniendo alta efectividad con una huella espacial manejable. Estas cualidades resaltan la idoneidad de ColBERT para su despliegue en entornos de producción donde tanto el rendimiento como la eficiencia computacional son primordiales.
4. Escalabilidad con el Tamaño de la Colección de Documentos: Quizás la conclusión más sorprendente es la escalabilidad y eficiencia de ColBERT en el manejo de colecciones de documentos a gran escala. La arquitectura permite el pre-cálculo de embeddings de documentos y aprovecha el procesamiento eficiente por lotes para la interacción consulta-documento, permitiendo que el sistema escale efectivamente con el tamaño de la colección de documentos. Esta escalabilidad es contraintuitiva cuando se considera la complejidad y profundidad de comprensión requerida para una recuperación efectiva de documentos, mostrando el enfoque innovador de ColBERT para equilibrar la eficiencia computacional con la efectividad de recuperación.

tagUsando jina-colbert-v1-en: un modelo ColBERTv2 de longitud 8192

Jina-ColBERT está diseñado tanto para recuperación rápida como precisa, soportando longitudes de contexto más largas de hasta 8192, aprovechando los avances de JinaBERT, que permite el procesamiento de secuencias más largas debido a sus mejoras en la arquitectura.

jinaai/jina-colbert-v1-en · Hugging Face

We're on a journey to advance and democratize artificial intelligence through open source and open science.

tagMejoras de Jina sobre el ColBERT original

El principal avance de Jina-ColBERT es su columna vertebral, jina-bert-v2-base-en, que permite procesar contextos significativamente más largos (hasta 8192 tokens) comparado con el ColBERT original que usa bert-base-uncased. Esta capacidad es crucial para manejar documentos con contenido extenso, proporcionando resultados de búsqueda más detallados y contextuales.

tagComparación de rendimiento de jina-colbert-v1-en vs. ColBERTv2

Evaluamos jina-colbert-v1-en en datasets BEIR y el nuevo benchmark LoCo que favorece el contexto largo, probándolo contra la implementación original de ColBERTv2 y basada en no-interacciónmodelo jina-embeddings-v2-base-en.

Esta tabla demuestra el rendimiento superior de jina-colbert-v1-en, especialmente en escenarios que requieren longitudes de contexto más largas frente al ColBERTv2 original. Ten en cuenta que jina-embeddings-v2-base-en utiliza más datos de entrenamiento, mientras que jina-colbert-v1-en solo utiliza MSMARCO, lo que puede justificar el buen rendimiento de jina-embeddings-v2-base-en en algunas tareas.

tagEjemplo de uso de jina-colbert-v1-en

Este fragmento describe el proceso de indexación con Jina-ColBERT, mostrando su soporte para documentos largos.

from colbert import Indexer
from colbert.infra import Run, RunConfig, ColBERTConfig

n_gpu: int = 1  # Set your number of available GPUs
experiment: str = ""  # Name of the folder where the logs and created indices will be stored
index_name: str = ""  # The name of your index, i.e. the name of your vector database

if __name__ == "__main__":
    with Run().context(RunConfig(nranks=n_gpu, experiment=experiment)):
        config = ColBERTConfig(
          doc_maxlen=8192  # Our model supports 8k context length for indexing long documents
        )
        indexer = Indexer(
          checkpoint="jinaai/jina-colbert-v1-en",
          config=config,
        )
        documents = [
          "ColBERT is an efficient and effective passage retrieval model.",
          "Jina-ColBERT is a ColBERT-style model but based on JinaBERT so it can support both 8k context length.",
          "JinaBERT is a BERT architecture that supports the symmetric bidirectional variant of ALiBi to allow longer sequence length.",
          "Jina-ColBERT model is trained on MSMARCO passage ranking dataset, following a very similar training procedure with ColBERTv2.",
          "Jina-ColBERT achieves the competitive retrieval performance with ColBERTv2.",
          "Jina is an easier way to build neural search systems.",
          "You can use Jina-ColBERT to build neural search systems with ease.",
          # Add more documents here to ensure the clustering work correctly
        ]
        indexer.index(name=index_name, collection=documents)
Copiar

tagUso de jina-colbert-v1-en en RAGatouille

RAGatouille es una nueva biblioteca de Python que facilita el uso de métodos avanzados de recuperación dentro de pipelines RAG. Está diseñada para ser modular y de fácil integración, permitiendo a los usuarios aprovechar la investigación de vanguardia sin problemas. El objetivo principal de RAGatouille es simplificar la aplicación de modelos complejos como ColBERT en pipelines RAG, haciendo accesible para los desarrolladores utilizar estos métodos sin necesitar experiencia profunda en la investigación subyacente. Gracias a Benjamin Clavié, ahora puedes usar jina-colbert-v1-en fácilmente:

from ragatouille import RAGPretrainedModel

# Get your model & collection of big documents ready
RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")
my_documents = [
    "very long document1",
    "very long document2",
    # ... more documents
]

# And create an index with them at full length!
RAG.index(collection=my_documents,
          index_name="the_biggest_index",
          max_document_length=8190,)

# or encode them in-memory with no truncation, up to your model's max length
RAG.encode(my_documents)
Copiar

Para obtener información más detallada y explorar más a fondo Jina-ColBERT, puedes visitar la página de Hugging Face.

tagConclusión

ColBERT representa un avance significativo en el campo de la recuperación de información. Al permitir longitudes de contexto más largas con Jina-ColBERT y mantener la compatibilidad con el enfoque de interacción tardía de ColBERT, ofrece una poderosa alternativa para los desarrolladores que buscan implementar funcionalidades de búsqueda de última generación.

Junto con la biblioteca RAGatouille, que simplifica la integración de modelos complejos de recuperación en pipelines RAG, los desarrolladores pueden ahora aprovechar el poder de la recuperación avanzada con facilidad, optimizando sus flujos de trabajo y mejorando sus aplicaciones. La sinergia entre Jina-ColBERT y RAGatouille ilustra un notable avance en hacer que los modelos avanzados de búsqueda con IA sean accesibles y eficientes para uso práctico.