OPRD: On-Policy Representation Distillation (arxiv.org)

Researchers have introduced On-Policy Representation Distillation (OPRD), a novel method in machine learning that enhances the efficiency of knowledge transfer between teacher and student models. Unlike traditional on-policy distillation methods that focus solely on output probabilities, OPRD aligns the hidden states of student and teacher models across selected layers during training. This innovative approach aims to mitigate the drawbacks of previous methods, such as significant sampling variance from Monte Carlo estimates and a lack of structural insights from hidden states, providing a more comprehensive learning mechanism. OPRD not only demonstrates theoretical advantages by reducing sampling variance but also shows empirical success, narrowing the performance gap between student and teacher models in benchmarks like AIME 2024/2025 and AIMO. Additionally, it boasts a 1.44x increase in training speed and a 54% reduction in memory usage compared to top-k OPD methods. This advancement is significant for the AI/ML community as it promises faster model training and more effective distillation strategies, which could lead to more powerful and efficient AI systems in various applications.