AsyncOPD: Quão Desatualizada Pode Ser a Destilação On-Policy?

Resumo

Destilação on-policy (OPD) treina um estudante em suas próprias trajetórias guiadas pelo feedback do professor e está se tornando cada vez mais importante para o pós-treinamento de modelos de linguagem de grande escala (LLMs). Assim como o aprendizado por reforço (RL), no entanto, a OPD enfrenta um gargalo de sistemas on-policy, pois as trajetórias podem dominar o tempo de treinamento para cargas de raciocínio. Pipelines de treinamento assíncrono podem aliviar esse gargalo ao desacoplar a geração de trajetórias das atualizações do aprendiz, mas isso introduz dados de política desatualizada. Embora trabalhos anteriores tenham estudado dados desatualizados em RL assíncrono, seus efeitos na OPD permanecem pouco explorados. Apresentamos o primeiro estudo sistemático da desatualização em OPD assíncrona, focando em um cenário prático onde o feedback do professor é implementado por meio de perdas KL locais e os logits do professor de vocabulário completo são muito caros para armazenar ou transferir, exigindo caches finitos de pontuações do professor. Primeiro, mostramos que a direção da KL altera o problema de dados desatualizados: a KL direta ponderada pelo professor é mais robusta a trajetórias desatualizadas, enquanto a KL reversa ponderada pelo estudante é vulnerável. Segundo, para este caso vulnerável de KL reversa, estudamos se métodos projetados para estabilizar RL assíncrono podem mitigar a desatualização da OPD. Em nossos experimentos, eles não melhoram em relação a um substituto mais simples específico da OPD: recomputar o sinal da KL reversa sob o estudante atual no momento do aprendiz. Terceiro, analisamos como caches finitos de pontuações do professor criam um compromisso entre viés e variância para estimadores de OPD de KL reversa esparsos e amostrados. Isso motiva o Monte Carlo (MC) multi-amostra, que preserva a corrigibilidade do MC enquanto reduz a variância de uma amostra. Finalmente, apresentamos e disponibilizamos como código aberto o AsyncOPD, um pipeline de treinamento OPD totalmente assíncrono construído a partir dessas escolhas de estimador. Experimentos mostram que o AsyncOPD melhora o throughput de treinamento em 1,6 a 3,8 vezes em relação ao treinamento síncrono estrito, alcançando precisão comparável.

English

On-policy distillation (OPD) trains a student on its own rollouts guided by teacher feedback and is becoming increasingly important for large language model (LLM) post-training. Like reinforcement learning (RL), however, OPD faces an on-policy systems bottleneck, as rollouts can dominate training time for reasoning workloads. Asynchronous training pipelines can alleviate this bottleneck by decoupling rollout generation from learner updates, but doing so introduces stale-policy data. While prior work has studied stale data in asynchronous RL, its effects in OPD remain underexplored. We present the first systematic study of staleness in asynchronous OPD, focusing on a practical setting where teacher feedback is implemented through local KL losses and full-vocabulary teacher logits are too expensive to store or transfer, necessitating finite teacher-score caches. We first show that KL direction changes the stale-data problem: teacher-weighted forward KL is more robust to stale rollouts, whereas student-weighted reverse KL is vulnerable. Second, for this vulnerable reverse-KL case, we study whether methods designed to stabilize asynchronous RL can mitigate OPD staleness. In our experiments, they do not improve over a simpler OPD-specific surrogate: recomputing the reverse-KL signal under the current student at learner time. Third, we analyze how finite teacher-score caches create a bias-variance tradeoff for sparse and sampled reverse-KL OPD estimators. This motivates multi-sample Monte Carlo (MC), which preserves MC correctability while reducing one-sample variance. Finally, we present and open-source AsyncOPD, a fully asynchronous OPD training pipeline built from these estimator choices. Experiments show that AsyncOPD improves training throughput by 1.6times to 3.8times over strict synchronous training while reaching comparable accuracy.