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

Abstract: On-policy distillation (OPD) supervises a student LLM 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.

• The paper introduces CROP, a selective on-policy distillation method that ranks tokens by counterfactual prompt sensitivity calibrated against paraphrases, while keeping student rollouts, teacher targets, and the underlying update unchanged.
• CROP improves aggregate performance by 3.11 points over Pure OPD and 1.92 points over TIP for Qwen3-4B to Qwen3-1.7B, while CROP-ent reaches a 57.48 score for Qwen3-8B to Qwen3-4B at a 10% supervision budget.
• Controlled experiments show CROP’s ranking—not sparsity alone—drives gains, outperforming random selection by 2.00 points and lowest-score selection by 3.34 points, although results remain limited to mathematical tasks and Qwen-family models.

Motivation: two dimensions of token supervision value

On-policy distillation (OPD) supervises a student LLM on trajectories sampled from its own current policy, querying a stronger teacher on the prefixes actually visited. Standard OPD, however, treats every response token as an equally suitable imitation target. The authors argue that existing selective OPD methods—entropy-based selectors such as Entropy and TIP (Jin et al., 7 Mar 2026, Xu et al., 15 Apr 2026), teachability- and trust-region-based methods such as TA-OPD and TrOPD (Wang et al., 26 May 2026, Xing et al., 31 May 2026), and position- or training-dynamics-based approaches (Liu et al., 20 May 2026, Xie et al., 21 Jun 2026, Jiang et al., 10 May 2026)—predominantly measure optimization need: whether a token is uncertain, discrepant from the teacher, or locally correctable. A complementary dimension, task relevance—whether the supervision is tied to the semantic content of the specific input—is largely unaddressed. The closest prior work, CREDIT (Shen et al., 12 May 2026), rescores a fixed response under unrelated batch queries, but those queries vary task semantics and nuisance factors jointly, so the resulting contrast conflates semantic dependence with other input differences.

Method

CROP (Counterfactual Relevance for On-Policy Distillation) operationalizes task relevance through a paraphrase-calibrated counterfactual sensitivity margin. For each source prompt $x$, an offline pipeline constructs a validated triplet $(x, x^{\mathrm{para}}, x^{\mathrm{cf}})$, where the paraphrase preserves all task-relevant meaning and the counterfactual changes exactly one material condition while preserving task type and answer format. An LLM generator produces candidate rewrites and an independent critic validates them against criteria of semantic preservation, single-condition intervention, determinacy, and answer-type consistency; only strictly validated triplets enter training.

At training time, given a student rollout $y \sim \pi_{\bar\theta}(\cdot \mid x)$ sampled from the rollout-time snapshot, CROP holds $y$ and all prefixes fixed and rescoring the same response under all three prompts. At each position it computes top-$K$ Jensen–Shannon divergence ($K=64$, with a residual probability bucket for omitted vocabulary mass) between the original-prompt distribution and both the counterfactual and paraphrase distributions:

$$s_t^{\mathrm{CROP}} = d_t^{\mathrm{sem}} - d_t^{\mathrm{surf}}$$

where $$d_t^{\mathrm{sem}} = \operatorname{JSD}_K(P_t^o, P_t^c)$$ and $$d_t^{\mathrm{surf}} = \operatorname{JSD}_K(P_t^o, P_t^p)$$. Holding the realized rollout fixed prevents differences among independently sampled trajectories from being mistaken for prompt sensitivity. The signed difference is explicitly characterized as a model-internal, contrast-specific ranking heuristic—not a measure of correctness or an unrestricted causal effect—and because JSD is nonlinear, it is not an additive decomposition of semantic and surface effects.

Scores are converted into a batch-global binary mask under a nominal supervised-token budget $\rho$ (10% in main experiments), with a per-sample minimum-retention safeguard of one token; the realized retention can therefore slightly exceed the nominal budget. Crucially, CROP changes only which tokens enter the masked sampled-token OPD loss—the rollout distribution, teacher target, and underlying clipped PPO-style update are unchanged. An optional variant, CROP-ent, combines the relevance score with normalized student entropy via a Soft-OR to test uncertainty–relevance complementarity.

The triplet-construction funnel is reported transparently: from 17,398 deduplicated DAPO-Math-17K prompts, 16,148 pass strict validation in the first round and 446 more after feedback-based repair (73.48% repair success), yielding 16,594 validated triplets (95.38% of sources). Critic-assigned confidences are self-reported diagnostics rather than calibrated correctness estimates, a caveat the authors state plainly.

Main results

Experiments cover Qwen3-4B → Qwen3-1.7B and Qwen3-8B (GRPO checkpoint) → Qwen3-4B, evaluated on AIME24, AIME25, MATH-500, GPQA-Diamond, HumanEval, and IFEval at a nominal 10% budget:

Setting	Pure OPD	TIP (best non-CROP)	CROP	CROP-ent
Qwen3-4B → Qwen3-1.7B	44.80	46.06	47.98	46.58
Qwen3-8B → Qwen3-4B	55.25	54.08	57.13	57.48

In the first setting, CROP improves over Pure OPD by 3.11 points and over the strongest selective baseline (TIP) by 1.92 points. In the second, CROP-ent leads with 57.48 and CROP reaches 57.13, both exceeding Pure OPD by more than 1.8 points. The reversal between settings indicates that entropy complements counterfactual relevance for some teacher–student pairs but is not uniformly beneficial—an honest inconsistency the paper does not smooth over.

Selector behavior and ablations

An entropy-rank analysis shows CROP is not reducible to an uncertainty proxy: 10.0% of its retained tokens fall in the lowest-entropy decile, versus 3.6% for CREDIT and nearly zero for entropy-oriented selectors. Budget sweeps show CROP wins at 10%, 15%, and 20% budgets but performance varies non-monotonically; TA-OPD is strongest at 5% and Entropy at 30–40%, so the preferred selector depends on supervision density, with CROP most effective in the sparse-to-moderate regime.

Component ablations isolate the signal's ingredients: CROP (47.98) exceeds CS-OPD (raw observed-token counterfactual log-probability change, 46.00) by 1.98 points, PC-OPD (paraphrase sensitivity alone, 43.09) by 4.89 points, and CROP-Teacher (teacher-side scoring, 44.73) by 3.25 points. This establishes counterfactual sensitivity as the primary signal and paraphrase calibration as a meaningful refinement beyond raw counterfactual change.

Matched selection controls directly test whether the ranking selects useful tokens: at the identical 10% budget, highest-scoring CROP selection (47.98) beats random selection (45.98) by 2.00 points and lowest-scoring selection (44.64) by 3.34 points. The gain over random shows the score carries information; the gain over bottom-selection shows the ordering itself is meaningful. Notably, Random-10% achieves the highest MATH-500 score (80.20 vs. 78.60 for CROP), though CROP leads the aggregate and five of six columns.

A contamination audit finds seven exact MATH-500 overlaps with training prompts plus one condition-level near duplicate; on the decontaminated MATH-500-Clean subset, CROP retains its lead (79.07 clean accuracy, 48.06 clean average), and the maximum possible distortion (0.27 points) is far smaller than the reported margins.

Limitations and open questions

The evidence base has clear boundaries that the authors acknowledge. All experiments use mathematical prompts and Qwen-family models, so transfer across model families and domains remains untested. The triplet-generation prompts are instantiated only for mathematics, although the construction is claimed to extend to coding and open-domain QA. The validation pipeline controls semantic preservation, single-condition intervention, determinacy, and answer-type consistency, but does not independently control other contrast properties such as problem difficulty—a high margin may reflect overreaction to residual prompt differences rather than genuine semantic dependence. Additionally, the top-$(x, x^{\mathrm{para}}, x^{\mathrm{cf}})$0 JSD score is an approximation rather than exact full-vocabulary divergence, and the critic's confidence values are self-reported and uncalibrated. Open questions include whether multiple valid interventions per prompt or difficulty-controlled counterfactuals improve selection, and whether the low-entropy tail concentration translates into gains outside mathematics.

Conclusion

CROP contributes a complementary criterion for selective on-policy distillation: task relevance, measured as paraphrase-calibrated counterfactual sensitivity on a fixed student rollout. Implemented as a hard batch-global token mask over an otherwise unchanged sampled-token OPD objective, it consistently outperforms full-token OPD and optimization-need selectors at matched budgets, with controlled selection experiments confirming that the ranking—not merely sparsity—drives the gains. The results support matched semantic interventions as a practical basis for allocating token-level distillation supervision, within the demonstrated scope of mathematical tasks and Qwen-based teacher–student pairs.