Resumo Técnico: Como Você Começa é Como Você Raciocina: Impulsionando a Exploração em RLVR via Priors Ajustados por Prefixo

1. Declaração do Problema

A Aprendizagem por Reforço com Recompensas Verificáveis (RLVR) emergiu como um paradigma poderoso para aprimorar as capacidades de raciocínio de Modelos de Linguagem de Grande Escala (LLMs) em tarefas como resolução de problemas matemáticos e prova de teoremas. No entanto, os autores identificam uma limitação crítica nas abordagens atuais de RLVR: o colapso de entropia.

Embora o RLVR frequentemente melhore a precisão em uma única execução (Pass@1), ele frequentemente falha em expandir a cobertura de trajetórias de raciocínio bem-sucedidas. Isso se manifesta como uma redução na diversidade de saída, onde o modelo converge para um conjunto restrito de caminhos de raciocínio. Técnicas de exploração existentes, como regularização de entropia ou recompensas intrínsecas, frequentemente tratam a distribuição do modelo base como fixa e tentam modificar o objetivo de RL ou o procedimento de amostragem. O artigo argumenta que esses métodos são insuficientes porque o RL é fortemente moldado pela distribuição a priori do modelo base. Além disso, técnicas de exploração passiva frequentemente degradam a qualidade da geração, resultando em execuções ruidosas que não contribuem para um raciocínio eficaz.

O desafio central é incentivar a exploração eficaz — descobrir comportamentos de raciocínio diversos, relevantes para a tarefa e de alta qualidade — sem colapsar a política em trajetórias redundantes ou de baixa qualidade, e sem o custo computacional de atualizar todos os parâmetros do modelo.

2. Metodologia: Framework IMAX

Os autores propõem o Aumento de Exploração Maximizador de Informação (IMAX), um framework leve que desloca o alvo de otimização dos parâmetros do LLM para o ajuste de prompt suave (soft prompt tuning). Em vez de modificar o processo de RL sobre um prior fixo, o IMAX aprende um conjunto de prefixos suaves que remodelam ativamente o prior do modelo sobre trajetórias de raciocínio.

2.1 Raciocínio Condicionado a Prefixo

O mecanismo central envolve um conjunto de $C$ prefixos suaves (tokens virtuais contínuos) $\mathcal{E} = \{e_1, \dots, e_C\}$. Para um prompt de entrada dado $x$, um prefixo $e_z$ é adicionado ao início da sequência de embeddings de entrada. O LLM de base permanece congelado, e a política $\pi_E(y|x, z)$ é condicionada tanto ao prompt quanto à identidade do prefixo selecionado $z$.

• Inicialização: Os prefixos são inicializados a partir de um conjunto de prompts de texto semanticamente significativos (por exemplo, "Introduza variáveis explicitamente", "Particione em casos") gerados por um LLM externo. Isso garante que os prefixos comecem com estratégias de raciocínio distintas.
• Controle: Cada prefixo atua como um "botão de controle", induzindo uma distribuição de execução distinta e direcionando o modelo para diferentes regiões de seu espaço de raciocínio (por exemplo, soluções algébricas vs. diagramáticas vs. semelhantes a código).

2.2 Recompensa Maximizadora de Informação (InfoMax)

Para evitar que o conjunto de prefixos aprendido colapse em controles redundantes (onde diferentes prefixos induzem trajetórias idênticas), os autores introduzem uma recompensa InfoMax. O objetivo é maximizar a informação mútua (MI) entre a identidade do prefixo $Z$ e a trajetória gerada $Y$, condicionado ao prompt $X$:
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
Intuitivamente, isso incentiva que diferentes prefixos produzam comportamentos de raciocínio válidos e distinguíveis para a mesma pergunta.

2.3 Limite Inferior Variacional e Treinamento

Calcular diretamente a informação mútua é intratável. Os autores derivam um limite inferior variacional usando um posterior variacional $q_\phi(z | x, y)$, que é um classificador leve treinado para prever qual prefixo gerou uma resposta específica $y$ dado o prompt $x$.

• Forma Reversa: Eles utilizam a decomposição reversa da MI: $I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$.
• Formulação da Recompensa: A recompensa InfoMax é definida como $\log q_\phi(z | x, y)$. O objetivo total para o conjunto de prefixos torna-se:
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  onde $R(x, y)$ é a recompensa verificável padrão (correção) e $\beta$ é um hiperparâmetro.

2.4 Otimização Alternada

O procedimento de treinamento segue um esquema de otimização alternada (Inferência Variacional de Ascensão Coordenada):

1. Atualizar Posterior ($q_\phi$): Com o conjunto de prefixos fixo, atualizar o classificador $q_\phi$ para prever a identidade do prefixo a partir do par prompt-resposta (gerado sem o prefixo suave para evitar vazamento de informação).
2. Atualizar Conjunto de Prefixos ($E$): Com $q_\phi$ fixo, atualizar os prefixos suaves usando um algoritmo de Otimização de Política Relativa em Grupo (GRPO) aumentado com a recompensa InfoMax.
3. Amostragem Estratificada: Durante o treinamento, cada prompt é emparelhado com identidades de prefixo distintas para garantir uma cobertura uniforme do espaço de prefixos dentro de cada grupo GRPO.

3. Contribuições Principais

1. Paradigma Leve de RLVR: O artigo introduz um método para melhorar o raciocínio otimizando um conjunto de prompts suaves em vez de todos os parâmetros do modelo. Isso remodela a distribuição de geração condicional de um modelo congelado sem a sobrecarga computacional do ajuste fino completo.
2. Recompensa InfoMax para Diversidade: Os autores propõem uma recompensa InfoMax em nível de resposta que incentiva explicitamente que diferentes prefixos induzam trajetórias de raciocínio distinguíveis e relevantes para a tarefa, prevenindo o colapso do conjunto de prefixos em controles redundantes.
3. Validação Empírica: Extensos experimentos em benchmarks de raciocínio matemático (GSM8K, MATH-500, Minerva, AMC) em três escalas de modelo (Qwen2.5-1.5B, Qwen3-4B, Qwen3-8B) demonstram que o IMAX melhora tanto a qualidade quanto a diversidade do raciocínio.

4. Resultados Experimentais

Os autores avaliaram o IMAX contra várias linhas de base, incluindo GRPO de política completa, inicialização por prompt de texto e outros métodos de RLVR orientados à diversidade (DIVER, MERCI).

• Ganhos de Desempenho: O IMAX superou consistentemente as linhas de base padrão de RLVR ajustadas por prefixo.
  • No Qwen3-4B, o IMAX melhorou o Pass@4 em até 11,60% e o Avg@4 em 10,57% em comparação com linhas de base padrão em benchmarks específicos.
  • No Qwen3-8B, o DAPO+IMAX mostrou melhorias significativas sobre o DAPO-prefixo no GSM8K, MATH-500 e Minerva.
• Diversidade e Cobertura:
  • Análise t-SNE: Visualizações dos embeddings de resposta mostraram que os prefixos induzidos pelo IMAX resultaram em clusters mais separados em comparação com outros métodos, indicando maior diversidade intra-grupo.
  • Componentes de Raciocínio: Mapas de calor revelaram que diferentes prefixos induziram com sucesso padrões de raciocínio distintos (por exemplo, o Prefixo 1 favorecendo manipulação algébrica, o Prefixo 2 favorecendo partição de casos e verificação), alinhando-se com sua inicialização semântica.
  • Desempenho Oráculo: O desempenho combinado de múltiplos prefixos (Oráculo) superou o melhor prefixo único, confirmando que os prefixos cobrem regiões bem-sucedidas complementares do espaço de problemas.
• Eficiência: O IMAX requer significativamente menos custo de treinamento (3,0x–6,0x menos horas de GPU) em comparação com o GRPO de parâmetros completos, pois atualiza apenas os embeddings de prefixo leves e a cabeça do classificador.

5. Significado e Alegações

O artigo afirma que como um modelo inicia um processo de raciocínio (o prior) fundamentalmente dita como ele raciocina. Ao tratar prompts suaves como variáveis de controle reutilizáveis, o IMAX oferece uma solução escalável para o problema de exploração no RLVR.

• Agnóstico de Algoritmo: O framework é projetado para ser integrado perfeitamente em pipelines existentes de RLVR (por exemplo, GRPO, DAPO).
• Diversidade Estruturada: Ao contrário de métodos que simplesmente aumentam a aleatoriedade, o IMAX incentiva a diversidade estruturada de raciocínio, garantindo que a exploração leve a caminhos semanticamente distintos e de alta qualidade.
• Limitações Modestas: Os autores reconhecem que o método é projetado principalmente para tarefas de raciocínio, mostrando ganhos limitados em tarefas gerais de seguimento de instruções. Além disso, a eficácia depende da qualidade dos prompts de texto iniciais, que podem exigir design manual ou reajuste entre domínios. Os prefixos aprendidos também introduzem uma sobrecarga persistente, embora pequena, durante a inferência.

Em resumo, o IMAX demonstra que otimizar o "início" do processo de geração via prefixos suaves é uma estratégia altamente eficaz para impulsionar a exploração e melhorar as capacidades de raciocínio em LLMs, oferecendo uma alternativa computacionalmente eficiente à aprendizagem por reforço de política completa.