UniDoc-RL: Coarse-to-Fine Visual RAG with Hierarchical Actions and Dense Rewards

1. Problema: Limitaciones del RAG Visual Existente

El artículo identifica que los sistemas actuales de Generación Aumentada por Recuperación (RAG) Visual, diseñados para modelos de lenguaje visión-idioma grandes (LVLMs), sufren de tres deficiencias críticas al manejar documentos visuales densos (como gráficos, informes escaneados y tablas):

• Recuperación Genérica: Los métodos existentes dependen de señales de recuperación genéricas (similitud de embeddings) que a menudo fallan en capturar la semántica específica de la tarea necesaria para el razonamiento complejo, introduciendo ruido visual irrelevante.
• Utilización Visual Pasiva: La mayoría de los enfoques actuales codizan imágenes completas de manera pasiva en el contexto del modelo. Esto ignora la naturaleza jerárquica de la comprensión visual, retiene gran cantidad de ruido de fondo y desperdicia la capacidad del contexto.
• Recompensas Escasas (Sparse Rewards): Los métodos de Aprendizaje por Refuerzo (RL) previos suelen utilizar recompensas basadas únicamente en el resultado final. Esto crea un problema de asignación de crédito, donde el modelo no recibe retroalimentación explícita sobre decisiones intermedias críticas como la búsqueda, la selección de imágenes o el recorte activo, haciendo que la optimización sea ineficaz.

2. Metodología: El Marco UniDoc-RL

Los autores proponen UniDoc-RL, un marco unificado de aprendizaje por refuerzo donde un agente LVLM realiza conjuntamente la recuperación, la reordenación, la percepción visual activa y el razonamiento. La metodología se basa en tres pilares principales:

A. Espacio de Acciones Jerárquico (Coarse-to-Fine)

El proceso de adquisición de información visual se formula como un problema de toma de decisiones secuencial con un espacio de acciones jerárquico que refina progresivamente la evidencia:

1. Búsqueda de Imágenes (Search): Un primer paso de recuperación gruesa que utiliza herramientas externas para obtener un conjunto amplio de candidatos basados en una consulta generada por el modelo.
2. Selección Precisa (Precise Selection): Un mecanismo impulsado por el LVLM que evalúa semánticamente los candidatos recuperados y filtra las imágenes irrelevantes. Esto actúa como un puente entre la recuperación genérica y el razonamiento específico.
3. Percepción Visual Activa (Visual Perception): En lugar de procesar la imagen completa, el agente aprende a realizar operaciones de recorte y zoom (<bbox>) para localizar y extraer regiones de alta densidad de información. Esto imita la atención visual humana, reduciendo la redundancia y preservando detalles de alta resolución.

B. Sistema de Recompensas Densas y Multi-etapa

Para superar el problema de las recompensas escasas, UniDoc-RL introduce un esquema de recompensas densas que supervisa cada etapa del proceso de decisión:

• Recompensa de Búsqueda ($r_{ir}$): Basada en el NDCG (Normalized Discounted Cumulative Gain) para evaluar la calidad de la recuperación inicial.
• Recompensa de Selección ($r_{sel}$): Evalúa si la imagen seleccionada por el agente es relevante. Incluye una estrategia de "pseudo-supervisión" para casos donde no hay imágenes de referencia en el conjunto recuperado.
• Recompensa de Percepción ($r_{crop}$): Basada en la IoU (Intersección sobre Unión) entre los recortes predichos y las regiones de interés reales (ground-truth), incentivando la precisión en la localización.
• Recompensa de Patrón ($r_{pat}$): Basada en reglas para asegurar que el modelo siga el formato correcto de las acciones (uso de etiquetas XML como <search>, <select>, <bbox>).
• Recompensa de Resultado ($r_{ans}$): Una recompensa basada en modelo que evalúa la corrección de la respuesta final.

C. Entrenamiento con GRPO

El modelo se entrena utilizando Optimización de Política Relativa de Grupo (GRPO), un algoritmo que alinea el comportamiento del agente con múltiples objetivos sin necesidad de una red de valores separada, lo que mejora la eficiencia del entrenamiento.

3. Contribuciones Clave

1. Marco UniDoc-RL: Un enfoque unificado que modela la recuperación, la reordenación, la percepción visual activa y el razonamiento en un único proceso de decisión, superando las arquitecturas desacopladas anteriores.
2. Nuevo Dataset de Alta Calidad: Los autores han curado y liberado un conjunto de datos exhaustivo con 12,621 muestras para SFT (Fine-Tuning Supervisado) y 5,537 muestras para RL, que incluyen trayectorias de razonamiento de alta calidad con anotaciones de acciones granulares (búsquedas, selecciones y recortes).
3. Mecanismo de Recompensas Densas: La introducción de un sistema de recompensas específico para cada etapa (búsqueda, selección, recorte) que resuelve el problema de asignación de crédito y permite una optimización end-to-end efectiva.

4. Resultados Experimentales

Los experimentos se realizaron en tres benchmarks desafiantes: SlideVQA, ViDoSeek y MMLongBench.

• Rendimiento Superior: UniDoc-RL supera consistentemente a los métodos más avanzados (SOTA), incluyendo baselines de RAG tradicional, ReAct y métodos basados en RL como VRAG-RL.
• Mejoras Cuantitativas: Se observan ganancias de hasta un 17.7% sobre métodos anteriores basados en RL. Por ejemplo, en el modelo de 3B, UniDoc-RL alcanza un 82.2% de precisión en SlideVQA frente al 65.3% de VRAG-RL.
• Análisis de Ablación:
  • La acción de Selección Precisa mejora significativamente la recuperación de imágenes relevantes, cerrando la brecha semántica entre el recuperador y el razonador.
  • La acción de Percepción Visual es crucial para tareas que requieren discriminación visual fina (como MMLongBench), demostrando que el recorte activo es superior a la lectura pasiva de imágenes completas.
  • El entrenamiento con RL transforma el comportamiento del modelo: mientras que el modelo SFT tiende a ser pasivo y recortar imágenes completas o irrelevantes, el modelo RL aprende a realizar recortes precisos y activos solo cuando es necesario.

5. Significado e Impacto

El trabajo de UniDoc-RL representa un avance significativo en la comprensión de documentos visuales complejos. Al integrar la percepción visual activa (capacidad de "mirar más de cerca") dentro del ciclo de razonamiento y optimizar este proceso mediante recompensas densas, el modelo logra una eficiencia y precisión superiores.

Esto demuestra que los LVLMs no solo deben "ver" imágenes, sino aprender a interactuar con ellas de manera dinámica (buscar, filtrar, recortar) para resolver problemas complejos. La liberación del dataset y el código facilita la investigación futura en RAG visual y agentes multimodales, estableciendo un nuevo estándar para sistemas que requieren razonamiento profundo sobre documentos visuales densos.