ColBERT

ColBERT (interacción tardía contextualizada sobre BERT) es un modelo neuronal de recuperación que representa una consulta y cada documento con un embedding por token y puntúa un documento mediante interacción tardía: cada token de la consulta busca su token de documento más parecido y esas similitudes máximas, denominadas MaxSim, se suman para obtener la puntuación de relevancia. Omar Khattab y Matei Zaharia, de la Universidad de Stanford, lo presentaron en el artículo de SIGIR de 2020 ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. Los embeddings de los tokens de los documentos se calculan una sola vez, sin conexión, por lo que en el momento de la consulta basta con un cálculo de similitud barato en lugar de ejecutar BERT por completo para cada documento candidato BERT.

Dónde encaja la interacción tardía#

Los recuperadores neuronales se diferencian en el momento en que se encuentran la consulta y el documento. Los bi-encoders comprimen cada texto en un único vector, el método habitual para la búsqueda vectorial: los vectores de los documentos se precalculan y la búsqueda es rápida, pero un solo vector difumina las diferencias entre términos concretos. Los cross-encoders procesan juntos la consulta y el documento mediante un Transformer, de modo que la atención compara todos los pares de tokens. Son los más precisos y deben ejecutarse una vez por cada pareja de consulta y documento, lo que los limita al reranking de listas cortas de candidatos.

ColBERT codifica la consulta y el documento por separado, como hace un bi-encoder, pero conserva un vector de 128 dimensiones por token y permite que ambas partes interactúen solo en el paso MaxSim final. Las consultas se rellenan con el token de máscara de BERT hasta alcanzar una longitud fija de 32, lo que da al modelo posiciones adicionales para ampliar la consulta con términos relacionados. En los benchmarks de búsqueda de pasajes MS MARCO y TREC CAR, ColBERT fue competitivo con los rankers basados en BERT, a la vez que era cien veces más rápido y usaba diez mil veces menos FLOPs por consulta.

Cálculo de una puntuación MaxSim#

El operador MaxSim se implementa en unas pocas líneas de PyTorch. Este ejemplo lo aplica a embeddings contextuales de tokens procedentes de un codificador pequeño de la familia BERT:

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
query = "who wrote the novel Dune"
docs = [
    "Frank Herbert published the science fiction novel Dune in 1965.",
    "A dune is a hill of sand shaped by the wind.",
]

# One L2-normalized contextual embedding per query token
q = F.normalize(encoder.encode(query, output_value="token_embeddings"), dim=-1)
for doc in docs:
    d = F.normalize(encoder.encode(doc, output_value="token_embeddings"), dim=-1)

    # MaxSim: each query token keeps its best-matching document token, and the maxima are summed
    score = (q @ d.T).max(dim=1).values.sum()
    print(f"{score.item():.2f}  {doc}")

El pasaje sobre la novela obtiene una puntuación de 5,90 y el pasaje sobre la duna de arena, de 4,33. Ambos contienen «duna», que encaja muy bien en los dos casos, pero en el primer pasaje el token de consulta «escribió» encuentra «publicó» con una similitud coseno de 0,72 y «novela» se encuentra a sí mismo, mientras que en el segundo pasaje la mejor coincidencia para «escribió» solo llega a 0,29.

ColBERT frente a términos relacionados#

• BERT en sí: el codificador preentrenado que ColBERT ajusta para la recuperación, al que añade una capa lineal que proyecta cada token a 128 dimensiones.
• Fragmentación tardía: también codifica el texto en contexto antes de dividirlo y, después, promedia los tokens en un vector por fragmento, mientras que ColBERT conserva todos los vectores de tokens.
• Búsqueda semántica: la tarea que realiza ColBERT, que ordena los pasajes por significado en lugar de por palabras clave compartidas.
• Búsqueda híbrida: combina resultados por palabras clave y por vector único para reunir la coincidencia exacta de términos y el significado. La coincidencia a nivel de token de ColBERT captura parte de esa exactitud en un único modelo neuronal.

Interacción tardía para imágenes#

La interacción tardía también funciona con imágenes. ColPali (2024) representa la captura de pantalla de cada página de un documento como una cuadrícula de vectores de parches mediante un modelo de visión y lenguaje y las puntúa frente a los tokens de la consulta con MaxSim, llevando la interacción tardía a la generación aumentada por recuperación multimodal sobre archivos PDF. La detección de vocabulario abierto mantiene las características de imagen en un nivel detallado por un motivo relacionado: Ultralytics YOLOE almacena un embedding de texto por cada clase indicada en la instrucción, como «autobús de dos pisos», y lo compara con las características de región en cada ubicación de sus mapas de características, de modo que una instrucción breve puede coincidir con un objeto pequeño en lugar de con un único vector que resume toda la imagen.

ColBERTv2 y PLAID#

Almacenar un vector por cada token hace que un índice ColBERT sea aproximadamente diez veces mayor que un índice de vector único, un coste importante cuando ColBERT actúa como recuperador de primera etapa o reranker en un flujo de generación aumentada por recuperación (RAG). El repositorio ColBERT de Stanford implementa dos soluciones. PLAID es un motor de búsqueda que puntúa los candidatos frente a centroides de clústeres antes de ejecutar MaxSim completo. ColBERTv2 (NAACL 2022) almacena cada vector de token como el ID de su centroide más cercano, junto con un residuo cuantizado a 1 o 2 bits por dimensión, y entrena con supervisión depurada destilada de un cross-encoder; su compresión de residuos redujo entre 6 y 10 veces el espacio ocupado por la interacción tardía.