AsyncOPD: ¿Qué tan desactualizada puede ser la destilación on-policy?

Resumen

La destilación on-policy (OPD) entrena a un estudiante utilizando sus propias trayectorias simuladas, guiado por la retroalimentación del profesor, y se está volviendo cada vez más importante para el post-entrenamiento de modelos de lenguaje de gran escala (LLM). Sin embargo, al igual que el aprendizaje por refuerzo (RL), la OPD enfrenta un cuello de botella sistémico on-policy, ya que las trayectorias pueden dominar el tiempo de entrenamiento en tareas de razonamiento. Los pipelines de entrenamiento asíncrono pueden aliviar este cuello de botella al desacoplar la generación de trayectorias de las actualizaciones del aprendiz, pero esto introduce datos de política desactualizada. Si bien trabajos previos han estudiado datos desactualizados en RL asíncrono, sus efectos en OPD siguen siendo poco explorados. Presentamos el primer estudio sistemático sobre la desactualización en OPD asíncrona, enfocándonos en un escenario práctico donde la retroalimentación del profesor se implementa mediante pérdidas KL locales y los logits completos del profesor son demasiado costosos de almacenar o transferir, lo que requiere cachés finitos de puntuaciones del profesor. Primero, mostramos que la dirección de la KL modifica el problema de los datos desactualizados: la KL directa ponderada por el profesor es más robusta a trayectorias desactualizadas, mientras que la KL inversa ponderada por el estudiante es vulnerable. Segundo, para este caso vulnerable de KL inversa, estudiamos si los métodos diseñados para estabilizar el RL asíncrono pueden mitigar la desactualización en OPD. En nuestros experimentos, no mejoran respecto a un sustituto más simple específico de OPD: recalcular la señal de KL inversa bajo el estudiante actual en el momento del aprendiz. Tercero, analizamos cómo los cachés finitos de puntuaciones del profesor generan un compromiso entre sesgo y varianza para estimadores de OPD con KL inversa dispersa y muestreada. Esto motiva el uso de Monte Carlo (MC) de múltiples muestras, que preserva la corregibilidad de MC mientras reduce la varianza de una sola muestra. Finalmente, presentamos y publicamos como código abierto AsyncOPD, un pipeline de entrenamiento OPD completamente asíncrono construido a partir de estas elecciones de estimadores. Los experimentos muestran que AsyncOPD mejora el rendimiento de entrenamiento entre 1,6 y 3,8 veces en comparación con el entrenamiento estrictamente síncrono, al tiempo que alcanza una precisión comparable.

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.