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

Summary (Overview)

• Novel training objective: The paper introduces Best-of-Depth (BoD), which trains looped Transformers by minimizing the minimum prediction loss across supervised depths, treating the recurrent trajectory as a candidate set of solutions rather than supervising only the final output.
• Theoretical grounding: Under exact-match loss, BoD is proven to maximize solution coverage—the fraction of all valid solutions recovered along the trajectory—connecting the training objective directly to solution discovery.
• Strong empirical gains on combinatorial tasks: On N-Queens completion, BoD achieves 91.03% coverage vs. 45.75% for final-depth training; on graph coloring, it reaches 72.16% coverage vs. 44.39% for the best baseline (Soft-Depth), with success nearly saturating while coverage continues to grow.
• Discovery beyond training data: BoD recovers solutions withheld from teacher pools on familiar inputs and finds new valid solutions beyond the training depth, demonstrating genuine exploratory dynamics.
• Language-model validation: Fine-tuning the 3B-parameter Huginn looped LM with BoD yields higher pooled code-generation success (37.8% HumanEval, 44.2% MBPP) than final-depth training (34.8%, 39.2%).

---

Introduction and Theoretical Foundation

Background and Motivation

Looped Transformers extend inference-time computation by repeatedly applying a shared transformation to a latent state, increasing computational depth without adding parameters. While greater recurrent depth can improve reasoning, the learning objective determines how successive updates contribute to problem-solving.

The key insight: a single recurrent trajectory produces multiple predictions at successive depths, which can be treated as a candidate set—analogous to Best-of-N sampling but generated sequentially along one trajectory rather than in parallel.

Problem with Existing Objectives

• Last-Depth training (Eq. 2): Only supervises the final prediction, so intermediate states receive no direct learning signal for producing useful alternatives.
• Mean/Random-Depth: Supervising every depth against the same target forces all readouts to reproduce that target, discouraging diversity.
• Fixed-point convergence approaches promote stability but do not determine which alternatives a trajectory discovers.

Theoretical Foundation

The paper formalizes the connection between BoD and solution coverage. For exact-match loss, define:

$$\mathcal{L}_{\mathrm{BoD}}^{0-1} = \min_{d \in \mathcal{D}} \mathbf{1}[\hat{y}_d \neq y]$$

Proposition 1 (First-hit interpretation): Let $\tau_y = \min\{d \in \mathcal{D} : \hat{y}_d = y\}$, with $\tau_y = \infty$ if no readout matches $y$. Then:

$$\mathbb{E}[\mathcal{L}_{\mathrm{BoD}}^{0-1}] = \Pr(\tau_y > D)$$

BoD evaluates whether the target is reached anywhere within the depth budget.

Proposition 2 (Connection to solution coverage): Fix $x$ and $S_D(x) = \{\hat{y}_d(x) : d \in \mathcal{D}\}$ independently of target $Y \sim \text{Uniform}(\mathcal{Y}^\star(x))$. Then:

$$\mathbb{E}_Y[\mathcal{L}_{\mathrm{BoD}}^{0-1} \mid x] = 1 - \frac{|S_D(x) \cap \mathcal{Y}^\star(x)|}{|\mathcal{Y}^\star(x)|} = 1 - \text{Coverage@Depth}(D; x)$$

Minimizing BoD risk maximizes solution coverage; repeated candidates do not increase coverage.

Proposition 3 (Separation in solution discovery): There exists a multi-solution task where, under cross-entropy training with uniformly resampled valid targets, greedy coverage satisfies:

$$\lim_{n \to \infty} \text{Coverage@Depth}(D; \theta_A^{(n)}) = \begin{cases} 0, & A \in \{\text{Last}, \text{Mean}, \text{Random}, \text{Soft}\}, \\ 1, & A = \text{BoD}, \end{cases}$$

almost surely. BoD breaks symmetry by selecting one loss-minimizing depth, allowing different targets to reinforce the same latent solution through different recurrent phases.

---

Methodology

Looped Transformer Formulation

A looped Transformer encodes input $x$ and repeatedly applies a shared transformation:

$$h_0 = \mathrm{Enc}_\theta(x), \qquad h_d = F_\theta(h_{d-1}; x), \qquad p_d = \mathrm{Dec}_\theta(h_d), \quad d = 1, \ldots, D$$

One application of $F_\theta$ is a loop; $d$ denotes recurrent depth. A shared decoder reads prediction distribution $p_d$ from each state.

The Best-of-Depth Objective

At training depth $D = D_{\text{train}}$, with supervised depths $\mathcal{D} \subseteq \{1, \ldots, D\}$ (including $D$), BoD minimizes:

$$\mathcal{L}_{\mathrm{BoD}}(x, y) = \min_{d \in \mathcal{D}} L_d(x, y)$$

For autoregressive targets, the teacher-forced loss is:

$$L_d(x, y) = -\frac{1}{T} \sum_{t \in \mathcal{T}} \log p_d(y_t \mid x, y_{<t}), \qquad T = |\mathcal{T}|$$

BoD selects one depth for the entire target and trains on its loss, breaking exact ties at random.

Comparison Objectives

Evaluation Metrics

$$\text{Success@Depth}(D; x) = \mathbf{1}[S_D(x) \cap \mathcal{Y}^\star(x) \neq \varnothing]$$ $$\text{Coverage@Depth}(D; x) = \frac{|S_D(x) \cap \mathcal{Y}^\star(x)|}{|\mathcal{Y}^\star(x)|}$$

Coverage is the primary metric for controlled tasks, measuring the fraction of all valid solutions recovered along the trajectory.

---

Empirical Validation / Results

N-Queens Completion

Setup: 10×10 board with 5 fixed queens; model completes the board validly. Trained on 24,320 problems; evaluated on 1,000 multi-solution test problems with all valid completions enumerated.

Results: BoD achieves 91.03% coverage and 99.03% success at depth 64, versus 45.75% coverage for Last-Depth. Gains are largest near the training depth, with further discoveries beyond it.

Graph Coloring

Setup: 3-color assignment on connected 10-vertex graphs; 4,096 training graphs, 512 test graphs. Each graph has 2–128 enumerated solution families.

Results: BoD achieves 72.16% coverage at depth 64 and solves nearly all test graphs. Success nearly saturates at training depth while coverage continues increasing—additional depth mainly recovers alternatives for already-solved graphs. Soft-Depth achieves high success but only 44.39% coverage.

Code Generation (Huginn 3B)

Fine-tuned on 20,000 Python examples from CodeFeedback at depth 32; evaluated on HumanEval (164 problems) and MBPP (500 problems) with greedy decoding at fixed depths {4, 8, 12, 16, 24, 32}.

Table 1: Code-generation results (%)

BoD achieves the highest Success@32 on both benchmarks, though other objectives perform better at shallow depths.

Ablation Studies

Fixed recurrent-step budget ($B = ND = 64$): BoD peaks at $(N, D) = (4, 16)$ on both tasks, favoring combining depth with multiple trajectories.

Training depth extrapolation: All BoD settings show further discoveries beyond training depth, with diminishing gains. Deeper training reaches higher coverage by depth 64 but can lag at shallower depths.

Discovery beyond teacher pools (Table 2): On graph coloring with restricted teacher pools:

BoD has the highest unseen coverage in every restricted-pool condition, recovering solutions withheld from training targets on familiar inputs.

Trajectory Analysis

• BoD discovers new valid solutions near and just beyond training depth, where Last-Depth makes few discoveries.
• PCA projections show BoD moving among states decoding to different valid solutions, with greater relative state displacement near and beyond training depth.
• Loss-minimizing depths span all 16 supervised depths, providing best predictions for different input–target pairs.

---

Theoretical and Practical Implications

Theoretical Contributions

1. First-hit interpretation (Prop. 1): BoD directly optimizes the probability that a target is reached within the depth budget—an inference-aware objective.
2. Coverage equivalence (Prop. 2): Minimizing expected BoD loss under exact-match is equivalent to maximizing solution coverage, providing a principled training objective for candidate-set evaluation.
3. Separation result (Prop. 3): Demonstrates that BoD uniquely breaks symmetry in a cyclic recurrence, where other objectives (Last, Mean, Random, Soft) converge to symmetric predictors with invalid greedy outputs.

Practical Implications

• Training objectives as design choices: The paper shows that how we supervise recurrent computation fundamentally shapes what a trajectory discovers, not just its final accuracy.
• Candidate-set training without architectural changes: BoD requires no additional parameters, depth-specific labels, or prescribed solution ordering—making it broadly applicable.
• Complementarity with test-time compute: Fixed-budget results suggest combining recurrent depth with multiple trajectories (Best-of-N style) yields the best coverage, pointing toward hybrid inference strategies.
• Applicability to language models: BoD improves pooled code-generation success, demonstrating utility beyond small combinatorial tasks.

---

Conclusion

Best-of-Depth trains looped Transformers on the lowest-loss prediction along their trajectory, allowing different depths to match different valid targets without architectural changes. Key findings:

1. Improved coverage: On N-Queens (91.03% vs. 45.75%) and graph coloring (72.16% vs. 44.39%), BoD recovers substantially more distinct valid solutions along a single trajectory.
2. Discovery beyond training: BoD finds solutions withheld from teacher pools and discovers new valid outputs beyond training depth.
3. Exploratory dynamics: Trajectory analyses associate broader discovery with sustained latent movement and use of multiple intermediate depths.
4. Language-model gains: Higher pooled code-generation success when pooling programs from different fixed depths.

Limitations: Evidence covers small combinatorial tasks and one pretrained language model; discovery saturates with depth. Future directions include cross-depth sampling, branching, restarts, verification, adaptive stopping, and testing on larger tasks and additional models.

The paper's central thesis: training and evaluating recurrent trajectories as candidate sets of complementary valid solutions across depth is a promising paradigm for making additional inference-time computation genuinely useful for solution discovery.