Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Resumo Técnico: Unindo Trajetórias de Raciocínio na Destilação On-Policy via Orientação de Futuro Próximo

Definição do Problema
A Destilação On-Policy (OPD) é um pilar fundamental para o pós-treinamento moderno de Grandes Modelos de Linguagem (LLMs), projetada para melhorar o raciocínio treinando um modelo estudante em trajetórias amostradas de sua própria política sob supervisão de um professor. A OPD padrão opera em um mecanismo de "amostragem de trajetória, mas aprendizado de token": ela identifica desvios via tokens de alta perda e os repara usando a correção de divergência KL reversa local.

Os autores identificam um desajuste fundamental nesta abordagem, argumentando que a supervisão ao nível de token é estruturalmente desalinhada com os requisitos de correção de raciocínio ao nível de trajetória. Esse desalinhamento se manifesta em dois modos críticos de falha:

1. Falsos Alarmes Ruidosos: Tokens de alta perda nem sempre correspondem a uma divergência genuína de raciocínio. A análise empírica revela que aproximadamente 30% dos tokens de alta perda caem em um regime de baixa divergência, representando discrepâncias de forma superficial (ex: preferências estilísticas como "e" vs. "então") em vez de bifurcações lógicas. Tratar esses casos como erros críticos fornece uma supervisão ruidosa que pode interferir ativamente na correção do raciocínio.
2. A Armadilha do Aprendizado Token a Token: Mesmo quando pontos de divergência genuínos são identificados, a correção KL reversa de um único token é frequentemente insuficiente para redirecionar toda a trajetória de raciocínio. Falhas de raciocínio tipicamente evoluem como deriva de distribuição de curto horizonte. Corrigir um token localmente não garante que o modelo se alinhará com a trajetória futura do professor; o modelo frequentemente deriva de volta para ramos errôneos imediatamente após a correção local.

Metodologia: OPD Consciente de Trajetória (TOPD)
Para abordar essas limitações, os autores propõem a TOPD (Trajectory-aware OPD), que une a lacuna entre a correção de token local e a evolução da trajetória global ao incorporar informações de trajetória de futuro próximo. O método consiste em duas etapas primárias:

1. Detecção de Divergência:

   • Para posições candidatas de alta perda, o sistema retrocede ao token precedente para estabelecer um prefixo compartilhado.
   • Tanto o professor quanto o estudante geram continuações de janela curta (comprimento $K=50$) a partir deste prefixo.
   • O sistema mede a discrepância de nível de trajetória usando a distância de Transporte Ótimo (OT) entre as continuações do professor e do estudante.
   • Pontos divergentes reais são identificados como posições que exibem tanto alta perda ao nível de token quanto alta divergência de OT em janela curta. Isso filtra discrepâncias superficiais benignas.

2. Injeção de Sinal do Professor ao Nível de Trajetória:

   • Em vez de otimizar apenas o token atual, a TOPD estende a supervisão para toda a trajetória de janela curta.
   • Usando o plano de transporte OT ($\gamma$), o método constrói um "alvo suave consciente de trajetória" para cada posição futura do estudante. Este alvo alinha os tokens futuros do estudante com a trajetória do professor com base no mapeamento de transporte ótimo.
   • O estudante é treinado para corresponder a esses alvos suaves alinhados por OT através da janela, aprendendo efetivamente como guiar seu caminho de raciocínio futuro próximo em direção ao caminho do professor, e não apenas como corrigir o token atual.
   • O objetivo final combina a perda OPD padrão com um termo de perda de trajetória ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Principais Contribuições

• Identificação de Desajuste: Este trabalho estabelece que a supervisão ao nível de token é um "proxy impreciso" para a correção de raciocínio ao nível de trajetória, falhando em distinguir entre bifurcações lógicas genuínas e variações superficiais benignas.
• Caracterização Empírica de Modos de Falha: Os autores fornecem evidências de que ~30% dos tokens de alta perda são falsos alarmes de baixa divergência e demonstram a "armadilha do aprendizado token a token", onde correções locais falham em restaurar a trajetória de raciocínio geral.
• Princípio de Orientação ao Nível de Trajetória: O artigo introduz um princípio onde a informação de trajetória de futuro próximo é distribuída por múltiplos tokens. Ao usar o alinhamento OT de janela curta, a TOPD injeta as discrepâncias de caminho professor-estudante diretamente na função de perda, permitindo que o estudante aprenda mudanças de trajetória em vez de edições de tokens isolados.

Resultados Experimentais
Experimentos foram conduzidos utilizando modelos Qwen3 (destilando um professor de 30B para um estudante de 4B) em benchmarks competitivos de matemática (AIME24, AIME25, HMMT25-Feb).

• Ganhos de Desempenho: A OPD padrão melhorou a precisão média do modelo de aquecimento (warm-start) de 38,9% para 47,8%. A TOPD aumentou isso para 52,2%.
• Benchmarks Específicos: A TOPD melhorou a precisão do AIME24 de 60,0% para 63,3% e do AIME25 de 46,7% para 53,3%.
• Estudos de Ablação:
  • Suprimir tokens de alta perda e baixa OT (falsos alarmes) na OPD padrão resultou em um ganho modesto (+0,4%), confirmando que esses tokens são ruidosos.
  • Suprimir tokens de alta OT e alta perda (divergências reais) causou uma queda significativa no desempenho, confirmando sua necessidade.
  • A análise de correção de trajetória local mostrou que, embora a OPD padrão reduzisse a distância de OT de janela curta parcialmente (de 0,538 para 0,350), a TOPD alcançou uma redução substancial (para 0,204), demonstrando uma capacidade superior de redirecionar trajetórias de futuro próximo.
• Eficiência: A TOPD introduz um overhead computacional modesto, com o tempo de treinamento sendo 1,41$\times$ o da OPD padrão, pois o OT é computado apenas em janelas selecionadas de alta perda.

Significância e Alegações
O artigo afirma que as falhas de raciocínio na OPD são fundamentalmente fenômenos de nível de trajetória e não apenas incompatibilidades de tokens isolados. Ao mudar de edições de tokens pontuais para um alinhamento consciente de transição global, a TOPD permite um alinhamento de raciocínio mais eficaz. Os autores posicionam este trabalho como um passo em direção a objetivos de destilação mais conscientes de trajetória, argumentando que a correção de raciocínio bem-sucedida exige que o estudante aprenda não apenas qual token corrigir, mas como navegar pelo caminho de raciocínio subsequente para permanecer consistente com o professor.

Limitações
Os autores reconhecem que a TOPD incorre em custos computacionais adicionais devido aos cálculos de OT, foca principalmente em divergências de curto horizonte (podendo perder erros de contexto longo) e é validada atualmente principalmente em benchmarks de raciocínio matemático, e não em tarefas de código ou de escrita aberta. Além disso, a distância OT mede a proximidade da trajetória, mas pode não capturar a equivalência semântica onde caminhos diferentes levam à mesma resposta correta.