OPRD: Destilación de Representaciones On-Policy

Resumen

La destilación on-policy (OPD) supervisa al estudiante únicamente en el espacio de salida, igualando las probabilidades del siguiente token. Este paradigma centrado exclusivamente en la salida presenta dos limitaciones: (1) la varianza de muestreo derivada de las estimaciones de KL de Monte Carlo sobre vocabularios extensos (p. ej., los ~150k tokens de Qwen) persiste durante todo el entrenamiento, y (2) trata al profesor como una caja negra, descartando todos los estados ocultos intermedios tras la cabeza del LM. Proponemos la Destilación de Representaciones On-Policy (OPRD), que eleva la destilación al espacio de los estados ocultos al alinear las representaciones del estudiante y del profesor en capas seleccionadas a lo largo de los mismos despliegues, evitando por completo la cabeza del LM. Teóricamente, OPRD elimina la varianza de muestreo y proporciona información estructural más rica por capa. Empíricamente, OPRD cierra la brecha estudiante-profesor en AIME 2024/2025 y AIMO, mientras que las líneas base de OPD en el espacio de salida se estancan por debajo del profesor. Además, OPRD entrena 1,44 veces más rápido y utiliza un 54 % menos de memoria que la OPD top-k. Código: https://github.com/ShenzhiYang2000/OPRD.

English

On-policy distillation (OPD) supervises the student only in output space by matching next-token probabilities. This output-only paradigm has two limits: (1) sampling variance from Monte Carlo KL estimates over large vocabularies (e.g., Qwen's ~150k tokens) persists throughout training, and (2) it treats the teacher as a black-box, discarding all intermediate hidden states after the LM head. We propose On-Policy Representation Distillation (OPRD), which lifts distillation into hidden-state space by aligning student and teacher representations across selected layers on the same rollouts, bypassing the LM head entirely. Theoretically, OPRD eliminates sampling variance and provides richer per-layer structural information. Empirically, OPRD closes the student-teacher gap on AIME 2024/2025 and AIMO, while output-space OPD baselines plateau below the teacher. OPRD also trains 1.44x faster and uses 54% less memory than top-k OPD. Code: https://github.com/ShenzhiYang2000/OPRD.