SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models

Título: SPS: Redirecionamento do Efeito de Espremimento de Probabilidade para Melhor Exploração em Aprendizado por Reforço para Grandes Modelos de Linguagem

1. O Problema: A Limitação da Exploração no RL para Raciocínio

O Aprendizado por Reforço (RL) tem sido fundamental para treinar modelos de linguagem (LLMs) focados em raciocínio, utilizando sinais de recompensa baseados em regras (como validadores de respostas matemáticas). No entanto, o artigo identifica uma limitação crítica:

• Foco em Pass@1 vs. Pass@k: O RL tradicional tende a melhorar drasticamente a taxa de sucesso de uma única amostra (Pass@1), mas falha em melhorar o desempenho com múltiplas amostras (Pass@k), que depende da capacidade de explorar trajetórias de raciocínio diversas.
• O Efeito de Espremimento (Squeezing Effect): A análise preliminar dos autores revela que o RL causa um "efeito de espremimento" sistemático. Quando o modelo recebe gradientes negativos para respostas de baixa probabilidade (erradas), a massa de probabilidade removida não é redistribuída uniformemente para outras opções. Em vez disso, devido à normalização da função softmax, essa massa é "espremida" e absorvida desproporcionalmente pela resposta dominante (a mais provável).
• Consequência: Isso leva a uma concentração excessiva da distribuição de saída em um subconjunto estreito de trajetórias de alta recompensa, reduzindo a entropia e limitando a descoberta de novas e melhores trajetórias de raciocínio. O modelo converge prematuramente para comportamentos "gananciosos" (greedy).

2. Metodologia: Steering Probability Squeezing (SPS)

Para mitigar esse efeito e promover uma exploração genuína, os autores propõem o SPS, um paradigma de treinamento que intercala Aprendizado por Reforço (RL) convencional com Aprendizado por Reforço Inverso (IRL).

Mecanismo Central:

O SPS trata os rollouts (amostras geradas) do próprio modelo durante a fase de RL como demonstrações para uma fase subsequente de IRL.

1. Fase de RL (Exploração Inicial): O modelo é treinado com RL padrão (usando GRPO - Group Relative Policy Optimization) para gerar trajetórias e receber feedback de recompensa.
2. Fase de IRL (Redirecionamento): Em vez de apenas maximizar a recompensa, o IRL é usado para remodelar explicitamente a distribuição de trajetórias induzida.
   • O objetivo é realocar a massa de probabilidade que foi "espremida" para as respostas dominantes de volta para trajetórias subexploradas, mas potencialmente valiosas.
   • Utiliza-se uma função de perda baseada em Forward-KL para fazer com que a política atual se alinhe com a distribuição empírica dos rollouts gerados, sem introduzir supervisão externa ou conhecimento prévio.
3. Estratégia Iterativa: O processo alterna entre RL e IRL. O IRL "suaviza" a distribuição antes que ela colapse, permitindo que o próximo ciclo de RL explore novas regiões do espaço de soluções.

Componente Adicional: Ênfase em Trajetórias de Baixa Probabilidade (L2TE)

Para combater ainda mais o efeito de espremimento, o SPS emprega uma estratégia chamada Low-Likelihood Trajectory Emphasis (L2TE). Durante a fase de IRL, o modelo prioriza a amostragem de rollouts que possuem probabilidade relativamente baixa sob a política atual. Isso amplifica o sinal de aprendizado para soluções subexploradas, incentivando a diversidade.

3. Contribuições Chave

1. Análise Teórica e Empírica do Efeito de Espremimento: Os autores formalizam matematicamente como a redistribuição de massa de probabilidade sob objetivos padrão de RL leva a um afunilamento (sharpening) da distribuição, estabelecendo um limite superior empírico para o Pass@k.
2. Proposta do SPS: Introdução de um novo framework que utiliza IRL para redirecionar ativamente o efeito de espremimento, transformando-o de um limitador de exploração em um mecanismo de redistribuição de probabilidade.
3. Resultados Empíricos Robustos: Demonstração de que o SPS supera consistentemente os métodos de RL padrão (como GRPO, DAPO e GSPO) em benchmarks de raciocínio matemático de nível olímpico, melhorando significativamente o Pass@k sem sacrificar o Pass@1.

4. Resultados Experimentais

Os experimentos foram conduzidos em cinco benchmarks matemáticos de alto nível (AIME, BRUMO, HMMT) utilizando modelos da série Qwen2.5-Math (1.5B e 7B parâmetros).

• Desempenho em Pass@128: O SPS demonstrou melhorias substanciais no Pass@128, indicando uma capacidade de exploração expandida.
  • Exemplo Notável: No modelo Qwen2.5-Math-1.5B no benchmark BRUMO, o SPS alcançou uma pontuação de 63.33 (Pass@128), uma melhoria de +10.00 pontos em comparação com o GRPO padrão (53.33).
• Comparação com Baselines: O SPS superou consistentemente o GRPO, DAPO e GSPO na métrica Pass@128, mantendo desempenho comparável ou superior no Pass@1 (Avg@128).
• Análise de Diversidade: Métricas de similaridade de trajetória mostraram que o SPS gera respostas mais diversas (menor similaridade) do que o GRPO, confirmando que a melhoria não vem apenas de uma melhor amostragem, mas de uma exploração real de novos caminhos de raciocínio.
• Custo Computacional: A fase de IRL adiciona apenas cerca de 3% ao tempo total de treinamento por iteração, tornando o método altamente eficiente.

5. Significado e Conclusão

O trabalho oferece uma mudança de paradigma na forma como entendemos e otimizamos o RL para LLMs de raciocínio:

• Superação do Colapso de Entropia: O SPS demonstra que é possível mitigar o colapso de entropia (onde o modelo para de explorar) não apenas adicionando regularização de entropia, mas reestruturando a dinâmica de aprendizado através do IRL.
• Exploração Intrínseca: A metodologia prova que é possível melhorar a exploração utilizando apenas os dados gerados pelo próprio modelo (auto-supervisão), sem necessidade de dados externos ou recompensas complexas.
• Implicações Futuras: O estudo sugere que a alternância entre RL e IRL é um caminho eficaz para estender os limites de exploração de modelos de linguagem, permitindo que eles descubram trajetórias de raciocínio mais complexas e diversas, essenciais para tarefas de raciocínio avançado e descoberta científica.

Em resumo, o SPS resolve o problema fundamental de concentração excessiva de probabilidade no RL, transformando um defeito intrínseco em uma oportunidade para expandir as fronteiras do raciocínio dos LLMs.