IG-Search: Step-Level Information Gain Rewards for Search-Augmented Reasoning

Título: IG-Search: Recompensas de Ganho de Informação em Nível de Passo para Raciocínio Aumentado por Busca

1. O Problema

O aprendizado por reforço (RL) tem sido eficaz para treinar Grandes Modelos de Linguagem (LLMs) em tarefas de raciocínio aumentado por busca (Search-Augmented Reasoning). No entanto, as abordagens existentes enfrentam duas limitações estruturais críticas ao utilizar recompensas baseadas em trajetória (nível de episódio):

1. Incapacidade de Distinguir a Qualidade da Consulta: Em um grupo de trajetórias amostradas para uma mesma pergunta, se duas trajetórias gerarem a resposta final correta, elas recebem a mesma recompensa. Isso ocorre mesmo que uma tenha feito uma busca precisa e a outra uma busca vaga ou redundante. O modelo não consegue aprender qual formulação de consulta foi mais eficaz.
2. Sinal de Gradiente Desaparecente em Cenários de Falha Total: Quando todas as trajetórias amostradas falham em produzir a resposta correta, as vantagens calculadas no nível da trajetória tendem a zero. Isso resulta em um sinal de gradiente quase nulo, impedindo o aprendizado em momentos cruciais (treinamento inicial ou perguntas difíceis de múltiplos saltos), exatamente quando a supervisão é mais necessária.

2. Metodologia: IG-Search

O IG-Search propõe um framework de RL que introduz uma recompensa baseada em Ganho de Informação (IG) em nível de passo. A ideia central é avaliar a qualidade de cada consulta de busca individualmente, independentemente do resultado final da resposta.

A. Definição do Ganho de Informação (IG)

Para cada passo de busca $t$ em uma trajetória, o IG mede o quanto os documentos recuperados realmente melhoram a confiança do modelo na resposta correta ($a^*$) em comparação com um baseline contrafactual.

• Contexto Real ($C_{real}$): O contexto completo até o passo $t$, incluindo os documentos recuperados pela consulta real.
• Contexto Contrafactual ($C_{rand}$): O mesmo contexto, mas substituindo os documentos e refinamentos do passo $t$ por pares aleatórios de documentos e refinamentos de outras perguntas no mesmo batch de treinamento.
• Cálculo: O IG é a diferença entre a probabilidade logarítmica da resposta correta no contexto real e a média das probabilidades nos $N$ contextos aleatórios:
  $$IG_t = \log \pi_{\theta}(a^* | C_{real}) - \frac{1}{N} \sum_{j=1}^{N} \log \pi_{\theta}(a^* | C_{rand, j})$$

B. Modulação de Vantagem por Token (dentro do GRPO)

O algoritmo utiliza o Group Relative Policy Optimization (GRPO). Diferente do GRPO padrão, onde todos os tokens de uma trajetória compartilham a mesma vantagem, o IG-Search modula a vantagem especificamente nos tokens da consulta de busca:

• A vantagem de cada token na consulta é ajustada adicionando um termo proporcional ao IG daquele passo, dividido pelo número de tokens na consulta (para evitar recompensas baseadas no comprimento da frase).
• Isso permite uma atribuição de crédito granular: consultas melhores recebem gradientes positivos, enquanto consultas ruins recebem negativos ou zero, mesmo dentro da mesma trajetória.

C. Mecanismos de Estabilização

Para garantir que o treinamento seja robusto, o IG bruto é processado através de três mecanismos:

1. Filtro de Zona Morta (Dead Zone Filtering): Se $|IG_t| < \delta$ (padrão 0.5), o valor é zerado. Isso evita penalizar ou recompensar consultas em perguntas que o modelo já sabe responder de memória (onde os documentos não trazem nova informação).
2. Escala Assimétrica Negativa: Valores negativos de IG (quando a busca real é pior que a aleatória) são reduzidos por um fator $\lambda < 1$ (padrão 0.1). Isso impede que o modelo aprenda a evitar completamente a busca para evitar penalidades.
3. Recorte Suave (Soft Clipping): Valores extremos de IG são limitados para evitar instabilidade no gradiente.

3. Principais Contribuições

1. Recompensa de Processo Aumentada por Busca: O IG-Search é uma recompensa de processo fundamentada na recuperação, derivada das próprias probabilidades de geração do modelo, sem necessidade de modelos de recompensa externos ou anotações intermediárias.
2. Sinal de Aprendizado em Cenários de Falha: Ao contrário de métodos baseados em trajetória, o IG-Search fornece um sinal de gradiente significativo mesmo quando todas as trajetórias amostradas falham, permitindo que o modelo aprenda a formular melhores consultas durante o processo de raciocínio.
3. Atribuição de Crédito Granular: Permite que o modelo diferencie a qualidade de múltiplos passos de busca dentro de uma única trajetória, algo impossível com recompensas de nível de episódio.

4. Resultados Experimentais

Os autores avaliaram o IG-Search em 7 benchmarks de Perguntas e Respostas (QA), cobrindo cenários de salto único (ex: NQ, TriviaQA) e múltiplos saltos (ex: HotpotQA, Musique).

• Desempenho Geral: Utilizando o modelo base Qwen2.5-3B, o IG-Search alcançou uma precisão média de Exact Match (EM) de 0.430.
  • Superou a melhor linha de base de nível de trajetória (MR-Search) em 1.6 pontos.
  • Superou o método de nível de passo anterior (GiGPO) em 0.9 pontos em média.
• Raciocínio de Múltiplos Saltos: Os ganhos foram particularmente pronunciados em tarefas de múltiplos saltos, onde a precisão de cada passo de busca é crítica.
• Escalabilidade: Os resultados foram consistentes e superiores também na escala de 7B (Qwen2.5-7B).
• Eficiência: A introdução do sinal denso de nível de passo aumentou o tempo de treinamento por passo em apenas ~6.4%, sem alterar a latência de inferência.

5. Significado e Conclusão

O IG-Search representa um avanço significativo no treinamento de LLMs para raciocínio complexo com busca. Ao transformar um sinal de recompensa esparsa (nível de trajetória) em um sinal denso e fundamentado na recuperação (nível de passo), o método permite que os modelos aprendam não apenas se devem buscar, mas como formular consultas eficazes em cada etapa do raciocínio.

A capacidade de fornecer sinais de aprendizado mesmo quando o modelo falha em todas as tentativas de resposta é uma inovação crucial para o treinamento robusto em tarefas difíceis. O trabalho demonstra que a confiança do modelo na resposta correta, condicionada ao contexto recuperado, é um proxy suficiente e prático para a supervisão de nível de passo, eliminando a necessidade de anotações humanas intermediárias complexas.