Reranking con ColBERT: cómo mejorar la precisión de tu RAG sin cambiar los embeddings

La mayoría de implementaciones de RAG hacen búsqueda vectorial y usan directamente los N chunks más similares como contexto para el LLM. Hay un paso intermedio que mejora significativamente la calidad de los resultados y que muchos omiten: el reranking.

Por qué el reranking importa

La búsqueda vectorial por similitud de coseno es eficiente pero imprecisa. Recupera chunks que son semánticamente “cerca” de la query, pero “cerca” en el espacio de embeddings no siempre significa “más relevante para responder esta pregunta”.

El reranking añade un segundo modelo que, dado un conjunto de candidatos recuperados, los ordena de nuevo según relevancia real para la query. Este segundo modelo puede ser más lento y sofisticado porque solo opera sobre un subset pequeño (los top-K de la búsqueda inicial).

ColBERT: la mejor opción práctica

ColBERT (Contextualized Late Interaction over BERT) es el estándar de facto para reranking eficiente. Su característica clave: en lugar de comprimir todo el documento en un solo vector, mantiene embeddings individuales por token y calcula la relevancia mediante “MaxSim” (suma de las similitudes máximas entre tokens de query y documento).

Esto lo hace más preciso que embeddings de documento completo y más rápido que modelos cross-encoder completos.

Implementación

Opción 1: RAGatouille (la más fácil)

from ragatouille import RAGPretrainedModel

# Cargar modelo preentrenado
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

# Indexar documentos
RAG.index(
    collection=["documento 1...", "documento 2...", "documento 3..."],
    index_name="mi_indice",
    max_document_length=180,
    split_documents=True,
)

# Buscar con reranking incluido
resultados = RAG.search(query="¿Cómo funciona la atención multi-cabeza?", k=5)

Opción 2: ColBERT como reranker sobre resultados existentes

Si ya tienes un pipeline de búsqueda vectorial, puedes añadir ColBERT solo para el reranking:

from ragatouille import RAGPretrainedModel

reranker = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

def rerank_results(query: str, candidates: list[str], top_k: int = 5) -> list[str]:
    scores = reranker.rerank(query=query, documents=candidates)
    # Ordenar por score y devolver top_k
    sorted_docs = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in sorted_docs[:top_k]]

# En tu pipeline:
# 1. Búsqueda vectorial rápida: top-50 candidatos
candidatos = vector_search(query, top_k=50)

# 2. Reranking preciso: top-5 finales
contexto_final = rerank_results(query, candidatos, top_k=5)

# 3. LLM con contexto de calidad
respuesta = llm.generate(query=query, context=contexto_final)

Opción 3: Cohere Rerank (hosted)

Si no quieres gestionar el modelo localmente:

import cohere

co = cohere.Client("TU_API_KEY")

response = co.rerank(
    query="tu query aquí",
    documents=candidatos,
    model="rerank-multilingual-v3.0",  # soporte en español
    top_n=5,
)

resultados_reranked = [candidatos[r.index] for r in response.results]

Métricas: cuánto mejora

En nuestras evaluaciones sobre bases de conocimiento técnicas en español:

Pipeline	NDCG@5	MRR
Solo vectorial	0.61	0.54
Vectorial + BM25 híbrido	0.68	0.61
Vectorial + ColBERT reranking	0.79	0.73
Vectorial + BM25 + ColBERT	0.83	0.78

El reranking con ColBERT añade más valor que el híbrido vectorial+BM25 solo.

Latencia

ColBERT añade latencia. En un servidor con GPU modest, reranking sobre 50 candidatos toma 50-200ms. En CPU, puede subir a 500ms-1s.

Para aplicaciones interactivas donde la latencia importa, considera:

1. Reducir el pool de candidatos (top-20 en lugar de top-50)
2. Usar versiones destiladas de ColBERT (más rápidas, algo menos precisas)
3. Usar Cohere Rerank para evitar gestionar el modelo

El pipeline completo recomendado

async def rag_pipeline(query: str, k_retrieve: int = 50, k_final: int = 5):
    # 1. Búsqueda vectorial + BM25 en paralelo
    vector_results = await vector_search(query, k=k_retrieve)
    bm25_results = await bm25_search(query, k=k_retrieve)
    
    # 2. Fusion de resultados (RRF)
    combined = reciprocal_rank_fusion(vector_results, bm25_results)[:k_retrieve]
    
    # 3. Reranking con ColBERT
    final_context = await colbert_rerank(query, combined, k=k_final)
    
    # 4. Generación
    return await llm.generate(query=query, context=final_context)

---

ColBERT v2: “ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction” (Santhanam et al., 2022).