Best-of-Depth: Learning Exploratory Dynamics in Looped Transformers
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
| 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 |
| Last-Depth | Last-Depth + BoN | Best-of-Depth | ||||
|---|---|---|---|---|---|---|
| Training solutions | Seen Coverage | Unseen Coverage | Seen Coverage | Unseen Coverage | Seen Coverage | Unseen Coverage |
| 1 | ||||||
| 2 | ||||||
| 3 | ||||||
| 4 | ||||||
| All | N/A | N/A | 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 | 16 |
| Optimizer | AdamAtan2 |
| Learning rate / batch size | / 768 |
| 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 |
| Method | @1 | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|---|
| Last-Depth | 0.00 0.00 | 0.00 0.00 | 0.00 0.00 | 22.27 1.27 | 43.29 3.12 | 45.75 3.93 |
| Mean-Depth | 0.00 0.00 | 2.51 0.18 | 20.42 1.20 | 36.52 2.21 | 44.52 1.66 | 49.94 1.05 |
| Random-Depth | 0.00 0.00 | 2.46 0.11 | 19.93 0.53 | 35.42 0.91 | 44.15 1.24 | 49.63 1.34 |
| Soft-Depth | 0.00 0.00 | 2.63 0.37 | 21.53 0.82 | 38.99 0.79 | 48.26 0.83 | 53.01 1.47 |
| Best-of-Depth | 0.00 0.00 | 0.10 0.09 | 17.11 1.25 | 85.92 0.93 | 90.97 0.98 | 91.03 1.05 |
| Method | @1 | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|---|
| Last-Depth | 0.00 0.00 | 0.00 0.00 | 0.00 0.00 | 44.80 2.52 | 78.23 4.48 | 81.13 4.94 |
| Mean-Depth | 0.00 0.00 | 5.10 0.40 | 40.37 1.83 | 65.97 3.25 | 75.37 2.50 | 81.97 2.46 |
| Random-Depth | 0.00 0.00 | 5.00 0.20 | 39.43 1.34 | 64.50 1.61 | 75.27 1.72 | 82.30 2.35 |
| Soft-Depth | 0.00 0.00 | 5.30 0.70 | 42.37 1.11 | 68.67 1.69 | 79.13 2.10 | 84.20 2.49 |
| Best-of-Depth | 0.00 0.00 | 0.20 0.17 | 33.03 2.20 | 97.13 0.55 | 99.00 0.10 | 99.03 0.12 |
| Method | @1 | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|---|
| Last-Depth | 0.05 0.01 | 2.01 0.10 | 8.55 0.77 | 15.79 0.88 | 18.30 0.88 | 18.62 0.76 |
| Mean-Depth | 2.67 0.09 | 11.99 0.25 | 14.11 0.05 | 14.84 0.03 | 15.09 0.07 | 15.15 0.07 |
| Random-Depth | 2.38 0.20 | 11.51 0.15 | 13.87 0.29 | 14.42 0.25 | 14.68 0.15 | 14.75 0.16 |
| Soft-Depth | 0.88 0.25 | 15.17 0.59 | 24.36 0.66 | 36.64 0.53 | 42.18 0.77 | 44.39 0.80 |
| Best-of-Depth | 0.71 0.42 | 14.18 0.53 | 30.76 0.70 | 54.95 0.80 | 66.79 1.89 | 72.16 3.62 |
| Method | @1 | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|---|
| Last-Depth | 0.59 0.34 | 20.70 1.47 | 50.98 2.92 | 70.31 1.86 | 73.76 1.39 | 74.09 1.30 |
| Mean-Depth | 23.37 0.96 | 61.72 1.22 | 67.38 0.20 | 68.42 0.11 | 68.82 0.11 | 68.88 0.11 |
| Random-Depth | 21.88 1.28 | 60.16 0.52 | 66.28 0.63 | 67.12 0.30 | 67.32 0.11 | 67.38 0.00 |
| Soft-Depth | 10.74 1.79 | 79.88 2.34 | 94.60 1.58 | 97.98 0.11 | 98.44 0.20 | 98.63 0.20 |
| Best-of-Depth | 9.18 4.04 | 74.48 1.47 | 94.21 0.74 | 99.67 0.30 | 99.87 0.23 | 99.93 0.11 |
| Depth | Trajectories | Last-Depth | Best-of-Depth |
|---|---|---|---|
| 1 | 64 | 0.00 0.00 | 0.00 0.00 |
| 2 | 32 | 0.00 0.00 | 0.00 0.00 |
| 4 | 16 | 0.00 0.00 | 0.50 0.30 |
| 8 | 8 | 0.00 0.00 | 31.78 1.81 |
| 16 | 4 | 48.49 2.06 | 94.13 0.56 |
| 32 | 2 | 56.25 2.40 | 93.45 0.92 |
| Depth | Trajectories | Last-Depth | Best-of-Depth |
|---|---|---|---|
| 1 | 64 | 2.25 1.25 | 33.85 10.34 |
| 2 | 32 | 2.26 1.35 | 62.86 2.63 |
| 4 | 16 | 20.61 1.82 | 77.28 0.67 |
| 8 | 8 | 36.39 0.65 | 82.80 0.20 |
| 16 | 4 | 37.50 0.51 | 85.86 0.49 |
| 32 | 2 | 28.54 0.52 | 82.67 1.44 |
| Training depth | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|
| 4 | 6.48 3.88 | 8.11 3.96 | 8.21 4.01 | 8.21 4.01 | 8.21 4.01 |
| 8 | 6.61 2.73 | 60.21 11.72 | 75.15 11.99 | 75.66 11.90 | 75.70 11.95 |
| 16 | 0.10 0.09 | 17.11 1.25 | 85.92 0.93 | 90.97 0.98 | 91.03 1.05 |
| Training depth | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|
| 4 | 21.09 0.50 | 27.64 0.83 | 31.53 0.54 | 33.32 0.64 | 34.13 0.69 |
| 8 | 17.02 0.55 | 34.69 0.92 | 44.86 0.90 | 50.38 1.10 | 53.11 1.57 |
| 16 | 14.18 0.53 | 30.76 0.70 | 54.95 0.80 | 66.79 1.89 | 72.16 3.62 |
| Training depth | @4 | @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|
| N-Queens | |||||
| 4 | 12.67 7.31 | 15.90 7.41 | 16.07 7.51 | 16.07 7.51 | 16.07 7.51 |
| 8 | 13.73 5.78 | 83.27 8.42 | 94.40 4.54 | 94.77 4.37 | 94.80 4.39 |
| 16 | 0.20 0.17 | 33.03 2.20 | 97.13 0.55 | 99.00 0.10 | 99.03 0.12 |
| Graph Coloring | |||||
| 4 | 88.09 0.85 | 93.29 0.92 | 94.86 1.71 | 95.51 1.28 | 95.57 1.30 |
| N-Queens | Graph Coloring | |||
| Method | Success | Coverage | Success | Coverage |
| Transformer (8 layers) | 9.8 1.3 | 4.8 0.6 | 39.6 1.8 | 5.2 0.5 |
| Transformer (32 layers) | 10.9 1.1 | 5.4 0.6 | 44.5 1.1 | 6.3 0.3 |
| Autoregressive Transformer | 57.4 3.0 | 27.4 1.4 | 67.9 1.5 | 8.2 0.4 |
| Masked Diffusion | 34.2 4.1 | 16.3 2.0 | 76.1 1.5 | 10.3 0.3 |
| HRM | 98.6 0.5 | 48.0 0.6 | 67.4 0.9 | 13.7 0.2 |
| Training objective | 0 | 2 | 16 | 256 | 4,096 |
|---|---|---|---|---|---|
| Last-Depth | 0.0 | 15.6 | 1.6 | 0.0 | 0.0 |
| Mean-Depth | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| Random-Depth | 0.0 | 18.8 | 1.6 | 0.0 | 0.0 |
| Soft-Depth ( ) | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| Soft-Depth ( ) | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| Soft-Depth ( ) | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |