OPRD: On-Policy дистилляция представлений

Аннотация

Политико-ориентированная дистилляция (On-policy distillation, OPD) контролирует студента только в пространстве выходных данных путем согласования вероятностей следующего токена. Эта парадигма, ограниченная только выходными данными, имеет два недостатка: (1) вариативность выборки из-за оценок Монте-Карло расхождения Кульбака–Лейблера (KL) для большого словарного запаса (например, ~150 тыс. токенов у Qwen) сохраняется на протяжении всего обучения, и (2) она рассматривает учителя как «черный ящик», отбрасывая все промежуточные скрытые состояния после головы языковой модели (LM head). Мы предлагаем политико-ориентированную дистилляцию представлений (On-Policy Representation Distillation, OPRD), которая переносит дистилляцию в пространство скрытых состояний путем выравнивания представлений студента и учителя на выбранных слоях на одних и тех же роллаутах, полностью минуя голову LM. Теоретически OPRD устраняет вариативность выборки и предоставляет более богатую структурную информацию на уровне слоев. Эмпирически OPRD сокращает разрыв между студентом и учителем на AIME 2024/2025 и AIMO, в то время как базовые методы OPD в пространстве выходных данных достигают плато ниже уровня учителя. OPRD также обучается в 1,44 раза быстрее и использует на 54% меньше памяти, чем top-k OPD. Код: 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.