Resumo Técnico: POISE (Otimização de Política com Estimativa de Valor de Estado Interno)

Declaração do Problema

A aprendizagem por reforço com recompensas verificáveis (RLVR) é um motor central para a melhoria de Modelos de Raciocínio de Grande Escala (LRMs) em tarefas complexas como o raciocínio matemático. No entanto, uma RLVR estável e eficiente depende fortemente de uma estimativa precisa da linha de base para reduzir a variância dos gradientes da política. As abordagens existentes incorrem em custos computacionais significativos:

• PPO (Otimização de Política Próxima): Exige o treinamento de um modelo crítico separado, na escala de LLM, que processa a sequência gerada completa a cada atualização, aproximadamente dobrando o consumo de memória e a complexidade de otimização.
• GRPO (Otimização de Política Relativa de Grupo): Evita um crítico separado estimando linhas de base a partir da média de recompensas de um grupo de trajetórias amostradas por prompt. Esta abordagem "troca parâmetros por amostras", exigindo múltiplas trajetórias por prompt para manter uma linha de base estável. Sob um orçamento de computação fixo, isso reduz a diversidade de prompts em um lote, inflando a variância do gradiente. Além disso, uma computação substancial é desperdiçada em prompts "degenerados" onde todas as trajetórias recebem recompensas idênticas (vantagem zero).

A questão central abordada por este trabalho é: É possível extrair uma linha de base eficaz dos cálculos já realizados durante o treinamento da política, eliminando a necessidade de recursos extras?

Metodologia: POISE

Os autores propõem o POISE (Otimização de Política com Estimativa de Valor de Estado Interno), um método que aproveita as próprias representações internas do modelo de política para estimar a função de valor $V^\pi(x)$ a custo negligenciável.

1. Estimativa de Valor de Estado Interno

Em vez de treinar um crítico separado, o POISE treina uma sonda leve para prever a recompensa verificável esperada diretamente a partir dos estados ocultos do modelo de política e das estatísticas de tokens coletadas durante a passagem frontal.

• Recursos de Entrada: A sonda utiliza três sinais complementares:
  1. Recursos no nível do prompt: Estados ocultos dos tokens finais do prompt, capturando a representação do modelo do prompt e a dificuldade antecipada.
  2. Recursos no nível da trajetória: Estados ocultos do final da trajetória de raciocínio, expondo informações relevantes para o valor não presentes apenas no prompt.
  3. Estatísticas de entropia no nível do token: Recursos leves de incerteza.
• Objetivo de Treinamento: A sonda é treinada para prever o valor no nível do prompt $V^\pi(x)$. Para evitar que o alvo vaze informações sobre a trajetória específica sendo avaliada, os autores utilizam um alvo de Monte Carlo leave-one-out. Para um prompt $x$ com $K$ trajetórias amostradas, o alvo para a trajetória $i$ é a recompensa média das outras $K-1$ trajetórias.

2. Construção Cruzada de Trajetórias para Imparcialidade

Um desafio teórico crítico é garantir que a linha de base permaneça condicionalmente independente da ação (a resposta amostrada) para preservar a imparcialidade do gradiente da política. Se a linha de base depender da mesma trajetória que avalia, introduz viés.

• Solução: O POISE emprega uma construção cruzada de trajetórias. Para cada prompt, duas trajetórias independentes, $y^{(1)}$ e $y^{(2)}$, são amostradas.
  • A linha de base para $y^{(1)}$ é prevista usando os estados internos de $y^{(2)}$.
  • A linha de base para $y^{(2)}$ é prevista usando os estados internos de $y^{(1)}$.
• Isso garante que a linha de base usada para atualizar uma trajetória específica seja derivada de uma amostra independente, satisfazendo a condição $E[b(x) \nabla \log \pi(y|x)] = 0$.

3. Integração na Otimização de Política

O POISE integra este estimador em um loop de otimização estilo PPO:

• Treinamento Online: O estimador de valor é treinado conjuntamente com a política em um buffer deslizante de trajetórias recentes. Como os recursos são extraídos da passagem frontal já exigida para atualizações da política, o sobrecarga computacional é negligenciável.
• Cálculo de Vantagem: As vantagens são calculadas como $A^{(i)}(x) = R(x, y^{(i)}) - b^{(i)}(x)$, onde $b^{(i)}(x)$ é a linha de base cruzada de trajetórias.
• Eficiência: Ao contrário do GRPO, que exige um grande grupo (por exemplo, 8 trajetórias) para estimar uma linha de base, o POISE requer apenas um par de trajetórias. Isso permite que o orçamento de computação fixo seja alocado para um número maior de prompts distintos por lote, reduzindo teoricamente a variância do gradiente.

Contribuições Principais

1. Estimativa de Linha de Base Inovadora: Introduz um método para derivar linhas de base a partir dos estados internos do modelo de política, removendo a necessidade de um crítico separado na escala de LLM ou de grandes grupos de trajetórias.
2. Imparcialidade Teórica: Propõe uma arquitetura cruzada de trajetórias que garante a independência condicional necessária para gradientes de política imparciais, enquanto utiliza recursos condicionados à trajetória.
3. Eficiência Computacional: Demonstra que uma sonda leve (regressão linear) pode alcançar desempenho de estimativa de valor comparável ou superior a um crítico separado, exigindo significativamente menos computação.
4. Generalização: Mostra que os sinais de estado interno generalizam entre diferentes domínios (matemática, codificação, chamada de ferramentas, seguimento de instruções) e escalas de modelo.

Resultados Experimentais

O método foi avaliado no Qwen3-4B e no DeepSeek-R1-Distill-Qwen-1.5B em benchmarks de raciocínio matemático de nível olímpico (AMC, AIME, HMMT, BRUMO).

• Desempenho: O POISE alcança desempenho comparável ao DAPO (um algoritmo baseado em GRPO de última geração).
  • No Qwen3-4B, o POISE alcançou uma pontuação média Avg@32 de 0.500, correspondendo de perto aos 0.508 do DAPO, enquanto superou-o em benchmarks específicos como AMC23 e HMMT25.
  • No DeepSeek-R1-Distill-Qwen-1.5B, o POISE melhorou a pontuação média para 0.303, comparado aos 0.296 do DAPO.
• Eficiência de Treinamento: O POISE requer substancialmente menos tempo de relógio. No DeepSeek-R1-Distill-Qwen-1.5B, atingir desempenho comparável levou 18 horas com POISE versus 24 horas com DAPO em uma única GPU B200. No Qwen3-4B, o tempo foi de 36 horas versus 49 horas.
• Estabilidade: As normas de gradiente no POISE foram consistentemente mais baixas do que no DAPO, indicando variância de gradiente reduzida e otimização mais estável.
• Qualidade de Estimativa de Valor: A sonda de estado interno alcançou uma correlação de Pearson de 0.870 com recompensas empíricas, superando um crítico separado treinado na escala de LLM (0.676) na mesma tarefa.
• Generalização: O estimador manteve alto desempenho em tarefas de codificação, chamada de ferramentas e seguimento de instruções, frequentemente superando críticos separados nesses domínios.

Significado e Alegações

O artigo afirma que o POISE fornece um caminho eficiente em computação para uma RLVR estável e escalável para modelos de raciocínio de grande escala. Ao aproveitar as próprias representações internas do modelo, o método:

• Elimina o sobrecarga de memória e computação do treinamento de um modelo crítico separado.
• Remove a ineficiência de amostragem de linhas de base relativas a grupos (GRPO), permitindo maior diversidade de prompts dentro de um orçamento fixo.
• Evita o desperdício de computação em prompts degenerados que resultam em vantagem zero.
• Demonstra que estados internos, anteriormente usados principalmente como ferramentas de diagnóstico, podem servir como sinais de otimização práticos e de alta fidelidade para aprendizagem por reforço.

Os autores concluem que o POISE permite uma otimização de política mais estável e eficiente sem a necessidade de linhas de base relativas a grupos ou modelos de valor auxiliares, tornando-o uma alternativa viável para o treinamento de modelos de raciocínio de grande escala sob restrições de recursos.