Hydra: Unifying Document Retrieval and Generation in a Single Vision-Language Model

1. El Problema

Los sistemas actuales de Inteligencia Artificial para documentos (Document AI) enfrentan un dilema de eficiencia y complejidad: deben resolver dos tareas fundamentalmente diferentes que tradicionalmente requieren modelos separados:

1. Recuperación (Retrieval): Encontrar páginas relevantes dadas consultas de texto (ej. usando modelos tipo ColBERT con interacción tardía).
2. Comprensión y Generación (Generation): Extraer e interpretar información de esas páginas para responder preguntas (ej. usando Modelos de Lenguaje Visuales o VLMs autoregresivos).

Ineficiencias del estado actual:

• Duplicidad de recursos: Los enfoques modernos cargan dos modelos distintos (uno para recuperación y otro para generación) en la memoria GPU, duplicando los requisitos de memoria y la complejidad del despliegue.
• Pérdida de capacidades: Los modelos de recuperación (como ColPali o ColQwen2) son VLMs preentrenados que se han ajustado finamente (fine-tuning) con atención bidireccional y cabezas de proyección personalizadas. Este proceso suele degradar o sacrificar la capacidad de generación autoregresiva del modelo base.
• Desperdicio de arquitectura: Ambos modelos comparten la misma arquitectura base (transformador visión-lenguaje), pero se ejecutan de forma independiente.

2. Metodología: La Arquitectura Hydra

Hydra propone un enfoque de cabeza dual que permite que un único Modelo Visión-Lenguaje (VLM) realice ambas tareas mediante el uso de un adaptador LoRA (Low-Rank Adaptation) que se activa o desactiva en tiempo de inferencia.

Componentes Clave:

• Modelo Base: Se utiliza un VLM congelado (ej. Qwen3.5-4B) con un codificador de visión congelado.
• Adaptador LoRA Único: Se entrena un solo adaptador LoRA ($r=16, \alpha=64$) exclusivamente para la tarea de recuperación. Este adaptador se aplica a las capas de proyección del lenguaje y a una nueva cabeza de proyección personalizada (custom_text_proj).
• Dos Modos de Operación (Switching):
  1. Modo Recuperación (LoRA ON): Se habilita el adaptador y se cambia la atención de las capas completas a atención bidireccional. La salida pasa por custom_text_proj para generar embeddings multivectoriales (320 dimensiones) compatibles con ColBERT.
  2. Modo Generación (LoRA OFF): Se deshabilita el adaptador, restaurando exactamente los pesos del modelo base ($W_{base}$). Se vuelve a la atención causal (autoregresiva) y se utiliza la cabeza de lenguaje original (lm_head) para generar texto.

Requisitos de Ingeniería Críticos (Contribución Técnica)

El artículo identifica tres requisitos de ingeniería que, si se omiten, rompen silenciosamente la capacidad de generación al desactivar el LoRA, a pesar de que la teoría sugiere que los pesos se restauran perfectamente:

1. Restauración del modo de atención: Las capas de atención deben volver explícitamente a la máscara causal. Si se mantienen bidireccionales, la generación falla porque el modelo puede "ver" tokens futuros.
2. Preservación de lm_head: La cabeza de lenguaje (lm_head) debe cargarse desde una instancia separada del modelo base. Durante el entrenamiento, incluso si no es objetivo de LoRA, los gradientes pueden corromperla debido a pesos atados (tied weights) con las embeddings de entrada o a la sincronización de gradientes en DDP (Distributed Data Parallel).
3. Soporte para KV-Cache: Para que la generación sea viable, se debe implementar un manejo manual del caché de claves y valores (KV-cache), ya que los wrappers estándar de recuperación no lo soportan. Sin esto, la generación es ~38 veces más lenta.

3. Contribuciones Principales

1. Unificación con un solo adaptador: Demuestra que un solo adaptador LoRA, entrenado solo para recuperación, es suficiente para habilitar tanto la recuperación ColBERT como la generación autoregresiva de alta calidad. No se requiere entrenamiento para generación.
2. Análisis de Fallos Silenciosos: Identifica y resuelve los tres mecanismos de fallo técnico (atención, lm_head, KV-cache) que impiden que este enfoque funcione en pipelines estándar.
3. Validación de la Toggling (Conmutación): Prueba empíricamente que la conmutación LoRA ON/OFF recupera los pesos base de manera byte-identica (100% de coincidencia en 10,500 muestras), eliminando la necesidad de entrenamiento conjunto (joint training).
4. Reducción de Memoria: Logra una reducción del 41% en el uso de memoria VRAM pico al eliminar la necesidad de cargar dos modelos.

4. Resultados Experimentales

Rendimiento de Recuperación (ViDoRe)

• ViDoRe V1: Hydra (4B) alcanza un nDCG@5 promedio de 0.8842, quedando a menos de 1 punto porcentual de una línea base controlada de un solo cabezal (0.8892) entrenada bajo las mismas condiciones.
• ViDoRe V2 y V3: Hydra supera a la línea base en tareas más difíciles (V2: +0.7 pp, V3: +4.7 pp), aunque los autores notan que esto requiere verificación con múltiples semillas debido a que se basan en una sola ejecución de entrenamiento.
• Conclusión: No hay un costo significativo en la recuperación al unificar el modelo.

Calidad de Generación

• Equivalencia Exacta: En 10,500 muestras (DocVQA y TextVQA), la salida de Hydra con LoRA desactivado es byte-identica a la del modelo base original (Qwen3.5) en un 100% de los casos bajo decodificación greedy.
• Pruebas ANLS: En cuatro benchmarks de VQA (DocVQA, ChartQA, InfoVQA, TextVQA), la diferencia máxima absoluta en ANLS es de 0.0044, sin degradación estadísticamente significativa.
• Ablación de Entrenamiento Conjunto: Se comparó Hydra con un enfoque de "entrenamiento conjunto" estilo GritLM (alternando lotes de recuperación y generación).
  • Resultado: El entrenamiento conjunto no ofrece beneficios medibles.
  • Fallo Crítico: El modelo entrenado conjuntamente falla catastróficamente al intentar generar texto con el LoRA activado (produce solo el token "The" independientemente de la imagen). Esto confirma que, dentro del régimen LoRA de rango bajo ($r=16$), la conmutación es estructuralmente necesaria; el espacio de subespacio de bajo rango no puede soportar simultáneamente ambos modos de atención.

Eficiencia

• Memoria: Reducción de 17,913 MB (dos modelos) a 10,496 MB (un modelo), un ahorro del 41%.
• Latencia: El cambio de modo (retrieval $\to$ generation) introduce una latencia de ~5.9 ms, lo cual es insignificante (1.8%) respecto al tiempo de inferencia.
• Aislamiento de Estado: Se verificó que no hay fugas de estado (KV-cache) entre modos; los embeddings y las generaciones son idénticos en ciclos repetidos.

Extensión Omnimodal (Prueba de Concepto)

El mecanismo se aplicó sin entrenamiento adicional a Qwen2.5-Omni-3B (que soporta audio y video).

• Logró recuperación de imágenes, audio y video (zero-shot para audio/video).
• Mantuvo la generación de texto y voz.
• Esto demuestra que la arquitectura es generalizable a diferentes familias de modelos y modalidades.

5. Significado e Impacto

• Eficiencia Operativa: Hydra resuelve el problema de la "duplicidad" en RAG (Retrieval-Augmented Generation) visual, permitiendo desplegar sistemas de documentos complejos con la mitad de la memoria de GPU.
• Cambio de Paradigma en LoRA: El artículo redefine el uso de LoRA no solo como una herramienta de entrenamiento eficiente, sino como un interruptor de modo en tiempo de inferencia.
• Simplicidad vs. Complejidad: Demuestra que los enfoques complejos de entrenamiento conjunto (GritLM) son innecesarios para la unificación de tareas en modelos de visión-lenguaje cuando se utiliza LoRA, siempre que se resuelvan los problemas de ingeniería de implementación.
• Limitaciones: Actualmente probado en arquitecturas Qwen. La evaluación de recuperación de video y la evaluación de extremo a extremo (recuperar + generar respuesta final) son áreas de trabajo futuro.

En resumen, Hydra es una arquitectura práctica y eficiente que unifica la recuperación de documentos visuales y la generación de respuestas en un solo modelo, reduciendo drásticamente los costos de infraestructura sin sacrificar la precisión, siempre que se implementen los requisitos de ingeniería específicos identificados por los autores.