Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Resumo Técnico: GRPO Regularizado por Future-KL

Declaração do Problema
A Otimização de Política Relativa por Grupo (GRPO) emergiu como um paradigma dominante para aprendizado por reforço (RL) sem crítico no pós-treinamento de Modelos de Linguagem de Grande Escala (LLM), particularmente para raciocínio matemático. A GRPO elimina a necessidade de uma função de valor aprendida ao normalizar recompensas de resultado dentro de um grupo de respostas amostradas. No entanto, implementações padrão de regularização KL na GRPO tipicamente tratam a penalidade como uma penalidade local de token do lado da perda (por exemplo, usando o estimador $k_3$).

Os autores identificam uma falha fundamental nessa abordagem: ela ignora a atribuição de crédito em nível de processo induzida pela geração autoregressiva. Em uma trajetória autoregressiva, um token inicial não apenas incorre em seu próprio custo de regularização local, mas também altera a distribuição de futuros prefixos, afetando assim os custos KL de todos os tokens subsequentes. Termos de perda locais padrão diferenciam apenas a penalidade do token amostrado, omitindo o "retorno futuro de regularização causal" (future-regularization return-to-go).

Além disso, o artigo destaca uma incompatibilidade estrutural entre a normalização por grupo da GRPO e formulações padrão de KL-em-recompensa. Diferentemente de objetivos de recompensa lineares (por exemplo, RL de MaxEnt), onde recompensa e regularizador podem ser fundidos, a normalização por grupo da GRPO induz uma utilidade não linear em nível de prompt (especificamente, uma transformação arcsin para recompensas binárias). Essa não linearidade significa que adicionar penalidades KL às recompensas antes da normalização distorce o objetivo implícito, colapsando informações KL em nível de token em um único escalar em nível de trajetória e falhando em preservar o sinal de gradiente correto.

Metodologia
O artigo deriva o gradiente de política on-policy para objetivos estilo GRPO sob regularização geral de divergência $f$ token a token. A principal insight teórica é que o gradiente se decompõe em dois termos distintos com propriedades de escalonamento diferentes:

1. O Termo de Recompensa: Herdado da utilidade não linear $\ell(p)$, este termo carrega um ganho dependente do prompt $\ell'(p)$. Para recompensas binárias, isso recupera a vantagem GRPO padronizada.
2. O Termo de Regularizador: Este termo é independente de $\ell'$ e decompõe-se em uma penalidade local de token mais um retorno futuro de regularização causal.

Com base nessa derivação, os autores propõem a Otimização de Política Regularizada por Future-KL (FRPO). O algoritmo modifica a atualização padrão da GRPO de maneira desacoplada:

1. Calcular a vantagem de resultado normalizada por grupo ($\hat{A}^{GRPO}$) estritamente a partir das recompensas originais do verificador, preservando a utilidade não linear implícita.
2. Construir um termo de correção future-KL em nível de token calculando a soma cumulativa reversa das razões logarítmicas por token ($\delta_{i,j} = \log \frac{\pi_\theta}{\pi_{ref}}$) do token atual até o final da trajetória.
3. Adicionar este termo future-KL à vantagem após a construção da vantagem, em vez de moldar a recompensa antes da normalização.

A regra de atualização para um token $t$ na resposta $i$ é:
$$\hat{A}^{FRPO}_{i,t} = \text{sg}\left( \hat{A}^{GRPO}_i - \beta \sum_{j=t}^{T_i} \delta_{i,j} \right)$$
onde $\text{sg}$ denota stop-gradient. Esta abordagem não requer crítico, nenhuma passagem extra de modelo e mantém a mesma sobrecarga computacional da GRPO padrão.

Principais Contribuições

1. Identificação de Utilidade Não Linear: O artigo demonstra que a normalização por grupo da GRPO induz uma utilidade não linear ($\ell(p) = 2 \arcsin \sqrt{p}$ para recompensas binárias), tornando recompensa e KL não fundíveis. Adicionar KL às recompensas antes da normalização distorce estruturalmente o objetivo.
2. Derivação do Gradiente On-Policy: Os autores derivam o gradiente on-policy correto para regularização de divergência $f$ token a token, mostrando que isso exige um termo de retorno futuro de regularização ausente em implementações locais do lado da perda.
3. Algoritmo FRPO: Eles propõem o FRPO, uma atualização simples e desacoplada que preserva a vantagem de resultado da GRPO enquanto injeta a atribuição de crédito future-KL necessária. Este método separa o ranqueamento em nível de trajetória da regularização KL em nível de processo.
4. Interpretação de Optimalidade: A estrutura future-KL é mostrada como alinhada com as condições de optimalidade de RL regularizado por entropia, onde a política ótima satisfaz uma forma de recompensa de processo envolvendo custos KL futuros.

Resultados Experimentais
Os autores avaliam o FRPO em tarefas de raciocínio matemático usando modelos base Qwen3 (1.7B, 4B e 30B-A3B) treinados no conjunto de dados DAPO-Math-17k.

• Desempenho: No modelo 30B-A3B, o FRPO melhora o pass@16 em quase 5% de ganho absoluto sobre a melhor linha de base (GRPO padrão com KL do lado da perda) no benchmark MATH500. Também supera as linhas de base no AIME24, AIME25 e outros benchmarks de nível olímpico.
• Estabilidade e Deriva: O FRPO mantém entropia final mais alta e deriva de política menor (medida por PPO-KL) em comparação com linhas de base de KL do lado da perda. Isso sugere que o método elicita padrões de raciocínio de Cadeia de Pensamento (CoT) longos sem desvio excessivo da política de referência.
• Ablação: Experimentos confirmam que o KL reverso (busca de modo) performa melhor que o KL direto (cobertura de modo) para esta tarefa. O artigo também observa que o uso do estimador $k_1$ (razão logarítmica) é preferível ao $k_3$ para o termo future-KL, pois a não negatividade do $k_3$ cria um sinal de penalidade unilateral que prejudica a atribuição de crédito.
• Modos de Falha: O artigo observa que abordagens "KL-em-recompensa" (adicionar KL às recompensas antes da normalização) levam ao colapso do modelo ou degradação de desempenho, validando a afirmação teórica de que tal fusão é estruturalmente incorreta sob a utilidade não linear da GRPO.

Significado
O artigo argumenta que o significado deste trabalho reside na correção de um erro de implementação sutil, mas crítico, em um framework de RL amplamente utilizado. Ao mudar a perspectiva de "penalidade de perda local" para "atribuição de crédito em nível de processo", os autores mostram que a GRPO pode ser melhorada sem introduzir funções de valor complexas ou custos computacionais adicionais. O método FRPO proposto fornece uma modificação teoricamente fundamentada e leve que aprimora as capacidades de raciocínio enquanto preserva as propriedades de estabilidade e generalização do modelo base. O trabalho sublinha que, em configurações autoregressivas, a regularização é inerentemente um sinal em nível de processo, e ignorar o custo futuro leva a uma atribuição de crédito subótima.