Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Resumo Técnico: Além da Normalização: Repensando a Função de Partição como um Agendador de Dificuldade para RLVR

Declaração do Problema
A Aprendizagem por Reforço com Recompensas Verificáveis (RLVR) demonstrou ser eficaz para aprimorar as capacidades de raciocínio de Modelos de Linguagem de Grande Escala (LLMs). No entanto, métodos padrão de maximização de recompensa, como PPO e GRPO, frequentemente levam a distribuições de saída excessivamente afiadas, reduzindo a diversidade de geração. Embora abordagens recentes que utilizam GFlowNets abordem isso treinando políticas para corresponder a uma distribuição-alvo induzida por recompensa, elas tipicamente tratam a função de partição aprendida ($Z$) apenas como uma constante de normalização necessária para a estabilidade do treinamento. Este artigo argumenta que essa perspectiva subutiliza a informação codificada dentro de $Z$, perdendo uma oportunidade de melhorar a eficiência de amostragem sem incorrer em custos computacionais adicionais.

Metodologia: PACED-RL
Os autores propõem o RL Guiado pela Função de Partição (PACED-RL), um framework de pós-treinamento que reinterpreta a função de partição do GFlowNet não apenas como um normalizador, mas como um sinal de precisão online por prompt. A metodologia consiste em três componentes principais:

1. Reinterpretação Teórica da Função de Partição:
   O artigo estabelece uma ligação teórica entre a função de partição ótima $Z^*(x)$ e a precisão online $p_{old}(x)$ de um prompt de pergunta $x$ sob a política pré-atualização. Ao modificar a perda de Equilíbrio de Trajetória (TB) para ancorar a divergência KL na política atual ($\pi_{old}$) em vez de um modelo de referência fixo, os autores derivam que $Z^*(x)$ codifica diretamente informações de precisão. Especificamente, $p_{old}(x) \approx \beta \log Z^*(x)$, onde o termo de divergência KL é mostrado empiricamente como negligenciável durante a dinâmica padrão de treinamento. Isso permite que a função de partição, já sendo aprendida para normalização, sirva como um estimador de precisão online gratuito.

2. Seleção Adaptativa de Prompts:
   Aproveitando as estimativas de precisão derivadas de $Z_\phi(x)$, o PACED-RL implementa uma estratégia de seleção de prompts sensível à dificuldade. Em vez de amostrar prompts uniformemente, o método seleciona greedy os top-$m$ prompts cujas precisões estimadas estão mais próximas de um valor-alvo $\tau$ (tipicamente 0,5). Isso foca o treinamento em amostras de "dificuldade intermediária", que trabalhos anteriores sugerem oferecerem os sinais de aprendizado mais informativos, maximizando assim a eficiência de amostragem.

3. Replay Priorizado por Erro de Estimativa de Precisão:
   Para aprimorar ainda mais a eficiência e calibrar a função de partição, o framework introduz um mecanismo de replay. Um buffer armazena pares prompt-saída com recompensas corretas ($r=1$). Dentro deste buffer, as amostras são priorizadas com base na magnitude de seu erro de estimativa de precisão: $| \frac{N_{correto}}{N} - \hat{p}_{old}(x) |$. Isso garante que o modelo revise amostras onde a previsão da função de partição diverge mais da realidade observada, refinando as estimativas de precisão e a normalização da distribuição-alvo.

Crucialmente, tanto a seleção de prompts quanto a priorização de replay reutilizam informações geradas durante o treinamento padrão de GFlowNet, amortizando a sobrecarga computacional no processo de otimização existente.

Principais Contribuições

• Insight Teórico: O artigo demonstra que a função de partição do GFlowNet para pós-treinamento de LLMs está intrinsecamente ligada às precisões online, permitindo seu uso direto como estimador de precisão sem mecanismos auxiliares.
• Framework PACED-RL: A introdução de um novo framework de treinamento que utiliza essas estimativas de precisão para seleção adaptativa de prompts e replay priorizado por erro, melhorando a eficiência de amostragem enquanto preserva os benefícios de diversidade dos GFlowNets.
• Validação Empírica: Avaliações extensas em benchmarks de raciocínio matemático e geração de código mostram que o PACED-RL alcança ganhos de desempenho significativos sobre bases fortes (GRPO, FlowRL) e métodos de amostragem adaptativa (DS, LILO, MoPPS), sem a sobrecarga computacional associada à superamostragem.

Resultados Experimentais
Os experimentos foram conduzidos em geração de código (conjunto de dados DeepCoder, DeepSeek-R1-Distill-Qwen-1.5B) e raciocínio matemático (conjunto de dados DeepScaleR, Qwen2.5-Math-1.5B/7B).

• Ganhos de Desempenho: No benchmark AIME 24, o PACED-RL melhorou o desempenho médio pass@1 em 29,1% sobre o GRPO e em 40,0% sobre o FlowRL usando o modelo Qwen2.5-Math-7B.
• Eficiência de Amostragem: O PACED-RL alcançou o melhor desempenho do GRPO e do FlowRL em significativamente menos tempo (por exemplo, 0,42× o tempo de treinamento no LiveCodeBench). Diferentemente de métodos que exigem superamostragem (DS, LILO), o PACED-RL manteve uma razão de rollout de 1,0× em relação ao treinamento padrão, evitando as penalidades de latência associadas à geração de amostras extras para estimativa de precisão.
• Diversidade: Medida via métricas pass@k, o PACED-RL consistentemente superou as bases, mostrando melhorias de até 14,2% sobre o GRPO e 9,1% sobre o FlowRL, confirmando que o método preserva a diversidade de saída e a capacidade de exploração.
• Estudos de Ablação: Os autores verificaram que omitir o termo de divergência KL no estimador de precisão tem impacto negligenciável na dinâmica de amostragem, enquanto economiza computação significativa. Eles também mostraram que o método naturalmente desloca o treinamento para perguntas mais difíceis à medida que o modelo melhora, prevenindo o colapso para um subconjunto estreito do conjunto de dados.

Significado e Alegações
O artigo alega que o PACED-RL oferece uma maneira "principlada e prática" de aprimorar o treinamento de correspondência de distribuição para LLMs. Ao reinterpretar a função de partição como um sinal de precisão online, o método permite um treinamento mais eficiente em termos de amostragem, sem o ônus computacional de estimadores de precisão auxiliares ou geração excessiva de rollouts. Os autores posicionam este trabalho como um passo em direção a um pós-treinamento RLVR mais eficiente, sugerindo que o papel da função de partição se estende além da normalização para atuar como um agendador de dificuldade dinâmico. Trabalhos futuros são sugeridos para explorar estratégias no momento da inferência (por exemplo, autoconsistência adaptativa) e estender o framework para configurações de GFlowNet assíncronas.