Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

Resumo Técnico: Hidden States Know Where Reasoning Diverges

Título Original: Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

1. O Problema: Atribuição de Crédito em RLVR

O aprendizado por reforço com recompensas verificáveis (RLVR) é fundamental para melhorar o raciocínio em Grandes Modelos de Linguagem (LLMs). O algoritmo padrão, GRPO (Group Relative Policy Optimization), utiliza uma recompensa baseada apenas no resultado final (correto ou incorreto).

O problema central é a atribuição de crédito (credit assignment): o GRPO atribui o mesmo valor de vantagem (advantage) a todos os tokens de uma trajetória. Se um raciocínio de 500 tokens for correto, mas contiver um erro lógico no meio que foi "corrigido" por sorte ou que não afetou o resultado, o modelo ainda assim reforça todos os tokens. Identificar exatamente onde o raciocínio divergiu (o erro local) é crucial, mas métodos atuais como Modelos de Recompensa de Processo (PRMs) exigem anotações humanas caríssimas passo a passo.

2. Metodologia: O Método SHEAR

Os autores propõem o SHEAR (Span-level Hidden state Enabled Advantage Reweighting), uma abordagem de auto-supervisão que extrai sinais de erro diretamente das representações internas (estados ocultos) do modelo, sem necessidade de modelos de recompensa externos ou anotações adicionais.

A intuição fundamental: As distribuições dos estados ocultos (hidden states) de trajetórias corretas e incorretas tornam-se sistematicamente diferentes nos pontos onde o raciocínio diverge.

O processo técnico consiste em quatro etapas:

1. Decomposição em Spans (Intervalos): A trajetória é dividida em janelas sobrepostas (sliding windows) de estados ocultos.
2. Cálculo da Distância de Wasserstein: Para cada intervalo (span), calcula-se a Distância de Wasserstein ($W_1$) entre a distribuição de seus estados ocultos e a distribuição dos estados ocultos do grupo oposto (ex: comparar o intervalo de uma resposta correta com os intervalos de respostas incorretas no mesmo grupo). A distância de Wasserstein é escolhida por capturar diferenças na forma e na variância da distribuição, não apenas no deslocamento da média.
3. Agregação via Max-Pooling: Os valores de distância calculados para os intervalos são projetados de volta para o nível de token individual usando max-pooling, garantindo que o token receba o sinal de discrepância mais forte de qualquer intervalo que o contenha.
4. Reajuste da Vantagem (Reweighting): A vantagem do GRPO é multiplicada por esse peso derivado da distância. Tokens em regiões onde as trajetórias corretas e incorretas são muito distintas recebem atualizações de gradiente mais intensas (amplificando o aprendizado onde o erro é detectável).

3. Contribuições Principais

• Descoberta Empírica: Demonstraram que a distância de Wasserstein entre estados ocultos é um proxy eficaz para a qualidade do raciocínio local, correlacionando-se fortemente com a probabilidade de sucesso de continuação.
• Fundamentação Teórica: Provaram um Teorema de Separação, estabelecendo matematicamente que, sob certas condições, os intervalos após o erro (post-divergence) apresentam distâncias de Wasserstein significativamente maiores do que os intervalos antes do erro (pre-divergence), desde que o sinal de divergência supere o ruído de amostragem finita.
• Algoritmo SHEAR: Um método simples que exige mudanças mínimas no pipeline de treinamento e não requer modelos auxiliares.

4. Resultados e Validação

Os experimentos foram realizados em 5 benchmarks de matemática e 5 de geração de código, utilizando backbones como Qwen2.5 e Llama 3.1.

• Desempenho Superior: O SHEAR superou consistentemente o GRPO padrão em todos os testes.
• Competitividade com PRMs: Surpreendentemente, o SHEAR superou métodos baseados em Modelos de Recompensa de Processo (PRMs) treinados, pois o sinal do SHEAR é "nativo" ao modelo e não sofre com o descasamento de distribuição que ocorre quando um PRM externo tenta avaliar um modelo que está evoluindo.
• Eficiência: O custo computacional adicional é baixo (entre 7% e 16% de aumento no tempo de treinamento), tornando-o viável para escala industrial.
• Robustez: O método funciona melhor em cadeias de raciocínio longas, onde a atribuição de crédito granular é mais necessária.

5. Significância

Este trabalho é significativo porque resolve um dos maiores gargalos do RL para raciocínio: a necessidade de supervisão densa (passo a passo). Ao provar que o próprio modelo "sabe" onde errou através de suas representações latentes, os autores oferecem um caminho para treinar modelos de raciocínio complexo de forma muito mais eficiente, utilizando apenas o feedback binário (certo/errado) que já é comum em tarefas matemáticas e de programação.