REWIND: Uncertainty-Guided Trajectory Repair for Test-Time Reasoning

Yi Ding; Mengyu Wang; Min Zhou; Menglin Yang

cs.CLOct 14, 2026

REWIND: Uncertainty-Guided Trajectory Repair for Test-Time Reasoning

Authors: Yi Ding, Mengyu Wang, Min Zhou, Menglin Yang

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · University of Edinburgh · Huawei

Abstract

An early mistake in the generated reasoning can constrain all subsequent reasoning. Repeated-generation methods, like MCMC-based approaches, attempt to improve output quality by sampling multiple generations, but their effectiveness can be limited by randomly selecting where to regenerate. On MATH500 with Qwen2.5-Math-7B, we find that incorrect responses have substantially higher mean token entropy than correct responses, with an average within-question relative increase of approximately 112%, suggesting that entropy can signal reasoning errors. Motivated by this finding, we introduce REWIND, a training-free controller that uses entropy dynamics to propose targeted suffix regenerations. Across three model families and five benchmarks, REWIND outperforms all baselines, with gains of 0.51-10.00 percentage points over the strongest baseline. These results demonstrate the effectiveness of targeted, uncertainty-guided trajectory repair across models and tasks. Our code is available at https://anonymous.4open.science/r/rewind-D310.

Figures & tables

Figure 1: Incorrect-response entropy on MATH500. The x-axis groups questions by the number of correct responses out of 20. Yellow bars show equal-question means of incorrect-response entropy (Equation 9 ). Numbers above bars give entropy in nats; numbers below bin labels count questions in each group.

Figure 2: REWIND proposal policy. The entropy peak $t^{*}$ anchors gradient backtracking to $t_{\text{start}}$ . Boundary alignment gives the suffix start $t_{\text{aligned}}$ ; regeneration and score-based acceptance complete the proposal.

Qwen2.5-Math-7B (Specialist)
Methods	GSM8K	MATH500	GPQA	AIME2025	BBH
Vanilla	$61.41$	$48.13$	$26.77$	$2.22$	$63.73$
Temperature	$75.86$	$66.80$	$34.68$	$7.78$	$71.20$
Typical-p	$78.26$	$67.40$	$31.31$	$7.78$	$74.40$
Self-refine	$77.63$	$68.40$	$32.82$	$6.67$	$69.73$
DoLa	$77.86$	$70.20$	$34.84$	$10.00$	$73.60$

Table 2: Main results. Accuracy (%) averaged over three runs. Best and second-best are highlighted. Vanilla uses

T=1

Variant	Site / alignment	Acc.
MCMC	Global random / none	72.60
Random delimiter	Global random / full	73.20
Peak revision	Peak / full	74.20
No boundary	Local grad. / none	75.00
Random in window	Local random / full	75.20
Restricted boundary	Local grad. / . and newline	76.80

Table 3: Proposal ablation on Qwen2.5-Math-7B/MATH500. Accuracy (%) with

K=10

random MCMC steps per block. Revision variants add a suffix proposal. Local variants use the peak window; “grad.” uses its largest entropy increase.

Signal	Peak	Tracing	$\Delta$
Shannon entropy	74.20	77.6	+3.4
Temporal KL	71.80	75.0	+3.2
Top- $K$ similarity ( $K=5$ )	73.40	74.20	+0.8
Top- $K$ similarity ( $K=3$ )	72.60	73.40	+0.8
Skewness	73.20	73.40	+0.2
Kurtosis	72.00	73.00	+1.0

Table 4: Signal choice vs. intervention strategy. MATH500 accuracy (%) on Qwen2.5-Math-7B for six signal configurations.

\Delta

is tracing minus peak revision.

Trigger signal	GSM8K	MATH	AIME
MCMC	79.45	72.60	10.00
Shannon entropy	84.31	77.60	16.67
Skewness	80.97	73.40	12.22
Top- $K$ similarity ( $K=3$ )	76.19	73.40	12.22
Temporal KL	77.86	75.00	10.00
Kurtosis	80.59	73.00	8.89

Table 5: Trigger signals across benchmarks. Accuracy (%) on Qwen2.5-Math-7B.

HumanEval
Model / benchmark	Temp.	MCMC	Rewind
Qwen2.5-Math-7B	52.4	56.1	58.5
Tülu-8B	58.5	62.8	66.5
8-shot CoT
MATH500	68.4	73.2	78.8
GSM8K	85.5	86.3	88.4

Table 6: Transfer across tasks and prompts. Accuracy (%) on HumanEval and under 8-shot CoT math prompting. The latter uses Qwen2.5-Math-7B.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Metric	Naive	Tracing	$\Delta$ Gain
MCMC	$72.6$	-	-
Shannon Entropy	$74.2$	77.6	+3.4
KL Divergence (Temporal)	$71.8$	$75.0$	$+3.2$
Top-K Similarity ( $K=5$ )	$73.4$	$74.2$	$+0.8$
Top-K Similarity ( $K=3$ )	$72.6$	$73.4$	$+0.8$
Skewness	$73.2$	$73.4$	$+0.2$

Appendix

Table 9: Peak vs. gradient-based intervention localization. Accuracy (%) on MATH500 under naive peak revision and gradient tracing.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 931 papers

Inject a Turn, Not a Token: Learning Where to Redirect LLM Reasoning for Self-Correction
Oct 9, 2026Jiwon Jeon, Myungsik Cho, Jeonghye Kim +1LLM Self-Correction LLM Reasoning
Korea Advanced Institute of Science and Technology (KAIST)
LLM Reasoning as Trajectories: Step-Specific Representation Geometry and Correctness Signals
Apr 7, 2026Lihao Sun, Hang Dong, Bo Qiao +3LLM Reasoning Representation Geometry in Language Models
Microsoft
Reasoning Can Be Restored by Correcting a Few Decision Tokens
May 16, 2026Changshuo Shen, Leheng Sheng, Yuxin Chen +2Efficient Language Model Reasoning Inference-Time Compute Allocation
University of Science and Technology of China, Hefei, China · National University of Singapore, Singapore

Oct 9, 2026J/K move · Enter open · S save

Inject a Turn, Not a Token: Learning Where to Redirect LLM Reasoning for Self-Correction

Jiwon Jeon, Myungsik Cho, Jeonghye Kim, Youngchul Sung

Korea Advanced Institute of Science and Technology (KAIST)

LLM Self-Correction LLM Reasoning LLM Prompting

Open paper AskPDF arXiv

Large language models struggle to reliably identify and correct errors in their own reasoning. Unlike prior training-based approaches that fine-tune an LLM to generate revisions or feedback, we learn where to intervene in failed rollouts and determine how to redirect them. To choose where to intervene, we train a 130K-parameter MLP head on a frozen encoder using only final-answer correctness. For redirection, we find that under a fixed generation budget, the tested inline injections fail to reliably redirect reasoning, whereas presenting the same instruction in a new user turn changes the model's course. The resulting correction framework, UPTURN (User-Prompt TURN injection at a learned position), requires no LLM fine-tuning: it retains a selected prefix, injects an instruction as a new user turn, and regenerates the suffix. Across four reasoners and six math benchmarks, UPTURN improves average accuracy in every setting over a 32-rollout baseline. It outperforms self-correction and reflection-token injection baselines under matched correction budgets, transfers across reasoners and to GPQA-Diamond without retraining, and complements GRPO and self-correction methods.

Abstract

LLM Reasoning

Explore similar work

Inject a Turn, Not a Token: Learning Where to Redirect LLM Reasoning for Self-Correction

LLM Reasoning as Trajectories: Step-Specific Representation Geometry and Correctness Signals

Reasoning Can Be Restored by Correcting a Few Decision Tokens