COSEARCH: Joint Training of Reasoning and Document Ranking via Reinforcement Learning for Agentic Search

Resumen Técnico: COSEARCH

1. El Problema: El Cuello de Botella de la Recuperación

El campo de la búsqueda agéntica (agentic search) ha avanzado significativamente mediante el uso de Aprendizaje por Refuerzo (RL) para entrenar agentes que razonan, generan consultas y sintetizan información para responder preguntas complejas. Sin embargo, un enfoque común en trabajos recientes (como Search-R1) es tratar el sistema de recuperación (el motor de búsqueda o el modelo de recuperación) como una herramienta fija e inmutable.

Los autores identifican que esto introduce un cuello de botella crítico:

• Análisis de errores: La mayoría de los fallos en agentes entrenados con RL no provienen de la incapacidad de razonamiento del agente, sino de que el sistema de recuperación falla en devolver documentos relevantes para las consultas generadas.
• Brecha del Oráculo: Experimentos preliminares muestran que, si se utiliza un sistema de recuperación "oráculo" (que siempre coloca los documentos correctos en la parte superior), el rendimiento del agente mejora drásticamente.
  • Para un agente de 7B parámetros, la mejora es del +14.7% en F1 relativo.
  • Para un agente de 3B parámetros, la mejora es del +26.8% en F1 relativo.
• Conclusión: El sistema de recuperación es un limitante fundamental para escalar el rendimiento de la búsqueda agéntica, independientemente de la capacidad de razonamiento del agente.

2. Metodología: COSEARCH

Para abordar este problema, los autores proponen COSEARCH, un marco de trabajo que entrena conjuntamente un agente de razonamiento principal y un modelo generativo de clasificación de documentos (ranker) mediante Aprendizaje por Refuerzo.

Arquitectura del Sistema:
El sistema de recuperación se descompone en dos etapas:

1. Recuperador denso (Fijo): Un recuperador de primera etapa (ej. E5) que devuelve un conjunto candidato de $N$ documentos.
2. Ranker Generativo (Entrenable): Un modelo (basado en LLM) que recibe las consultas del agente y los documentos candidatos, razona sobre su relevancia y selecciona/reranea un subconjunto de $K$ documentos para que el agente los observe.

Optimización Conjunta con GRPO:
Ambos componentes (Agente Principal y Ranker) se optimizan simultáneamente utilizando Group Relative Policy Optimization (GRPO), un algoritmo de RL libre de valores (value-free) que es más eficiente que PPO.

Desafíos Técnicos y Soluciones:

• Desafío 1: Agrupación para el Ranker (Semantic Grouping)

  • Problema: GRPO requiere agrupar múltiples respuestas generadas desde el mismo prompt de entrada para calcular ventajas normalizadas. Para el agente principal, esto es natural (misma consulta de usuario). Sin embargo, el ranker recibe sub-consultas que varían entre diferentes trayectorias de razonamiento, impidiendo una agrupación directa.
  • Solución: Los autores introducen una estrategia de agrupación semántica. Agrupan las sub-consultas basándose en la similitud F1 a nivel de tokens. Si varias sub-consultas de diferentes trayectorias son semánticamente equivalentes (aunque difieran en la superficie), se agrupan para formar un grupo válido de GRPO. Esto evita la necesidad de realizar rollouts adicionales costosos para el ranker.

• Desafío 2: Diseño de Recompensa Compuesta

  • Problema: La recompensa basada solo en la respuesta final es ruidosa para el ranker (un buen documento puede ser ignorado por un mal razonamiento, y viceversa). Las recompensas de relevancia pura (Hit@k) no capturan si el documento contribuye a la respuesta final.
  • Solución: Se diseña una recompensa compuesta que combina:
    1. Recompensa de Relevancia ($r_{rel}$): Basada en métricas de calidad de ranking (Hit@k) usando etiquetas de pseudo-relevancia (documentos que contienen la respuesta correcta).
    2. Recompensa de Trayectoria ($r_{main}$): Basada en la corrección de la respuesta final del agente.
  • Lógica: Si hay documentos relevantes disponibles pero mal clasificados, se prioriza $r_{rel}$. Si el ranking ya es bueno, se añade $r_{main}$ para dar crédito a largo plazo. Si no hay documentos relevantes en el candidato, se usa solo $r_{main}$.

3. Contribuciones Clave

1. Identificación del Cuello de Botella: Demostración empírica de que la recuperación fija limita severamente el potencial de los agentes de búsqueda, con brechas de rendimiento de hasta +26.8% frente a un oráculo.
2. Marco COSEARCH: Propuesta de un marco de entrenamiento conjunto RL para el agente y el ranker, superando la limitación de tratar la recuperación como una herramienta estática.
3. Innovaciones en GRPO:
   • Estrategia de agrupación semántica para permitir el entrenamiento de rankers con entradas variables sin coste adicional de muestreo.
   • Diseño de recompensa compuesta que equilibra la precisión de recuperación inmediata con la utilidad a largo plazo de la respuesta.

4. Resultados Experimentales

El método se evaluó en 7 benchmarks de preguntas y respuestas (QA), cubriendo tanto saltos únicos (single-hop) como múltiples (multi-hop): PopQA, NQ, TriviaQA, HotpotQA, 2WikiMultiHopQA, Musique y Bamboogle.

• Rendimiento General:

  • Con un agente principal de 7B (Qwen2.5-7B), COSEARCH supera a Search-R1 (el estado del arte anterior) en un +6.6% de F1 relativo.
  • Con un agente principal de 3B, la mejora es aún más notable: +10.8% sobre Search-R1.
  • COSEARCH cierra aproximadamente el 27% de la brecha de rendimiento respecto al límite superior (Oráculo) para el agente de 7B.

• Análisis de Ablación:

  • La eliminación de la agrupación semántica causa la mayor caída de rendimiento (-3.3% relativo), confirmando su importancia crítica.
  • La recompensa compuesta es superior a usar solo recompensas de relevancia o solo de respuesta final.
  • El uso de un ranker de 3B causó inestabilidad en el entrenamiento, sugiriendo que se necesita cierta capacidad del modelo para manejar la salida de permutaciones complejas.

• Eficiencia:

  • COSEARCH reduce el número promedio de vueltas de búsqueda necesarias para resolver una pregunta en comparación con sistemas de ranker fijos, indicando que la recuperación de mayor calidad permite al agente llegar a la respuesta más rápido.

5. Significado e Impacto

El trabajo COSEARCH marca un cambio de paradigma en la construcción de agentes de búsqueda. Demuestra que:

1. La recuperación y el razonamiento deben co-evolucionar: Optimizar solo el agente de razonamiento mientras se mantiene el sistema de recuperación estático es subóptimo.
2. La viabilidad del entrenamiento conjunto: Es posible y altamente efectivo entrenar el sistema de recuperación (ranker) y el agente de razonamiento de forma simultánea mediante RL, utilizando recompensas derivadas de la corrección de la respuesta final.
3. Escalabilidad: La mejora es más pronunciada en modelos más pequeños (3B), lo que sugiere que un sistema de recuperación adaptativo es crucial para desbloquear el potencial de modelos de lenguaje más eficientes en tareas de búsqueda complejas.

En resumen, COSEARCH establece que el entrenamiento conjunto de la recuperación y el razonamiento es un ingrediente clave para la próxima generación de agentes de búsqueda autónomos.