Do Halting Heads Really Save Compute in Looped Language Models?

Tsung-Han Wu; Minsuk Heo; Kenny Tan

cs.CLOct 23, 2026

Do Halting Heads Really Save Compute in Looped Language Models?

Authors: Tsung-Han Wu, Minsuk Heo, Kenny Tan

Organizations: Apple

Abstract

Looped language models run each token through the same layers several times, and each run is one iteration. To save compute, an extra layer, the halting head, gives after each iteration the probability that the token stops there, so that easy tokens can stop early. Pretraining models from scratch, up to 1.3B parameters, we find two problems. First, heads trained on a mix of all iterations predict that tokens need under half the iterations, yet under the standard stopping rule tokens run almost every iteration, because this training never scores the single iteration where a token stops. Scoring every iteration instead fixes this, but in our runs it works only with a strong penalty favoring early stops. Second, savings are commonly reported as an average number of iterations per token, but when sequences are generated in batches that run until their slowest token stops, a model without a head is almost as good at that cost. The reported average can therefore overstate the saving. We recommend reporting the iterations tokens actually run, with the stopping threshold and batch size, and give a test, run before training a head, that bounds its gain.

Figures & tables

Figure 1: A head can predict

3.7

iterations and still run all eight. (a) A transformer stacks eight different blocks of layers; a looped model runs one block up to eight times. (b) After each iteration, a halting head gives the probability that the token stops there. Its average, the expected halting steps, is

1{\times}0.33+2{\times}0.15+\dots+8{\times}0.18=3.7

iterations. An exit rule stops the token once the running total reaches

90\%

, which for this illustrative token happens only at iteration 8.

Figure 2: A batched pass costs more than the average token. (a) Four tokens processed together: the pass keeps running until its deepest token stops. (b) Two 370M-parameter models: the predicted average (what the head’s probabilities predict), spent (what its exit rule runs) and paid (what a batched pass over

32

sequences runs). Mixture: trained on a mixture of all iterations; per-step: every iteration scored. Thresholds as in Section 2.2 .

	Iterations			$\Delta$ BPB $\downarrow$
						decode
Scale	$\mathbb{E}[N]$	$\bar{N}\downarrow$	$R\,\bar{N}\downarrow$	at $\bar{N}$	at $R\,\bar{N}$	hidden	reused
Mixture, $\varepsilon=0.1$
140M	3.81	7.86	–	$\mathbf{-0.0093}$	–	–	–
370M	3.74	7.76	8.0	$-0.0063$	–	–	–
770M	3.63 †	7.67	–	$-0.0065$	–	–	–

Table 1: What the head predicts is not what a batched pass pays. The three Iterations columns are the three numbers of Section 2.2 : how many iterations the head predicts a token needs (

\mathbb{E}[N]

), how many a token spends under the exit rule (

\bar{N}

), and how many a batched pass over

32

sequences pays (

R\,\bar{N}

\Delta

BPB is our model’s BPB minus that of the sampled-depth baseline run for the same number of iterations, so a negative value means our model is better. That baseline runs

\bar{N}

iterations in the column at

\bar{N}

(teacher-forced) and

R\,\bar{N}

in the column at

R\,\bar{N}

(a batched real decode). The mixture-loss

R\,\bar{N}

is derived from

\bar{N}

(Appendix E.1 ). The two decode columns simulate generating one token at a time, and differ in what a stopped token leaves at the iterations it skipped: nothing ( hidden ), or its keys and values from the iteration where it stopped ( reused ). –: not measured. Bold : beyond noise. † One seed.

Method	Training depth	What the coda reads	Loss term on depth
Dense	no sharing, $n_{p}+n_{r}+n_{c}$ layers	final layer	—
Sampled-depth	$T\sim$ Poisson, mean $8$	$h_{i,T}$	—
ACT ( Graves, 2016 )	fixed $T{=}8$	mixture up to its halting iteration	ponder cost, weight $\tau_{p}{=}10^{-4}$
Mixture	fixed $T{=}8$	mixture $\sum_{t}p_{i,t}h_{i,t}$	KL to geometric $\kappa{=}4$ , weight $\beta$
Per-step	fixed $T{=}8$	each $h_{i,t}$ , weighted by $p_{i,t}$	KL to geometric $\kappa{=}4$ , weight $\beta$

Table 2: The methods, and how each treats depth. The sampled-depth baseline follows Prairie et al. (2026) . Dense has the same number of parameters as the others but runs each layer once, so it uses fewer floating-point operations. The ponder cost is ACT’s penalty on the number of iterations a token runs (Appendix E.8 varies

\tau_{p}

Figure 3: The cost in quality of exiting early. BPB against realized steps, the average number of iterations the exit rule actually runs. Each dot is one exit threshold of a 370M model. The dotted line is the sampled-depth baseline run for that many iterations; a curve below it is a lead at that cost. The mixture-loss model is above the baseline until

7.3

steps; the per-step model is on or below it at every threshold shown.

Mixture
Scale	Reads mixture $\downarrow$	Reads state at cap $\downarrow$	Readout gap $\downarrow$
140M	1.2058 $\pm$ .0006	1.4091 $\pm$ .040	$+0.203$ $\pm$ .040
370M	1.0807 $\pm$ .0001	1.2559 $\pm$ .015	$+0.175$ $\pm$ .015
770M	1.0090 $\pm$ .0012	1.1767 $\pm$ .006	$+0.168$ $\pm$ .007
1.3B	0.9643 $\pm$ .0017	1.1535 $\pm$ .016	$+0.189$ $\pm$ .015
Per-step

Table 3: Reading out one iteration costs the mixture-loss model about

0.2

BPB and the per-step model nothing. Reads mixture and Reads state at cap : BPB when the coda reads the mixture or only the state at the cap; Readout gap : the second minus the first. All eight iterations run, so only what is read out differs. Per-step:

\beta{=}0.01

(

0.1

at 770M; Appendix F.5 ).

Scale	Loss	BPB $\downarrow$	$0.1$	$1.0$
		$\beta{=}0.01$	change $\downarrow$ at $\beta{=}$
140M	Mixture	1.813	$+0.038$	$+0.052$
	Per-step	1.711	$\mathbf{-0.394}$	$\mathbf{-0.450}$
370M	Mixture	1.610	$-0.002$	$+0.031$
	Per-step	1.430	$\mathbf{-0.273}$	$\mathbf{-0.308}^{\dagger}$

Table 4: A stronger prior helps the per-step model by over

0.25

BPB and does not help the mixture-loss model. BPB on CC with every token forced to run one iteration. Bold : beyond noise. † One seed (Appendix F.1 ).

BPB $\downarrow$	Exit rule	Shuffled
Per-step, 140M	1.2202	$\mathbf{1.2285}$

Table 5: Shuffling the depths the exit rule chose costs

0.008

BPB at the same compute. At

\beta{=}0.1

and

\varepsilon{=}0.7

. Exit rule : each token runs the number of iterations the exit rule chose for it. Shuffled : we keep the same iteration counts but hand them to the tokens in a random order, so an easy token may get the many iterations chosen for a hard one. Bold : beyond noise.

Appendix figures & tables32 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Where	Meaning
Model and depth
$n_{p},n_{r},n_{c}$	2.1	layers in the prelude, the core and the coda; only the core repeats
$T$	2.1	the cap : the most iterations a token can run, set at inference
$T_{\text{train}}$	Tab. 7	the cap a model was trained at; $8$ unless stated
$t$	2.1	an iteration, $1\leq t\leq T$ ; one run of the core
$h_{i,t}$	2.1	the state of token $i$ after the core has run $t$ times

Appendix

Table 6: Notation, baselines and measures, with the section that defines each.

Scale	Params	$d_{\text{model}}$	$d_{\text{ffn}}$	$n_{p}/n_{r}/n_{c}$	Tokens
140M	144.3M	768	3072	2/2/2	11.2B
370M	388.0M	1024	4096	4/4/4	29.6B
770M	776.7M	1280	5120	6/6/6	61.6B
1.3B	1.339B	1536	6144	8/8/8	104B

Appendix

Table 7: Model configurations. Looped models train at

T_{\text{train}}=8

unless stated (Table 2 ); the dense models (140M–770M) run the same

n_{p}+n_{r}+n_{c}

layers once, without looping. Parameter counts are exact and include the sampled-depth baseline’s value embeddings (Appendix A ), 45–52% of parameters; output embeddings are tied. Scale names follow that baseline’s.

d_{\text{model}}

: the width of the hidden state;

d_{\text{ffn}}

: the width of the MLP’s hidden layer.

Scale	Workload	$T{=}1$	$T{=}2$	$T{=}3$	$T{=}4$	$T{=}5$	$T{=}6$	Fitted $f$	Largest error
140M	prefill, $b{=}32$	2.76	2.21	1.84	1.57	1.38	1.22	$0.729$	$0.005$
370M	prefill, $b{=}32$	3.05	2.36	1.92	1.62	1.40	1.24	$0.768$	$0.005$
770M	prefill, $b{=}32$	3.17	2.42	1.96	1.64	1.42	1.24	$0.782$	$0.005$
1.3B	prefill, $b{=}32$	3.22	2.45	1.97	1.65	1.42	1.25	$0.788$	$0.005$
140M	decode, $b{=}1$	3.16	2.42	1.95	1.65	1.41	1.24	$0.781$	$0.009$
370M	decode, $b{=}1$	3.24	2.43	1.92	1.61	1.41	1.24	$0.790$	$0.056$

Appendix

Table 8: The architectural speedup ceiling. Forced-depth speedup against

T=8

, median of nine repeats. Run-to-run variation is at most

0.070\times

(median

0.010\times

). Across the four scales at one

T

, speedups differ by under

0.2\times

in decode, and also in prefill for

T\geq 3

. Fitted

f

: the cost model’s one parameter (Section 8 ), fitted to the row’s

T{=}1

speedup; Largest error : its largest error in predicting the row’s

T=2

6

speedups (Appendix C.2 ).

Figure 4: A batched pass costs the maximum, not the mean. Measured speedup over

T{=}8

at 370M: the exit rule (filled) beside the same weights at a fixed depth equal to their realized steps (open, interpolated). Thresholds differ between points (Table 9 ): compare filled with open at one

x

only. Per-step prefill at

b{=}32

was not timed at 370M. The

x

-axis is tokens per batched pass (sequences

\times

positions). This count, not the number of sequences, decides how deep the deepest token goes.

Model	Workload	$\varepsilon$	$\bar{N}$	$\max_{i}N_{i}$	Measured	Fixed at $\bar{N}$
370M, both models, nine repeats per cell
Mixture	prefill, $b{=}1$	0.3	5.44	8	$1.00\times$	$1.33\times$
Mixture	prefill, $b{=}32$	0.6	2.17	5	$1.40\times$	$2.29\times$
Mixture	decode, $b{=}1$	0.5	2.86	5	$1.93\times$	$2.04\times$
Mixture	decode, $b{=}32$	0.5	2.90	5	$1.65\times$	$2.03\times$
Per-step, $\beta{=}0.1$	prefill, $b{=}1$	0.7	2.93	6	$1.18\times$	$1.96\times$

Appendix

Table 9: Speedup against

T=8

. Measured : timed with the exit rule. Fixed at

\bar{N}

: the same weights run at a fixed depth equal to their realized steps (interpolated between integer depths).

\bar{N}

and

\max_{i}N_{i}

: the mean and the largest realized steps over the timed tokens. The 370M block is timed on one H200 and the 770M block on the device named in its header; inputs are validation tokens.

Scale	Steps	BPB	at the cap	falls	to fall
		Mixture	Baseline	Curve	Would have
140M	3.020	1.4493	1.2159	0.0256	$0.2334$
370M	3.017	1.2759	1.0875	0.0202	$0.1884$
770M	2.885	1.1881	1.0157	0.0224	$0.1724$
1.3B	2.925	1.1353	0.9693	0.0217	$0.1660$

Appendix

Table 10: How far the depth curve would have to fall for the quality condition to hold: the mixture-loss model’s BPB at about three iterations (

\varepsilon{=}0.5

) minus the depth curve at the cap, beside the fall the curve actually shows from the head’s steps to the cap. Steps : realized steps at

\varepsilon{=}0.5

; the other columns are BPB.

		forced depth $t$
Model	Seed	$1$	$2$	$3$	$4$	$6$	$8$	$12$	$16$	$24$	$32$
Sampled-depth	42	7.2937	7.4076	5.0295	3.8793	2.7472	1.7466	1.3038	1.2311	1.2098	1.2081
Sampled-depth	7	5.9709	5.8090	4.6748	3.9846	2.5963	1.6305	1.2904	1.2285	1.2105	1.2092
Per-step, $\beta{=}0.1$	42	1.3337	1.2532	1.2290	1.2206	1.2163	1.2156	1.2154	1.2154	1.2154	1.2154
Per-step, $\beta{=}0.1$	7	1.3386	1.2654	1.2430	1.2352	1.2313	1.2307	1.2305	1.2305	1.2305	1.2305

Appendix

Table 11: Forced-depth BPB at

T_{\text{train}}{=}32

, 140M, CC,

2/2/2

11.2

B tokens, weights frozen and only the forced depth

t

changing. Two models, two seeds each; the anchor is

t{=}12

, at

t/T=3/8

. The baseline’s cells below

t{=}8

are not a usable depth curve: its training rarely drew such depths .

		forced depth $t$
Model	Seed	$1$	$2$	$3$	$4$	$5$	$6$	$8$	$12$	$16$
Sampled-depth	42	1.7147	1.4347	1.3293	1.2763	1.2477	1.2311	1.2155	1.2078	1.2068
Sampled-depth	7	1.6856	1.4331	1.3310	1.2782	1.2496	1.2332	1.2177	1.2098	1.2087
Per-step, $\beta{=}0.1$	42	1.3324	1.2519	1.2274	1.2188	1.2156	1.2144	1.2137	1.2136	1.2136
Per-step, $\beta{=}0.1$	7	1.3236	1.2490	1.2257	1.2176	1.2147	1.2136	1.2129	1.2127	1.2127
Specialist $T{=}6$	42	—	—	—	—	—	1.2077	—	—	—

Appendix

Table 12: Forced-depth BPB at

T_{\text{train}}{=}16

, 140M, CC,

2/2/2

11.2

B tokens, weights frozen and only the forced depth

t

changing. The per-step model reaches all sixteen iterations with gradient, the sampled-depth rows eight . The specialist is shown at its trained depth.

\mathbb{E}[N]

: expected halting steps with the cap set to

t

$\varepsilon$	Seed	Steps	BPB	Specialist	Model $-$ specialist
$0$	42	15.993	1.2132	1.2047	$+0.0085$
$0$	7	16.000	1.2124	1.2047	$+0.0077$
$0.3$	42	6.531	1.2128	1.2065	$+0.0063$
$0.3$	7	6.491	1.2119	1.2066	$+0.0053$
$0.5$	42	4.798	1.2139	1.2114	$+0.0025$
$0.5$	7	4.755	1.2130	1.2116	$+0.0014$

Appendix

Table 13: Threshold early exit at

T_{\text{train}}{=}16

, 140M, CC, two seeds, exit-state readout. Steps : realized steps, not

\mathbb{E}[N]

. Specialist : the specialist’s BPB at the row’s realized steps, interpolated linearly between the specialists of Appendix D.2 (trained at

2

5

and

8

iterations) and the six-iteration specialist trained at this cap, and held at the eight-iteration value beyond eight. At matched realized steps the two are inside the evaluation noise floor for

\varepsilon\geq 0.5

, the same no-advantage verdict as at a cap of eight.

		CC			FW
Scale	Method	$n$	Core-Ext $\uparrow$	BPB $\downarrow$	$n$	Core-Ext $\uparrow$	BPB $\downarrow$
140M	Dense	2	0.0488 $\pm$ .0021	1.2576 $\pm$ .0010	–	–	–
140M	Sampled-depth	3	0.0641 $\pm$ .0010	1.2193 $\pm$ .0026	3	0.0906 $\pm$ .0111	0.9054 $\pm$ .0012
140M	ACT	3	0.0724 $\pm$ .0021	1.2027 $\pm$ .0007	–	–	–
140M	Mixture	3	0.0676 $\pm$ .0037	1.2052 $\pm$ .0007	3	0.1014 $\pm$ .0051	0.8983 $\pm$ .0005
370M	Dense	2	0.0765 $\pm$ .0027	1.1208 $\pm$ .0011	–	–	–

Appendix

Table 14: Quality at the training depth

T=8

. Core-Ext is the DCLM CORE-extended aggregate; BPB is measured on the validation split of the training corpus. Mean

\pm

s.d. over

n

seeds, best per column-block in bold, with any value within

0.0005

of it. Dense models were trained on CC only, at 140M to 770M. Core-Ext differences below

0.0042

are unresolved. ‡ BPB over

n=2

. The mixture-loss BPB values differ from those of Tables 3 and 17 by at most

0.0006

: a different evaluation session, within the evaluation noise floor.

Model	$t{=}2$	$t{=}3$	$t{=}4$	$t{=}5$	$t{=}8$
Specialist $T{=}2$	1.2349	1.3124	1.3953	1.4353	1.4779
Specialist $T{=}3$	1.6913	1.2236	1.3212	1.3854	1.3434
Specialist $T{=}4$	1.8764	1.7231	1.2152	1.3088	1.3838
Specialist $T{=}5$	1.8973	1.8333	1.6849	1.2104	1.5053
Specialist $T{=}8$	1.8785	1.8578	1.8347	1.7676	1.2047
Per-step, $\beta{=}0.1$	1.2456	1.2251	1.2185	1.2164	1.2154

Appendix

Table 15: Specialists at 140M on CC, each trained at one number of iterations

T

with gradients through all of them. Bold : each specialist’s lowest BPB, at its own training depth. Entries off the diagonal are depths the model never saw in training. Entries are BPB with every token forced to run

t

iterations; lower is better. Specialists at

T{=}3,4

: three-seed means on the diagonal, seed

42

elsewhere; other specialists: seed

42

. Per-step rows:

\beta{=}0.1

seed

42

\beta{=}1.0

three-seed means.

$t$	Released	Mean $4$	Mean $2$	Specialist $T{=}3$	Per-step
1	1.3801 $\pm$ .0015	1.2808 $\pm$ .0006	1.2574 $\pm$ .0015	1.8779 $\pm$ .060	1.3166 $\pm$ .0044
2	1.2779 $\pm$ .0006	1.2388 $\pm$ .0005	1.2409 $\pm$ .0008	1.7251 $\pm$ .038	1.2457 $\pm$ .0019
3	1.2418 $\pm$ .0007	1.2264 $\pm$ .0007	1.2375 $\pm$ .0007	1.2236 $\pm$ .0007	1.2251 $\pm$ .0012
4	1.2269 $\pm$ .0006	1.2225 $\pm$ .0008	1.2367 $\pm$ .0006	1.3167 $\pm$ .0063	1.2185 $\pm$ .0010
8	1.2160 $\pm$ .0006	1.2208 $\pm$ .0008	1.2364 $\pm$ .0006	1.3985 $\pm$ .103	1.2155 $\pm$ .0008

Appendix

Table 16: Sampled-depth baselines trained at shallower depths. 140M, CC, three seeds each, one evaluation session. BPB with every token forced to run

t

iterations. Released : Poisson mean

8

, the baseline of the main text. Per-step:

\beta{=}0.1

\pm

: seed s.d.; lower is better.

Scale	$\varepsilon$	Steps $\downarrow$	Mixture $\downarrow$	Baseline $\downarrow$	$\Delta$ $\downarrow$
140M	0.0	8.000	1.2058 $\pm$ .0006	1.2159	$-0.0101$
140M	0.1	7.860	1.2067 $\pm$ .0011	1.2160	$-0.0093$
140M	0.2	6.993	1.2301 $\pm$ .0080	1.2169	$+0.0132$
140M	0.3	5.755	1.2855 $\pm$ .0148	1.2185	$+0.0670$
140M	0.5	3.020	1.4493 $\pm$ .0165	1.2415	$+0.2078$
370M	0.0	8.000	1.0808	1.0875	$-0.0067$

Appendix

Table 17: Realized compute under threshold early exit, CC. The mixture-loss column is a three-seed mean; the Baseline column uses three seeds at 770M, seed

42

at 370M, and three seeds at

t=1,3,5,6

at 1.3B. Larger

\varepsilon

exits earlier. “Baseline” is the sampled-depth baseline at the same number of iterations, interpolated between integer

T

on that scale’s depth curve; the curve is convex, so interpolation favors the mixture-loss model.

\varepsilon=0

runs all eight iterations, so its

\Delta

is the model’s margin at full depth. At 140M the baseline has three seeds of its own, the mixture-loss model uses the mixture readout, and

\pm

is the seed s.d. All values are BPB; lower is better.

\Delta

is the mixture-loss model minus the baseline, so a negative value favors the model.

Scale	$\varepsilon$	$\bar{N}$	Mixture	Exit state	Over own
140M	0.5	4.08	$-0.0095$	$-0.0106$	$-0.0025$
	0.7	3.11	$-0.0199$	$-0.0205$	$-0.0048$
	0.95	1.88	$-0.0503$	$-0.0504$	$-0.0143$
370M	0.5	3.79	$-0.0100$	$-0.0108$	$-0.0025$
	0.7	2.83	$-0.0214$	$-0.0219$	$-0.0050$
	0.95	1.63	$-0.0589$	$-0.0590$	$-0.0174$

Appendix

Table 18: Per-step

\beta{=}0.1

model, both readouts, CC. Mixture , Exit state : model minus sampled-depth baseline at

\bar{N}

. Over own : the exit-state margin minus the margin that the model’s own forced-depth curve gives at the same

\bar{N}

(every token run at that depth). Each model and each baseline uses three seeds. Seed s.d. of each margin at most

0.0011

, except

0.0027

at 140M,

\varepsilon{=}0.95

. At 140M and 370M the baseline’s three seeds were measured in the same session as the model, so margins differ slightly from Table 19 .

Model	$\varepsilon$	Steps $\downarrow$	BPB $\downarrow$	Baseline	$\Delta$ $\downarrow$	Depth equivalent / steps [floor interval]
140M; $\beta{=}0.1$ rows three seeds, $\beta{=}1.0$ rows seed $42$
$\beta{=}0.1$	0.0	7.997	1.2143	1.2155	$-0.0012$	–
$\beta{=}0.1$	0.5	4.074	1.2159	1.2261	$-0.0102$	$1.86\times$ [ $1.22$ , open]
$\beta{=}0.1$	0.7	3.113	1.2196	1.2398	$-0.0202$	$1.67\times$ [ $1.41$ , open]
$\beta{=}0.1$	0.9	2.160	1.2317	1.2721	$-0.0404$	$1.69\times$ [ $1.56$ , $1.83$ ]
$\beta{=}0.1$	0.95	1.879	1.2399	1.2905	$-0.0506$	$1.65\times$ [ $1.56$ , $1.81$ ]

Appendix

Table 19: Threshold early exit on the per-step model under a strong prior, CC; larger

\varepsilon

exits earlier. “Baseline” is the sampled-depth baseline interpolated to the same realized steps. The depth equivalent is the depth that baseline needs to match the model’s BPB, so

1.86\times

means the baseline needs

1.86

times as many iterations as the model runs. The bracketed interval recomputes it after moving the model’s BPB by the evaluation noise floor and is the figure to read; it includes only the evaluation noise floor, so the true uncertainty is at least this large. “Open” intervals have no upper end (see below the table). The mixture-loss rows are the prior without the per-step loss, against the sampled-depth baseline’s three-seed curve;

\pm

is the seed s.d.

Gate chose	Share	$t{=}1$	$t{=}2$	$t{=}3$	$t{=}4$
$1$	$18.5\%$	0.6665	0.6677	0.6613	0.6601
$2$	$30.6\%$	4.0175	2.3504	2.3113	2.3033
$3$	$38.8\%$	7.2121	4.9394	3.3773	3.3300
$4$	$12.1\%$	8.3433	7.4025	5.0786	3.9426
Corpus BPB		2.4907	1.7641	1.3296	1.2531

Appendix

Table 20: Ouro 2.6B at

\tau_{\text{Ouro}}=0.05

: mean per-token loss (nats) at every fixed readout depth

t

, bucketed by the step the gate chose (bold).

Figure 5: The per-step loss with a strong prior, and what pricing the pass does to it. 370M, CC; model three seeds, baseline seed

42

. Against the sampled-depth baseline read at

\bar{N}

the model wins everywhere. Read at the depth a batched pass runs,

R\,\bar{N}

, the baseline improves as the pass widens until it straddles the model, and further left, the baseline at the widest pass is ahead outright. Legend: the baseline read at

R\,\bar{N}

for a pass of

32

tokens (decoding

32

sequences;

R

counted in a real decode) and of

2048

tokens (one prompt;

R

from the timing table). Here “static” in the legend is the sampled-depth baseline. Dashed line:

\varepsilon=0.7

, the threshold used. The point at

1.1689

lies above the axis. Model BPB is teacher-forced. The band is the evaluation noise floor, either side.

Tokens/pass	$R$	$R\,\bar{N}$	Sampled-depth there	Per-step $-$ sampled-depth	Verdict
370M, $\varepsilon{=}0.7$ , $\bar{N}=2.82$ , BPB $1.0908$ — same scale, model and threshold as $R$
$1$ (decode, $b{=}1$ )	1.06	2.99	1.1082	$\bf-0.0174$	model ahead
$2048$ (prefill, $b{=}1$ )	2.19	6.17	1.0889	$+0.0019$	inside floor
$65{,}536$ (prefill, $b{=}32$ )	2.32 †	6.54	1.0886	$+0.0022$	inside floor
370M, $\varepsilon{=}0.95$ , $\bar{N}=1.62$ , BPB $1.1086$ — $R$ carried from $\varepsilon{=}0.7$
$1$ (decode, $b{=}1$ )	1.06	1.72	1.1605	$\bf-0.0519$	model ahead

Appendix

Table 21: The

R

-priced comparison behind Figure 5 , at the timed widths; the

32

-sequence pass of Figure 5 is in Tables 1 and 24 . Sampled-depth there is the sampled-depth baseline read at

R\,\bar{N}

R

: recovered from Table 9 ’s per-step rows at 370M and

\varepsilon{=}0.7

(Appendix C.2 ); block headers state what they borrow. Verdict and bold: “model ahead” or “baseline ahead” if the margin exceeds twice the evaluation noise floor, “inside floor” otherwise. All margins are teacher-forced; with stopped tokens’ keys and values hidden (simulated decode) the one-token rows change sign (Appendix E.5 ). †

R

borrowed from a mixture-loss row.

Figure 6: The speedup is largest when one token is generated at a time and is gone in large batches. Horizontal axis: how many tokens share one batched pass;

1

and

32

are decoding one or

32

sequences at once, and

2048

and

65{,}536

are processing one or

32

prompts of

2048

tokens. Vertical axis: how many times faster the 370M per-step model at

\varepsilon{=}0.95

is than the sampled-depth baseline at

T{=}3

, whose teacher-forced quality it matches; above

1

is faster. Iterations : the speedup if only the core cost compute (

3

divided by the iterations each pass runs). Throughput : the speedup of the whole pass, as predicted by the cost model of Section 8 ; it is smaller because the prelude and coda run once however many iterations the core runs (at one token per pass,

1.33\times

against

1.75\times

). In a real decode with reused keys and values, the per-step model is slightly worse than that baseline (Appendix E.6 ).

Scale	$\varepsilon$	teacher-forced	keep-first	teacher-forced	keep-first	drop	512-token prompt
		Steps		Model $-$ baseline at the model’s own steps
140M	0.7	3.113	3.112	$-0.0202$	$\mathbf{+0.0824}$ $\pm$ .014	$+0.1542$	$\mathbf{+0.0433}$ $\pm$ .0082
	0.95	1.877	1.889	$-0.0506$	$\mathbf{+0.0764}$ $\pm$ .015	$+0.1233$	$\mathbf{+0.0279}$ $\pm$ .0087
370M	0.7	2.827	2.833	$-0.0221$	$\mathbf{+0.0626}$ $\pm$ .004	$+0.1238$	$\mathbf{+0.0348}$ $\pm$ .0002
	0.95	1.628	1.651	$-0.0597$	$\mathbf{+0.0439}$ $\pm$ .004	$+0.0643$	$\mathbf{+0.0125}$ $\pm$ .0006
770M	0.7	2.752	2.767	$-0.0246$	$\mathbf{+0.0576}$ $\pm$ .002	$+0.1955$	$\mathbf{+0.0276}$ $\pm$ .0006

Appendix

Table 22: The per-step

\beta{=}0.1

model, CC, three seeds per cell (mean, and seed s.d. where shown), against the sampled-depth baseline at the model’s own realized steps. The last column keeps the first

512

positions fully computed and scores both only after them. Bold (keep-first and prompt columns): over twice the evaluation noise floor. The margin is positive in every simulated decode cell and beyond twice the evaluation noise floor in all but one (770M,

\varepsilon=0.95

, 512-token prompt). Margins are in BPB; a positive margin means the model is worse than the baseline.

Scale	Prompt	real	emulated	real	emulated	drop
		Steps		Model $-$ baseline
140M	512	3.006	3.006	$+0.0701$	$+0.0666$	n/a
	1	2.979	3.005	$+0.1941$	$+0.1155$	$+0.1940$
370M	512	2.734	2.734	$+0.0342$	$+0.0319$	n/a
	1	2.638	2.746	$+0.1272$	$+0.0581$	$+0.1276$
770M	512	2.651	2.650	$+0.0223$	$+0.0198$	n/a

Appendix

Table 23: Real decode against the emulation with keys and values hidden: per-step

\beta{=}0.1

model,

\varepsilon{=}0.7

, three seeds, first

64

validation sequences (

98

k scored positions behind

512

tokens,

131

k behind one). Emulated : prompt keys and values kept at every depth. Drop : none kept. Margins against the sampled-depth baseline on the same positions; teacher-forced behind

512

tokens,

-0.018

-0.021

-0.026

. Full-split margins: Table 22 . Margins are in BPB; a positive margin means the model is worse than the baseline.

		Pass		Model $-$ baseline at
Scale	$\bar{N}$	depth	$R$	$\bar{N}$	pass depth	TF
$\varepsilon=0.7$
140M	3.13	4.71	1.51	$\mathbf{-0.0182}$	$-0.0004$ $\pm$ .0013	$-0.0030$
370M	2.84	4.40	1.55	$\mathbf{-0.0197}$	$-0.0010$ $\pm$ .0005	$-0.0034$
770M	2.76	4.37	1.58	$\mathbf{-0.0224}$	$\mathbf{-0.0030}$ $\pm$ .0002	$-0.0053$
$\varepsilon=0.95$

Appendix

Table 24: The per-step

\beta{=}0.1

model under a batched real decode: the full split,

32

sequences per pass behind a

512

-token prompt, three seeds. Pass depth : mean iterations per pass,

R\,\bar{N}

measured. Margins against the sampled-depth baseline on the same sequences; TF : teacher-forced at the pass depth.

\pm

: seed s.d. (at most

0.0033

\bar{N}

0.0017

for TF); bold (

\bar{N}

and pass-depth columns): over twice it.

Model	Variant	Steps	Change vs teacher-forced	Condition
Ouro 2.6B	teacher-forced	2.444	—	$\mathbf{-0.300}$
	keep-first	2.475	$+0.111$	$\mathbf{-0.176}$
	drop	2.943	$+0.264$	$+0.181$
Ouro 1.4B	teacher-forced	2.669	—	$\mathbf{-0.278}$
	keep-first	2.654	$+0.075$	$\mathbf{-0.211}$

Appendix

Table 25: Ouro at

\tau_{\text{Ouro}}=0.05

and one token per pass: the quality condition, the model minus the depth curve at the realized steps (negative: the model passes). Keep-first : the first token keeps its deeper keys and values; drop : it hides them like any exited token. On the 1.4B model drop and keep-first coincide.

			emulated, full split			real, $64$ seq.
Scale	$\varepsilon$	TF	hidden	copy	reuse	reuse
140M	0.7	$-0.021$	$+0.043$	$+0.015$	$-0.019$	$-0.016$
	0.95	$-0.050$	$+0.028$	$+0.012$	$-0.039$	$-0.044$
370M	0.7	$-0.022$	$+0.035$	$+0.098$	$-0.020$	$-0.019$
	0.95	$-0.059$	$+0.013$	$+0.209$	$-0.050$	$-0.050$
770M	0.7	$-0.025$	$+0.028$	$+0.237$	$-0.022$	$-0.025$

Appendix

Table 26: Filling a stopped token’s deeper keys and values: the per-step

\beta{=}0.1

model minus the sampled-depth baseline at the model’s own realized steps, after a

512

-token prompt, three seeds. TF : teacher-forced. Seed s.d. at most

0.0087

(hidden),

0.098

(copy),

0.0039

(reuse). Margins are in BPB; a positive margin means the model is worse than the baseline.

Model	Readout	Steps	BPB $\downarrow$	Baseline	$\Delta$ $\downarrow$
Sampled rule: realized steps $=\mathbb{E}[N]$ by construction
140M mixture-loss	mixture	3.830	1.4633	1.2291	$+0.2342$
140M mixture-loss	exit state	3.830	1.5262	1.2291	$+0.2971$
370M mixture-loss	mixture	3.762	1.2966	1.0990	$+0.1976$
370M mixture-loss	exit state	3.762	1.3439	1.0990	$+0.2449$
140M per-step, $\beta{=}0.1$	exit state	4.335	1.2198	1.2244	$-0.0046$

Appendix

Table 27: What realizing the average costs. PonderNet’s sampled rule, whose realized steps equal

\mathbb{E}[N]

by construction, against the sampled-depth baseline at the same realized steps (“Baseline”); the threshold rule at 140M is included for contrast. BPB; three seeds, three random draws of the stopping iteration per seed. Variation across seeds is

0.0070

–

0.0149

on the mixture-loss models and

0.0003

–

0.0008

on the per-step models, against the mixture-loss effects of

0.198

and above in the upper block.

$\tau_{p}$	$\varepsilon_{\text{ACT}}$	Steps	BPB $\downarrow$	Baseline
$10^{-4}$ (reported) ‡	$[0,0.5]$	7.99	1.2027	1.2159
$3\times 10^{-3}$ ‡	$[0,0.3]$	7.95–7.98	1.2032	1.2159
$3\times 10^{-3}$ ‡	0.5	7.862	1.2049	1.2160
$10^{-2}$ ‡	0	7.352	1.2048	1.2163
$10^{-2}$ ‡	0.1	7.182	1.2049	1.2164
$10^{-2}$ ‡	0.2	6.977	1.2056	1.2166

Appendix

Table 28: ACT’s ponder weight selects a depth, not a policy. Realized steps and BPB at 140M on CC as the ponder weight

\tau_{p}

and the eval-time threshold

\varepsilon_{\text{ACT}}

vary; “Baseline” is the sampled-depth baseline at the same realized steps. ‡ Three seeds, Baseline from the three-seed curve; other rows one seed. A row spanning an

\varepsilon_{\text{ACT}}

interval moves by at most

0.03

steps and

0.0001

BPB within it.

140M, BPB at forced $T=1$
Loss	$\beta{=}0.01$	$\beta{=}0.1$	$\beta{=}1.0$
Mixture	1.8127 $\pm$ .017	1.8508 $\pm$ .044 $(+0.038)$	1.8648 $\pm$ .044 $(+0.052)$
Per-step	1.7111 $\pm$ .003	1.3166 $\pm$ .004 $\bf(-0.394)$	1.2608 $\pm$ .000 $\bf(-0.450)$
370M, BPB at forced $T=1$
Mixture	1.6104 $\pm$ .011	1.6083 $\pm$ .027 $(-0.002)$	1.6415 $\pm$ .026 $(+0.031)$
Per-step	1.4302 $\pm$ .011	1.1576 $\pm$ .001 $\bf(-0.273)$	1.1225 $\bf(-0.308)$

Appendix

Table 29: Forced-depth BPB at

T{=}1

under both losses as the prior weight varies — the cells underlying Table 4 (Section 6.2 ). Parenthesized values are the change from

\beta=0.01

. 140M is three seeds per cell. At 370M every printed change is the difference of the two printed means, so the

\beta{=}1.0

change compares one seed against a three-seed base; the within-seed figure is in the text. Bold : beyond twice the seed s.d.

Model	$T{=}1$	$T{=}2$	$T{=}3$	$T{=}4$	$T{=}5$	$T{=}6$	$T{=}8$
Sampled-depth baseline	1.3818	1.2779	1.2415	1.2266	1.2201	1.2174	1.2155
Per-step, $\beta{=}0.01$	1.7139	1.5612	1.5015	1.4544	1.4000	1.3327	1.2039
Per-step, $\beta{=}0.1$	1.3169	1.2456	1.2250	1.2185	1.2164	1.2157	1.2154
Per-step, $\beta{=}1.0$	1.2608	1.2344	1.2278	1.2261	1.2256	1.2255	1.2255
$\Delta$ , $\beta{=}0.01$	$+0.3321$	$+0.2833$	$+0.2600$	$+0.2278$	$+0.1799$	$+0.1153$	$\bf-0.0116$
$\Delta$ , $\beta{=}0.1$	$\bf-0.0649$	$\bf-0.0323$	$\bf-0.0165$	$-0.0081$	$-0.0037$	$-0.0017$	$-0.0001$

Appendix

Table 30: Forced-depth BPB at 140M on CC as the prior weight

\beta

varies under the per-step loss, and the difference against the sampled-depth baseline at each depth. Weights are unchanged between columns; only the forced depth

T

changes. Seed

42

for every row except

\beta{=}1.0

(three-seed means) and the baseline (three-seed mean). Bold : the model ahead by more than twice the evaluation noise floor, the criterion since most rows are single-seed. The resolved band at

\beta{=}0.1

T\leq 3

Model	$T{=}1$	$T{=}3$	$T{=}5$	$T{=}8$
Sampled-depth baseline	1.2214	1.1079	1.0912	1.0875
Per-step, $\beta{=}0.01$	1.4302	1.2505	1.1817	1.0924
Per-step, $\beta{=}0.1$	1.1576	1.0935	1.0880	1.0874
Per-step, $\beta{=}1.0$	1.1225	1.0955	1.0940	1.0939
Mixture, $\beta{=}0.01$	1.6104	1.2765	1.1844	1.0807
Mixture, $\beta{=}0.1$	1.6083	1.2764	1.1588	1.0813

Appendix

Table 31: Forced-depth BPB at 370M under both losses and three prior weights. Model rows are three seeds, except per-step

\beta{=}1.0

(seed

42

); the baseline is seed

42

Model	$T{=}1$	$T{=}2$	$T{=}3$	$T{=}4$	$T{=}8$
Sampled-depth	–	0.9439	–	0.9122	0.9054
Mixture $0.01$	1.7119	1.1462	1.0577	1.0155	0.8997
Mixture $0.1$	1.3168 $(-.395)$	1.1487 $(+.003)$	1.0601 $(+.002)$	1.0072 $(-.008)$	0.8995 $(-.000)$
Per-step $0.01$	1.3446	1.2675	1.2464	1.2241	0.8969
Per-step $0.1$	0.9766 $(-.368)$	0.9263 $(-.341)$	0.9108 $(-.336)$	0.9057 $(-.318)$	0.9035 $(+.007)$

Appendix

Table 32: Loss

\times

prior weight at 140M on FW, forced-depth BPB, seed

42

. Row labels give

\beta

. Parenthesized values are the change from

\beta=0.01

at that depth. At

T=8

the changes are at most

0.007

Seed	Per-step	Specialist $T{=}8$	$\Delta$
42	1.2039	1.2047	$-0.0008$
7	1.2056	1.2035	$+0.0021$
123	1.2055	1.2065	$-0.0010$
Mean	1.2050 $\pm$ .0010	1.2049 $\pm$ .0015	$+0.0001$ $\pm$ .0017

Appendix

Table 33: The per-step model (

\beta{=}0.01

) against a specialist trained at

T=8

, 140M, three seeds each, same data and schedule.

\Delta

is the model minus the specialist (negative favors the model); mean

\pm

s.d. over seeds.

		Quintile of loss at the cap
Scale		1	2	3	4	5
140M	step	2.28	3.08	3.33	3.40	3.48
	improvement	0.06	0.29	0.34	0.24	0.12
370M	step	2.06	2.69	3.06	3.11	3.21
	improvement	0.01	0.15	0.28	0.20	0.09
770M	step	2.03	2.54	2.98	3.06	3.16

Appendix

Table 34: Where the per-step model (

\beta{=}0.1

\varepsilon{=}0.7

) spends depth (quintile 1: easiest). Step : realized steps. Improvement : the loss drop from one iteration to the cap, in nats per token. CC; three-seed means.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,080 papers

Sep 30, 2026

Closing the Loop: Practical Training Recipes for Looped Language Models

Looped language models increase effective depth by repeatedly applying a shared block of layers, but existing large-scale recipes require multi-stage training over trillions of tokens, while the benefits of recurrence remain difficult to separate from differences in data and training. In this work, we establish practical training recipes for looped language models, with three main results. (1) We develop a compute-efficient from-scratch pipeline that reduces the training budget from 7.7T tokens in Ouro to 310B tokens while retaining strong reasoning performance. Pretraining followed by high-quality mid-training, together with learning-rate warmup and stronger exit-gate regularization, enables stable recurrent training without prior multi-stage schedules. (2) Under controlled comparisons, our 1.4B LoopLM outperforms a parameter-matched dense model trained on the same data and token budget on all 12 evaluated benchmarks, including +14 points on GSM8K, +10 on MATH, and +22 on DROP. At matched inference compute, it approaches a 3.9B dense model on mathematical reasoning and reading comprehension while using only 36% as many parameters. (3) We introduce a minimal recipe for converting pretrained dense models into looped ones: a single learned input-mixing scalar and a smoothed exit loss, with no step-specific parameters. Applied to Qwen3-1.7B-Base, Looped Qwen improves over an identically continued dense baseline on every evaluated benchmark across two data regimes, with statistically clear gains on GSM8K, MATH, and MMLU-Pro on the curated mixture. Together, these results make looped language models substantially cheaper to train from scratch and practical to introduce into existing pretrained checkpoints, while isolating the gains due to recurrence itself.

Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin +5

RND NLP, DAIMLD, Russian Federation

Efficient Language Model Training LLM Training Recursive Language Models Recurrent Language Models

Ask PDF

Oct 6, 2026

Hybrid Latent Attention for Looped Language Models

Looped language models apply the same stack of layers T times to each token, which deepens the model without adding parameters but multiplies its key-value (KV) cache by T. The larger cache limits how many sequences a GPU can decode at once and slows each decoding step, which reads the whole cache. We propose Hybrid Latent Attention (HLA), which keeps exact keys and values within a sliding window of W recent tokens and stores each older token as a compact latent that the query of each loop reads directly, without reconstructing keys and values. We uptrain HLA on Ouro looped models (T=4) with 1.4B and 2.6B parameters, keeping the pretrained weights frozen and training only the added parameters to reproduce the original attention. The cache shrinks by 10.7x per token, fitting 4.0-8.8x as many concurrent sequences per GPU, and decoding throughput improves by 2.5x at 1K-token contexts and by up to 7.4x at 16K. HLA retains over 97% of the original accuracy on math, knowledge and reasoning benchmarks, and 96-100% on long-context retrieval up to 16K tokens. After supervised fine-tuning, it performs on par with the fine-tuned original model on competition-level math.

Yuhan Chen, Siyuan Zhang, Nan Wang +2

Virginia Tech · Independent Researcher

Long-Context Language Model Inference LLM Inference Acceleration KV-Cache Compression Hybrid Attention Recursive Language Models

Ask PDF

Oct 19, 2026

Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation

Looped transformers reuse a shared block of layers to increase computational depth without increasing the number of trainable parameters. They are trained with backpropagation through time (BPTT). To save training compute, BPTT is often truncated to the last few loops, but truncation yields biased gradients that may weaken the benefit of looping. In this work, we study the value of truncation in looped transformer training in a systematic way: we jointly vary model size, training tokens, and forward and backward loop counts across nearly 600 training runs and fit a unified scaling law that captures the effects of looping and truncation through effective parameters. Despite the training-compute savings of truncation, our fitted law favors full BPTT when forward and backward loop counts are jointly optimized. In matched-compute runs, full BPTT through fewer loops consistently outperforms truncated settings of the same compute. The law further shows that at fixed model size, the compute-optimal loop count and training-token budget both grow with compute. When model size is also free, the law identifies two fully trained loops as the best looped setting. Experimentally, this setting achieves lower validation loss than the compute-optimal non-looped model with about 27% fewer non-embedding parameters and only 8% more training compute.

Xuyuan Xiong, Siyu Chen, Zihan Shao +5

MMLab, The University of Hong Kong · Yale University

Truncated Backpropagation Through Time Efficient Language Model Training Recurrent Transformers Language Model Scaling Laws Recurrent Neural Networks

Ask PDF

Abstract

LLM Inference Acceleration

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models

Hybrid Latent Attention for Looped Language Models

Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation