Resumo Técnico: Rethinking On-Policy Distillation (OPD)

1. Problema e Contexto

A Distilação em Política (On-Policy Distillation - OPD) emergiu como uma técnica central no pós-treinamento de Grandes Modelos de Linguagem (LLMs), sendo adotada por grandes players industriais (como Qwen3, MiMo e GLM-5). Diferente da distilação off-policy (que treina o aluno em sequências fixas geradas pelo professor e sofre de viés de exposição), a OPD faz com que o aluno gere suas próprias trajetórias e utiliza os log-probabilidades do professor como um sinal de recompensa denso (token a token) para refinar o comportamento.

Apesar do sucesso prático, a OPD é frágil e mal compreendida. O artigo identifica um modo de falha surpreendente: um professor mais forte (com melhores benchmarks) pode falhar completamente em melhorar o aluno, enquanto um professor mais fraco pode ter sucesso. O trabalho busca responder:

1. Sob quais condições a OPD succeed ou falha?
2. Qual é o mecanismo em nível de token que governa esse processo?
3. Como recuperar configurações de OPD que falham?

2. Metodologia

Os autores realizaram uma investigação sistemática combinando análise fenomenológica (empírica), análise de mecanismos em nível de token e proposição de receitas práticas.

• Configurações Experimentais: Utilizaram famílias de modelos como Qwen e DeepSeek, variando tamanhos (1.5B, 4B, 7B) e estados de treinamento (Base, Pós-RL, com/sem raciocínio).
• Métricas Dinâmicas: Introduziram métricas para monitorar o treinamento em tempo real:
  • Taxa de Sobreposição (Overlap Ratio): Proporção de tokens que aparecem simultaneamente nos top-$k$ do aluno e do professor.
  • Vantagem de Token de Sobreposição (Overlap-Token Advantage): Mede o alinhamento das distribuições dentro dos tokens compartilhados.
  • Gap de Entropia: Diferença na incerteza entre as políticas do aluno e do professor.
• Abordagem de "Distilação Reversa": Testaram distilar um modelo pós-RL (forte) de volta para seu checkpoint pré-RL (fraco) para isolar o efeito do padrão de pensamento versus a capacidade bruta.
• Ablação: Testaram otimizar apenas em tokens de sobreposição versus tokens não sobrepostos para identificar a fonte do gradiente.

3. Contribuições Principais e Descobertas

A. Fenomenologia: Condições para o Sucesso

O estudo identifica duas condições fundamentais que governam o sucesso da OPD:

1. Consistência do Padrão de Pensamento (Thinking-Pattern Consistency):

   • O aluno e o professor devem compartilhar padrões de raciocínio compatíveis.
   • Descoberta Chave: Um professor com pontuação de benchmark mais alta não garante sucesso se o padrão de pensamento for incompatível. Exemplo: Um professor "sem raciocínio" (non-thinking) falha em ensinar um aluno que precisa aprender raciocínio, mesmo que o professor tenha alta precisão. A sobreposição inicial baixa impede a recuperação do sinal de distilação.
   • Evidência: Em experimentos de distilação reversa, um aluno pós-RL (JustRL-1.5B) regrediu para o nível pré-RL ao ser distilado de um professor do mesmo tipo (R1-Distill-1.5B), ignorando a superioridade de um professor maior (R1-Distill-7B) que tinha o mesmo padrão de pensamento.

2. Novo Conhecimento vs. Escala (Higher Scores $\neq$ New Knowledge):

   • Mesmo com padrões de pensamento consistentes, a OPD falha se o professor não oferecer conhecimento genuinamente novo além do que o aluno já viu durante o treinamento.
   • Descoberta Chave: Professores de mesma família e pipeline de treinamento (ex: 1.5B vs 7B) tendem a convergir para distribuições similares. Um professor 7B não traz "novas capacidades" transferíveis se o aluno 1.5B já foi treinado nos mesmos dados. O ganho só ocorre quando o professor possui capacidades adquiridas via pós-treinamento (ex: RL) que o aluno não possui.

B. Mecanismo: Alinhamento Progressivo em Tokens de Alta Probabilidade

A análise em nível de token revela como a OPD funciona quando bem-sucedida:

• Alinhamento Progressivo: O sucesso é caracterizado por um aumento steady na Taxa de Sobreposição entre os tokens de alta probabilidade do aluno e do professor (subindo de ~72% para ~91%).
• Suficiência da Sobreposição: A otimização ocorre quase exclusivamente nos tokens compartilhados (sobreposição).
  • Tokens de sobreposição concentram 97%-99% da massa de probabilidade combinada.
  • Experimentos de ablação mostraram que otimizar apenas nos tokens de sobreposição (Overlap Top-$k$) recupera quase todo o benefício da OPD padrão, enquanto otimizar em tokens não sobrepostos é ineficaz.
• Dinâmica de Auto-Reforço: Uma vez que um token entra na região de alta probabilidade compartilhada, as atualizações de KL reverso concentram mais massa nele, empurrando tokens concorrentes para fora, criando um ciclo virtuoso de alinhamento.

C. Receita Prática: Como Recuperar OPD Falha

Para configurações onde a OPD falha (devido a lacunas no padrão de pensamento), os autores propõem duas estratégias:

1. Início Frio Off-Policy (Off-Policy Cold Start):
   • Realizar uma fase de Supervised Fine-Tuning (SFT) no aluno usando rollouts gerados pelo professor antes de iniciar a OPD.
   • Resultado: Isso aumenta a taxa de sobreposição inicial, fechando a lacuna de padrão de pensamento e permitindo que a OPD subsequente seja estável e eficaz.
2. Seleção de Prompts Alinhados ao Professor:
   • Utilizar prompts que correspondam ao formato e conteúdo dos dados de pós-treinamento do professor.
   • Resultado: Melhora o alinhamento em tokens de alta probabilidade. No entanto, os autores alertam que usar apenas esses prompts pode causar colapso de entropia (redução excessiva da exploração), sugerindo a necessidade de misturar com prompts fora da distribuição do professor.

D. Limitações e Custos

• Degradação da Recompensa com Profundidade: A qualidade do sinal de recompensa (token a token) degrada-se conforme a profundidade da trajetória aumenta. Em sequências longas (>10k tokens), a instabilidade surge no final da resposta e propaga-se para trás, sugerindo que a OPD pura pode não escalar bem para raciocínio de longo horizonte ou agentes multi-turno.
• Geometria de Otimização: Mesmo quando a recompensa global é informativa (correlacionada com a correção), a OPD pode falhar devido à geometria local do gradiente (anisotropia), onde sinais de vantagem por token se cancelam em atualizações de gradiente.

4. Significado e Impacto

Este trabalho é fundamental porque:

1. Desmistifica a OPD: Mostra que a força bruta do professor (benchmark score) não é o fator determinante; a compatibilidade de "pensamento" e a novidade do conhecimento são cruciais.
2. Revela o Mecanismo de Aprendizado: Demonstra que a OPD aprende essencialmente "padrões de pensamento" através do alinhamento progressivo em tokens de alta probabilidade, e não apenas a memorização de respostas.
3. Oferece Soluções Práticas: As estratégias de "Cold Start" e seleção de prompts fornecem um roteiro claro para engenheiros de ML que enfrentam falhas na distilação on-policy.
4. Alerta sobre Escalabilidade: Identifica um teto prático imposto pela degradação do sinal em trajetórias longas, apontando para a necessidade de abordagens híbridas (combinação de recompensa densa e esparsa) para tarefas complexas futuras.

Em suma, o paper redefine a OPD não como uma simples transferência de conhecimento, mas como um processo delicado de alinhamento de distribuições condicionadas a padrões de pensamento específicos, onde a qualidade do sinal local é tão importante quanto a capacidade global do professor.