Resumo Técnico: PAIR (Modelo de Recompensa Interna Consciente de Prefixo para Otimização de Agentes Multi-turno)

1. Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) atuais têm dificuldade com tarefas complexas de agentes multi-etapa que exigem sequências de ações (por exemplo, chamadas de API, invocações de ferramentas) e feedback ambiental. Embora a Otimização de Política Relativa em Grupo (GRPO) seja um algoritmo de aprendizado por reforço (RL) líder para essas tarefas devido à sua simplicidade e eficiência, ela sofre de recompensas de resultado esparsas. No GRPO padrão, uma única recompensa escalar é atribuída apenas após a conclusão da trajetória, tornando difícil realizar a atribuição de crédito entre etapas intermediárias. Essa esparsidade leva a uma convergência lenta e a uma baixa eficiência de amostragem.

As tentativas existentes de densificar recompensas enfrentam limitações significativas:

• Rollouts Completos: A propagação de recompensas de resultado para trás exige a conclusão de trajetórias completas, gerando rollouts desperdiçados a partir de caminhos já falhos.
• Juízes Externos: O uso de LLMs externos (por exemplo, GPT-4) para avaliar cada etapa incorre em altos custos de API e latência, tornando-o impraticável para RL online em grande escala.
• Recompensas Intrínsecas: Métodos que dependem de respostas de verdade fundamental (por exemplo, log-probabilidades forçadas por professor) exigem acesso em tempo de execução a rótulos de ouro, o que frequentemente não está disponível ou é operacionalmente custoso.
• Modelos de Recompensa Externos: Modelos de Recompensa de Processo (PRMs) treinados separadamente introduzem sobrecarga no treinamento e na manutenção.

O desafio central abordado por este trabalho é: Podemos gerar sinais de recompensa densos, em nível de etapa, usando apenas os estados internos do agente, sem modelos externos, dependências de verdade fundamental ou rollouts de trajetória completa?

2. Metodologia

2.1 Motivação Empírica: A Falha da Sondagem Ingênua

Os autores investigam primeiro se sondas de correção internas (treinadas para prever a correção da etapa a partir de estados ocultos) podem servir como recompensas. Eles introduzem o conceito de contaminação de prefixo, onde um agente gera uma etapa intermediária errônea, corrompendo o prefixo para etapas subsequentes.

Através de uma análise sistemática na base Qwen-2.5-7B-Instruct, eles descobrem três achados críticos:

1. Degradação do Estado Oculto: Sondas de estado oculto (por exemplo, Último-Token, Média-Pool) degradam-se severamente sob contaminação de prefixo. Elas rastreiam consistência de crença (coerência com o prefixo corrompido) em vez de correção fundamentada (utilidade real da tarefa). Consequentemente, elas frequentemente recompensam turnos que perpetuam erros e penalizam turnos de reparo.
2. Robustez da Atenção: Características baseadas em atenção permanecem relativamente robustas à contaminação. Elas capturam padrões de comportamento estrutural (por exemplo, se o modelo atende a tokens relevantes para a tarefa) em vez de coerência semântica com o prefixo.
3. Complementaridade: Os estados ocultos se destacam em prefixos limpos, mas falham em prefixos contaminados; as características de atenção são robustas, mas têm desempenho inferior em prefixos limpos. Uma concatenação ingênua de ambos os sinais falha porque o sinal invertido do estado oculto cancela o sinal correto da atenção.

2.2 Método Proposto: PAIR

Para aproveitar a natureza complementar desses sinais, os autores propõem o PAIR (Modelo de Recompensa Interna Consciente de Prefixo), uma arquitetura de dois estágios:

• Estágio 1: Estimador de Consistência de Crença (Congelado): Uma sonda de regressão logística pré-treinada mapeia características de estado oculto (especificamente o estado oculto do último-token) para uma pontuação de consistência de crença ($s_{bc}$). Esta sonda é treinada em dados limpos e congelada durante o treinamento de RL. Ela captura a capacidade padrão da sonda de detectar correção quando o prefixo está limpo.
• Estágio 2: Correção Baseada em Atenção: Um cabeçalho leve de regressão logística toma as características de atenção (estatísticas como magnitude de pico, dispersão, razão de prefixo e razão própria) concatenadas com $s_{bc}$ como entrada. Ele prevê uma pontuação final ($s_{final}$).
  • Mecanismo: Em prefixos limpos, o cabeçalho de correção aprende a produzir $s_{final} \approx s_{bc}$. Em prefixos contaminados, as características de atenção detectam anomalias estruturais, e o cabeçalho aprende a ajustar $s_{final}$ longe da $s_{bc} enviesada em direção à correção fundamentada.
• Integração de Recompensa: A pontuação final é transformada via uma sigmoide com limite de temperatura e combinada com um mecanismo de recompensa baseada em momento. Isso contrasta a pontuação da etapa atual com a média móvel da trajetória para recuperar a variância relativa ao grupo, o que é essencial para a estimativa de vantagem do GRPO.

3. Contribuições Principais

1. Descoberta Empírica: O artigo fornece o primeiro estudo sistemático demonstrando que sondas de correção internas falham sob contaminação de prefixo em configurações multi-etapa, rastreando especificamente a consistência de crença em vez da correção fundamentada.
2. Insight Mecanístico: Estabelece que os estados ocultos codificam consistência de crença, enquanto os padrões de atenção codificam comportamento de roteamento estrutural. Esses sinais são complementares, explicando por que a concatenação simples falha e orientando o design de arquiteturas adaptativas.
3. Inovação Metodológica: O PAIR é proposto como um modelo de dois estágios que combina adaptativamente sinais de estado oculto e de atenção. Ele preserva alto desempenho em trajetórias limpas enquanto corrige viés sob contaminação.
4. Eficiência Prática: O PAIR opera sem chamadas de LLM externas, acesso a verdade fundamental em tempo de execução ou treinamento separado de modelo de recompensa. Ele incorre em custo de inferência negligenciável (nível de sonda) reutilizando ativações já computadas durante a passagem direta da política.

4. Resultados Experimentais

Os autores avaliam o PAIR em dois benchmarks de agentes multi-etapa: GTA (Agentes de Ferramentas Gerais) e ToolBench, usando o Qwen-2.5-7B-Instruct como modelo de política.

• Previsão em Nível de Etapa: O PAIR alcança a maior AUROC (0,9217 no GTA, 0,8129 no ToolBench) entre todas as bases de sondagem internas, superando sondas de estado oculto, sondas de atenção e métodos não supervisionados (por exemplo, Lente de Retroanálise, Entropia de Cabeçalho).
• Desempenho de RL: No treinamento GRPO, o PAIR alcança as maiores taxas de sucesso em ambos os conjuntos de dados (0,2489 no GTA, 0,4498 no ToolBench), superando a segunda melhor base (LLM-como-Juiz) no ToolBench e igualando-a no GTA.
• Transferência entre Domínios: Quando treinado no GTA e avaliado no HotpotQA (sem ajuste fino), o PAIR ocupa o primeiro lugar entre 16 métodos, demonstrando transferibilidade superior em comparação com métodos baseados em rollout ou dependentes de verdade fundamental.
• Eficiência: O PAIR opera com latência de nível de sonda (~0,2 ms/etapa), o que é aproximadamente 2.500× mais rápido que chamadas de API de LLM-como-Juiz e ordens de magnitude mais rápido que métodos de busca em árvore.

5. Significado e Alegações

O artigo afirma que o PAIR aborda as quatro limitações primárias dos mecanismos de recompensa densa existentes simultaneamente:

1. Sem Rollouts Completos: As recompensas são computadas por etapa sem esperar pela conclusão da trajetória.
2. Sem Chamadas Externas: Ele depende exclusivamente das representações internas do agente, eliminando custos de API e latência.
3. Sem Dependência de Verdade Fundamental: A sonda é treinada offline e aplicada em tempo de execução sem necessidade de respostas de ouro para cada etapa.
4. Sem Modelo de Recompensa Separado: Ele não exige o treinamento de uma rede de recompensa separada e em grande escala.

Os autores posicionam o PAIR como uma solução prática para permitir sinais de recompensa densos, em nível de etapa, no treinamento de RL de agentes multi-turno. Ao resolver o modo de falha das sondas de estado oculto sob contaminação de prefixo através de uma arquitetura adaptativa e corrigida por atenção, o PAIR torna a análise de estado interno uma alternativa viável e eficiente a fontes de recompensa externas para otimizar comportamentos complexos de agentes.