PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Depth in Looped Transformers: Diagnosing Learned Halting Gates and Trajectory Readouts

Andrei Cristian Popescu, Haitz Sáez de Ocáriz Borde, Pietro Liò

arXiv:2607.20519v1cs.LG

TL;DR

Adaptive depth is hard to diagnose because learned halting can both select exits and shape the recurrent states being selected. The paper separates trajectory formation from exit readout across controlled tasks and Ouro checkpoints, finding that simple readouts on well-shaped trajectories can match or outperform learned gates.

  • Problem

    Learned halting objectives use the exit distribution both to choose inference-time stopping depths and to weight training losses, entangling readout quality with trajectory formation.

    Source: §1 · p. 2

  • Method

    The study separates trajectory formation from exit selection by comparing fixed-prior and learned-gate training, multiple post-hoc readouts, forced exits, and frozen-trajectory gate fitting across controlled tasks and Ouro checkpoints.

    Source: §1 · p. 2, §5.4 · p. 10

  • Results

    Across MANO, parity, and Ouro evaluations, simple readouts on well-shaped trajectories match or outperform learned gates; geometric-prior MANO with logit margin reaches 95%–99% accuracy at 1.4–1.6 average loops.

    Source: §5.2 · p. 9, §6 · p. 12

  • Takeaways & Limitations

    Adaptive depth in looped Transformers is a joint trajectory-formation and exit-readout problem, and gate failures can reflect the trajectory induced by joint training rather than limited gate expressivity.

    Source: §5.4 · p. 11, §6 · p. 12

Abstract

from arXiv · show

Looped Transformers increase test-time computation by repeatedly applying a shared recurrent block. Learned halting objectives in looped Transformers typically use a single exit distribution both as the inference-time stopping rule and as the training-time weighting of per-depth losses. This entangles exit selection with trajectory formation: the gate not only chooses which recurrent state to use, but also determines how strongly each intermediate state is supervised. Consequently, poor adaptive-compute performance can arise from the readout, the induced trajectory, or their interaction. We study adaptive depth in looped Transformers through this trajectory--readout lens, across controlled synthetic tasks (modular arithmetic and binary parity) and large-scale Ouro-1.4B and 2.6B checkpoints. We find that fixed-prior depth supervision, which shapes the trajectory without an input-dependent halting policy, produces difficulty-aware trajectories whose intermediate states expose useful stopping signals, and that simple post-hoc confidence readouts often match or outperform learned linear and MLP gates. Fitting gates on frozen trajectories localizes the failure: it appears to stem mainly from the trajectory induced by joint gate training rather than from limited gate expressivity. The same pattern is present in Ouro evaluations, where pretrained ponder gates are competitive but not uniformly Pareto-optimal, and measured latency confirms that the resulting reductions in average exit depth translate into practical inference-time savings. Our systematic diagnostic evaluation reframes adaptive depth in looped Transformers as a joint problem of trajectory formation and exit readout, rather than gate learning alone, highlighting a distinction that prior learned-halting work has often left implicit.

1 Introduction

Learned halting couples inference-time exit selection with training-time depth supervision, making adaptive-compute failures difficult to diagnose. The paper separates trajectory formation from exit readout and finds that fixed-prior supervision can produce useful trajectories for simple post-hoc exits.

1 Introduction

  • Learned exit distributions both select inference-time stopping depths and weight training losses across recurrent depths.This coupling means poor tradeoffs may arise from the stopping rule, the induced trajectory, or their interaction.Source: §1 · p. 2
  • Fixed-prior depth supervision can produce difficulty-aware trajectories whose intermediate states support strong post-hoc early exits.Simple confidence readouts match or outperform learned gates in compute–accuracy tradeoffs in several settings.Source: §1 · p. 2
  • The study evaluates training objectives, gate parameterizations, post-hoc readouts, forced exits, model scales, and latency.Its evaluations cover controlled MANO and parity tasks, plus Ouro-1.4B and Ouro-2.6B checkpoints.Source: §1 · p. 2

2 Related Work

Prior work uses looped Transformers to add recurrent computation and adaptive methods to allocate computation dynamically. This paper asks how training objectives shape recurrent trajectories and how readouts turn them into compute–quality tradeoffs.

2.1 Looped Transformers

  • Looped Transformers repeat shared layers, decoupling effective computational depth from parameter count.Prior examples include the Universal Transformer, Deep Equilibrium Models, and ALBERT.Source: §2.1 · p. 2
  • Recent work treats recurrence as latent computation, refining hidden representations at each step rather than extending the input context.The paper focuses on recurrent-depth computation and its exit readouts.Source: §2.1 · pp. 2–3

2.2 Adaptive Computation

  • Adaptive-computation methods allocate different amounts of computation to different tokens using learned halting policies, routers, or other rules.Looped Transformers provide intermediate states from which computation can exit early or continue refining predictions.Source: §2.2 · p. 3
  • Confidence, entropy, and prediction-stability signals have been used as stopping criteria in early-exit models and can be measured across recurrent iterations.Related signals include changes in hidden states, logits, and predictive distributions.Source: §2.2 · p. 3
  • This study asks which component explains adaptive-depth success or failure, separating trajectory formation from exit readout.A halting mechanism also determines which recurrent depths receive supervision during training, shaping the trajectory used for later exits.Source: §2.2 · p. 3

3 Formalizing Adaptive Depth as Trajectory and Readout

The section formalizes adaptive depth as a recurrent trajectory plus a rule for selecting an exit. It contrasts learned gates, which shape trajectories while selecting exits, with fixed-prior training and post-hoc readouts that separate those roles.

3.1 Looped Transformer Trajectories

  • A looped Transformer trajectory is generated by repeatedly applying a shared block to successive hidden states.The model produces hidden states across recurrent depths, which can be unembedded for predictions.Source: §3.1 · p. 3

3.2 Gating Mechanism for Adaptive Computation

  • A learned gate predicts conditional halt probabilities that induce a distribution over exit depths.The remaining probability mass is assigned to the final depth, making the exit probabilities a valid distribution.Source: §3.2 · p. 4
  • At inference, the learned exit distribution can stop computation when its cumulative probability exceeds threshold α.The resulting early-exit depth is defined by a cumulative-threshold rule.Source: §3.2 · p. 4
  • During training, the learned exit distribution weights per-depth losses, so the gate changes both the readout and the trajectory it later reads.The trajectory depends on the halting objective and gate parameters because the learned weights determine each depth’s contribution to training loss.Source: §3.2 · p. 4

3.3 Fixed-Prior Depth Supervision

  • Fixed-prior supervision uses input-independent depth weights rather than learning an input-dependent stopping policy.This controls trajectory shaping without simultaneously learning an adaptive readout.Source: §3.3 · p. 4
  • The study considers uniform and truncated geometric priors, with the latter placing greater loss weight on earlier depths.The geometric prior assigns remaining probability mass to the final recurrent depth.Source: §3.3 · p. 4
  • Fixed-prior models require a separate inference-time readout, allowing tests of whether their intermediate states contain useful stopping signals.Effective post-hoc exits from these trajectories would show that adaptive exit need not rely on a learned gate.Source: §3.3 · p. 4

3.4 Trajectory Readouts

  • A readout maps a trained trajectory to an exit depth; post-hoc confidence or convergence rules use thresholded signals without shaping training.Confidence signals include entropy, top-1 confidence, and logit margin; convergence signals include predictive KL, logit change, and hidden-state movement.Source: §3.4 · p. 5

4 Methodology

The experiments separate trajectory quality from exit-readout quality across controlled MANO models and pretrained Ouro checkpoints. They compare training objectives and gate types, then assess adaptive exits using validation-selected thresholds and held-out compute–quality outcomes.

4.2 Training Objectives and Gate Parameterizations

  • The controlled study trains fixed-prior and learned-gate models, then evaluates multiple readouts on their resulting trajectories.Fixed-prior models use uniform or truncated geometric priors; learned gates use linear projections or two-layer MLPs.Source: §4 · p. 5, §4.2 · pp. 5–6

4.1 Controlled MANO Experiments

  • MANO tests synthetic modular arithmetic over integers modulo 23 using prefix expressions of up to 10 operations, without an input length token.Examples are grouped by operation count to analyze difficulty; the shared-backbone models use maximum recurrent depth T = 6.Source: §4.1 · p. 5

4.3 Trajectory Readouts and Threshold Selection

  • Readouts include native learned-gate cumulative-threshold rules and post-hoc confidence or convergence signals swept over validation-calibrated thresholds.Confidence signals include entropy, top-1 confidence, and logit margin; convergence signals include predictive KL, logit change, and hidden-state displacement.Source: §4.3 · p. 6

4.4 Evaluation Protocol and Metrics

  • The evaluation separates trajectory quality, measured with forced exits at every recurrent depth, from adaptive stopping measured by threshold-swept compute–quality curves.Computation is mean selected depth; MANO frontiers use D@X, while Ouro reports held-out accuracy at validation-selected operating points.Source: §4.4 · p. 6
  • Thresholds are selected on validation data and evaluated on held-out data, with MANO using D@X and Ouro reporting held-out accuracy at selected operating points.MANO uses separate seeded train, validation, and test sets; the common compute measure is average exit depth.Source: §4.1 · p. 5, §4.4 · p. 6

4.5 Large-Scale Looped Language-Model Evaluation

  • Ouro-1.4B and Ouro-2.6B are evaluated as released, comparing pretrained ponder gates with post-hoc readouts on their recurrent trajectories.The models use up to four recurrent steps; six standard benchmarks are evaluated under benchmark-specific prompting and held-out protocols.Source: §4.5 · p. 6

5 Results

Across MANO and Ouro, the results separate the quality of intermediate recurrent states from the readout that chooses exits. Fixed-prior trajectories support strong early exits, while alternative readouts remain competitive on pretrained Ouro and reduced exit depth corresponds to latency savings.

5.1 Fixed-prior training produces difficulty-aware trajectories

  • Fixed-prior supervision produces difficulty-aware trajectories, with confidence readouts reaching 99% accuracy at roughly 1.5 average loops under the geometric prior.The same structure persists beyond the training range and appears on binary parity, where diagnostics remain ordered by bit length.Source: §5.1 · p. 7, Figure 2 · p. 7

5.2 Fixed-prior readouts can match or outperform learned gates

  • Across three seeds, geometric fixed-prior trajectories with logit margin reach 95%–99% accuracy at 1.4–1.6 average loops, while learned gates require greater depth.The table comparison uses one fixed post-hoc readout for fixed-prior models but selects the highest-performing β for learned gates.Source: Figure 4 · p. 8, §5.2 · pp. 8–9

5.3 The training objective shapes trajectories through depth-weighted losses

  • Forced exits show that fixed-prior training yields useful early intermediate predictions, while learned exit allocation alone cannot ensure accurate states.The same learned exit distribution both weights depth-specific losses and defines inference-time stopping; fixed-prior supervision separates these roles.Source: §5.3 · pp. 9–10, Figure 5 · p. 9

5.4 Post-hoc fitting of a learned gate over a pretrained trajectory

  • On frozen jointly learned trajectories, new gates do not improve the frontier, pointing to the induced trajectory rather than gate expressivity as the main constraint.On fixed-prior trajectories, post-hoc gates recover strong early exits, while simple readouts remain competitive.Source: Figure 6 · p. 10, §5.4 · pp. 10–11

5.5 From controlled tasks to large-scale checkpoints

  • Across Ouro-1.4B and Ouro-2.6B, selected post-hoc readouts beat the pretrained gate in held-out accuracy in 5 of 12 pairs, use fewer loops in 7, and improve both in 3.Several accuracy differences are small and reported on single runs; the best readout varies by model size and benchmark.Source: Table 3 · p. 11, §5.5 · p. 11, Figure 7 · p. 11

5.6 Average exit depth corresponds to latency savings

  • On MANO, top-1 confidence and logit margin preserve 100.0% accuracy while reducing average depth from 6.0 to 1.26 loops, with 1.36× and 1.29× speedups.Across Ouro benchmarks, the pretrained gate's estimated speedups range from 1.24× to 1.58×; post-hoc readouts yield similar reductions at comparable average loops.Source: §5.5 · p. 12, §5.6 · p. 12, Table 4 · p. 12

6 Conclusion

Adaptive depth is a trajectory–readout problem: training shapes which recurrent states are useful, while a separate readout selects when to stop. The conclusion calls for future work on alternative pretraining objectives and looped Transformer architectures.

6 Conclusion

  • The study frames adaptive depth as a trajectory–readout problem rather than solely a gate-learning problem.It identifies the training objective and the stopping readout as distinct influences on adaptive-compute performance.Source: §6 · p. 12
  • Across MANO, parity, and Ouro-1.4B and 2.6B, simple post-hoc readouts on well-shaped trajectories match or outperform learned gates.The conclusion attributes gate failures to the induced trajectory rather than gate expressivity.Source: §6 · p. 12

A Model Background

  • The appendix provides background on the looped Transformer formulation and the Ouro architecture and training procedure.Source: A Model Background · p. 15

A.1 Looped Transformer Architecture

Looped Transformers repeatedly apply a shared block to form recurrent states, which can be read at different stopping depths. The paper uses this framework to compare halting and readout approaches across synthetic tasks and pretrained Ouro models.

A.1 Looped Transformer Architecture

  • A looped Transformer iterates a shared block over embedded inputs, reusing parameters across recurrent depths.The recurrent states yield per-depth logits and losses, allowing predictions at intermediate loops.Source: §A.1 · p. 15
  • Adaptive depth selects a stopping depth and returns that recurrent state’s prediction; fixed-depth, learned-gate, and post-hoc methods differ in how they select it.Post-hoc readouts use trajectory signals such as entropy, maximum probability, logit margin, predictive change, or hidden-state displacement.Source: §A.1 · p. 15
  • The paper’s synthetic experiments train from scratch at six loops, while Ouro evaluations keep pretrained trajectories fixed and vary the exit readout.This common setup compares learned gates, fixed-depth exits, and post-hoc readouts within one framework.Source: §A.1 · p. 16

A.2 The Ouro Model Family

  • Ouro applies a shared stack of decoder layers recurrently; its 1.4B and 2.6B checkpoints were pretrained for four recurrent steps on 7.7T tokens.Stage I jointly trains the gate and backbone with an entropy-regularized objective; Stage II freezes the backbone and trains the gate using marginal utility.Source: §A.2 · p. 16
  • Ouro’s Stage-II loss improvement becomes a soft continue/exit label and binary-cross-entropy target for gate training.This objective directly inspires the paper’s marginal-utility post-hoc gate on frozen trajectories.Source: §A.2 · p. 16
  • Ouro’s authors report that the 1.4B and 2.6B checkpoints match dense Transformer baselines 2–3× larger across general, math, and coding benchmarks.The paper uses the released base checkpoints without modifying or retraining them.Source: §A.2 · p. 16

B.1 MANO Task

  • MANO uses modular arithmetic expressions modulo 23, with up to 10 operations and no expression-length input token.Operation count serves as the difficulty variable, avoiding a length-token shortcut for compute allocation.Source: §B.1 · p. 17

B.2 MANO Architecture

  • The MANO backbone has four Transformer layers, hidden size 512, and maximum recurrent depth six; the appendix also summarizes experimental protocols.The backbone ties input embeddings and language-model head weights.Source: B Experimental Details · p. 16, §B.2 · p. 17

B.3 MANO Training

The appendix specifies MANO optimization and model variants, threshold selection, readout definitions, and Ouro benchmark protocols. These details define how the paper trains and evaluates its adaptive-depth comparisons.

B.3 MANO Training

  • MANO variants use AdamW with learning rate 5 × 10−4 and train for 40,000 steps with bf16 mixed precision.Training uses batch size 1024, gradient clipping at 1.0, and a warmup-cosine schedule.Source: §B.3 · p. 17

B.4 MANO Variants

  • Fixed-prior models use uniform or truncated geometric priors, while learned linear and MLP gates sweep specified KL regularization strengths.The geometric-prior sweep uses λ ∈{0.2, 0.3, 0.5, 0.7}.Source: §B.4 · p. 17

B.5 Threshold Selection

  • Post-hoc readouts use a 51-point validation-quantile threshold grid; learned gates use cumulative-probability thresholds from 0.01 to 1.Fixed-depth baselines require no threshold selection.Source: §B.5 · p. 17

B.6 Post-Hoc Readout Definitions

  • Confidence readouts exit when top-1 confidence or logit margin reaches a threshold, while convergence readouts exit when a change signal falls below one.Readouts default to maximum depth when their threshold is never reached.Source: §B.6 · p. 18

B.7 Ouro Models and Benchmarks

  • The Ouro evaluation uses six standard benchmarks, including MMLU, ARC-Challenge, ARC-Easy, OpenBookQA, CommonsenseQA, and HellaSwag.Prompting ranges from zero-shot OpenBookQA to 25-shot ARC-Challenge.Source: §B.7 · p. 18

B.8 Ouro Threshold Selection

  • Ouro thresholds are selected on validation data and evaluated on test data when available; otherwise, labeled validation data is split for selection and held-out evaluation.CommonsenseQA and HellaSwag use a 20% validation subset and an 80% held-out test subset.Source: §B.8 · p. 18

C Fixed-Prior Trajectory Analysis

Fixed-prior supervision shapes trajectories whose confidence signals reflect input difficulty, while overly aggressive early weighting can degrade later states. Across linear and MLP comparisons, simple readouts on fixed-prior trajectories retain a strong low-depth frontier.

C.1 Pareto Curves for the Geometric Prior over Multiple λ Values

  • Moderate geometric-prior early weighting, particularly λ = 0.3, yields useful early predictions, while aggressive weighting degrades later states and the accuracy ceiling.Uniform supervision produces useful early exits only at higher average depth.Source: §C.1 · p. 18, Figure 8 · p. 19

C.2 Uniform Prior Difficulty Convergence

  • Under uniform supervision, entropy, logit margin, and predictive KL show clearer operation-count difficulty ordering than hidden-state displacement.Harder expressions remain higher-entropy and lower-margin for more loops; Table 5 summarizes signed Spearman correlations across seeds.Source: §C.2 · p. 18, Table 5 · p. 19, Figure 9 · p. 20

C.3 Geometric Prior Difficulty Convergence

  • For geometric prior λ = 0.3, difficulty ordering is preserved across four signals, with entropy and logit margin at ρ = 0.98 and predictive KL at ρ = 0.87.Hidden-state displacement is ρ = 0.69 ± 0.34, roughly double the uniform-prior value.Source: §C.3 · p. 19, Table 6 · p. 20

C.4 Difficulty Extrapolation on MANO

  • Beyond the training range, harder expressions remain uncertain longer, develop logit margin more slowly, and tend to show larger hidden-state changes under both priors.The experiment is a trajectory diagnostic, not a length-generalization benchmark; the geometric trajectory separates operation counts more clearly than the uniform one.Source: §C.4 · pp. 20, 22, Figure 11 · p. 21, Figure 12 · p. 22

D Learned Gate Comparison

  • Across gate classes and prior settings, fixed-prior trajectories with simple confidence readouts provide a strong low-depth frontier, and added gate capacity narrows but does not close the gap.The full comparison covers learned linear and MLP gates.Source: D Learned Gate Comparison · p. 22

D.1 Linear Gate

  • Learned linear gates do not dominate fixed-prior confidence readouts; large β pulls gate exit mass toward the prior and can collapse it to near-fixed depth.With aggressive geometric priors λ = 0.5 or 0.7, both gate and trajectory saturate at a lower accuracy ceiling.Source: Figure 13 · p. 22, §D.1 · p. 23

D.2 MLP Gate

  • A two-layer MLP improves on the linear-gate baseline in several panels but does not displace the fixed-prior low-depth frontier.The comparison attributes the remaining limitation to the jointly induced trajectory rather than halt-head form.Source: §D.2 · p. 23, Figure 14 · p. 23

E Trajectory Diagnostics

Forced-exit accuracy reveals that fixed-prior supervision shapes more useful intermediate predictions than joint learned-gate training, while learned gates still allocate exits by difficulty.

E Trajectory Diagnostics

  • Forced exits measure per-depth trajectory quality without applying an adaptive stopping rule, while learned exit distributions show how gate models allocate exits by difficulty.Source: E Trajectory Diagnostics · pp. 23–24
  • Fixed-prior trajectories support strong early exits, especially on harder examples, while learned-gate trajectories often gain accuracy only at later loops.The contrast indicates that the gate objective changes which recurrent states become accurate, not only how the model selects among them.Source: E Trajectory Diagnostics · p. 24, Figure 15 · p. 24
  • Learned gates shift exit mass toward later depths for harder examples, but those later exits help only when the corresponding recurrent states are accurate.Forced-exit curves show that jointly trained trajectories do not always provide accurate states at those later depths.Source: E Trajectory Diagnostics · p. 24

F Parity Experiments

In a lightweight parity experiment, fixed-prior confidence readouts reach near-ceiling accuracy at low average depth, and longer inputs show difficulty-aware trajectories.

F Parity Experiments

  • The parity experiment uses binary sequences of lengths 1–40 as a lightweight robustness check and does not repeat the full MANO benchmark ablation.Models use six recurrent loops; learned gates use β = 0.1.Source: F Parity Experiments · pp. 24–25
  • Fixed-prior confidence readouts reach near-ceiling accuracy at low average depth across the three-seed parity diagnostic.The compute–accuracy curves show average depths comparable to, and often lower than, learned gates as settings approach saturation.Source: F Parity Experiments · pp. 25–26, Table 7 · p. 25
  • On parity, longer bit strings remain higher-entropy for more loops and develop logit margin more slowly than shorter strings.The trajectory diagnostics group examples by bit length, used as a simple difficulty variable.Source: F Parity Experiments · p. 26, Figure 18 · p. 26

G Additional Ouro Results

Across Ouro-1.4B and 2.6B, post-hoc confidence and convergence readouts track the pretrained ponder gate and sometimes match or improve its compute–accuracy tradeoff.

G Additional Ouro Results

  • Across both Ouro model sizes, post-hoc readouts match or improve on the pretrained ponder gate’s compute–accuracy tradeoff in several benchmarks.The pretrained gate remains competitive and close to the frontier, but is not uniformly Pareto-optimal; the best readout varies by benchmark and model size.Source: G Additional Ouro Results · pp. 27–28, Figure 19 · p. 27, Figure 20 · p. 27
  • Static-depth baselines are generally dominated by adaptive readouts, supporting input-dependent stopping rather than uniform depth reduction as the source of the gains.Source: G Additional Ouro Results · p. 28
  • The oracle earliest-correct marker shows the accuracy ceiling if exit always occurs at the first correct loop, while the gap to realized readouts marks remaining room to improve stopping rules.Source: G Additional Ouro Results · p. 28

H Frozen Trajectory Gate Fitting

Frozen-trajectory gate fitting separates readout capacity from trajectory quality: fitted gates work well on fixed-prior trajectories but do not improve jointly trained ones. Reasoning-tuned Ouro results likewise leave simple heuristics competitive.

H Frozen Trajectory Gate Fitting

  • Post-hoc linear and MLP gates on frozen fixed-prior trajectories recover strong early exits and near-final-loop accuracy at much lower average depth.Confidence readouts remain competitive, indicating that useful stopping signals are already exposed in the prediction trajectory.Source: H Frozen Trajectory Gate Fitting · p. 28, Figure 21 · p. 28
  • On trajectories shaped by joint learned-gate training, newly fitted gates do not improve the frontier and remain bounded by the underlying trajectory across gate classes.The PonderNet-style objective outperforms the marginal-utility objective, which exits later without improving the frontier.Source: Figure 22 · p. 29, Figure 23 · p. 29, H Frozen Trajectory Gate Fitting · p. 30

I Reasoning-Tuned Ouro

  • Across 12 reasoning-tuned Ouro settings, a simple post-hoc heuristic is best in 7, while the pretrained gate remains useful but not necessarily best.The evaluation uses Ouro-1.4B-Thinking and 2.6B-Thinking on the same six benchmarks as the main Ouro experiments.Source: I Reasoning-Tuned Ouro · p. 30, Table 8 · p. 30

Related papers

Loading 2607.20519v1…