# Adaptive-Depth Is a Finishing Stage

> Fixed-first training cuts adaptive-depth transformer training compute by 30-45% with negligible quality loss by reserving full-depth updates for a brief final phase.

- **Source:** [arXiv](https://arxiv.org/abs/2610.14252)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/xG2wkK
- **Whiteboard:** https://picx.dev/p/xG2wkK/image

## Summary

## 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

| Setting | 200M | 1B |
|---------|------|-----|
| Data | UltraData-Math-L3 | Ultra-FineWeb |
| Updates | 4,800 (2.52B tokens) | 20,000 (10.49B tokens) |
| Serving depths evaluated | $K = 1, \dots, 16$ | $K = 1, \dots, 16$ |
| Fixed-first variants | $K_0 \in \{1, 2, 4\}$ | $K_0 = 4$ |
| Switch points | 67% or 80% of training | 67% or 80% of training |
| Finish distributions | Uniform or Poisson-lognormal | Uniform |

### 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

| Schedule | FLOP savings | Wall-clock savings | Terminal NLL diff |
|----------|-------------|-------------------|-------------------|
| $K_0=4$, 67% switch | 30.8% | 32.1% | within 0.8% |
| $K_0=4$, 80% switch | 36.9% | 38.5% | within 0.8% |

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

---

_Markdown view of https://picx.dev/p/xG2wkK, served by PicX — AI-generated visual whiteboard summaries of research papers._