Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Each language version is independently generated for its own context, not a direct translation.

Imagine que você está ensinando um aluno muito inteligente, mas inexperiente (o Estudante) a resolver problemas complexos, usando um professor brilhante (o Mestre) como referência.

O objetivo é fazer o aluno aprender não apenas copiando o que o professor já escreveu, mas pensando e gerando suas próprias respostas. Isso é chamado de Distilação On-Policy.

No entanto, os pesquisadores descobriram que o método atual de ensinar tem um grande defeito: é como tentar ensinar alguém a dirigir olhando apenas para uma única roda do carro a cada segundo. Se a roda estiver torta, o aluno fica confuso. Se o carro começar a sair da pista, o professor não sabe como corrigir porque está olhando apenas para aquele instante.

Aqui está a explicação simples do que o artigo propõe, usando analogias do dia a dia:

1. O Problema: O "Olhar de Águia" vs. O "Olhar de Formiga"

O método antigo (chamado de Sampled-token OPD) funciona assim:
O aluno escreve uma palavra. O professor olha apenas para aquela palavra específica e diz: "Isso está certo" ou "Isso está errado".

Por que isso falha?

O Efeito Dominó: Em uma conversa longa ou um raciocínio complexo, uma palavra pode ser tecnicamente correta sozinha, mas levar a uma frase sem sentido depois. O método antigo ignora o futuro. É como um professor que elogia o aluno por escrever a palavra "banana", mesmo que o aluno esteja escrevendo "banana de sapato" e a frase inteira esteja errada.
O Aluno se Perde: Conforme o aluno começa a pensar de forma diferente do professor (criando caminhos novos), o professor pode não entender o contexto. Ele vê uma palavra que parece estranha para ele e pune o aluno, mesmo que o aluno esteja no caminho certo.
O Ruído da Tradução: Às vezes, o professor e o aluno usam "dicionários" diferentes (tokenizadores diferentes). O aluno escreve uma palavra que o professor vê como três pedaços estranhos. O professor pune o aluno por algo que não é um erro real, apenas uma diferença de "idioma".

Resultado: O aluno fica nervoso, aprende coisas erradas e começa a repetir frases sem sentido ou a "alucinar" (inventar coisas), porque o feedback é muito ruidoso e focado no detalhe errado.

2. A Solução: O "Círculo de Apoio" (Teacher Top-K Local Support Matching)

Os autores propõem uma mudança inteligente. Em vez de olhar apenas para a palavra que o aluno escolheu, o professor olha para um pequeno grupo de palavras prováveis que ele mesmo consideraria boas naquele momento.

A Analogia do Guia de Montanha:

Método Antigo: O guia diz: "Você pisou na pedra X. Está errado. Volte." (Mesmo que a pedra X fosse segura, mas o guia não a conhece).
Novo Método: O guia diz: "Olhe para este grupo de 5 pedras seguras que eu conheço aqui. Você está perto de uma delas? Ótimo. Vamos ajustar sua direção para ficar mais próximo desse grupo."

O novo método compara a distribuição de probabilidades do aluno com a do professor dentro desse "grupo seguro".

Se o aluno escolheu uma palavra que o professor também acha plausível, ele é recompensado.
Se o aluno escolheu algo que o professor acha improvável, ele é corrigido.
Mas, o professor não pune o aluno apenas porque ele escolheu uma palavra específica, e sim porque ele se afastou do "grupo de segurança".

3. As "Correções de Engenharia" (Os Truques Práticos)

Para que isso funcione na prática, eles adicionaram três "óculos de proteção":

Filtrar o Caos (Top-p Sampling): Eles impedem o aluno de gerar respostas totalmente aleatórias e loucas antes de receber o feedback. É como dizer ao aluno: "Só pense em ideias que tenham pelo menos 90% de chance de fazer sentido". Isso evita que o professor fique confuso com loucuras.
Máscara de Erros de Tradução (Special-token Masking): Eles ensinam o sistema a ignorar os erros de "dicionário" (tokenização). Se o professor e o aluno usam formatos diferentes para a mesma coisa, o sistema ignora essa diferença e foca no significado.
Normalização: Eles ajustam a matemática para garantir que, ao olhar apenas para um grupo pequeno de palavras, a comparação seja justa. É como pesar as frutas em uma balança que foi recalibrada para o tamanho da cesta, e não para o tamanho do armazém inteiro.

4. O Resultado: Um Aluno Mais Estável e Criativo

Com essa nova abordagem:

Menos Histeria: O aluno não fica nervoso com cada pequena mudança. O aprendizado é mais suave.
Melhor Raciocínio: Em tarefas de matemática e lógica complexa, o aluno consegue manter o foco por mais tempo sem "travar" ou repetir frases.
Aprendizado Real: O aluno aprende a pensar como o professor, mas mantendo sua própria capacidade de explorar novos caminhos, sem se perder no processo.

Resumo Final:
O artigo diz que ensinar uma IA olhando apenas para a palavra atual é como tentar montar um quebra-cabeça olhando apenas para uma peça de cada vez, sem ver a imagem geral. O novo método ensina a IA a olhar para o "grupo de peças" que fazem sentido juntas, criando um aprendizado mais estável, menos propenso a erros e muito mais eficiente para tarefas complexas e longas.

Each language version is independently generated for its own context, not a direct translation.

Resumo Técnico: Revisão da Distorção On-Policy (OPD)

1. O Problema

A Distorção On-Policy (OPD) é uma técnica atraente para o pós-treinamento de Grandes Modelos de Linguagem (LLMs), onde um modelo "aluno" é treinado em seus próprios rollouts (trajetórias geradas), enquanto recebe feedback de um modelo "professor" mais forte. Diferente da distilação offline (que usa traços fixos do professor), a OPD é essencial para cenários de longo horizonte e tarefas agênticas, onde o aluno rapidamente gera prefixos que são raros ou inexistentes nos dados de treinamento do professor.

No entanto, a implementação atual padrão da OPD, conhecida como comparação de token amostrado (sampled-token OPD), mostra-se frágil e instável em cenários de longo prazo. O artigo identifica três modos de falha principais:

Sinal desequilibrado de um único token: A atualização depende apenas da razão de log-probabilidades do token amostrado. Isso cria um sinal de aprendizado altamente desequilibrado, onde a maioria dos tokens recebe recompensas negativas e o aprendizado é impulsionado desproporcionalmente por uma pequena fração de tokens com vantagem positiva.
Guia não confiável do professor: Em prefixos gerados pelo aluno que são atípicos para o professor, a probabilidade que o professor atribui a um token pode não refletir a qualidade real da trajetória. O professor pode atribuir alta probabilidade a tokens que parecem plausíveis localmente, mas que levam a loops de repetição, reinicializações de raciocínio ou desvios da direção desejada.
Distorções por incompatibilidade de tokenização: Se o aluno e o professor usam tokenizadores diferentes, o mesmo texto bruto pode ser segmentado de formas distintas. Isso faz com que tokens semanticamente corretos gerados pelo aluno recebam baixa probabilidade do professor apenas por uma questão de alinhamento de tokens (ex: < vs think), distorcendo o sinal de recompensa.

Além disso, existe um trade-off teórico: métodos de nível de sequência (como KL reverso completo) capturam melhor a qualidade da trajetória, mas sofrem de variância de gradiente extremamente alta ( $O(T^4)$ ), tornando a otimização instável. A OPD de nível de token reduz a variância ( $O(T^2)$ ), mas introduz viés ao ignorar o acoplamento com recompensas futuras.

2. Metodologia Proposta

Os autores propõem uma nova abordagem chamada Correspondência de Suporte Local Top-K do Professor (Teacher Top-K Local Support Matching). Em vez de comparar apenas o token amostrado, o método compara as distribuições do aluno e do professor sobre um conjunto de suporte local definido pelo professor.

A implementação técnica envolve:

Definição do Suporte: Para cada prefixo, o conjunto de suporte $S$ é definido como os $K$ tokens de maior probabilidade segundo o professor ( $TopK_q$ ).
KL Reverso Truncado: O objetivo de treinamento é uma média do KL reverso truncado sobre esse conjunto de suporte, em vez de todo o vocabulário ou apenas um token. As distribuições do aluno e do professor são renormalizadas dentro desse suporte local.
Amostragem Top-p: Os rollouts são gerados usando amostragem top-p para manter as trajetórias próximas de continuação típicas, evitando prefixos onde o sinal do professor se torna pouco informativo.
Mascaramento de Tokens Especiais: Tokens especiais problemáticos (como marcadores de fim de sequência ou tokens de pensamento) são mascarados para evitar penalidades falsas devido a incompatibilidades de tokenização.

A fórmula do objetivo ( $L_{LSM}$ ) calcula a divergência entre as distribuições renormalizadas do aluno e do professor dentro do conjunto $S$ , proporcionando um sinal de atualização mais equilibrado e robusto, mantendo a eficiência computacional.

3. Contribuições Principais

Análise Teórica do Trade-off: Demonstraram que a OPD de nível de token é enviesada em relação ao objetivo de nível de sequência, mas possui um limite de variância muito mais apertado, o que é crucial para o treinamento de longo horizonte.
Identificação de Modos de Falha: Isolaram empiricamente três falhas críticas da OPD de token amostrado: sinal desequilibrado, guia não confiável em prefixos fora da distribuição do professor e distorções de tokenização.
Nova Metodologia e Resultados: Propuseram a correspondência de suporte local Top-K, implementada como KL reverso truncado. Demonstraram que essa abordagem supera a OPD padrão em termos de estabilidade de otimização e desempenho downstream.

4. Resultados Experimentais

Os experimentos foram conduzidos em dois cenários: raciocínio matemático de tarefa única e treinamento multi-tarefa (raciocínio matemático + tarefas agênticas).

Raciocínio Matemático (Tarefa Única):
- O método proposto ("Ours") alcançou uma pontuação média de 41.5 em benchmarks matemáticos (Math500, AIME, etc.), superando a OPD de token amostrado (36.4) e a OPD com apenas mascaramento (40.7).
- Isso indica que a melhoria não vem apenas da correção de tokenização, mas de um sinal de distilação local mais forte e equilibrado.
Treinamento Multi-tarefa (Matemática + Agente):
- O método proposto melhorou significativamente o desempenho matemático (aumentando a média de 36.6 para 41.7) enquanto mantinha um desempenho forte nas tarefas agênticas (ALFWorld), chegando a 97.7 de taxa de sucesso com mascaramento.
- A OPD padrão tendia a ter desempenho inferior na parte matemática quando alternada com tarefas agênticas.
Dinâmica de Treinamento:
- O método proposto exibiu menor variância de gradiente, normas de gradiente menores e menor fração de clipping (corte de gradiente) em comparação com a OPD padrão.
- A convergência foi mais estável, com ganhos consistentes ao longo do treinamento, e não apenas no checkpoint final.
- Ajustes de ablação mostraram que a renormalização e a amostragem top-p são componentes essenciais para a estabilidade; sem eles, o treinamento colapsa rapidamente.

5. Significado e Conclusão

Este trabalho é significativo porque oferece uma solução prática e simples para um problema fundamental no pós-treinamento de LLMs: a instabilidade da distilação on-policy em trajetórias longas.

Equilíbrio entre Viés e Variância: A proposta ocupa um "meio-termo" ideal, mantendo a supervisão local o suficiente para controlar a variância (evitando o colapso de otimização) enquanto substitui a estimativa pontual de um único token por uma comparação de distribuição local, mitigando o viés e a fragilidade.
Robustez: A abordagem é menos sensível a artefatos de tokenização e a desvios de distribuição entre professor e aluno.
Aplicabilidade: Os resultados sugerem que, para tarefas complexas de raciocínio e agentes autônomos, a supervisão baseada em distribuição local é superior à supervisão baseada em tokens individuais, permitindo um treinamento mais estável e desempenho superior em benchmarks desafiadores.

O artigo conclui que, embora a correspondência com o professor não seja uma solução perfeita para o sucesso da tarefa (devido a possíveis "hacks" de recompensa local), a correspondência de suporte local é um componente vital para estabilizar o treinamento on-policy e reduzir a lacuna de desempenho entre o aluno e o professor em cenários de longo horizonte.