Not Every Divergence Should Be Suppressed: Counterfactual Recoverability in On-Policy Distillation

2026-08-04

AI Summary

This work addresses a key limitation in policy distillation: conventional trajectory-divergence-based supervision cannot reliably determine whether an erroneous prefix is recoverable. To resolve this, the paper introduces counterfactual recoverability as a novel decision variable, classifying each erroneous state by comparing teacher rollouts from continuation versus rollback branches. Based on this classification, selective supervision—either retaining, rolling back, or applying standard supervision—is applied. Integrating budget-matched imitation policies with recoverability-guided labels, the method achieves a success rate of 0.578 on AIME2025 (baseline: 0.517), an average@32 of 0.3125 across AIME2024–2025, and 0.3070 on GPQA-Diamond, with a proxy recoverability metric attaining an AUC of 1.000—significantly outperforming divergence-only baselines.

Abstract

On-policy distillation (OPD) supervises student-visited trajectories, yet divergence-based rules cannot determine whether an erroneous prefix remains correctable. We formulate this decision as counterfactual recoverability and replay each error state through budget-matched teacher-continuation and rollback branches. Based on their relative success, states are categorized as recoverable, irreversible-but-avoidable, or ambiguous, and these labels guide whether training retains, rolls back, or conventionally supervises the corresponding trajectory. On AIME branch diagnostics, the mean continuation-minus-rollback effect is 0.185 for recoverable states and -1.000 for irreversible-but-avoidable states, demonstrating opposite intervention preferences. A branch-derived recoverability proxy achieves an AUC of 1.000, substantially outperforming divergence alone at 0.392. Across frozen evaluations, recoverability-aware control achieves the strongest recorded performance, reaching 0.578 success on held-out AIME2025 compared with 0.517 for the best baseline. It also improves AIME2024-2025 average@32 from 0.2656 to 0.3125 and GPQA-Diamond average@32 from 0.2702 to 0.3070. Component ablations further show that retaining teacher-correctable prefixes provides the largest individual contribution. These findings establish recoverability as an outcome-grounded decision variable for selective supervision in OPD.

Problem

• on-policy distillation
• counterfactual recoverability
• trajectory divergence
• selective supervision
• error correction

Innovation

• counterfactual recoverability
• on-policy distillation
• trajectory supervision
• teacher-student alignment
• selective intervention