CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding

1. El Problema

La Recuperación de Documentos Visuales (VDR, por sus siglas en inglés) es fundamental para sistemas modernos de búsqueda y RAG (Generación Aumentada por Recuperación). Aunque los Modelos de Lenguaje Multimodales (MLLM) han demostrado un gran potencial al generar embeddings multi-vector de alta calidad (como en ColPali), estos enfoques enfrentan un obstáculo crítico: la sobrecarga de almacenamiento.

• Ineficiencia de Representación: Los métodos actuales representan una página de documento como una secuencia de miles de tokens visuales (patches), lo que genera miles de vectores por página.
• Limitaciones de Escalabilidad: Esta gran cantidad de vectores hace que el despliegue en entornos de producción sea costoso en términos de memoria y latencia.
• Deficiencias de los Métodos de Compresión Actuales: Las estrategias existentes para reducir este tamaño, como el pruning (podado) basado en clustering o la agregación de tokens, a menudo resultan en una pérdida significativa de precisión de recuperación o requieren presupuestos de tokens fijos durante el entrenamiento, limitando su adaptabilidad.

2. Metodología: CausalEmbed

El paper propone CausalEmbed, un enfoque novedoso que transforma la generación de embeddings de documentos visuales de un proceso paralelo a uno autoregresivo en el espacio latente.

• Generación Autoregresiva: En lugar de codificar toda la página en paralelo, CausalEmbed utiliza un MLLM pre-entrenado (como Qwen2.5-VL o PaliGemma) para generar una secuencia de vectores latentes de manera secuencial (token a token).
  • Para un documento $I$, el modelo genera una secuencia $Z = [z_1, ..., z_L]$ donde cada token $z_k$ depende del contexto original y de los tokens previamente generados ($z_1, ..., z_{k-1}$).
  • Esto permite sintetizar representaciones compactas y estructuradas con solo decenas de tokens (ej. 32) en lugar de miles.
• Análisis de Flujo de Gradientes: El paper demuestra teóricamente que, bajo un objetivo de interacción tardía (Late Interaction) basado en MaxSim, la generación autoregresiva permite que los gradientes se propaguen a través de toda la cadena de generación. Esto mejora la eficiencia del entrenamiento en comparación con los métodos de parches paralelos, donde la mayoría de los tokens de documento rara vez reciben señales de gradiente directas.
• Función de Pérdida (Training Objective): Para evitar el colapso de la representación y fomentar la diversidad, se introduce un objetivo compuesto:
  1. Pérdida de Alineación Contrastiva ($L_m$): Estándar para emparejar consultas y documentos relevantes.
  2. Pérdida de Refinamiento Progresivo ($L_d$): Penaliza si los tokens adicionales no mejoran la recuperación, incentivando que cada paso autoregresivo añada valor semántico.
  3. Regularización de Diversidad ($L_q$): Penaliza la similitud coseno entre tokens generados en la misma secuencia para evitar patrones repetitivos (homogeneización).

3. Contribuciones Clave

1. Nuevo Paradigma Autoregresivo: Introducen la primera arquitectura que genera embeddings multi-vector de documentos visuales de forma secuencial, alejándose de la codificación paralela basada en parches.
2. Alta Compresión sin Pérdida de Precisión: Logran una compresión de tokens de 30x a 155x (reduciendo de miles a ~32 tokens) manteniendo un rendimiento competitivo e incluso superior a los métodos de podado en escenarios de alta compresión.
3. Escalabilidad en Tiempo de Inferencia (Test-time Scaling): Descubren que la calidad de la recuperación se puede ajustar dinámicamente modificando el número de tokens generados durante la inferencia. Esto permite un equilibrio flexible entre latencia y precisión sin reentrenar el modelo.
4. Generalización Robusta: El método funciona eficazmente en diversas arquitecturas de base (backbones), mejorando significativamente modelos más débiles (como PaliGemma) y superando a los baselines en modelos potentes (Qwen2.5-VL).

4. Resultados Experimentales

Los experimentos se realizaron en los benchmarks ViDoRe V1, V2 y V3, que cubren diversos dominios (académico, industrial, multilingüe).

• Rendimiento Superior: CausalEmbed supera consistentemente a los métodos basados en clustering (K-Means, SemCluster) y pooling bajo restricciones de compresión estrictas.
  • En ViDoRe V2 y V3, la versión CausalQwen (32 tokens) supera a los mejores baselines de podado en un margen significativo.
  • En ViDoRe V1, alcanza un nDCG@5 de 0.811, muy cerca del mejor baseline.
• Mejora en Modelos Débiles: Al aplicarse a PaliGemma (un modelo más pequeño), CausalPali supera al modelo completo ColPali (con todos sus tokens) en un ~5% en todos los benchmarks, demostrando la capacidad del método para extraer semántica de alta calidad.
• Eficiencia de Latencia: Aunque la generación autoregresiva añade un pequeño costo en la codificación de consultas, reduce drásticamente el tiempo de la fase de interacción tardía (MaxSim) (aprox. 11.8x más rápido) debido a la reducción extrema de tokens, resultando en una latencia de extremo a extremo competitiva.
• Comportamiento de Escalado: Las curvas de escalado muestran que aumentar el número de tokens de 1 a 32 mejora la precisión de manera constante, acercándose al límite superior de los modelos de tokens completos.

5. Significado e Impacto

El trabajo de CausalEmbed representa un paso crucial hacia la viabilidad práctica de los embeddings multimodales en aplicaciones del mundo real:

• Viabilidad Económica: Al reducir drásticamente los requisitos de almacenamiento y memoria, hace posible desplegar sistemas de recuperación de documentos visuales de alta precisión en entornos con recursos limitados.
• Flexibilidad Operativa: La capacidad de ajustar la precisión en tiempo de ejecución (test-time scaling) ofrece una nueva flexibilidad para sistemas que requieren diferentes niveles de latencia según la demanda.
• Paradigma Generativo: Establece un nuevo estándar en la recuperación de documentos, demostrando que la generación autoregresiva puede ser más eficiente y expresiva que las representaciones espaciales densas tradicionales, abriendo la puerta a futuras investigaciones en embeddings generativos.

En resumen, CausalEmbed resuelve el cuello de botella de almacenamiento de los métodos multi-vector actuales mediante una generación inteligente y secuencial, logrando un equilibrio óptimo entre eficiencia, escalabilidad y precisión.