LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference

Resumen Técnico: LightRetriever

Declaración del Problema

Los sistemas actuales de recuperación de texto basados en Grandes Modelos de Lenguaje (LLM) emplean típicamente arquitecturas de codificador dual simétrico, donde tanto los documentos como las consultas son procesados por LLMs profundamente parametrizados para generar representaciones densas o dispersas. Si bien este enfoque mejora significamente la capacidad de recuperación mediante un modelado contextual profundo, introduce un cuello de botella crítico en el despliegue en línea. Los documentos pueden ser pre-codificados e indexados fuera de línea (offline), pero las consultas de búsqueda llegan en tiempo real y requieren una codificación inmediata. Servir un LLM de tamaño completo para la inferencia de consultas en servidores en línea (por ejemplo, GPUs A800) resulta en alta latencia, bajo rendimiento (throughput) y un consumo excesivo de recursos. El artículo postula una tensión fundamental: mientras que los documentos se benefician de un modelado pesado pre-calculado, los codificadores de consultas deben operar bajo estrictas restricciones en línea, sin embargo, los métodos existentes obligan a las consultas a someterse al mismo modelado profundo y computacionalmente costoso que los documentos.

Metodología: LightRetriever

Los autores proponen LightRetriever, una novedosa arquitectura de recuperación híbrida que rompe explícitamente la simetría entre los codificadores de documentos y de consultas. La idea central es que, mientras los documentos requieren el poder de modelado completo de los LLMs, las consultas no necesitan un procesamiento igualmente pesado para lograr un fuerte rendimiento de recuperación. La arquitectura traslada el mayor costo de modelado al lado del documento, reduciendo el lado de la consulta a una operación extremadamente ligera.

1. Arquitectura Asimétrica

• Lado del Documento: Mantiene un codificador LLM de tamaño completo ($Enc_d$) tanto para la codificación densa como para la dispersa. Los documentos se pre-calculan fuera de línea.
• Lado de la Consulta: Elimina la inferencia de Transformer profunda durante el servicio en línea. El codificador de la consulta ($Enc_q$) se reduce a un mecanismo simple de búsqueda de embeddings (embedding lookup).

2. Mecanismo de Recuperación Densa

Para la recuperación densa, LightRetriever emplea un proceso de tres etapas:

• Entrenamiento: Se utiliza un LLM de tamaño completo para entrenar el sistema. Para un token de consulta dado $t_i$ (precedido por una instrucción de tarea), el modelo genera una representación a nivel de token $v^{den}_{t_i} = Enc_q(Inst; t_i)$. El vector denso final de la consulta $v^{den}_q$ se calcula promediando estos vectores de token:
  $$v^{den}_q = \frac{1}{n} \sum_{i=0}^{n-1} v^{den}_{t_i}$$
  El sistema se optimiza utilizando una pérdida contrastiva (Ecuación 2) para alinear los vectores de la consulta con los vectores de los documentos relevantes.
• Almacenamiento en Caché (Caching): Dado que los tokens de las consultas no interactúan profundamente entre sí durante la inferencia, los vectores de los tokens son cacheables. Antes del servicio, el sistema precalcula los embeddings de los tokens para todo el vocabulario ($V$) del LLM y los almacena en una única matriz de búsqueda de embeddings $E$ (forma $[V, H]$).
• Servicio en Línea: En el momento de la inferencia, la consulta se tokeniza y el vector denso se genera mediante la búsqueda de los embeddings en caché para cada token y su posterior promedio. Esto reemplaza el paso hacia adelante (forward pass) de un Transformer grande con una operación de búsqueda de embedding y promedio, requiriendo únicamente una capa de embedding residente en la memoria RAM.

3. Mecanismo de Recuperación Dispersa (Sparse)

Para la recuperación dispersa, el lado de la consulta es aún más ligero:

• Codificación de la Consulta: El vector disperso de la consulta $v^{spr}_q$ es un vector basado en términos, no aprendible, derivado directamente de la tokenización, donde el valor para un token $t$ es simplemente su conteo en la consulta. No se requiere un codificador.
• Codificación del Documento: El documento es procesado por el LLM completo. Los estados ocultos de la última capa se proyectan al espacio del vocabulario a través de la cabeza de modelado de lenguaje del modelo. Para inducir la dispersión, se aplica una activación ReLU y una función de saturación logarítmica, seguida de un max-pooling sobre la longitud de la secuencia.
• Puntuación Híbrida: Las puntuaciones densas y dispersas se interpolan linealmente para producir el ranking de recuperación final.

Contribuciones Clave

1. Velocidad de Inferencia de Consulta Extrema: Al reemplazar el modelado profundo de la consulta con una búsqueda de embedding, LightRetriever logra una aceleración de 1000x en la codificación de consultas en comparación con servir un modelo Llama-8b completo en una GPU A800.
2. Mejora del Rendimiento (Throughput): La arquitectura permite un incremento de 10x en el rendimiento de recuperación de extremo a extremo (Consultas Por Segundo).
3. Retención de Rendimiento: A pesar de la drástica reducción en la computación del lado de la consulta, el método mantiene aproximadamente el 95% del rendimiento de recuperación de los recuperadores basados en LLM simétricos completos en diversos benchmarks.
4. Generalización: El enfoque es validado a través de varios modelos fundacionales (Llama-1B/3B/8B, Qwen-1.5B/3B/7B) y tareas de recuperación (BeIR, CMTEB), demostrando robustez sin requerir cambios arquitectónicos específicos para la tarea.

Resultados Experimentales

El artículo evalúa LightRetriever en 15 tareas en inglés de BeIR y 8 tareas en chino de CMTEB Retrieval.

• Velocidad: En un benchmark de 65,536 consultas frente a 1M de pasajes, LightRetriever redujo el tiempo total de recuperación de extremo a extremo de ~119s (Full-Llama8b) a ~9.36s, logrando un rendimiento de 6,999 QPS comparado con los 549 QPS del modelo completo.
• Efectividad:
  • LightRetriever-Llama3.1-8b logró un nDCG@10 de 54.4 en BeIR y 63.0 en CMTEB-R.
  • Esto representa una degradación modesta de -2.4 y -4.6 puntos respectivamente en comparación con la línea base simétrica completa de Llama3.1-8b (56.8 y 67.6), evitando al mismo tiempo el pesado costo de inferencia de la consulta.
  • El método supera a otros baselines eficientes como BGE-m3 (49.6 en BeIR) y se acerca al rendimiento de E5-Mistral y LLM2Vec.
• Estudios de Ablación:
  • Eliminar el codificador de documento completo (simétrico ligero) causó caídas severas en el rendimiento (por ejemplo, -13.8 en BeIR), confirmando la necesidad de un modelado profundo de documentos.
  • Usar solo un "embedding bag" para el entrenamiento de la consulta (sin el Transformer completo) también resultó en una pérdida significativa de rendimiento (-11.2 en BeIR), indicando que el modelado profundo es esencial durante la fase de entrenamiento para aprender representaciones efectivas, incluso si se descarta durante la inferencia.

Significado y Reivindicaciones

El artículo afirma que LightRetriever ofrece un paradigma escalable para escenarios de recuperación sensibles a la latencia y de alto rendimiento. Aborda la ineficiencia práctica de desplegar gigantescos LLMs para la codificación de consultas en tiempo real sin sacrificar los beneficios de comprensión semántica de la recuperación basada en LLM. Los autores enfatizan que su trabajo transfiere la carga computacional del lado de la consulta en línea al lado del documento fuera de línea, permitiendo que los sistemas aprovechen el poder de los modelos grandes manteniendo la capacidad de respuesta requerida para los motores de búsqueda de producción. El método se presenta como el primero en abordar explícitamente el cuello de botella de la eficiencia en línea en el lado de la consulta de los recuperadores basados en LLM mediante este desacoplamiento extremo.