Summary (Overview)

• Core contribution: Introduces fixed-first training, a new recipe for adaptive-depth looped transformers that trains at a fixed depth for most of the run and only varies depth during a short final "finishing stage." This reduces training compute by 30-45% while maintaining performance within 1% of uniform-depth training.
• Compute savings: At 200M parameters, using $K_0 = 2$ for 67% of training cuts modeled compute by 45.3% while keeping mean serving-depth NLL within 0.5% of uniform training. At 1B parameters, savings reach 31-37% in FLOPs and 32-39% in wall-clock time.
• Key insight: Fixed-depth training learns quickly but specializes near its training depth, while uniform training retains broader depth generalization but learns more slowly. Fixed-first combines the best of both: fast early learning plus a short adaptation phase.
• Diagnostic finding: A single deep update's advantage over a shallow update becomes statistically indistinguishable from zero after just 16 uniform-depth updates, explaining why a short finishing phase suffices.
• Practical significance: This makes adaptive-depth models substantially cheaper to train, potentially democratizing access to serving-depth-flexible models.

---

Introduction and Theoretical Foundation

Background: Looped Transformers and Adaptive Depth

A looped transformer reuses the same layer (with shared weights) multiple times per token. The paper uses the middle-loop architecture of Saunshi et al. [2025] and Geiping et al. [2025], where an untied prelude and coda surround a recurrent block repeated $K$ times. The key benefit: one model can serve at multiple depths—producing answers cheaply at small $K$ or more accurately at larger $K$ (the serving depth).

The Problem with Existing Adaptive-Depth Training

Training at a single depth produces a depth-specialized model. Prior adaptive methods therefore train across multiple depths:

• Uniform-depth training (Geiping et al. [2025], Hegazy et al. [2026]): samples $K$ uniformly from $\{1, \dots, 16\}$ per update
• Progressive curriculum (McLeish et al. [2025]): increases the mean of a Poisson-lognormal depth distribution over training

Both approaches pay the cost of expensive high-depth updates throughout training. Since recurrent compute grows with $K$, training at deep depths is disproportionately expensive.

The Fixed-First Proposal

Fixed-first reverses this cost structure: hold depth constant at a small $K_0$ for the first 67-80% of training, then switch to uniform $K = 1,\dots,16$ training for the remainder. Because small-$K$ updates are cheaper, total training FLOPs drop substantially. Importantly, fixed-first changes only the recurrent depths used during training—not the architecture, parameter count, or token budget.

---

Methodology

Synthetic Factual Recall Task

To motivate the approach, the authors use a controlled memorization task based on the bioS datasets from Allen-Zhu and Li [2023]. Fictional people have six attributes (birth city, day, month, university, major, employer) containing 33.39 bits of information per person. Each fact is tested via five fixed templates; a fact counts as recalled only if the model gives the exact answer on ≥4 templates. Recalled bits per person = $\sum_{\text{attributes}} \log_2(\text{number of possibilities})$.

Models and Training

Update Diagnostic

To understand why the finishing phase works, the authors compare single updates at $K=4$ and $K=16$ from the same checkpoint, batch, and optimizer state. They isolate batch effects, match the RMS of updates, and define a benefit metric:

$$\text{Benefit}(K) = \text{NLL after } K{=}4 \text{ update} - \text{NLL after } K{=}16 \text{ update}$$

Positive values favor the deeper $K=16$ update.

---

Empirical Validation / Results

Synthetic Recall Task Results

From Figure 1 (matched training FLOPs at 15%, 45%, 100% of budget):

• Fixed $K=4$: learns quickly near training depth but recall collapses at larger serving depths
• Uniform training: is more stable across depths but learns much slower (partly due to fewer optimizer steps at matched FLOPs, since $K=4$ updates are cheaper)
• Fixed-first (80% fixed, then uniform): achieves near-task-ceiling recall across serving depths through $K=64$ with the full budget, retaining recall even beyond the $K=1,\dots,16$ training range

200M Language Model Results

With a uniform finish:

• $K_0 = 2$ at 67% switch: 45.3% compute savings, mean serving NLL within 0.5% of uniform training
• $K_0 = 2$ at 80% switch: 54.3% compute savings, within 0.8% of uniform
• $K_0 = 1$ gives best $K=1$ NLL at 67% switch, but order reverses from $K=3$ onward (low $K_0$ favors cheap serving depths)
• Uniform finish > Poisson-lognormal finish for all six fixed-first schedules

vs. the McLeish et al. [2025] progressive curriculum:

• $K_0 = 2$ uniform-finish schedules reach lower mean NLL with 18.2% less compute (67% switch) and 31.7% less (80% switch)

vs. uniform training:

• $K_0 = 4$ at 67% switch nearly matches full uniform mean NLL: 0.68643 vs 0.68582, with 31.3% less compute (compute-matched uniform: 0.72089)

1B Language Model Results

Downstream evaluations (HellaSwag, PIQA, ARC-Easy) stay within observed seed variation. A loss spike occurs at the switch point (Figure 3) but recovers quickly.

Update Diagnostic Results

Figure 4 shows:

• At the start of the uniform phase: the $K=16$ update improves the model at larger serving depths, while the $K=4$ update is more useful near $K=4$
• After 16 uniform updates: the depth-averaged advantage of the $K=16$ update is no longer statistically distinguishable from zero

This suggests the finishing phase quickly corrects depth specialization once a strong low-depth computation exists, though deep updates remain important for flexibility during the transition.

---

Theoretical and Practical Implications

1. Reconceptualizing adaptive-depth training: The paper reframes adaptive-depth training as a finishing stage rather than a cost paid throughout training. This is analogous to how curriculum learning often front-loads easy examples and fine-tunes at the end—here, low depth plays the role of the "easy" regime.

2. Compute efficiency: At 1B scale, 31-37% FLOP savings with negligible quality loss suggests that most of the depth-adaptive capability is learned during the final phase. This dramatically lowers the barrier to training serving-depth-flexible models.

3. Design principle for $K_0$: The initial fixed depth acts as a control knob for the compute/quality trade-off and determines which serving depths the final model favors. Lower $K_0$ = cheaper training, favoring shallow serving depths; higher $K_0$ = better deep-serving performance at higher training cost.

4. Synergy with serving economics: Adaptive-depth models let practitioners trade quality for latency at inference time. Fixed-first makes such models more accessible by cutting their training cost substantially.

5. Benchmark for future work: The authors' diagnostic (update-pair comparison) provides a reusable tool for studying when multi-depth updates matter during training.

---

Conclusion

Fixed-first training demonstrates that training over the full range of serving depths can be reserved for a short finishing phase with only small performance changes. At 200M and 1B scales, it reduces modeled training compute by 30-45% and wall-clock time by 26-39% while keeping NLL within 1% and matching downstream performance. The key mechanism: fixed-depth training learns quickly (depth-specialized), and a brief uniform-depth phase corrects this specialization rapidly, since deep updates quickly lose their relative advantage.

Future directions implied by this work:

• Investigating whether even shorter finishing phases suffice with better switch-point scheduling
• Exploring whether the finishing phase could be made adaptive (e.g., triggered by loss plateau rather than fixed token budget)
• Extending fixed-first to other architectures (e.g., mixture-of-experts or early-exit models) that exhibit similar depth-compute trade-offs
• Studying how the optimal $K_0$ and switch fraction scale with model size, data, and serving-depth requirements