Adaptive-Depth Is a Finishing Stage

Andrei Piterbarg; Sean McLeish; Dhruv Pai; Benjamin Keigwin

cs.LGOct 9, 2026

Adaptive-Depth Is a Finishing Stage

Authors: Andrei Piterbarg, Sean McLeish, Dhruv Pai, Benjamin Keigwin

Organizations: University of Pennsylvania · University of Maryland

Abstract

Looped transformers reuse the same layer multiple times, allowing inference compute to be adjusted without changing the model. Recent adaptive-depth methods train the model across multiple loop depths from the beginning of training. We introduce fixed-first training, which varies model depth only during a short final phase of training. At 200M parameters, training a model of fixed-depth 2 for the first 67% of the run reduces modeled training compute by 45% while keeping mean negative log-likelihood (NLL) across inference depths within 0.5% of training across depths throughout. At 1B parameters, training initially at depth 4 reduces modeled compute by 31% to 37% and wall-clock time by 32% to 39%, while keeping final NLL within 0.8% and downstream differences within the observed seed variation. These results show that training over the full range of depths can be reserved for a short finishing phase with only small changes in performance.

Figures & tables

Figure 1: Recalled bits per person across serving depths for fixed K=4 , an 80% fixed-first schedule, and uniform K=1,...,16 training at three matched training budgets (15%, 45%, and 100% of the full FLOP budget). Panels compare complete runs: we train a new model of each schedule from scratch at each budget, so the fixed-first models have switched to uniform training in each panel. Lines show four-seed means, bands show 95% confidence intervals, dashed line marks the task ceiling.

Figure 2: Results at 200M. Rows show the finishing distribution. Left: mean serving-depth NLL against modeled training FLOPs; filled and open markers denote 67% and 80% fixed phases. Right: serving-depth NLL minus uniform-training NLL for the 67% schedules; positive values are worse than uniform. Lines and bands show the mean and range across three paired seeds.

Figure 3: Results at 1B. Top: full and terminal training NLL trajectories. The loss spike is where we switch from fixed to adaptive-depth training. Bottom: HellaSwag ( Zellers et al., 2019 ) , PIQA ( Bisk et al., 2020 ) , and ARC-Easy ( Clark et al., 2018 ) accuracy across serving depths. Lines are three-seed means and bands span seeds.

Figure 4: We define the benefit as the NLL after the K=4 update minus the NLL after the K=16 update; positive values favor K=16 . Left: the benefit across serving depths at the start of the uniform phase and after 16 uniform updates. Right: the benefit averaged across serving depths after 0, 16, 64, and 256 uniform updates; error bars show 95% confidence intervals.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Part	Token sequence
Person’s name	N17 N204 N6
Birth city	V3 V11
Employer	V6 V10

Appendix

Table 1: Example token sequences for one fictional person and two attributes. Names and attribute values use several tokens.

Stage	What the model sees
Training	N17 N204 N6 [employer] V6 V10 [birth city] V3 V11
Evaluation	N17 N204 N6 [birth city] ___ ___ [employer] V6 V10
Correct response	V3 V11

Appendix

Table 2: The same facts presented in different orders. Evaluation requires exact prediction of the two-token birth-city value.

Figure 5: Recall across serving depths for fixed-depth models trained at K=1 , 2 , 4 , or 8 , normalized by each model’s peak recall. Each model performs best near the depth used during training.

Figure 6: A standard transformer applies a sequence of different layers, while a looped transformer applies the same layer repeatedly. Here, the looped model applies layer “A” three times ( K=3 ). “ P ” is the input embedding and “ H ” is the output head. Adapted, with permission, from Schwarzschild et al. [2021] .

Figure 7: Uniform-depth training and an 80% fixed-first schedule. Uniform-depth training samples K from 1 through 16 throughout training. Fixed-first training uses K=4 for the first 80% of training, then switches to the same uniform distribution.

200M schedule	Terminal NLL	Mean serving NLL	Modeled FLOPs	Wall-clock reduction
Uniform	0.67297 $\pm$ 0.00058	0.68582 $\pm$ 0.00077	100%	–
Fixed 67%	0.67424 $\pm$ 0.00027	0.68643 $\pm$ 0.00028	68.7%	26.2%
Uniform, matched	0.71936 $\pm$ 0.00066	0.72089 $\pm$ 0.00072	68.7%	–
Fixed 80%	0.67677 $\pm$ 0.00037	0.68941 $\pm$ 0.00087	62.4%	31.6%

Appendix

Table 3: Cross-scale efficiency and NLL results. “Fixed 67%” and “Fixed 80%” use K=4 for that portion of training, then switch to uniform K=1,...,16 . Terminal NLL is measured over the final 200 updates, and 200M serving NLL is averaged over serving depths K=1,...,16 . The compute-matched uniform run is a separate 3,296-update run at the same modeled FLOPs as Fixed 67%, with its schedule re-tuned for the budget. Values are mean

\pm

SD over three seeds.

Figure 8: Next-token NLL across serving depths for the 200M models. The left panel shows the full serving range, and the right panel expands K=4,...,16 . Both fixed-first schedules remain close to uniform training. Lines show three-seed means, and bands show the range across seeds.

Depth-averaged task	Uniform	Fixed 67%	Fixed 80%
HellaSwag	50.72 $\pm$ 0.33	50.56 $\pm$ 0.19	50.27 $\pm$ 0.24
PIQA	70.48 $\pm$ 0.42	70.52 $\pm$ 0.55	70.86 $\pm$ 0.53
ARC-Easy	56.11 $\pm$ 1.46	56.76 $\pm$ 1.17	57.43 $\pm$ 1.12

Appendix

Table 4: Exact results for the three 1B evaluations shown in Figure 3 . The top table averages each task separately over serving depths K=1,...,16 ; the bottom table reports the conventional K=16 endpoint. There is no averaging across tasks. Values are mean

\pm

SD over three seeds.

Efficient Language Model Training

Efficient language model training studies how to reduce the data, compute, memory, time, energy, and cost needed to build capable language models.

See all 162 papers

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching
Aug 10, 2026Kristian Schwethelm, Daniel Rueckert, Georgios KaissisLLM Inference LLM Inference Scheduling
Chair for AI in Healthcare and Medicine, Technical University of Munich (TUM), Germany · Department of Computing, Imperial College London, UK · Munich Center for Machine Learning (MCML), Germany +1
Do Transformers Use their Depth Adaptively? Evidence from a Relational Reasoning Task
Apr 14, 2026Alicia Curth, Rachel Lawrence, Sushrut Karmalkar +1Adaptive Computation Transformer Interpretability
Microsoft Research Cambridge
Looped Transformers as Optimizers
Sep 29, 2026Yulong Huang, Chen Jiang, Zhanpeng Zhou +5Fast Weights Recurrent Transformers
HKUST (GZ) · StepFun · SJTU +2

Aug 10, 2026J/K move · Enter open · S save

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

Chair for AI in Healthcare and Medicine, Technical University of Munich (TUM), Germany · Department of Computing, Imperial College London, UK · Munich Center for Machine Learning (MCML), Germany+1

LLM Inference LLM Inference Scheduling LLM Inference Acceleration Language Modeling Recursive Language Models Inference-Time Adaptation of Language Models

Open paper AskPDF arXiv

A main promise of looped language models is depth-adaptive inference. By looping a block of shared layers a variable number of times, the model can use less compute for "easy" tokens and more for "hard" ones. However, tokens with different numbers of loops cannot share a uniform forward pass and therefore cannot be handled by standard batching systems such as vLLM. The practical value of depth-adaptive inference thus hinges on whether batching can be made efficient. We introduce the first efficient method for depth-adaptive looped LMs via continuous depth batching (CDB), which forms new batches between loop steps. Our method dynamically schedules looped and non-looped parts of the architecture, manages looped KV-caching, and predicts which tokens will exit the loop in advance so it can prepare batches asynchronously. Experiments on Ouro 1.4B and Huginn 3.5B show that fully looped architectures are best suited to depth-adaptive inference, as large non-looped layers outside the recurrent core (e.g., token embedding, LM head, and unshared transformer blocks) slow down and complicate scheduling. Overall, CDB realizes up to 99% of the estimated maximum speedup available, leaving further gains primarily dependent on model architecture and exit behavior.

Abstract

Efficient Language Model Training

Explore similar work

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

Do Transformers Use their Depth Adaptively? Evidence from a Relational Reasoning Task

Looped Transformers as Optimizers