Search papers, labs, and topics across Lattice.
This paper introduces the concept of counterfactual recoverability in on-policy distillation (OPD), addressing the limitations of divergence-based rules in determining the correctability of erroneous prefixes in student-visited trajectories. By employing budget-matched teacher-continuation and rollback branches to categorize states as recoverable, irreversible-but-avoidable, or ambiguous, the authors demonstrate that this approach significantly enhances training interventions. The results reveal that recoverability-aware control achieves superior performance metrics compared to traditional divergence methods, with a notable AUC of 1.000 and improved success rates on benchmark evaluations.
Counterfactual recoverability transforms how we approach error correction in on-policy distillation, leading to a staggering AUC of 1.000 compared to 0.392 with divergence alone.
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.
