Search, Fail, Recover: A Training Framework for Correction-Aware Reasoning

2026-07-08

AI Summary

Traditional reasoning methods exhibit limited performance in complex tasks requiring backtracking and correction, particularly lacking effective recovery mechanisms in delayed-failure scenarios. This work proposes Pyligent, a framework that models reasoning as a verification-guided search process. By employing a task-specific verifier to label the success or failure of reasoning paths, Pyligent converts the search tree into three types of supervision signals—“continue,” “complete,” and “backtrack”—to train the model. Notably, it introduces explicit supervision on failed branches for the first time, enabling the model to actively recognize errors and initiate backtracking, thereby transcending the conventional paradigm of merely imitating successful trajectories. Experiments demonstrate that Pyligent improves solution rates by 72.7 percentage points on Hidden Graph tasks and by 13–27 percentage points on Sudoku and Blocksworld, significantly enhancing the model’s error-correction and recovery capabilities.

Abstract

Many reasoning tasks are not well described by a single left-to-right chain: a solver may need to pursue a plausible branch, observe delayed failure, and return to the latest prefix that can still be completed. We introduce Pyligent, a training and inference framework inspired by the Diligent Learner formulation that represents reasoning as validated search over partial solution chains. A task validator labels generated continuations and failures, and the resulting search trees are converted into supervised targets for three actions: continue, finish, and backtrack, with optional traces that summarize abandoned branches. We evaluate Pyligent on a hidden directed graph task designed to isolate delayed-failure recovery, and on structured reasoning domains with exact validators, including $4{\times}4$ Sudoku, Sudoku with reasoning traces, and Blocksworld. Compared with gold-only supervised fine-tuning, Pyligent improves solve rate by $72.7$ percentage points on hidden graphs, by $17$ and $18$ points on mixed and expert Sudoku, by $27$ and $14$ points on mixed and expert Sudoku with reasoning traces, and by $13$ points on Blocksworld. These results suggest that explicit failed-branch supervision can teach useful recovery behavior beyond imitation of polished solution chains.

Problem

• reasoning
• backtracking
• failure recovery
• search
• validation

Innovation

• validated search
• backtrack supervision
• reasoning with failure
• correction-aware reasoning
• search tree learning