Self-Distilled Policy Gradient

Resumo Técnico: Self-Distilled Policy Gradient (SDPG)

Definição do Problema
O Aprendizado por Reforço com Recompensas Verificáveis (RLVR) tornou-se um padrão para o pós-treinamento de Grandes Modelos de Linguagem (LLMs) em tarefas de raciocínio, como matemática e geração de código. Algoritmos como o Group Relative Policy Optimization (GRPO) otimizam contra recompensas de resultado baseadas em regras (ex: respostas corretas/incorretas) sem exigir anotações de preferência humana. No entanto, o RLVR enfrenta duas limitações primárias:

1. Recompensas Esparsas: O sinal de recompensa é aplicado ao nível da sequência, tornando difícil atribuir crédito a passos de raciocínio específicos (tokens), o que leva a uma exploração ineficiente.
2. Instabilidade de Treinamento: Os estágios iniciais de treinamento frequentemente sofrem com vantagens negativas, causando instabilidade e potencial degradação das capacidades linguísticas fundamentais.

Embora a Destilação On-Policy (OPD) tenha sido proposta para fornecer supervisão densa ao nível do token, as abordagens tradicionais dependem de modelos professores externos e maiores, o que acarreta altos custos de memória e introduz desajustes de distribuição. Métodos recentes de "autodestilação", onde o modelo atua simultaneamente como aluno e professor utilizando contexto privilegiado (ex: soluções de verdade fundamental/ground-truth), abordam o problema de memória, mas correm o risco de reforçar tokens localmente plausíveis em trajetórias globalmente incorretas se não forem devidamente controlados. Além disso, a autodestilação pura pode levar ao colapso de modo (mode collapse) e à redução da exploração.

Metodologia: Estrutura SDPG
O artigo propõe o Self-Distilled Policy Gradient (SDPG), uma estrutura que integra a autodestilação on-policy de vocabulário completo exata com a otimização de política regularizada por KL. O objetivo central combina três componentes:

1. Gradiente de Política Baseado em Resultado ($L_{out}$):
   O SDPG mantém o sinal de recompensa binário esparso de um verificador. Ele utiliza uma vantagem relativa ao grupo ($A_{out}$) normalizada dentro de um grupo de rollouts para estabilizar o aprendizado, de forma semelhante ao GRPO, mas formulado como uma perda de substituição direta do tipo REINFORCE, sem o clipping de razão de importância do estilo PPO.

2. Destilação On-Policy de Vocabulário Completo Exata ($L_{OPD}$):
   Em vez de amostrar tokens para aproximar o sinal do professor, o SDPG calcula a divergência Kullback-Leibler (KL) reversa exata entre a distribuição do aluno ($p_t$) e uma distribuição professora privilegiada ($q_t$) condicionada a um contexto adicional $c$ (ex: soluções ou caminhos de raciocínio).

   • Identidade de Gradiente: Os autores derivam um resultado teórico fundamental (Proposição 3.1): o gradiente da perda de KL reversa de vocabulário completo em relação à política do aluno é localmente idêntico a uma atualização de gradiente de política de amostragem desconectada (detached-sampling). A "vantagem" nesta atualização é uma razão logarítmica centrada das probabilidades do professor e do aluno.
   • Mecanismo de Portão (Gating): Para evitar que o professor privilegiado reforce trajetórias incorretas, a perda de destilação é controlada pelo portão da vantagem de resultado. A destilação só é aplicada quando o rollout é verificado como correto ($A_{out} > 0$).

3. Regularização de Política de Referência ($L_K$):
   Para evitar que a política se desvie demais de suas capacidades iniciais ou sofra colapso, o SDPG ancora a política atualizada a uma política de referência fixa ($\pi_{ref}$). O artigo emprega a divergência Unnormalized KL (UKL) (forward ou reverse) para esta regularização, que oferece gradientes mais elegantes e simétricos em comparação com a KL normalizada padrão, particularmente ao lidar com medidas não normalizadas.

Estratégias de Estabilização
Para gerenciar o ruído do sinal privilegiado e o equilíbrio entre exploração e explotação, o SDPG introduz dois estabilizadores:

• Portão de Vantagem Positiva (Positive Advantage Gating): Garante que a destilação ocorra apenas em trajetórias endossadas pelo verificador.
• Cronograma de Warmup-Decay ($\beta(k)$): O coeficiente de destilação $\beta$ é aumentado gradualmente (warmup) para permitir que o modelo encontre trajetórias corretas antes de depender da destilação, e então decai (decay) em direção ao final do treinamento. Isso evita que o sinal privilegiado sobre-restrinja o modelo uma vez que este tenha internalizado o conhecimento necessário, incentivando a exploração na fase final.

Principais Contribuições

1. Derivação Teórica: O artigo deriva uma forma exata de gradiente de política local para a autodestilação de vocabulário completo de KL reversa, mostrando que é equivalente a uma atualização de amostragem desconectada com uma vantagem de razão logarítmica centrada.
2. Objetivo Inovador: Propõe um objetivo unificado combinando recompensas de resultado binárias, destilação privilegiada de vocabulário completo com portão e regularização KL não normalizada contra uma referência fixa.
3. Técnicas de Estabilização: A introdução do portão de vantagem positiva e de um cronograma de warmup-decay para o coeficiente de destilação para equilibrar a supervisão densa com a exploração.

Resultados Experimentais
Os autores avaliaram o SDPG em benchmarks de raciocínio matemático (AIME2024, AIME2025, AMC23) usando o Qwen3-4B.

• Desempenho: Ambas as variantes do SDPG (usando Unnormalized Forward KL e Unnormalized Reverse KL) superaram o baseline GRPO e o baseline de autodestilação RLSD em todos os benchmarks.
• Estabilidade: As variantes do SDPG atingiram patamares de alta recompensa significativamente mais rápido que o GRPO. Notavelmente, o SDPG manteve uma entropia de ator mais alta durante o treinamento em comparação ao RLSD, que exibiu colapso de entropia (um sinal de colapso de modo) em torno do passo 250.
• Eficiência: Os métodos SDPG produziram comprimentos de resposta que foram suficientes para o raciocínio de múltiplos passos, porém mais concisos que as saídas verbosas do GRPO.

Significância
O artigo afirma que o SDPG consegue preencher a lacuna entre o RLVR de recompensa esparsa e a autodestilação de supervisão densa. Ao tratar a autodestilação como uma atualização de gradiente de política e controlar rigorosamente o sinal privilegiado através de portões e cronogramas, o SDPG alcança estabilidade e desempenho superiores. Demonstra que a supervisão densa ao nível do token, derivada do próprio contexto privilegiado do modelo, pode guiar modelos de raciocínio de forma eficaz sem o custo de memória de professores externos ou a instabilidade da autodestilação pura.