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

1. Problema e Motivação

O campo de Busca Agêntica (Agentic Search) envolve agentes de IA que iterativamente raciocinam, emitem consultas de busca e sintetizam informações recuperadas para responder a perguntas complexas. Embora métodos recentes baseados em Aprendizado por Reforço (RL), como o Search-R1, tenham demonstrado progresso ao treinar o agente de raciocínio, eles assumem que o sistema de recuperação (o motor de busca ou o classificador de documentos) é uma ferramenta fixa e imutável.

Os autores identificam que essa suposição cria um gargalo crítico:

• Análise Preliminar: Experimentos mostram que, mesmo com um agente de raciocínio forte, a qualidade da recuperação limita severamente o desempenho final.
• Experimento Oráculo: Ao promover documentos que contêm a resposta correta para o topo da lista (simulando um sistema de recuperação perfeito), observa-se uma melhoria relativa de F1 de +14,7% para um agente de 7B parâmetros e +26,8% para um agente de 3B parâmetros em sete benchmarks de QA.
• Conclusão: A maior parte dos erros não vem do raciocínio do agente, mas da incapacidade do sistema de recuperação de retornar documentos relevantes, forçando o agente a tentar reformular consultas ou tirar conclusões erradas de informações ruidosas.

O objetivo do trabalho é superar esse gargalo treinando simultaneamente o agente de raciocínio e o sistema de recuperação.

---

2. Metodologia: COSEARCH

O COSEARCH é um framework que utiliza Otimização de Política Relativa em Grupo (GRPO) para treinar conjuntamente um agente de raciocínio principal e um ranker (classificador) de documentos generativo.

Arquitetura do Sistema

O sistema de recuperação é decomposto em duas etapas:

1. Recuperador Denso (Fixo): Um modelo de recuperação densa (ex: E5) que retorna um conjunto candidato inicial de $N$ documentos.
2. Ranker Generativo (Treinável): Um modelo de linguagem que recebe os documentos candidatos e as sub-consultas do agente, analisando a relevância e reordenando um subconjunto de $K$ documentos para o agente observar.

O agente principal opera em um loop estilo ReAct (Raciocinar-Agir), gerando pensamentos e consultas, enquanto o ranker seleciona os documentos mais úteis para essa etapa específica.

Desafios Técnicos e Soluções

Aplicar GRPO ao ranker apresenta dois desafios principais, que são resolvidos da seguinte forma:

A. Agrupamento Semântico (Semantic Grouping)

• Desafio: O GRPO requer que múltiplas respostas sejam geradas a partir da mesma entrada para calcular vantagens normalizadas. Para o agente principal, isso é natural (mesma pergunta do usuário). Para o ranker, as entradas variam porque diferentes trajetórias de raciocínio geram sub-consultas diferentes.
• Solução: Os autores introduzem uma estratégia de agrupamento semântico. As sub-consultas de múltiplas trajetórias são agrupadas com base na similaridade F1 a nível de token.
  • Consultas semanticamente equivalentes (mesmo que com formulações superficiais diferentes) são agrupadas.
  • Grupos com menos de um tamanho mínimo ($k_{min}$) são descartados para garantir estimativas de vantagem estáveis.
  • Isso permite treinar o ranker sem a necessidade de rollouts adicionais (custo zero), reutilizando as trajetórias já geradas pelo agente principal.

B. Recompensa Composta (Composite Reward)

• Desafio: A recompensa baseada apenas na correção da resposta final é ruidosa para o ranker (o ranker pode ter feito um bom trabalho, mas o agente raciocinou mal, ou vice-versa). Recompensas baseadas apenas em relevância (Hit@k) falham em supervisionar etapas intermediárias onde documentos úteis para o raciocínio não contêm a resposta final.
• Solução: Foi desenhada uma recompensa composta que combina:
  1. Recompensa de Relevância ($r_{rel}$): Baseada em métricas de qualidade de ranking (Hit@k) usando rótulos pseudo-relevantes (documentos que contêm a resposta correta).
  2. Recompensa de Trajetória ($r_{main}$): Baseada na correção da resposta final (F1).
  • Lógica de Combinação: Se documentos relevantes existem no conjunto candidato mas estão mal classificados, o sistema prioriza a recompensa de relevância. Se o ranking já é forte, adiciona-se a recompensa de trajetória para sinalizar se os documentos contribuíram para a resposta correta. Se nenhum documento relevante existe, usa-se apenas a correção da resposta.

---

3. Contribuições Principais

1. Identificação do Gargalo: Demonstração empírica de que a qualidade da recuperação é um limitador fundamental em sistemas de busca agêntica, com lacunas de desempenho significativas entre sistemas fixos e ideais.
2. Framework COSEARCH: Proposta de um método de RL que treina o agente de raciocínio e o ranker de documentos de forma conjunta e end-to-end.
3. Técnicas Inovadoras:
   • Agrupamento Semântico: Permite a aplicação eficiente do GRPO em componentes com entradas variáveis (sub-consultas) sem custo computacional extra de amostragem.
   • Recompensa Composta: Um sinal de aprendizado híbrido que fornece feedback imediato sobre a precisão da recuperação e feedback de longo prazo sobre a utilidade para a tarefa final.

---

4. Resultados Experimentais

Os experimentos foram realizados em 7 benchmarks de QA (incluindo single-hop como NQ, TriviaQA e multi-hop como HotpotQA, Musique).

• Desempenho Geral: O COSEARCH superou consistentemente todas as linhas de base fortes, incluindo Search-R1, ZeroSearch e variantes com ranker fixo.
  • Com um agente principal de 7B, houve uma melhoria relativa de +6,6% no F1 em relação ao Search-R1.
  • Com um agente principal de 3B (acoplado a um ranker de 7B), a melhoria foi de +10,8% (ou +10,0% conforme o resumo final) em relação ao Search-R1.
• Eficiência de Busca: O COSEARCH reduziu o número médio de voltas de busca (search turns) necessárias para resolver perguntas, aproximando-se da eficiência do sistema "Oráculo".
• Estudos de Ablação:
  • Remover o agrupamento semântico causou a maior queda de desempenho (-3,3% relativo), confirmando sua importância para a estabilidade do treino.
  • A recompensa composta é superior a recompensas puras de relevância ou de trajetória.
  • O uso de um backbone compartilhado entre agente e ranker degradou o desempenho devido a gradientes conflitantes, validando a necessidade de arquiteturas separadas.

---

5. Significado e Conclusão

O trabalho COSEARCH estabelece que o treinamento conjunto do sistema de raciocínio e do sistema de recuperação não é apenas viável, mas essencial para escalar o desempenho de agentes de busca.

• Mudança de Paradigma: Move-se da visão de "recuperação como ferramenta fixa" para "recuperação como componente adaptativo" que aprende a complementar o raciocínio do agente.
• Eficiência: A estratégia de agrupamento semântico resolve o problema de escalabilidade ao permitir o treinamento do ranker sem a explosão exponencial de rollouts necessária em abordagens ingênuas.
• Futuro: Os resultados sugerem que futuros agentes de busca devem ser projetados com mecanismos de otimização conjunta, onde a recuperação se adapta dinamicamente às necessidades de raciocínio do agente em tempo real.

O código do projeto está disponível publicamente, facilitando a reprodução e o avanço da pesquisa em agentes de busca autônomos.