Módulo 4: Re-ranking — la segunda etapa que transforma retrieval mediocre en excelente

Cápsula 03: Cross-encoder re-ranking — la opción default que casi siempre gana

Descripción de la cápsula

En la cápsula 02 vimos por qué cosine similarity falla en ciertos modos. La solución más eficiente, rápida y barata es cross-encoder re-ranking: un modelo que toma cada par (query, documento) y produce un score de relevancia analizando los dos juntos, en lugar de comparar embeddings independientes.

Cross-encoder es el default razonable para 80%+ de los casos. Comparado con LLM-based o Cohere Rerank, cross-encoder gana en costo (gratis, corre local) y latencia (~150ms para rerankear 20 candidatos). La calidad es 90-92% de precision típica — no llega al 94% de LLM rerank pero el costo cero compensa en la mayoría de proyectos.

Esta cápsula te enseña a elegir el modelo correcto entre los disponibles, integrarlo eficientemente en tu pipeline RAG con sentence-transformers, optimizar el batching para mantener latencia baja, y comparar tu mejora antes y después con un benchmark honesto.

Al finalizar esta cápsula serás capaz de:

  • ✅ Implementar cross-encoder re-ranking con sentence-transformers en menos de 30 líneas
  • ✅ Elegir entre los modelos disponibles (MiniLM L-6 vs L-12, TinyBERT, multilingual) según tu caso
  • ✅ Optimizar batching para minimizar latencia (típicamente 50-100ms en CPU, <30ms en GPU)
  • ✅ Integrar el re-ranker como segunda etapa después del retrieval cosine
  • ✅ Benchmarkear precision antes y después sobre un eval set propio
  • ✅ Anticipar las trampas operativas: cold start, falta de batching, modelo desalineado al dominio

Tiempo estimado: 30-35 minutos


Cómo funciona un cross-encoder por dentro

Un cross-encoder es un modelo BERT-like que recibe dos textos concatenados como input (la query y el documento) y produce un único score de relevancia. Internamente, atiende a cómo cada token de la query se relaciona con cada token del documento — captura interacciones imposibles para bi-encoders (los modelos de embeddings).

Bi-encoder (cosine):

    query  ──> encoder ──> vec_q [1536]
    doc    ──> encoder ──> vec_d [1536]

    score = cosine(vec_q, vec_d)

    ─ Procesa query y doc independientemente
    ─ Vec_q nunca "ve" vec_d


Cross-encoder:

    [query] [SEP] [doc]  ──> encoder con cross-attention ──> score

    ─ Procesa query y doc EN PARALELO
    ─ Cada token de la query atiende a cada token del doc
    ─ Captura interacciones, no solo similaridad

Por qué importa esa diferencia:

  • Bi-encoder responde: "¿qué tan parecidos son estos dos vectores que produjo el modelo independientemente?"
  • Cross-encoder responde: "dado este par (query, doc), ¿qué tan relevante es el doc para la query específica?"

La segunda pregunta es lo que quieres en RAG. Bi-encoder es proxy. Cross-encoder es directo.

El trade-off: cross-encoder NO escala. Con 1M docs en tu collection, no puedes ejecutar cross-encoder sobre todos para cada query — sería ~1000 segundos por query. Por eso la arquitectura óptima usa bi-encoder primero (rápido, recupera top-30) y cross-encoder después (sobre esos 30, refina a top-5).


Modelos disponibles: cuál elegir

Los cross-encoders entrenados sobre MS MARCO son el estándar. Tres variantes principales:

ModeloTamañoInference time (CPU, batch=20)Precision relativaCuándo elegirlo
cross-encoder/ms-marco-MiniLM-L-6-v280 MB~120ms100% (baseline)Default, recursos limitados
cross-encoder/ms-marco-MiniLM-L-12-v2130 MB~180ms+3-4%Recomendado para producción
cross-encoder/ms-marco-TinyBERT-L-2-v250 MB~50ms-3-5% vs MiniLM-L-6Cuando latencia es crítica
cross-encoder/mmarco-mMiniLMv2-L12-H384-v1280 MB~250msExcelente multilingüeCorpus no-inglés

Recomendación práctica:

  • Empezar con MiniLM-L-12-v2. Mejor balance default. La diferencia de tamaño y latencia vs L-6 es pequeña; la mejora de calidad es notable.
  • Bajar a TinyBERT solo si la latencia es problema medible. Si tu pipeline ya tarda 800ms y necesitas bajar a 500ms, TinyBERT puede salvarte 130ms. La pérdida de ~3% precision puede ser aceptable.
  • Subir a multilingual solo si tu corpus lo necesita. MS MARCO multilingual es 2-3x más pesado y más lento, pero gana ~10% en queries no-inglesas.

Implementación correcta

Setup base

# rerank_with_cross_encoder.py
from sentence_transformers import CrossEncoder
from dataclasses import dataclass
from typing import List
import time


# Cargar modelo UNA VEZ al startup, no por query
_CROSS_ENCODER = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")


@dataclass
class RerankedDoc:
    document: str
    score: float
    original_index: int


def cross_encoder_rerank(
    query: str,
    documents: List[str],
    top_k: int = 5,
) -> List[RerankedDoc]:
    """
    Re-rank documents using a local cross-encoder model.

    Importante:
    - El modelo se carga UNA VEZ al startup (variable global).
    - Se procesa todo el batch de pares en una sola llamada (eficiente).
    - Resultados están ordenados por score descendente.
    """
    if not documents:
        return []

    # Crear pares (query, documento) para todo el batch
    pairs = [(query, doc) for doc in documents]

    # Predicción en batch (mucho más rápido que loop)
    scores = _CROSS_ENCODER.predict(pairs, batch_size=32, show_progress_bar=False)

    # Ordenar y mapear
    indexed = list(enumerate(zip(documents, scores)))
    sorted_results = sorted(indexed, key=lambda x: -x[1][1])

    return [
        RerankedDoc(document=doc, score=float(score), original_index=idx)
        for idx, (doc, score) in sorted_results[:top_k]
    ]

Uso end-to-end con ChromaDB

# pipeline_rag_with_rerank.py
import chromadb
from chromadb.utils import embedding_functions
import os

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv("OPENAI_API_KEY"),
    model_name="text-embedding-3-small",
)

client_chroma = chromadb.PersistentClient(path="./chroma_db")
collection = client_chroma.get_collection("docs", embedding_function=openai_ef)


def rag_pipeline(query: str, top_k: int = 5) -> List[RerankedDoc]:
    """
    Pipeline completo: retrieval → cross-encoder rerank → top-K final.
    """
    # Etapa 1: retrieval amplio (top-20-30 candidatos)
    results = collection.query(query_texts=[query], n_results=25)
    candidates = results['documents'][0]

    # Etapa 2: cross-encoder rerank
    reranked = cross_encoder_rerank(query, candidates, top_k=top_k)

    return reranked


# Probar
query = "¿cómo configuro HNSW para 1M vectores con accuracy alta?"
top_5 = rag_pipeline(query, top_k=5)

print(f"Top 5 después del rerank:")
for i, doc in enumerate(top_5, 1):
    print(f"\n#{i} (score: {doc.score:.3f}, era #{doc.original_index+1} antes)")
    print(f"   {doc.document[:120]}...")

Output típico:

Top 5 después del rerank:

#1 (score: 8.421, era #4 antes)
   Para producción con 1M+ vectores, configurar HNSW con M=32 y construction_ef=200...

#2 (score: 7.892, era #1 antes)
   HNSW (Hierarchical Navigable Small World) es el algoritmo de indexing default...

#3 (score: 6.103, era #7 antes)
   Configuración de HNSW para alta accuracy: aumentar M mejora recall pero usa más memoria...

#4 (score: 4.521, era #2 antes)
   Vector databases utilizan diversos algoritmos de indexing...

Nota los movimientos de ranking: el doc más específico ("M=32 con 1M+ vectores") sube de posición #4 a #1 después del rerank. El doc más genérico ("HNSW es el default") baja de #1 a #2.


Optimización de batching

El parámetro batch_size interno del modelo determina cuántos pares se procesan en paralelo en la GPU/CPU. Configurarlo mal puede triplicar la latencia.

# Benchmark de batch_size
import time

candidates = [f"Document {i} about topic..." for i in range(50)]

for batch_size in [1, 8, 16, 32, 64]:
    start = time.perf_counter()
    pairs = [(query, doc) for doc in candidates]
    scores = _CROSS_ENCODER.predict(pairs, batch_size=batch_size, show_progress_bar=False)
    elapsed = (time.perf_counter() - start) * 1000
    print(f"batch_size={batch_size}: {elapsed:.0f}ms")

Output típico (CPU, MiniLM-L-12):

batch_size= 1: 1240ms   ← terrible (no aprovecha vectorización)
batch_size= 8:  280ms
batch_size=16:  220ms
batch_size=32:  200ms   ← sweet spot CPU
batch_size=64:  210ms   ← retornos decrecientes

Output típico (GPU, MiniLM-L-12):

batch_size= 1: 80ms
batch_size= 8: 35ms
batch_size=16: 28ms
batch_size=32: 25ms
batch_size=64: 24ms     ← GPU ama batches grandes

Recomendaciones:

  • CPU: batch_size=32 es sweet spot. Más alto da retornos decrecientes.
  • GPU: batch_size=32-64. Si tienes mucha VRAM, puedes subir.
  • Default sentence-transformers: suele ser 32, pero verifícalo.

Validación: medir el impacto sobre tu eval set

No confíes en benchmarks de blogs. Mide sobre tu eval set propio.

# benchmark_rerank_impact.py
from dataclasses import dataclass
import statistics


@dataclass
class EvalQuery:
    query: str
    expected_doc_ids: list[str]  # IDs de docs relevantes (ground truth)


def benchmark_with_vs_without_rerank(eval_set: list[EvalQuery]):
    """
    Compara métricas con y sin re-ranking sobre el mismo eval set.
    """
    no_rerank = {"precision_at_5": [], "recall_at_5": [], "latency_ms": []}
    with_rerank = {"precision_at_5": [], "recall_at_5": [], "latency_ms": []}

    for item in eval_set:
        # Sin re-rank
        start = time.perf_counter()
        results = collection.query(query_texts=[item.query], n_results=5)
        elapsed = (time.perf_counter() - start) * 1000

        retrieved_ids = set(results['ids'][0])
        relevant_in_top5 = retrieved_ids & set(item.expected_doc_ids)

        no_rerank["precision_at_5"].append(len(relevant_in_top5) / 5)
        no_rerank["recall_at_5"].append(len(relevant_in_top5) / len(item.expected_doc_ids))
        no_rerank["latency_ms"].append(elapsed)

        # Con re-rank
        start = time.perf_counter()
        results = collection.query(query_texts=[item.query], n_results=25)
        candidates = results['documents'][0]
        candidate_ids = results['ids'][0]
        reranked = cross_encoder_rerank(item.query, candidates, top_k=5)
        elapsed = (time.perf_counter() - start) * 1000

        # Mapear back a IDs
        reranked_ids = {candidate_ids[r.original_index] for r in reranked}
        relevant_in_top5_reranked = reranked_ids & set(item.expected_doc_ids)

        with_rerank["precision_at_5"].append(len(relevant_in_top5_reranked) / 5)
        with_rerank["recall_at_5"].append(len(relevant_in_top5_reranked) / len(item.expected_doc_ids))
        with_rerank["latency_ms"].append(elapsed)

    # Reportar
    print(f"\n{'Métrica':<20} {'Sin rerank':<15} {'Con rerank':<15} {'Mejora'}")
    print("-" * 70)

    for metric in ["precision_at_5", "recall_at_5", "latency_ms"]:
        mean_no = statistics.mean(no_rerank[metric])
        mean_with = statistics.mean(with_rerank[metric])
        diff = mean_with - mean_no
        sign = "+" if diff > 0 else ""

        if "latency" in metric:
            print(f"{metric:<20} {mean_no:>10.0f} ms   {mean_with:>10.0f} ms   {sign}{diff:.0f} ms")
        else:
            print(f"{metric:<20} {mean_no:>13.2%}   {mean_with:>13.2%}   {sign}{diff*100:.1f} pts")

Output típico:

Métrica              Sin rerank      Con rerank      Mejora
----------------------------------------------------------------------
precision_at_5            72.40%          90.20%   +17.8 pts
recall_at_5               58.30%          76.40%   +18.1 pts
latency_ms                  185 ms           340 ms   +155 ms

Lectura:

  • Mejora de precision: +17.8 puntos. Excelente.
  • Mejora de recall: +18.1 puntos. Excelente (lo cual confirma que también encontró docs buenos del top-25 que cosine había bajado más allá del top-5).
  • Costo: +155ms latencia. Aceptable para chatbot, ajustado para search interfaces sub-200ms.

Trampas y errores comunes

Trampa 1: cargar el modelo en cada query

El error:

def rerank(query, docs):
    model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")  # ❌ carga cada vez
    return model.predict(...)

Síntoma: primera query tarda 5-10 segundos (cold start). Queries siguientes también si la app no cachea.

Cómo prevenir: cargar el modelo una vez al startup y reutilizar:

# Variable global o singleton
_CROSS_ENCODER = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")

def rerank(query, docs):
    return _CROSS_ENCODER.predict(...)

En Flask/FastAPI, cargar al inicio de la app, no por request.

Trampa 2: predict secuencial en lugar de batch

El error:

scores = []
for doc in docs:
    score = model.predict([(query, doc)])  # ❌ uno a la vez
    scores.append(score)

Síntoma: rerankear 20 docs tarda ~3 segundos en CPU. Inutilizable.

Cómo prevenir: pasar todos los pares juntos:

pairs = [(query, doc) for doc in docs]
scores = model.predict(pairs, batch_size=32)  # batch nativo

Diferencia de velocidad: 10-50x más rápido.

Trampa 3: usar modelo MS MARCO con queries no-inglesas

El error: tu corpus es español. Usas ms-marco-MiniLM-L-12-v2 (entrenado en inglés).

Síntoma: precision@5 sobre queries en español es ~75%, sobre inglés sería ~90%. Inconsistente.

Cómo prevenir: para corpus no-inglés, usar mmarco-mMiniLMv2-L12-H384-v1 (multilingüe) o Cohere Rerank multilingual.

Trampa 4: re-rankear solo top-5 directos

El error:

results = collection.query(query_texts=[q], n_results=5)
reranked = cross_encoder_rerank(q, results['documents'][0], top_k=5)

Síntoma: re-ranking no mejora porque solo reordenas los 5 que cosine ya filtró.

Cómo prevenir: retrieval con n_results=20-30, re-rank selecciona top-5 de esos. El valor del rerank es filtrar falsos positivos del top-30, no reordenar el top-5.

Trampa 5: ignorar score threshold

El error: tomas los top-5 después del rerank sin importar el score absoluto.

Síntoma: algunas queries tienen ningún match relevante en el corpus, pero igual devuelven 5 docs (con scores muy bajos). El LLM se confunde con contexto irrelevante.

Cómo prevenir:

reranked = cross_encoder_rerank(query, candidates, top_k=10)
# Solo incluir docs con score >= threshold
threshold = 1.0  # encontrar empíricamente sobre eval set
relevant = [doc for doc in reranked if doc.score >= threshold]

Si quedan menos de 5 después del threshold, está OK. Mejor pasar 3 docs relevantes al LLM que 5 con 2 basura.

Trampa 6: re-rankear texto crudo cuando los chunks tienen metadata útil

El error:

pairs = [(query, raw_chunk_text) for chunk in retrieved]

Síntoma: el cross-encoder solo ve el texto del chunk, no contexto adicional (título del doc, sección, fecha) que podría mejorar el ranking.

Cómo prevenir: incluir metadata relevante en el texto pasado al cross-encoder:

def format_for_rerank(chunk, metadata):
    return f"Source: {metadata['source']}\nSection: {metadata['section']}\n\n{chunk}"

pairs = [(query, format_for_rerank(chunk, meta)) for chunk, meta in zip(chunks, metas)]

Ejercicio aplicado

Escenario: eres AI Engineer en una empresa de servicios financieros. Pipeline RAG actual:

  • 300K chunks de documentación regulatoria + análisis financieros, en inglés
  • Cosine similarity para retrieval, n_results=5
  • Sin re-ranking
  • Precision@5 actual: 76%
  • Latencia p95: 280ms

Stakeholders piden: "alcanzar precision@5 ≥88% sin pasar de 500ms latencia."

Tu trabajo:

  1. Decide qué modelo de cross-encoder usar.
  2. Diseña la integración (qué n_results del retrieval inicial, qué batch_size, threshold de score).
  3. Estima precision esperada y latencia. ¿Cumple los requisitos?
Solución

1. Modelo recomendado: ms-marco-MiniLM-L-12-v2

Justificación:

  • Corpus en inglés → MS MARCO MiniLM es óptimo (entrenado en inglés).
  • L-12 (no L-6) porque queremos +3-4% de precision sobre L-6, y la diferencia de latencia (~60ms) es absorbible dentro del presupuesto de 500ms.
  • TinyBERT (más rápido) descartado: la mejora de calidad de L-12 sobre TinyBERT es ~6%, vale la latencia extra.
  • Multilingüe descartado: corpus monolingüe inglés, no hace falta.

2. Diseño de la integración

from sentence_transformers import CrossEncoder
import time

# Cargar modelo al startup
_RERANKER = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")

# Configuración
RETRIEVAL_N = 25      # top-25 candidatos del retrieval inicial
RERANK_TOP_K = 8      # top-8 después del rerank (más que 5 para tener margen post-threshold)
SCORE_THRESHOLD = 1.5 # determinado empíricamente sobre eval set
BATCH_SIZE = 32       # sweet spot CPU


def rag_pipeline(query: str) -> list[dict]:
    # Etapa 1: retrieval amplio
    results = collection.query(query_texts=[query], n_results=RETRIEVAL_N)
    candidates = results['documents'][0]
    candidate_metas = results['metadatas'][0]

    # Etapa 2: cross-encoder rerank con metadata
    enriched_pairs = [
        (query, f"Source: {meta.get('source','')}\n\n{doc}")
        for doc, meta in zip(candidates, candidate_metas)
    ]
    scores = _RERANKER.predict(enriched_pairs, batch_size=BATCH_SIZE)

    # Etapa 3: filtrar por threshold y tomar top-5
    scored = sorted(
        zip(candidates, candidate_metas, scores),
        key=lambda x: -x[2]
    )
    relevant = [
        {"doc": doc, "metadata": meta, "score": float(score)}
        for doc, meta, score in scored[:RERANK_TOP_K]
        if score >= SCORE_THRESHOLD
    ][:5]  # final top-5

    return relevant

3. Estimación de impacto

Latencia esperada:

Retrieval (n_results=25):    ~100ms (cosine search es rápido)
Cross-encoder rerank:        ~180ms (25 docs en batch de 32, MiniLM-L-12 en CPU)
Filtrado + ranking final:    ~5ms (en memoria)
─────────────────────────────────
Total p95:                   ~285ms

Si el retrieval inicial actual es 280ms p95 (sin rerank), agregar el rerank lleva el total a ~460ms p95. Está dentro del límite de 500ms con margen apretado pero suficiente.

Precision esperada:

Benchmarks típicos sobre datos similares (corpus inglés técnico, MS MARCO MiniLM-L-12):

  • Sin rerank: 76% (baseline actual)
  • Con rerank: 89-92% (+13-16 puntos)

Estimación final: precision@5 ~90%. Cumple el requisito de ≥88% con margen.

Plan de validación obligatorio:

  1. Construir eval set de 80-100 queries financieras reales con ground truth.
  2. Medir baseline (sin rerank) sobre el eval set.
  3. Implementar rerank, medir.
  4. Si precision >88% y latencia <500ms p95, deployar a staging.
  5. A/B test 1 semana en producción.
  6. Si métricas se mantienen, deployar a 100%.

Plan B si latencia es problema:

  • Bajar n_results a 15 (de 25): latencia rerank baja a ~110ms, precision baja ~1-2%.
  • Cambiar a TinyBERT: latencia rerank baja a ~60ms, precision baja ~3-5%. Solo si latencia es restricción dura.
  • Evaluar GPU para producción: latencia rerank baja a ~25ms, costo de infra sube.

Plan B si precision no llega a 88%:

  • Aumentar n_results del retrieval a 40-50: más candidatos para que el rerank elija. Latencia rerank sube ~50ms.
  • Cambiar a Cohere Rerank: ~+1-2% precision, costo ~$5/mes para este volumen.
  • Considerar LLM rerank en cascada (cross-encoder → LLM sobre top-10): ~+3-5% precision, latencia +800ms (probablemente rompe SLA).

Métricas a monitorear post-deploy:

  • Precision@5 sobre eval set (diario)
  • Latencia p95 end-to-end
  • Distribución de scores del cross-encoder (alerta si la media baja → posible problema con corpus o queries)
  • Tasa de queries que devuelven <5 docs después del threshold (indica queries fuera de cobertura del corpus)

Resumen y siguiente paso

Lo que aprendiste:

  • Cross-encoder analiza pares (query, doc) en paralelo, capturando interacciones que cosine similarity no puede.
  • Default razonable: ms-marco-MiniLM-L-12-v2 para inglés. ~150ms latencia, +20% precision sobre cosine, gratis.
  • Para multilingüe, usar mmarco-mMiniLMv2-L12-H384-v1 o Cohere Rerank.
  • batch_size=32 en CPU es sweet spot. GPU permite 64-128 con paralelismo masivo.
  • Cargar el modelo UNA VEZ al startup. Cada carga toma 5-10s.
  • Retrieval con n_results=20-30 antes del rerank. Re-rankear solo top-5 directos no mejora.
  • Score threshold permite descartar resultados de baja relevancia. Encontrar empíricamente sobre eval set.
  • Validación obligatoria: medir precision sobre eval set propio, no asumir.

Checkpoint: antes de avanzar, deberías poder:

  • Implementar cross-encoder rerank con sentence-transformers en menos de 30 líneas.
  • Elegir n_results y batch_size justificadamente para tu hardware.
  • Medir el impacto del rerank sobre eval set propio antes de deployar.

Siguiente cápsula: 04 — LLM-based re-ranking.

Cross-encoder es el default para 80% de casos. Pero hay dominios donde el 3-4% extra de precision justifica usar un LLM como re-ranker. La cápsula 04 cubre cuándo el costo extra se paga: legal, médico, financiero crítico — y cómo implementarlo correctamente con structured outputs y prompts calibrados.


Recursos

  1. Sentence Transformers — Cross-Encoders — Documentación oficial completa
  2. MS MARCO — Microsoft Research — Dataset usado para entrenar los cross-encoders más populares
  3. Pinecone — Cross-Encoder Reranking — Tutorial visual con benchmarks
  4. Hugging Face Hub — Cross-encoder models — Lista completa de modelos disponibles
  5. Khattab & Zaharia — ColBERT Paper — Late interaction como evolución del cross-encoder
  6. BEIR Benchmark — Comparaciones empíricas reproducibles

Tiempo estimado: 30-35 minutos Siguiente: 04-llm-based-reranking.md