Mezcla de Comportamientos en Región de Confianza para Destilación On-Policy

Resumen

La destilación en política (On-Policy Distillation, OPD) entrena a un estudiante con prefijos muestreados de su propia política, mientras iguala a un profesor más fuerte. Esto aborda el desajuste de prefijos de la destilación fuera de línea, pero las simulaciones tempranas del estudiante aún pueden ser deficientes, colocando la supervisión del profesor sobre prefijos débiles o de baja calidad. Proponemos la Mezcla de Comportamiento con Región de Confianza (Trust-Region behavior Blending, TRB), un método de calentamiento que reemplaza la política de simulación temprana con la política de comportamiento más cercana al profesor dentro de una región de confianza KL centrada en el estudiante, mientras mantiene inalterada la pérdida OPD de KL inversa por prefijo. El presupuesto KL se atenúa hasta cero, de modo que el entrenamiento retorna a simulaciones puras del estudiante tras el calentamiento. En dos entornos de destilación de razonamiento matemático, TRB alcanza el promedio más fuerte entre los métodos comparados.

English

On-policy distillation (OPD) trains a student on prefixes sampled from its own policy while matching a stronger teacher. This addresses the prefix mismatch of offline distillation, but early student rollouts can still be poor, placing teacher supervision on weak or low-quality prefixes. We propose Trust-Region behavior Blending (TRB), a warmup method that replaces the early rollout policy with the closest-to-teacher behavior policy inside a student-centered KL trust region, while keeping the per-prefix reverse-KL OPD loss unchanged. The KL budget is annealed to zero, so training returns to pure student rollouts after warmup. Across two math-reasoning distillation settings, TRB attains the strongest average among the compared methods.