Enhancing Lexicon-Based Text Embeddings with Large Language Models

Resumen Técnico: LENS (Lexicon-based Embeddings from LLMs)

1. El Problema

Aunque los modelos de lenguaje grandes (LLMs) han demostrado un rendimiento excepcional en tareas de incrustación de texto (text embeddings) mediante enfoques densos, las incrustaciones basadas en léxico (que asignan una dimensión a cada token del vocabulario) han permanecido subexploradas en este contexto. La aplicación directa de LLMs a incrustaciones basadas en léxico enfrenta dos desafíos principales, ilustrados en la Figura 1 del artículo:

1. Redundancia y Ruido en el Tokenizador: Los tokenizadores modernos (subpalabras) fragmentan las palabras (ej. "educación" $\rightarrow$ "edu", "cation") y crean múltiples formas para el mismo significado semántico (ej. "what", "What", " what"). Esto introduce inconsistencias y ruido en el emparejamiento léxico, degradando el rendimiento si se usa el vocabulario original directamente.
2. Atención Unidireccional: La mayoría de los LLMs (modelos causales) utilizan atención unidireccional durante el pre-entrenamiento, donde un token solo observa a los anteriores. Esto impide que cada token aproveche el contexto completo, lo cual es crítico para las incrustaciones basadas en léxico que derivan su representación de la salida de todos los tokens.

2. Metodología: El Framework LENS

Los autores proponen LENS, un marco diseñado para generar incrustaciones basadas en léxico de baja dimensión utilizando LLMs, abordando los problemas anteriores mediante las siguientes técnicas:

• Agrupación de Tokens (Token Clustering):

  • En lugar de usar el vocabulario original (que puede tener decenas de miles de tokens), LENS aplica K-means a las incrustaciones de los tokens de la cabecera de modelado de lenguaje (LM head).
  • Los tokens se agrupan en $k$ clústeres basados en su similitud semántica o morfológica.
  • Las incrustaciones originales de los tokens se reemplazan por los centroides de estos clústeres. Esto reduce la dimensionalidad (ej. de ~32k a 4k u 8k) y elimina la redundancia, agrupando tokens como "rapid" y "quickly" o "Edu" y "education" en la misma dimensión.

• Atención Bidireccional:

  • Contrario a la creencia popular de que mantener la arquitectura original del LLM es óptimo para incrustaciones densas, LENS modifica el mecanismo de atención durante el ajuste fino (fine-tuning) para ser bidireccional. Esto permite que cada token acceda al contexto completo del texto, mejorando significativamente la calidad de la representación.

• Generación de Representación:

  • Se utiliza un enfoque de log-saturation y max-pooling sobre los logits de salida (similar a SPLADE), pero aplicado sobre los clústeres de tokens.
  • Se utiliza una estrategia de desplazamiento (shifting) de los logits para tener en cuenta la naturaleza autoregresiva del modelo.

• Entrenamiento:

  • Se utiliza una pérdida InfoNCE (contrastiva) estándar.
  • El modelo se entrena con datos públicos exclusivamente (siguiendo el recipe de BGE-en-ICL), sin necesidad de minería de negativos complejos o datos sintéticos, lo que garantiza una comparación justa.

3. Contribuciones Clave

1. Primera Incrustación Léxica con LLMs de Alto Rendimiento: LENS es el primer enfoque que logra un rendimiento competitivo (y a veces superior) a las incrustaciones densas utilizando un enfoque basado en léxico con LLMs.
2. Reducción de Dimensionalidad Natural: Al agrupar tokens, LENS produce representaciones compactas (ej. 4000 o 8000 dimensiones) comparables en tamaño a las incrustaciones densas, facilitando su integración en sistemas existentes (como FAISS) sin restricciones de dispersión estrictas.
3. Poda Natural (Top-K Pruning): Debido a su diseño basado en léxico, LENS permite la poda de dimensiones (retener solo los top-K valores más altos) en tiempo de inferencia sin necesidad de técnicas especiales de entrenamiento como Matryoshka Representation Learning, manteniendo un alto rendimiento.
4. Sinergia con Incrustaciones Densas: La combinación de LENS con incrustaciones densas logra el estado del arte (SOTA) en tareas de recuperación.

4. Resultados Experimentales

Los experimentos se realizaron en el MTEB (Massive Text Embedding Benchmark) y AIR-Bench.

• Rendimiento en MTEB:
  • LENS-8000 (8000 clústeres) logra el mejor rendimiento promedio entre todos los modelos entrenados exclusivamente con datos públicos, superando a BGE-en-ICL (su contraparte densa entrenada con los mismos datos) en 6 de 7 categorías de tareas.
  • Supera a modelos densos complejos como NV-Embed-v2 en tareas específicas, a pesar de usar un proceso de entrenamiento mucho más simple.
• Rendimiento en Recuperación (BEIR):
  • En el subconjunto de recuperación de MTEB (BEIR), la combinación de LENS-8000 + BGE-en-ICL alcanza un puntaje promedio de 63.00, superando a NV-Embed-v2 (62.65) y estableciendo un nuevo SOTA.
• Eficiencia y Poda:
  • Al podar LENS-4000 a solo 256 dimensiones (Top-K), el modelo mantiene un rendimiento muy alto (57.19 en BEIR), mostrando una caída de rendimiento menor que modelos densos entrenados explícitamente para soportar múltiples tamaños (como Gecko).
• Análisis Cualitativo:
  • Los ejemplos muestran que LENS agrupa correctamente sinónimos y variantes morfológicas. Además, demuestra una comprensión profunda del texto (ej. asociar "hipoxia" con "oxígeno" o "Lisboa" con "primavera"), yendo más allá del significado superficial de los tokens.

5. Significado e Impacto

Este trabajo es significativo porque:

• Desafía el Dominio Denso: Demuestra que las incrustaciones basadas en léxico, a menudo relegadas a tareas de recuperación tradicional, pueden ser altamente efectivas en una amplia gama de tareas (clasificación, agrupamiento, similitud) cuando se potencian con LLMs modernos.
• Transparencia e Interpretabilidad: A diferencia de las incrustaciones densas (cajas negras), las incrustaciones de LENS ofrecen transparencia, ya que cada dimensión corresponde a un clúster de palabras comprensible, permitiendo inspeccionar qué conceptos activan la representación.
• Eficiencia Computacional: Ofrece una vía para obtener representaciones ricas semánticamente con menor costo de almacenamiento y búsqueda gracias a la poda natural y la reducción de vocabulario.
• Simplicidad: Logra resultados SOTA sin la necesidad de pipelines de entrenamiento complejos, validando que la arquitectura (atención bidireccional) y el preprocesamiento (agrupamiento de tokens) son factores críticos.

Limitaciones: El trabajo se centra actualmente en el idioma inglés y utiliza Mistral-7B como modelo base. Además, el uso de LLMs como columna vertebral incrementa los costos computacionales en comparación con modelos más pequeños como los usados en SPLADE.