MetaOPD: Meta-Learned Token Weighting for On-Policy Distillation

2026-10-08

AI Summary

This study addresses the limitation of fixed token weights in existing online knowledge distillation, which fail to adaptively match the dynamic learning requirements of student models. To overcome this, we propose a joint training framework based on bilevel optimization that introduces a lightweight token weighting network co-optimized with the student model. The core innovation lies in leveraging meta-learning principles to differentiate through virtual updates, connecting weight decisions with post-update performance and enabling the weight mapping to dynamically adapt as the student model evolves. Experimental results demonstrate that on mathematical reasoning tasks, both 0.6B and 1.7B scale models achieve significantly superior Avg@8 and Pass@8 metrics compared to existing baseline methods.

Abstract

On-policy distillation (OPD) trains a student on its own generated responses using token-level teacher supervision. However, uniform weighting overlooks differences in token learning value, while existing weighting methods rely on predefined mappings from prediction signals to token weights. These mappings are not learned from the effectiveness of the resulting student updates, limiting their ability to adapt to evolving learning needs. In this paper, we propose MetaOPD, a bilevel optimization framework that jointly learns the student model and a lightweight token-weighting network. The inner objective updates the student through weighted OPD, while the outer objective optimizes the weighting network using validation loss on reference solutions after a virtual student update. Differentiating through this update connects weighting decisions to their effects on post-update performance, allowing the mapping from prediction signals to token weights to evolve alongside the student. Experiments on six mathematical reasoning and three out-of-domain datasets, covering two student scales and seven baselines, demonstrate the effectiveness of MetaOPD, with Avg@8/Pass@8 gains over OPD of 1.99/5.97 percentage points for the 0.6B student and 2.25/6.41 points for the 1.7B student.

Problem

• On-Policy Distillation
• Token Weighting
• Knowledge Distillation
• Bilevel Optimization

Innovation

• On-Policy Distillation
• Meta-Learning
• Bilevel Optimization
• Token Weighting
• Knowledge Distillation