Improving Model Distillation via Task Sensitivity

Summary & Methodology Analysis

On-policy distillation (OPD), which involves training a smaller student model to match the output distribution of a larger teacher model during inference, typically applies uniform weight to all generated tokens. This approach is suboptimal because it treats generic language patterns with the same priority as the specific reasoning signals required for task completion. The CROP method addresses this by calculating a score that identifies the most task-relevant positions in a response. This process begins by building prompt triplets consisting of an original input, a paraphrase, and a counterfactual version. By measuring the Jensen-Shannon divergence of the model output across these variants, the system calculates a sensitivity margin. This margin isolates the influence of the task instructions versus generic generation noise. The result is a score that determines which tokens receive the highest supervision priority, applied via a binary mask during the distillation update. The method allocates a batch-global supervised-token budget based on these scores, ensuring that only the most critical segments of the output contribute to the student model gradient updates. Evaluations demonstrate that applying this selection to Qwen3-4B distilling into Qwen3-1.7B provides a 1.92 point performance gain over the strongest alternative selector. Furthermore, the CROP-10 percent selection configuration improves aggregate average scores by 2.00 points over random token selection across benchmarks including MATH-500, AIME24, and GPQA-Diamond. Despite these gains, the current implementation has notable limitations. The empirical evidence is restricted to mathematical tasks and specific Qwen models, leaving its generalization to other domains unverified. Additionally, the CROP score functions as a heuristic ranking mechanism rather than an objective measurement of ground-truth relevance, meaning its effectiveness is tied to the quality of the sensitivity margin calculation.