Learning When to Abort in Transformer Planning: A Mechanistic Study

Shi Feng; Qimin Zhong; Siwei Wang; Yifei Shen; Shang-Hua Teng; Hao Liao; Wei Chen

cs.LGOct 18, 2026

Learning When to Abort in Transformer Planning: A Mechanistic Study

Authors: Shi Feng, Qimin Zhong, Siwei Wang, Yifei Shen, Shang-Hua Teng, Hao Liao, Wei Chen

Organizations: Harvard University · Shenzhen University · Microsoft Research Asia · University of Southern California

Abstract

Planning does not by itself imply the ability to recognise errors, but recognising errors can improve planning capacity. In this paper, we study this phenomenon in Transformer planning. To analyze the inner mechanism, we follow the approach of Wang et al. (2024) and abstract planning as path-finding in graphs. In this setting, we train the model to detect the error of invalid transitions, i.e., after the model generates a next node that is not adjacent with the current node, the model should then generate a special abort token x. For expressiveness, we construct a Transformer with feed-forward width $O(|V|)$ for any fixed graph with node set $V$ . It can detect all possible invalid transitions, and select a valid next node if the prefix is valid. We then propose a worst-case lower bound on feed-forward width, which makes the detector's linear width near-optimal up to logarithmic factors. In trained models, we demonstrate that invalid-transition detection increases previous-node attention, and most of the valid--invalid separation locates in the feed-forward network's output. This accords with our theoretical insights. Finally, we use this error detection to improve the model's performance: at inference, generating the abort token x will directly cause an external procedure to discard the attempt and restart from the original prompt. Experiments show that error detection combined with retry improves model's planning reliability on graphs, Countdown, and mazes.

Figures & tables

Figure 3: Graph domain over five seeds, each an independently drawn graph; solid is baseline, dashed $+\texttt{x}$ ; harness budget $K=100$ . (a) Single-shot validity. (b) Harness validity. (c) Composition of invalid outputs at $n=300$ for a single draw and for the harness output. (d) Mean draws per instance.

Figure 4: Countdown, $1/2/384$ , five seeds per $\rho$ . (a) Single-shot validity against $\rho$ , with the paired gain over the baseline. (b) Harness validity. (c) Composition of invalid outputs at $\rho=1$ for a single draw and for the harness output. (d) Mean draws per instance. In (a), (b) and (d) the baseline is a level band (mean $\pm$ s.e. over seeds) because its training data do not depend on $\rho$ .

Figure 5: Maze task, $1/2/384$ , four independently generated mazes. (a) Composition of invalid outputs for a single draw and the harness output. (b) Harness validity per maze, averaged over two runs. (c) Mean draws per instance.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

condition	all prefixes	target reachable
invalid prefix, no intervention	$0.997\pm 0.000$	$0.987\pm 0.003$
remove previous-node attention, main head	$0.594\pm 0.037$	$0.089\pm 0.086$
remove previous-node attention, other head	$0.949\pm 0.037$	$0.861\pm 0.097$
remove previous-node attention, both heads	$0.604\pm 0.007$	$0.008\pm 0.002$
remove target attention, both heads	$0.729\pm 0.081$	$0.911\pm 0.046$
remove current-node attention, both heads	$0.998\pm 0.001$	$0.997\pm 0.001$

Appendix

Table 1: Interventions after an injected invalid transition, six

n=300

graphs: mean

P(\texttt{x})

\pm

s.e. over all prefixes and over those whose target remains reachable. Removing target or current-node attention, and the valid prefix, are controls.

		to abort logit		to margin
Prefix	Abort logit	pre-FFN	FFN	pre-FFN	FFN
valid continuation	$3.0$	$-0.95$	$+4.0$	$-1.6$	$-8.1$
swap-valid, target reachable	$0.3$	$-1.15$	$+1.5$	$-1.7$	$-10.2$
invalid, target reachable	$29.0$	$-0.68$	$+29.6$	$-1.0$	$+22.4$
invalid, target unreachable	$28.8$	$-0.46$	$+29.2$	$-0.7$	$+22.0$

Appendix

Table 2: Direct logit attribution at the abort position, means over six

n=300

graphs (s.e. at most

0.4

). Contributions of the pre-FFN stream

X+\operatorname{MHA}(\operatorname{LN}_{1}(X))

and FFN output

M

to the abort logit and margin (abort logit minus highest node logit); “swap-valid” swaps the previous node to make the last transition valid.

Model	Mean $P(\texttt{x})$ after an invalid transition	Unjustified-abort mass fraction
Additive ( 13 )	$56.6\%$	$43.4\%$
Trained, same positions	$99.99\%$	$0.03\%$

Appendix

Table 3: The previous-node additive fit against the trained

1/2/384

Transformer on the same

2.78

million supervised positions of one

n=300

graph.

Figure 6: Planning representation probes over five seeds; shaded bands are s.e. (a) Edge–non-edge gap in

W^{M^{\prime}}

n=300

. (b)

W^{V^{\prime}}

by reachability group in baseline models. Architecture colours match (a).

Figure 7: Average per-head attention for baseline and $+\texttt{x}$ models, $1/2/384$ at $n=300$ . One representative seed is shown because head identities can swap across runs.

Figure 8: Blocksworld, Qwen3-0.6B-Base; mean $\pm$ s.e. over three seeds. (a) Validity for a single draw and under the harness. (b) Invalid outputs for a single draw and for the harness output. (c) Mean draws per instance.

Figure 9: Countdown, five seeds, mean $\pm$ s.e. (a), (b) Final validity against output tokens per instance under the likelihood and harness selection rules; budgets are set by Self-Backtracking at $N\in\{2,4,8,16,32\}$ . Dashed orange and dash-dotted green exchange the inference procedures. (c) Single-shot validity.

Language Model Error Detection

Language model error detection identifies unreliable answers, reasoning steps, and actions so systems can flag, verify, or repair them.

See all 91 papers

On the Ability of Transformers to Verify Plans
Mar 20, 2026Yash Sarrof, Yupei Du, Katharina Stein +3Transformer Symbolic Planning
Saarland University · Australian National University · University of Toulouse / LAAS-CNRS
Symmetry-Aware Transformer Training for Automated Planning
Aug 11, 2025Markus Fritzsche, Elliot Gestrin, Jendrik SeippContrastive Learning Transformer
Linköping University, Sweden
On Sample-Efficient Generalized Planning via Learned Transition Models
Feb 26, 2026Nitin Gupta, Vishal Pallagani, John A. Aydin +1Model-Based Planning OOD Generalization
University of South Carolina

Mar 20, 2026J/K move · Enter open · S save

On the Ability of Transformers to Verify Plans

Yash Sarrof, Yupei Du, Katharina Stein, Alexander Koller+2

Saarland University · Australian National University · University of Toulouse / LAAS-CNRS

Transformer Symbolic Planning Length Generalization

Open paper AskPDF arXiv

Transformers have shown inconsistent success in AI planning tasks, and theoretical understanding of when generalization should be expected has been limited. We take important steps towards addressing this gap by analyzing the ability of decoder-only models to verify whether a given plan correctly solves a given planning instance. To analyse the general setting where the number of objects -- and thus the effective input alphabet -- grows at test time, we introduce C*-RASP, an extension of C-RASP designed to establish length generalization guarantees for transformers under the simultaneous growth in sequence length and vocabulary size. Our results identify a large class of classical planning domains for which transformers can provably learn to verify long plans, and structural properties that significantly affects the learnability of length generalizable solutions. Empirical experiments corroborate our theory.

Abstract

Language Model Error Detection

Explore similar work

On the Ability of Transformers to Verify Plans

Symmetry-Aware Transformer Training for Automated Planning

On Sample-Efficient Generalized Planning via Learned Transition Models