Best-of-Depth: Learning Exploratory Dynamics in Looped Transformers

Shota Takashiro; Takeshi Kojima; Masanori Koyama; Kohei Hayashi; Yusuke Iwasawa

cs.LGNov 3, 2026

Best-of-Depth: Learning Exploratory Dynamics in Looped Transformers

Authors: Shota Takashiro, Takeshi Kojima, Masanori Koyama, Kohei Hayashi, Yusuke Iwasawa

Organizations: The University of Tokyo

Abstract

Looped Transformers extend inference-time computation through repeated latent-state updates, but final-depth training does not directly encourage useful intermediate alternatives. We introduce Best-of-Depth (BoD), which trains a recurrent trajectory as a candidate set by minimizing prediction loss across depths. Different valid targets can be matched at different depths without architectural changes or depth-specific labels. An exact-match analysis connects this objective to solution coverage. On N-Queens completion and graph coloring, BoD improves coverage over final-depth and intermediate-supervision objectives. N-Queens coverage reaches 91.03%, versus 45.75% for final-depth training. Further experiments show discovery beyond the training depth and recovery of solutions withheld from teacher pools on familiar inputs. Fine-tuning a pretrained looped language model with BoD yields higher pooled code-generation success than final-depth training when programs are generated separately at evaluated fixed depths. These results support training recurrent computation for solution discovery along its trajectory.

Figures & tables

Figure 2: Depth-wise solution discovery. Coverage@Depth (top) and Success@Depth (bottom) for greedy trajectories trained at depth 16. Bands indicate

\pm 1

SD.

	HumanEval					MBPP
	Depth				Success@Depth	Depth				Success@Depth
Method	4	8	16	32	32	4	8	16	32	32
Base	0.0	12.8	24.4	25.6	30.5	1.2	17.0	29.0	29.8	39.8
Last-Depth	0.0	13.4	27.4	26.8	34.8	1.8	21.2	30.2	31.4	39.2
Mean-Depth	6.7	20.1	27.4	27.4	31.7	10.4	26.0	28.0	29.0	37.2
Random-Depth	4.9	20.1	26.8	28.7	30.5	8.8	25.8	29.6	29.6	37.4

Table 1: Code-generation results (%). Base is pretrained; other rows denote fine-tuning objectives. Depth columns show accuracy for one greedy program per problem. Success@32 pools the six evaluated depths

\{4,8,12,16,24,32\}

. Bold marks each column’s best result.

Figure 3: Depth allocation and extrapolation. (a, b) Coverage pooling intermediate predictions from randomized trajectories under

ND=64

. (c, d) BoD with unperturbed greedy readouts; open markers denote training depth. Bands indicate

\pm 1

SD.

Figure 4: Exploratory recurrent dynamics on graph coloring. (a) Example outputs; NEW marks first appearances. (b) Fraction of graphs with a new valid solution. (c) Separate PCA projections for the same example and seed as (a); markers identify solutions or invalid outputs. (d) Relative state displacement. (e) Loss-minimizing depths. Bands (b, d) and error bars (e) indicate

\pm 1

SD. Dotted lines mark training depth 16.

	Last-Depth		Last-Depth + BoN		Best-of-Depth
Training solutions	Seen Coverage	Unseen Coverage	Seen Coverage	Unseen Coverage	Seen Coverage	Unseen Coverage
1	$\mathbf{99.8}\!\pm\!0.0$	$13.5\!\pm\!2.4$	$25.4\!\pm\!0.3$	$25.4\!\pm\!0.4$	$98.8\!\pm\!0.4$	$\mathbf{29.4}\!\pm\!0.4$
2	$46.6\!\pm\!0.6$	$12.9\!\pm\!1.2$	$22.8\!\pm\!0.7$	$22.9\!\pm\!0.9$	$\mathbf{94.2}\!\pm\!1.1$	$\mathbf{37.0}\!\pm\!0.8$
3	$37.1\!\pm\!1.4$	$13.8\!\pm\!1.3$	$23.6\!\pm\!0.2$	$23.7\!\pm\!0.3$	$\mathbf{91.3}\!\pm\!1.3$	$\mathbf{42.2}\!\pm\!0.5$
4	$28.4\!\pm\!0.3$	$13.8\!\pm\!0.4$	$24.4\!\pm\!0.6$	$24.4\!\pm\!0.6$	$\mathbf{87.7}\!\pm\!1.6$	$\mathbf{46.2}\!\pm\!1.2$
All	$11.8\!\pm\!0.2$	N/A	$25.0\!\pm\!0.4$	N/A	$\mathbf{67.4}\!\pm\!4.4$	N/A

Table 2: Discovery beyond the training targets on graph coloring. Seen Coverage and Unseen Coverage (%) on 3,034 training graphs under

ND=64

. All uses every valid solution as a training target candidate; Unseen Coverage is then undefined (N/A).

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Recurrent transformation	Two shared Transformer layers per loop
Width / heads / FFN width	128 / 8 / 384
Output normalization	StableMax
Training depth $D_{\mathrm{train}}$	16
Optimizer	AdamAtan2
Learning rate / batch size	$10^{-4}$ / 768

Appendix

Table 3: Training settings for the N-Queens and graph-coloring objective comparisons.

Setting	Value
Model	Huginn 3.565B, shared recurrent parameters
Data	CodeFeedback Python examples; 20,000 train / 1,000 validation
Maximum sequence length	1,024 tokens, including prompt and target
Training duration	1 epoch, 625 updates
Effective batch size	32
Maximum training depth	$D_{\mathrm{train}}=32$

Appendix

Table 4: Shared fine-tuning settings for the five code-generation objectives.

Method	@1	@4	@8	@16	@32	@64
Last-Depth	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00	22.27 $\pm$ 1.27	43.29 $\pm$ 3.12	45.75 $\pm$ 3.93
Mean-Depth	0.00 $\pm$ 0.00	2.51 $\pm$ 0.18	20.42 $\pm$ 1.20	36.52 $\pm$ 2.21	44.52 $\pm$ 1.66	49.94 $\pm$ 1.05
Random-Depth	0.00 $\pm$ 0.00	2.46 $\pm$ 0.11	19.93 $\pm$ 0.53	35.42 $\pm$ 0.91	44.15 $\pm$ 1.24	49.63 $\pm$ 1.34
Soft-Depth	0.00 $\pm$ 0.00	2.63 $\pm$ 0.37	21.53 $\pm$ 0.82	38.99 $\pm$ 0.79	48.26 $\pm$ 0.83	53.01 $\pm$ 1.47
Best-of-Depth	0.00 $\pm$ 0.00	0.10 $\pm$ 0.09	17.11 $\pm$ 1.25	85.92 $\pm$ 0.93	90.97 $\pm$ 0.98	91.03 $\pm$ 1.05

Appendix

Table 5: N-Queens Coverage@Depth (%) on 1,000 multi-solution test problems. Models are trained at

D_{\mathrm{train}}=16

Method	@1	@4	@8	@16	@32	@64
Last-Depth	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00	44.80 $\pm$ 2.52	78.23 $\pm$ 4.48	81.13 $\pm$ 4.94
Mean-Depth	0.00 $\pm$ 0.00	5.10 $\pm$ 0.40	40.37 $\pm$ 1.83	65.97 $\pm$ 3.25	75.37 $\pm$ 2.50	81.97 $\pm$ 2.46
Random-Depth	0.00 $\pm$ 0.00	5.00 $\pm$ 0.20	39.43 $\pm$ 1.34	64.50 $\pm$ 1.61	75.27 $\pm$ 1.72	82.30 $\pm$ 2.35
Soft-Depth	0.00 $\pm$ 0.00	5.30 $\pm$ 0.70	42.37 $\pm$ 1.11	68.67 $\pm$ 1.69	79.13 $\pm$ 2.10	84.20 $\pm$ 2.49
Best-of-Depth	0.00 $\pm$ 0.00	0.20 $\pm$ 0.17	33.03 $\pm$ 2.20	97.13 $\pm$ 0.55	99.00 $\pm$ 0.10	99.03 $\pm$ 0.12

Appendix

Table 6: N-Queens Success@Depth (%) on 1,000 multi-solution test problems. Models are trained at

D_{\mathrm{train}}=16

Method	@1	@4	@8	@16	@32	@64
Last-Depth	0.05 $\pm$ 0.01	2.01 $\pm$ 0.10	8.55 $\pm$ 0.77	15.79 $\pm$ 0.88	18.30 $\pm$ 0.88	18.62 $\pm$ 0.76
Mean-Depth	2.67 $\pm$ 0.09	11.99 $\pm$ 0.25	14.11 $\pm$ 0.05	14.84 $\pm$ 0.03	15.09 $\pm$ 0.07	15.15 $\pm$ 0.07
Random-Depth	2.38 $\pm$ 0.20	11.51 $\pm$ 0.15	13.87 $\pm$ 0.29	14.42 $\pm$ 0.25	14.68 $\pm$ 0.15	14.75 $\pm$ 0.16
Soft-Depth	0.88 $\pm$ 0.25	15.17 $\pm$ 0.59	24.36 $\pm$ 0.66	36.64 $\pm$ 0.53	42.18 $\pm$ 0.77	44.39 $\pm$ 0.80
Best-of-Depth	0.71 $\pm$ 0.42	14.18 $\pm$ 0.53	30.76 $\pm$ 0.70	54.95 $\pm$ 0.80	66.79 $\pm$ 1.89	72.16 $\pm$ 3.62

Appendix

Table 7: Graph-coloring Coverage@Depth (%) on 512 test graphs. Models are trained at

D_{\mathrm{train}}=16

Method	@1	@4	@8	@16	@32	@64
Last-Depth	0.59 $\pm$ 0.34	20.70 $\pm$ 1.47	50.98 $\pm$ 2.92	70.31 $\pm$ 1.86	73.76 $\pm$ 1.39	74.09 $\pm$ 1.30
Mean-Depth	23.37 $\pm$ 0.96	61.72 $\pm$ 1.22	67.38 $\pm$ 0.20	68.42 $\pm$ 0.11	68.82 $\pm$ 0.11	68.88 $\pm$ 0.11
Random-Depth	21.88 $\pm$ 1.28	60.16 $\pm$ 0.52	66.28 $\pm$ 0.63	67.12 $\pm$ 0.30	67.32 $\pm$ 0.11	67.38 $\pm$ 0.00
Soft-Depth	10.74 $\pm$ 1.79	79.88 $\pm$ 2.34	94.60 $\pm$ 1.58	97.98 $\pm$ 0.11	98.44 $\pm$ 0.20	98.63 $\pm$ 0.20
Best-of-Depth	9.18 $\pm$ 4.04	74.48 $\pm$ 1.47	94.21 $\pm$ 0.74	99.67 $\pm$ 0.30	99.87 $\pm$ 0.23	99.93 $\pm$ 0.11

Appendix

Table 8: Graph-coloring Success@Depth (%) on 512 test graphs. Models are trained at

D_{\mathrm{train}}=16

Depth $D$	Trajectories $N$	Last-Depth	Best-of-Depth
1	64	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00
2	32	0.00 $\pm$ 0.00	0.00 $\pm$ 0.00
4	16	0.00 $\pm$ 0.00	0.50 $\pm$ 0.30
8	8	0.00 $\pm$ 0.00	31.78 $\pm$ 1.81
16	4	48.49 $\pm$ 2.06	94.13 $\pm$ 0.56
32	2	56.25 $\pm$ 2.40	93.45 $\pm$ 0.92

Appendix

Table 9: N-Queens Coverage@Depth (%) at

ND=64

, pooling all

N

trajectories through depth

D

Depth $D$	Trajectories $N$	Last-Depth	Best-of-Depth
1	64	2.25 $\pm$ 1.25	33.85 $\pm$ 10.34
2	32	2.26 $\pm$ 1.35	62.86 $\pm$ 2.63
4	16	20.61 $\pm$ 1.82	77.28 $\pm$ 0.67
8	8	36.39 $\pm$ 0.65	82.80 $\pm$ 0.20
16	4	37.50 $\pm$ 0.51	85.86 $\pm$ 0.49
32	2	28.54 $\pm$ 0.52	82.67 $\pm$ 1.44

Appendix

Table 10: Graph-coloring Coverage@Depth (%) at

ND=64

, pooling all

N

trajectories through depth

D

Training depth	@4	@8	@16	@32	@64
4	6.48 $\pm$ 3.88	8.11 $\pm$ 3.96	8.21 $\pm$ 4.01	8.21 $\pm$ 4.01	8.21 $\pm$ 4.01
8	6.61 $\pm$ 2.73	60.21 $\pm$ 11.72	75.15 $\pm$ 11.99	75.66 $\pm$ 11.90	75.70 $\pm$ 11.95
16	0.10 $\pm$ 0.09	17.11 $\pm$ 1.25	85.92 $\pm$ 0.93	90.97 $\pm$ 0.98	91.03 $\pm$ 1.05

Appendix

Table 11: N-Queens BoD Coverage@Depth (%) across training and test depths.

Training depth	@4	@8	@16	@32	@64
4	21.09 $\pm$ 0.50	27.64 $\pm$ 0.83	31.53 $\pm$ 0.54	33.32 $\pm$ 0.64	34.13 $\pm$ 0.69
8	17.02 $\pm$ 0.55	34.69 $\pm$ 0.92	44.86 $\pm$ 0.90	50.38 $\pm$ 1.10	53.11 $\pm$ 1.57
16	14.18 $\pm$ 0.53	30.76 $\pm$ 0.70	54.95 $\pm$ 0.80	66.79 $\pm$ 1.89	72.16 $\pm$ 3.62

Appendix

Table 12: Graph-coloring BoD Coverage@Depth (%) across training and test depths.

Figure 5: Success across training and evaluation depths. Success@Depth for BoD trained at depths 4, 8, and 16, using the same checkpoints and evaluation problems as Figure 3 (c, d). Open markers indicate evaluation at the training depth; bands indicate

\pm 1

SD.

N-Queens
Training depth	@4	@8	@16	@32	@64
4	12.67 $\pm$ 7.31	15.90 $\pm$ 7.41	16.07 $\pm$ 7.51	16.07 $\pm$ 7.51	16.07 $\pm$ 7.51
8	13.73 $\pm$ 5.78	83.27 $\pm$ 8.42	94.40 $\pm$ 4.54	94.77 $\pm$ 4.37	94.80 $\pm$ 4.39
16	0.20 $\pm$ 0.17	33.03 $\pm$ 2.20	97.13 $\pm$ 0.55	99.00 $\pm$ 0.10	99.03 $\pm$ 0.12
Graph Coloring
4	88.09 $\pm$ 0.85	93.29 $\pm$ 0.92	94.86 $\pm$ 1.71	95.51 $\pm$ 1.28	95.57 $\pm$ 1.30

Appendix

Table 13: BoD Success@Depth (%) for Figure 5 .

Figure 6: Sensitivity to the Soft-Depth temperature on graph coloring. Test success and coverage use all greedy readouts through depth 64. The curve varies

\tau=0.1,0.2,\ldots,1.0

; horizontal lines show BoD and Mean-Depth under the same configuration. Bands indicate

\pm 1

SD.

	N-Queens		Graph Coloring
Method	Success	Coverage	Success	Coverage
Transformer (8 layers)	9.8 $\,\pm\,$ 1.3	4.8 $\,\pm\,$ 0.6	39.6 $\,\pm\,$ 1.8	5.2 $\,\pm\,$ 0.5
Transformer (32 layers)	10.9 $\,\pm\,$ 1.1	5.4 $\,\pm\,$ 0.6	44.5 $\,\pm\,$ 1.1	6.3 $\,\pm\,$ 0.3
Autoregressive Transformer	57.4 $\,\pm\,$ 3.0	27.4 $\,\pm\,$ 1.4	67.9 $\,\pm\,$ 1.5	8.2 $\,\pm\,$ 0.4
Masked Diffusion	34.2 $\,\pm\,$ 4.1	16.3 $\,\pm\,$ 2.0	76.1 $\,\pm\,$ 1.5	10.3 $\,\pm\,$ 0.3
HRM	98.6 $\,\pm\,$ 0.5	48.0 $\,\pm\,$ 0.6	67.4 $\,\pm\,$ 0.9	13.7 $\,\pm\,$ 0.2

Appendix

Table 14: Baseline comparison on N-Queens and graph coloring. Success and Coverage (%) summarize each method’s candidate set. Ouro uses an instance-level adaptation of its training objective. Compute and candidate counts differ across architectures; Appendix C.6 gives the protocol.

Training objective	2	16
Last-Depth	15.6	1.6
Mean-Depth	0.0	0.0
Random-Depth	18.8	1.6
Soft-Depth ( $\tau=0.1$ )	0.0	0.0
Soft-Depth ( $\tau=0.25$ )	0.0	0.0
Soft-Depth ( $\tau=1.0$ )	0.0	0.0

Appendix

Table 15: Numerical illustration of Proposition 3 . Greedy Coverage@Depth at

D=6

(%), averaged over 64 independent SGD runs on the same matching problem. Columns count training updates. All objectives use the same teacher stream within each run and the same initialization and step sizes.

Deep Supervision

See all 8 papers

Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks
Oct 16, 2026Magnus Sesodia, Juil Sock, Philip Torr +1Recurrent Transformers Algorithmic Reasoning
University of Oxford · BBC R&D
Beyond the Training Horizon: Mechanisms and Limits of Length Generalization in Looped Transformers
Sep 27, 2026Jia Liang, Xi Jin, Liangming PanLength Generalization Recurrent Transformers
Institute for Computational and Mathematical Engineering, Stanford University · MOE Key Laboratory of Computational Linguistics, School of Computer Science, Peking University
Looped Transformers as Optimizers
Sep 29, 2026Yulong Huang, Chen Jiang, Zhanpeng Zhou +5Fast Weights Recurrent Transformers
HKUST (GZ) · StepFun · SJTU +2

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

Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks

Magnus Sesodia, Juil Sock, Philip Torr, Christian Schroeder de Witt

University of Oxford · BBC R&D

Recurrent Transformers Algorithmic Reasoning

Open paper AskPDF arXiv

Solving algorithmic tasks with many sequential steps typically requires scaling up transformer depth or chain-of-thought length, incurring unfavorable parameter and inference-compute costs, respectively. Looped transformers, which reuse one block of layers across many iterations, are poised to address this, since they scale effective depth without adding parameters or tokens. Yet, it remains unclear what each iteration is actually computing, or whether the loop merely simulates a deeper untied stack. Here, we study directed permutation traversal, in which a model follows

H

hops from a starting node and reports the node it reaches, so every intermediate hop has a known ground truth. We find that a minimum effective depth is necessary to solve the task, beyond which the degree of weight tying becomes the determining factor. For instance, a looped transformer with a four-layer block iterated eight times solves the task in every seed, while an identically-trained untied transformer of the same effective depth never does, succeeding only with a lower learning rate or supervision of its intermediate layers. Probing and intervening on the residual stream, we find that the looped model advances through the intermediate nodes gradually, forming a staircase that adapts to its iteration budget, whereas the untied model, at that lower rate, exposes them only in its final layers. Taken together, our results recast looped transformers as iterative program learners, which allocate their iterations to successive steps of the algorithm and solve tasks that identically-trained untied models of the same effective depth do not. Code is available at https://github.com/magnus-sesodia/insidetheloop.

Abstract

Deep Supervision

Explore similar work

Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks

Beyond the Training Horizon: Mechanisms and Limits of Length Generalization in Looped Transformers

Looped Transformers as Optimizers