Resumo Técnico: OPD+: Repensando o Design da Vantagem para Destilação On-Policy

1. Definição do Problema

A destilação on-policy (OPD) é uma técnica utilizada para transferir capacidades de modelos de linguagem professores capacitados para modelos estudantes base, otimizando um sinal de supervisão denso ao nível de token derivado de rollouts gerados pelo estudante. Embora as implementações de OPD existentes (ex: Tinker) frequentemente utilizem a divergência de Kullback-Leibler (KL) como o objetivo primário, elas empregam uma escolha de design específica: interromper o gradiente (stopping the gradient) no termo de recompensa (a diferença de log-razão entre a verossimilhança do estudante e do professor).

Este design é motivado pelo desejo de estabilidade numérica, traçando paralelos com redes de alvo em Deep Q-Networks (DQN). No entanto, os autores identificam uma falha teórica crítica nesta abordagem:

1. Estimativas de Gradiente Enviesadas: Para $f$-divergências gerais, a função de recompensa depende dos parâmetros da política do estudante $\theta$. Interromper o gradiente neste termo de recompensa ignora a dependência direta da recompensa em relação a $\theta$, levando a estimativas de gradiente matematicamente incorretas.
2. Degradação de Desempenho: Este viés causa colapso catastrófico no treinamento para funções de divergência que não sejam a Reverse KL (ex: Forward KL e Divergência de Jensen-Shannon), tornando-as destreináveis sob configurações padrão de OPD. Mesmo para a Reverse KL, embora o viés teoricamente desapareça devido à identidade do score, resultados empíricos sugerem que o design de interrupção de gradiente ainda subotimiza o desempenho.

2. Metodologia: O Framework OPD+

O artigo propõe o OPD+, um framework de otimização corrigido que reavalia o design da vantagem para destilação on-policy baseada em $f$-divergência.

2.1 Fundamentação Teórica

Os autores formulam o objetivo de OPD como a minimização da $f$-divergência entre a política do estudante $p_\theta$ e a política do professor $q$. A função objetivo é definida como:
$$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$$
onde $u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ é a razão de densidade ao nível de token, e $f$ é uma função convexa com $f(1)=0$.

2.2 A Correção

O artigo deriva o gradiente real deste objetivo. O OPD padrão utiliza um estimador de gradiente de política de "interrupção de gradiente" onde a recompensa $r = -f(u_{\theta,n})$ é tratada como fixa. Os autores mostram que o gradiente real contém dois termos:

1. O termo padrão de função de score: $\mathbb{E}[r \nabla_\theta \log p_\theta]$.
2. Um termo de gradiente de recompensa direta: $\mathbb{E}[\nabla_\theta r]$.

Para uma $f$-divergência geral, o termo de gradiente de recompensa direta não desaparece. Especificamente, $\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$.
Consequentemente, a função de vantagem correta não é simplesmente $-f(u)$, mas um termo ponderado:
$$w_f(u) = -f(u) + u f'(u)$$

2.3 Caso Especial: Reverse KL

Para a Reverse KL, $f(u) = -\log u$. O termo de correção torna-se:
$$w(u) = -(-\log u) + u(-1/u) = \log u + 1$$
O termo $+1$ adicional é uma constante de baseline. Como $\mathbb{E}[\nabla_\theta \log p_\theta] = 0$ (identidade do score), este termo teoricamente desaparece em expectativa. No entanto, os autores argumentam que, na prática, este baseline constante atua como um mecanismo de redução de variância, melhorando a velocidade de convergência e a estabilidade.

2.4 Implementação

O OPD+ requer apenas uma mudança de uma única linha de código nas implementações existentes de OPD: substituir a recompensa de divergência negativa bruta $-f(u)$ pelo peso corrigido $w_f(u)$. Este framework suporta várias $f$-divergências, incluindo Forward KL, Reverse KL e Divergência de Jensen-Shannon (JSD).

3. Contribuições Principais

1. Framework de Otimização Genérico: O artigo fornece um framework unificado para destilação on-policy baseado em $f$-divergências gerais, derivando o gradiente e as funções de perda corretos que consideram a dependência da recompensa em relação aos parâmetros da política.
2. Prova Teórica de Viés: Prova matematicamente que a operação padrão de interrupção de gradiente leva a estimativas de gradiente enviesadas para $f$-divergências gerais, invalidando a prática comum de tratar a recompensa como fixa.
3. Algoritmo OPD+: Os autores introduzem o OPD+, o algoritmo corrigido que elimina o colapso de treinamento observado na Forward KL e JSD e melhora o desempenho mesmo para a Reverse KL.
4. Validação Empírica: O trabalho desafia a crença prevalecente de que a Reverse KL é a métrica de distância superior ou a única viável para destilação on-policy, demonstrando que outras divergências podem superar a Reverse KL quando otimizadas corretamente.

4. Resultados Experimentais

Os autores avaliaram o OPD+ na família de modelos QWEN3 (1.7B e 8B parâmetros) em dois domínios distintos: Matemática de Competição (DeepMath, AIME 24/25, HMMT 25) e Uso de Ferramentas Agênticas (Agentic Tool Use).

4.1 Resgate de Divergências "Intreináveis"

Sob o design padrão de interrupção de gradiente (OPD), a Forward KL e a JSD exibiram falha catastrófica, caindo para 0% de precisão em tarefas de uso de ferramentas e falhando em igualar a linha de base SFT-inicializada em tarefas de matemática.

• Correção OPD+: Ao aplicar a vantagem corrigida $w(u)$, o OPD+ eliminou completamente este modo de falha.
• Desempenho JSD: A JSD+ não apenas estabilizou, mas superou a linha de base padrão de Reverse KL. No benchmark desafiador AIME 25, a JSD+ alcançou 50,42% de precisão comparado aos 48,54% da Reverse KL. Na tarefa de Uso de Ferramentas (Tool Use), a JSD+ atingiu 60,29%, superando a linha de base de Reverse KL de 59,97%.

4.2 Melhoria na Reverse KL

Mesmo dentro do domínio onde a Reverse KL é tradicionalmente dominante, o OPD+ apresentou melhorias empíricas estritas:

• AIME 24: A Reverse KL+ melhorou de 64,27% para 64,69%.
• Tool Use: A Reverse KL+ melhorou de 59,97% para 61,44%.
• Convergência: A dinâmica de treinamento (Figura 1) mostra que as variantes do OPD+ aceleram a convergência precoce e evitam platôs de desempenho ou degradações observadas nas linhas de base padrão.

4.3 Autodestilação (OPSD)

O método também foi validado em um cenário de Autodestilação On-Policy (OPSD), onde o modelo atua tanto como estudante quanto como professor (com o professor condicionado a soluções de verdade fundamental privilegiadas). O OPD+ consistentemente superou as versões vanilla de interrupção de gradiente em todos os tipos de divergência, com a Reverse KL corrigida mostrando o melhor desempenho entre as divergências.

5. Significância e Alegações

O artigo alega que o baixo desempenho histórico de outras $f$-divergências (Forward KL, JSD) na destilação on-policy não é uma limitação intrínseca das divergências em si, mas sim um efeito colateral de um coeficiente de gradiente mal especificado causado pelo design de interrupção de gradiente.

Ao corrigir a função de vantagem para ser matematicamente fiel ao objetivo da $f$-divergência, o OPD+:

1. Expande o Espaço de Design: Permite o uso de uma variedade mais ampla de distâncias estatísticas para destilação, indo além do padrão da indústria da Reverse KL.
2. Melhora o Desempenho: Estabelece novos picos de desempenho estado-da-arte em benchmarks de raciocínio e uso de ferramentas.
3. Desafia Convenções: Questiona a crença comum de que a Reverse KL é inerentemente superior devido à estabilidade, sugerindo que a estabilidade era um artefato do método de otimização e não da escolha da divergência.

Os autores concluem que seu trabalho fornece um algoritmo matematicamente mais correto e empiricamente mais eficaz, potencialmente desencadeando mais pesquisas no design de funções de vantagem para métodos de destilação on-policy. Eles reconhecem limitações relativas à análise de escala de modelos e à comparação entre destilação de verossimilhança de token único e de logits, apontando-as como áreas para trabalhos futuros.