CROP: Task Relevance via Counterfactuals for Selective On-Policy Distillation

2026-08-13

AI Summary

This work addresses the limitation of existing selective policy distillation methods, which lack explicit modeling of task relevance and thus struggle to accurately identify supervision signals truly useful for the semantics of a given input. To remedy this, the paper introduces task relevance as an independent dimension within the distillation framework and proposes a token-level supervision selection mechanism based on counterfactual sensitivity and paraphrase calibration. Specifically, by constructing original–paraphrase–counterfactual triplets and fixing the student’s trajectory, the method quantifies task relevance at each response position through paraphrase-calibrated counterfactual sensitivity, enabling fine-grained, contrastively specific supervision allocation. Experiments demonstrate that, under two teacher–student settings, the proposed approach outperforms the strongest non-CROP selector by 1.92 and 2.96 performance points, respectively, validating the efficacy of task relevance as a complementary selection criterion.

Abstract

On-policy distillation (OPD) supervises a student language model on trajectories sampled from its current policy, but assigns equal credit to response tokens with unequal supervision value. Selective OPD addresses this limitation by allocating supervision non-uniformly across response tokens according to their estimated training value. Most existing criteria, however, focus primarily on optimization need, such as uncertainty or teacher-student disagreement, while task relevance, namely whether the supervision is tied to the semantic content of the current input, remains less directly characterized as a complementary dimension. To address this gap, we introduce Counterfactual Relevance for On-Policy Distillation (CROP), which operationalizes task relevance through a paraphrase-calibrated counterfactual sensitivity margin. For each source prompt, CROP constructs a validated original-paraphrase-counterfactual triplet, holds the student rollout fixed, and measures each response position by its sensitivity to a task-relevant condition change calibrated by its sensitivity to a meaning-preserving rewrite. Matched selection controls show that CROP identifies more useful supervision positions than random or lowest-relevance selection, while component comparisons confirm the value of both counterfactual sensitivity and paraphrase calibration. Across two teacher-student settings, CROP improves aggregate performance by 1.92 and 2.96 points over the strongest non-CROP selector. These results support task relevance as a complementary criterion for selective OPD and establish CROP as a model-internal, contrast-specific method for allocating token-level supervision.

Problem

• on-policy distillation
• task relevance
• selective supervision
• counterfactual reasoning
• token-level supervision

Innovation

• counterfactual relevance
• on-policy distillation
• task relevance
• paraphrase calibration
• selective supervision