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

Xuyuan Xiong; Siyu Chen; Zihan Shao; Yue Su; Weiyang Jin; Ping Luo; Zhuoran Yang; Xihui Liu

cs.LGOct 19, 2026

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

Authors: Xuyuan Xiong, Siyu Chen, Zihan Shao, Yue Su, Weiyang Jin, Ping Luo, Zhuoran Yang, Xihui Liu

Organizations: MMLab, The University of Hong Kong · Yale University

Abstract

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.

Figures & tables

Figure 1: Overview of looped transformer scaling with truncated BPTT . (a) Step for

(T,B)=(4,2)

(

T

forward loops, last

B

trained). (b) At

8D_{0}

tokens, loops add effective parameters while truncation discounts them ( Section 4 ; lines: law, markers: runs). (c) At fixed size (120M), optimal

T

grows with compute under full BPTT ( Section 6.1 ). (d)

(2,2)

(344M) beats

(1,1)

(471M,

2.1\times 10^{19}

FLOPs) in loss (

2.36

vs.

2.37

nats) with

27\%

fewer parameters and

8\%

more compute ( Section 6.3 ).

Figure 2: No looping is computationally optimal; two fully trained loops saves most parameters with little extra compute. Each marker is a loop setting

(T,B)

, with

T

forward loops and the last

B

trained, at the parameter and token counts our scaling law finds compute-optimal. Compared with the model without looping at the same loss, the horizontal axis shows the training compute it needs and the vertical axis the fraction of non-embedding parameters it saves. Filled markers use full BPTT (

B=T

), and the dashed curve is the law’s prediction for full BPTT at any

T

from 2 to 16. Every truncated setting lies below this curve, so full BPTT forms the frontier and saves the most parameters at each compute. Dotted lines join the settings with the same

T

( Section 6.2 ).

		Scaling axes
Study	Arch.	$N$	$D$	$T$	$B$	Loop gain	BPTT	Regime
Schwethelm et al. (2026)	dense	✓	✓	✓	✗	power law ( $T^{\phi}$ )	$T$ ; $\lceil T/2\rceil$	fixed depth
Parcae ( Prairie et al., 2026 )	dense	✓	✓	✓	✗	linear ( $T$ )	$\lceil T/2\rceil$	fixed $N$
Chen et al. (2026b)	dense, MoE	✓	✓	✓	✗	bounded exp. ( $T$ )	full ( $T$ )	free $N$ , mem.
SMELT ( Wang et al., 2026b )	MoE	✓	✓	✗	✗	per-arch. fit	/	matched $N$ , KV
Lee et al. (2026)	dense, MoE	✓	✗	✗	✗	per-arch. fit	/	free $N$

Table 1: Scaling laws for looped language models. Earlier laws either fit a separate curve per architecture or model looping through a single loop count

T

, with backward loops

B

fixed or tied to

T

; ours models

T

and

B

independently with a bounded effective-parameter gain. /: not discussed; † appendix scaling fit.

Figure 3: The looped transformer, drawn for

(T,B)=(4,2)

. (a) One training step: a stop gradient on

s_{T-B}

, so the first

T-B

loops (dashed) run forward only. (b) One loop, ( 3.1 ); all loops share its weights. (c) Training FLOPs per parameter and token of each part ( 5.1 ).

Table 5

Figure 4: Training curves of the scaling suite for no looping, $(2,1)$ , and $(2,2)$ . Each curve is one run, colored by model size, against training compute ( Section D.1 ); each token budget is a separate run. Training loss of every tenth step, averaged over at most $\pm 250$ steps, not the validation loss we report. Highlighted line and dots: the frontier of run ends, the runs that end below every run with less compute. Appendix H shows all fifteen loop settings.

$E$	$A$	$H$	$\alpha$	$\beta$	$\kappa$	$c$	$p$	$q$
0.29865	6.48695	7.63343	0.30568	0.39281	0.32510	1.11198	0.23677	1.02181

Table 4: Fitted parameters on the 230-configuration fitting set. The reference scales are fixed at

N_{0}=2.4\times 10^{8}

parameters and

D_{0}=2^{30}

tokens.

Figure 7: More loops with full BPTT lower the loss at every size and budget and one law tracks all of them. Validation loss against training tokens for the four loop settings with full BPTT . Dots: runs, open where outside the fitting set; lines: the full-data fitted law ( Section 4 ).

Figure 8: Forward and backward loops at fixed size and budget (240M,

2D_{0}

). Markers show measured runs and solid lines show the fitted law. (a)

B=T

: the gain levels off; dashed, the power-law gain of prior work ( Schwethelm et al., 2026 ) keeps growing. (b) Fixed

T

: training more of the loops lowers the loss. (c) Fixed

B

: more forward loops help little. Gains level off, and training more loops (

B

) lowers the loss more than running more loops (

T

Figure 9: What a loop setting is worth in effective size. (a) Full BPTT , the looping coefficient

g(T,T)

, approaches its fitted ceiling (dotted); dashed beyond the fitted loop counts. (b) The looping coefficient

g(T,B)

against the trained fraction

B/T

. (c) Total effective size

N_{\mathrm{eff}}/N

, grouped by

B

. Lines, bars and numbers: the law; markers and error bars: run estimates, median and interquartile range over sizes and budgets ( Section C.2 ). Effective size is set mainly by

B

: within each group of equal

B

, neither the law nor the runs gain much from more forward loops.

$T$	$B^{*}$	$\chi(T,B^{*})$	$\chi(T,T)$	$N^{*}$	$D^{*}$
1	1	1.000	1.000	1	1
2	2	1.031	1.031	1	1
4	2	1.204	1.430	1.27	1.08
8	2	1.580	2.376	1.51	1.20
16	2	2.343	4.339	1.70	1.31

Table 5: With

T

prescribed, the law selects

B^{*}=2

and so truncates from

T=4

on. For each forward-loop count

T

B^{*}

is the backward-loop count that minimizes

\chi(T,B)=N_{\mathrm{comp}}/N_{\mathrm{eff}}

, the compute per effective parameter, over the integers

1\leq B\leq T

, with

N_{\mathrm{rec}}/N\simeq 0.507

and

N,D

optimized under leading FLOPs.

N^{*}

and

D^{*}

are the optimal model size and token count at

B^{*}

relative to full BPTT at the same

T

and compute

C

. These ratios do not depend on

C

. The gap between

\chi(T,B^{*})

and

\chi(T,T)

, and hence the compute that truncation saves, grows with

T

Figure 10: Separately fitted laws confirm the predicted truncation. Iso-loss contours (labels: loss in nats) and compute-optimal frontiers of laws fitted separately to

B=T

(dashed, grey) and to the predicted

B^{*}=2

(solid), at

T=4

(a) and

T=8

(b). Markers are measured runs; faded segments lie outside the runs of a setting. At every loss level, the

B^{*}=2

contour lies left of the

B=T

contour (shaded gap), so truncation reaches the same loss with less compute: 5–11% less at

T=4

and 33–35% less at

T=8

, where both compute-optimal points lie within the runs. The saving grows with

T

because further backpropagated loops add little effective size at full cost. The compute that truncation frees goes partly into a larger model.

Table 7: Fewer fully trained loops improve the compute-optimal frontier. Iso-loss contours of separately fitted laws for

(T,B)=(4,4)

(solid) and

(8,2)

(dashed), both with

T+2B=12

. Thick lines: compute-optimal frontiers; faded segments lie outside the runs of a setting. At equal loss, solid contours lie left of dashed ones:

(4,4)

uses 8–14% less compute, so once

T

can be chosen, full BPTT beats truncation.

Figure 11: How the compute-optimal allocation scales with training compute. Top: fixed model size (240M). (a) Loss against the loop count at five computes, with the optimum at each compute marked by a star. (b, c) The optimal loop count

T^{*}

and token count

D^{*}

of the stars, and the law’s optimum (line), whose exponents approach ( 6.4 ). Bottom: free model size. (d) Loss against model size at four computes for

(1,1)

(dashed) and

(2,2)

(solid), with their optima (stars) and the compute-optimal pairs of Section 6.3 (diamonds). (e, f)

N^{*}

and

D^{*}

of the stars and the law’s optimum ( 6.6 ), which places the optimum of

(1,1)

at a larger model than its runs do. Colours mark compute within each row, and compute is

6N_{\mathrm{comp}}D

( 5.1 ). At fixed size, added compute goes mostly to tokens while the loop count grows slowly. With free size,

(2,2)

tracks

(1,1)

at a smaller model.

Figure 12: At a fixed model size, the best loop count grows with training compute: no looping, then two, then four trained loops. Loss minus that of the model without looping at the same compute, which spends it on more tokens; markers are runs, with a smooth curve through each setting;

(4,2)

trains half of its four loops. The bold curve is the best setting, named in the strip by its loop count

T^{*}

(faded where extrapolated). Within each stretch the best model spends added compute on tokens, and at each switch it trades tokens for loops.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Evaluation	Fits	RMSE	MAE	Bias	Max. error
In-sample	1	$6.14\times 10^{-3}$	$4.91\times 10^{-3}$	$-3.73\times 10^{-16}$	$2.04\times 10^{-2}$
One configuration out	230	$6.40\times 10^{-3}$	$5.11\times 10^{-3}$	$-3.30\times 10^{-6}$	$2.10\times 10^{-2}$
One model size $N$ out	7	$7.48\times 10^{-3}$	$5.89\times 10^{-3}$	$-2.16\times 10^{-4}$	$2.58\times 10^{-2}$
One token budget $D$ out	11	$6.42\times 10^{-3}$	$5.12\times 10^{-3}$	$-3.42\times 10^{-5}$	$2.08\times 10^{-2}$
One $(T,B)$ setting out	15	$8.67\times 10^{-3}$	$6.74\times 10^{-3}$	$8.86\times 10^{-4}$	$2.63\times 10^{-2}$

Appendix

Table 8: Pooled errors in nats/token over all 230 configurations; all 264 fits converge (one at a bound).

Table 19

Best setting changes	60M	120M	240M	480M	1B
$(1,1)\to(2,2)$	16	17	16	15	14
$(2,2)\to(4,4)$	400	413	374 †	345 †	312 †
$(4,4)\to(8,8)$	9,622 †	9,648 †	8,605 †	7,833 †	7,025 †
$(8,8)\to(16,16)$	210,087 ‡	204,851 ‡	182,247 ‡	164,533 ‡	146,710 ‡

Appendix

Table 11: Model size fixed: the training compute at which the best loop setting changes, given as the tokens per parameter that the model without looping trains on at that compute. Among all fifteen trained settings, four of them with

T=16

, only full BPTT reaches the frontier, apart from

(4,3)

, which leads briefly between

(2,2)

and

(4,4)

by at most 0.002 nats. † Beyond the compute of the runs of that size. ‡ Involves

T=16

, trained at three sizes and short budgets only.

Figure 15: Truncation can be favored at a prescribed

T

, but not when

T

and

B

are chosen jointly. The law’s

\chi=N_{\mathrm{comp}}/N_{\mathrm{eff}}

; lower

\chi

means lower compute-optimal loss. (a) Stars: the best

B

at each

T

. (b) Percentages: compute saved by the best

B

over

B=T

at equal loss.

	$(1,1)$		$(2,2)$
Compute (FLOPs)	Parameters	Loss	Parameters	Loss	Gap
$1.6\times 10^{18}$	115M	2.7098	79M	2.7177	$(7.9\pm 1.9)\times 10^{-3}$
$6.0\times 10^{18}$	239M	2.5215	162M	2.5253	$(3.8\pm 1.8)\times 10^{-3}$
$2.1\times 10^{19}$	471M	2.3663	344M	2.3677	$(1.4\pm 1.7)\times 10^{-3}$
$2.2\times 10^{19}$ (+8%)	471M	2.3663	344M	2.3590	$(-7.3\pm 1.8)\times 10^{-3}$

Appendix

Table 12: Compute-optimal pairs. Each model at the size that minimizes its loss for the compute; loss: mean of three seeds; gap:

(2,2)

minus

(1,1)

\pm

standard error; parameters: non-embedding counts (the models without looping are the 120M, 240M and 480M models of Table 3 ). Compute includes attention, output, and optimizer updates ( D.3 ); the first three rows are matched in the compute of the forward and backward passes (

1.5\times 10^{18}

5.9\times 10^{18}

, and

2.0\times 10^{19}

FLOPs, or

9.6\times 10^{17}

4.1\times 10^{18}

, and

1.5\times 10^{19}

under the leading-FLOP count

6N_{\mathrm{comp}}D

of Figure 11 d; counting the optimizer update,

(1,1)

spends 0.7% more, which corresponds to about

10^{-3}

nats); the last row trains the 344M

(2,2)

model on 8% more tokens (+8% compute) against the

2.1\times 10^{19}

FLOP

(1,1)

reference.

Size	Model	Width	Parameters	Compute	Loss	$-$ untied
60M	untied	768 (6 $\times$ 128)	87M	1.00	2.7843
	$(2,2)$	768 (6 $\times$ 128)	58M	1.00	2.7931	$(+8.8\pm 3.0)\times 10^{-3}$
	$(2,2)$ , wider	840 (7 $\times$ 120)	69M	1.15	2.7668	$(-1.75\pm 0.19)\times 10^{-2}$
120M	untied	1024 (8 $\times$ 128)	180M	1.00	2.5809
	$(2,2)$	1024 (8 $\times$ 128)	115M	1.00	2.5982	$(+1.73\pm 0.14)\times 10^{-2}$
	$(2,2)$ , wider	1152 (9 $\times$ 128)	146M	1.22	2.5681	$(-1.28\pm 0.23)\times 10^{-2}$

Appendix

Table 13: Comparison with the untied model. At each size: the untied model,

(2,2)

at its width, and

(2,2)

widened to about 80% of its non-embedding parameters. Width with heads

\times

head dimension; parameters: non-embedding; compute: compute per token of the forward and backward passes relative to the untied model, the same ratio as the training compute since all models of a size train on the same tokens; loss: mean of three seeds;

-

untied: paired by seed, mean

\pm

standard error.

Budget $D/D_{0}$	$\tfrac{1}{2}$	1	2	4	8	16	32
Batch (sequences of 2,048 tokens)	48	64	88	128	176	256	352
Steps	5,464	8,192	11,920	16,384	23,832	32,768	47,664
Peak learning rate ( $\times 10^{-3}$ )	10.7	9.25	8.11	8.01	7.03	6.36	6.09

Appendix

Table 14: Batch size, number of steps and peak MuonH learning rate of the hidden weight matrices at each budget. The learning rate is ( F.1 ) with

\gamma=0.40

, with its ratio to the

\gamma=0.32

value rounded to a power of

2^{1/8}

; the AdamW learning rates are multiplied by the same ratio. Warmup takes 2% of the steps.

Figure 16: Learning-rate rule (( F.1 ),

\gamma=0.40

) against the measured optimum. Colors give the model size; filled markers:

(4,4)

, open markers:

(4,1)

. Solid gray lines: the rule; dashed lines: the alternative it is tested against; shading: the tolerance (0.25%). (a) Loss above the fitted optimum, in percent of the optimum loss, against the learning rate relative to the rule, for the ten token-exponent sweeps, with one quadratic in

\log\eta

per sweep. (b) Fitted optima relative to ( F.1 ) at

\gamma=0.32

, with the learning rates within the tolerance (bars); dashed:

\gamma=0.32

. (c) As (a), at

D_{0}

, for four sizes with

(4,4)

and for 60M with

(2,2)

and

(8,8)

. (d) Optima of these sweeps against the number of executed layers; dashed: the trend HyperP measures without residual scaling (the optimum halves from 8 to 24 layers), drawn through the rule at 20 layers.

Figure 17: Batch rule ( F.2 ) against the measured optimum. Colors give the model size. (a) Loss-optimal batch for

(4,4)

, the vertex of a quadratic in log batch fitted to each sweep, with the batches within the tolerance (bars). Lines: the rule (solid) and one ladder step lower (dashed). (b) Loss above the best batch, in percent of the best loss, at up to two ladder steps from the rule, over the 33 sweeps with at least three batch sizes (filled:

B=T

; open:

B<T

;

n

: sweeps run at that batch). Thick line: worst sweep; thin line: mean; shading: the tolerance.

Figure 18: The law across the scaling suite. Validation loss of every configuration in the five-size grid against training tokens (

D_{0}=2^{30}

), one panel per loop setting

(T,B)

with a row per forward-loop count, colored by model size. Filled dots: configurations the law is fitted on; open dots: the other configurations; lines: the law ( Section 4 ) over each size’s measured budgets.

Figure 19: Training curves of every run in the five-size grid. Training loss against compute ( Section D.1 ), one panel per loop setting, with a row per forward-loop count and backward loops increasing to the right, colored by model size; every tenth step, averaged over at most

\pm 250

steps. Highlighted line and dots: the frontier of run ends, the runs that end below every run with less compute.

Truncated Backpropagation Through Time

See all 8 papers

Oct 20, 2026

How Much Gradient Does Truncated Backprop Lose in Looped Transformers?

Looped transformers reuse one block of layers many times and are often trained by backpropagating through only the last k loops. We measure how much of the weight gradient this throws away, exactly and for every k from one backward pass, on nine looped language models, and report three findings. First, on a trained model this lost gradient is set near the output: once the loop count is well above k, it barely changes with the loop count (within about 0.03 from 32 to 64 loops on the three models we tested, as a pre-registered test predicted for Parcae). Where the backward signal shrinks,

ρ^k

, with

ρ

the signal's fitted shrink per loop, approximates it at moderate loop counts but overstates it at long ones, because the fitted

ρ

drifts up as the fit takes in more loops. On Ouro the signal grows, and keeping two of four loops still loses 81% of the final output's gradient. Second, the lost gradient falls during fine-tuning, by 0.07-0.22 within 300 steps on three of four models, on held-out problems too. The fall is not just the read-out's share of the gradient; training only the layers outside the loop reproduced up to half of it, and it also falls under truncated training. So a trained model's value overstates what truncation removes later. Third, we could not show that it predicts the training cost of truncation better than rules that need no measurement: on two new models it tied Parcae's cost at the same k, in a test too small to separate them.

Aur Shalev Merin

Independent researcher

Truncated Backpropagation Through Time Backpropagation Recurrent Transformers Transformer Training Dynamics

Ask PDF

Sep 28, 2026

Improving Test-Time Scaling with Adaptive Looped Transformers

Looped transformers have demonstrated promising parameter efficiency by reusing layers for latent computation. Prior studies compare looped and non-looped models at matched parameters or per-token FLOPs. However, to the best of our knowledge, whether looping improves test-time scaling as outputs grow longer remains underexplored. Through post-training looped transformers, we study the accuracy-compute slope, measured as the accuracy gain per doubling of test-time decoding FLOPs. We find that existing looped transformers often yield steeper slopes than their non-looped baseline, yet underperform it at matched compute. While fixed-depth looping spends extra iterations on every token, our analysis shows that many tokens do not benefit from extra iterations. We therefore propose TaH2, which enables the model to focus extra iterations on the tokens that benefit from looping. It jointly post-trains the backbone and an iteration decider through lookahead depth supervision, which uses online labels indicating whether further iteration improves the prediction. TaH2 improves both the efficiency and attainable accuracy of test-time scaling. On challenging AIME benchmarks, TaH2 improves the accuracy-compute slope by 53% (2.74 vs. 1.79) over the non-looped baseline, exceeding the baseline's peak accuracy by about 3.4 points at matched test-time compute. As the maximum iteration depth increases, existing looped models largely plateau, while TaH2's gain over the non-looped baseline continues to grow from +2.8 points at depth 2 to +3.9 points at depth 8. Our code is available at https://github.com/thu-nics/TaH.

Yichen You, Tianyu Fu, Aosong Feng +4

Tsinghua University · Yale University

Test-Time Scaling Recurrent Transformers LLM Inference Acceleration Adaptive Computation

Ask PDF

Jun 16, 2026

Loop- $μ$ P: Residual Scaling and Hyperparameter Transfer in Looped Transformers

Looped Transformers have drawn growing attention because looping improves multi-step reasoning and increases the executed depth of a model without adding parameters. They achieve this by reusing a stack of

L

layers for

N

loops. However, increasing the loop count

N

not only destabilizes training but also shifts the best hyperparameters, so they must be searched again for every loop count. Inspired by Depth-

μ

P (Yang et al., 2024), which multiplies each residual branch by

L^{-1/2}

for depth

L

so that a learning rate tuned on a shallow untied network transfers to a deep one, we study how the residual branches and the learning rate should scale with the loop count

N

. We find that loops cannot simply be treated as extra depth, because the loops share weights. At initialization, distinct layers add nearly uncorrelated increments to the residual stream, so their sum grows as

\sqrt{L}

, whereas a shared layer adds positively correlated increments at every loop, so their sum grows linearly in

N

. For a single layer looped

N

times, we show in the infinite-width limit that multiplying the residual branch by

1/N

and keeping the learning rate independent of

N

is the only parameterization with both a stable forward pass and maximal residual-stream updates. We then generalize the analysis to a stack of

L

distinct layers looped

N

times to obtain a parameterization we call Loop-

μ

P, in which the loop factor

1/N

combines with the depth factor

L^{-1/2}

. Empirically, we find in language modeling that the base learning rate tuned on the smallest model stays best under Loop-

μ

P for loop counts up to

8

and executed depths up to

384

layers. These results make Loop-

μ

P a practical recipe for scaling looped Transformers.

Shaowen Wang, Bingrui Li, Tansheng Zhu +4

Tsinghua University · ByteDance Seed · M-A-P

Residual Learning Hyperparameter Transfer Recurrent Transformers Neural Network Training Dynamics Transformer Training Dynamics

Ask PDF

Abstract

Truncated Backpropagation Through Time

Explore similar work

How Much Gradient Does Truncated Backprop Lose in Looped Transformers?

Improving Test-Time Scaling with Adaptive Looped Transformers

Loop-μμμP: Residual Scaling and Hyperparameter Transfer in Looped Transformers

Loop- $μ$ P: Residual Scaling and Hyperparameter Transfer in Looped Transformers