PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.26030v2cs.LGstat.ML

TL;DR

Looped-transformer training must balance the savings of truncated backpropagation against the loop gains it discounts, but evidence on jointly choosing forward and backward loops has been limited. The paper varies model size, training tokens, and both loop counts across nearly 600 runs and fits a unified scaling law. The law favors full BPTT when both counts are optimized, while two fully trained loops achieve lower validation loss than the compute-optimal non-looped model with fewer parameters and slightly more compute.

  • Problem

    Prior studies examine prescribed backpropagation recipes rather than jointly modeling forward and backward loop counts, leaving their joint choice across model sizes and token budgets unresolved.

    Source: Introduction · p. 2

  • Method

    The authors vary model size, training tokens, and forward and backward loop counts across nearly 600 runs, then fit a unified scaling law using effective parameters.

    Source: Introduction · pp. 2–3

  • Results

    When both loop counts are chosen, the fitted law favors full BPTT; with 27% fewer parameters, two fully trained loops achieve lower loss with 8% more compute than the compute-optimal non-looped model.

    Source: Conclusion · p. 18

  • Takeaways & Limitations

    At fixed model size, most additional compute should go to data while fully trained loop count grows only about C^0.24; with model size free, two fully trained loops come close to the compute-optimal non-looped model.

    Source: Conclusion · p. 18

    The study measures validation loss only, does not establish downstream-task or reasoning gains, and caps experiments at T ≤ 16, leaving deeper-loop behavior extrapolative.

    Source: Conclusion · p. 18

Abstract

from arXiv · show

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.

Introduction

The paper asks how model size, training tokens, and forward and backward loops should be allocated under a fixed compute budget, and whether BPTT should be truncated. It fits a unified scaling law and finds that full BPTT is favored when both loop counts can be chosen, while two fully trained loops offer a parameter-efficient setting when model size is free.

Introduction

  • Existing scaling studies fit prescribed backpropagation recipes or leave joint forward- and backward-loop optimization open.The paper frames its question around jointly choosing model size, training tokens, and both loop counts under fixed compute.Source: Introduction · p. 2
  • When both loop counts can be chosen, the fitted law favors full BPTT: reducing forward loops frees compute for backward loops until all remaining loops are trained.At a prescribed forward-loop count, truncation can instead pay because saved backward compute supports a larger model and more training tokens.Source: Introduction · p. 2
  • At fixed model size, the law allocates most additional compute to training tokens and increases loop count only slowly.When model size is free, the law identifies two fully trained loops as the best looped setting.Source: Introduction · pp. 2–3
  • At equal training compute, models with two fully trained loops nearly match compute-optimal non-looped models while using substantially fewer non-embedding parameters; with a little more training compute, the smaller model surpasses them.Source: Introduction · p. 3
  • The unified scaling law expresses loop gains and truncation discounts through effective parameters while jointly modeling model size, training tokens, and both loop counts.The study varies these quantities across nearly 600 training runs.Source: Introduction · pp. 2–3

Related work

Prior work establishes looping as a way to improve parameter efficiency and commonly truncates backpropagation to save compute. This study instead separates forward and backward loop counts to quantify truncation’s trade-off and fit a unified scaling law.

  • Existing studies show looping improves parameter efficiency across architectures, while this work fixes the architecture and studies loop training and compute allocation.Source: Related work · p. 4
  • Truncated BPTT saves activation memory and backward compute but introduces biased gradients; looped language models commonly backpropagate through only the last few loops or half of them.Source: Related work · p. 4
  • Earlier studies of truncation compare recipes at fixed forward depth, whereas this work varies backward depth independently and weighs saved compute against loop gain.Source: Related work · p. 4
  • Prior looped-model scaling laws either fit architectures separately or model looping with one count; this law treats forward and backward counts separately using bounded effective-parameter gains.Source: Related work · p. 4
  • Iso-Depth finds truncation weakens loop gain but can lower loss at fixed forward depth by freeing compute, while Parcae leaves the joint choice of forward and backward loops open.Source: Related work · p. 4

Preliminaries

The model applies a shared recurrent block repeatedly between a prelude and coda, while truncated BPTT determines how many of those loops receive gradients. Its scaling law represents looping through an effective parameter count.

  • A prelude maps input tokens to hidden states, a shared recurrent block refines them over T loops, and a coda produces next-token predictions.Source: Preliminaries · p. 4
  • Unlike prior laws summarized here, the proposed law models forward and backward loop counts independently with a bounded gain.Source: Preliminaries · p. 5
  • Each loop injects the initial state alongside the previous state, then applies the recurrent block and RMSNorm before passing the result onward.Source: Preliminaries · p. 5
  • With setting (T, B), all T loops run forward but only the last B receive gradients; B = T is full BPTT, while B < T is truncated BPTT.Source: Preliminaries · p. 6
  • Truncation detaches the state before the last B loops, saving activations and backward compute while giving recurrent weights and the prelude a biased gradient; the coda gradient remains exact.Source: Preliminaries · p. 6
  • The scaling law counts the prelude and coda once and scales the recurrent block through an effective parameter count.Source: Preliminaries · p. 6

A scaling law for forward and backward loops

The section develops and tests a scaling law that separates forward and backward loop counts, then interprets their effects through effective parameters. It finds diminishing gains from adding loops and a stronger contribution from training more loops than merely running more.

Fitting and validating the scaling law

Illustrating the scaling law with effective parameters

Truncate or not?

The fitted law favors truncation when forward depth is prescribed, but favors full BPTT when both loop counts are tunable. Matched-compute runs support full BPTT, while integer loop counts qualify the continuous optimum.

Fixed forward loopstruncation can be optimal

Free forward-loop countfull BPTT is optimal

  • When both loop counts are tunable, the fitted law predicts full BPTT is optimal at every compute because its fitted exponents satisfy q > 2p.The continuous optimum is T = B = S/3 for every S ≥ 3.Source: Free forward-loop countfull BPTT… · p. 12
  • The continuous optimum need not be an integer setting; at loop compute T + 2B = 10, (4, 3) leads both neighboring full-BPTT settings by at most 0.002 nats.The compared fully trained counterparts at equal loop compute are also non-integer settings.Source: Free forward-loop countfull BPTT… · p. 12
  • At D = 16D0, (4, 4) achieves lower loss than (8, 2) at every size in matched-compute runs.Separately fitted laws also predict that (4, 4) reaches the same loss with less compute when both optimal allocations lie within the run ranges.Source: Free forward-loop countfull BPTT… · p. 13

Allocating compute among loopsdataand model size

The law favors no looping when model size is free, though two fully trained loops come close with fewer parameters. At fixed size, optimal loop and token budgets both grow with compute, and experiments show the two-loop model can surpass the non-looped model with modestly more compute.

Fixed model sizeloops versus data

  • At fixed model size, optimal loops and tokens both grow with compute, asymptotically as C^0.24 and C^0.76 respectively; runs show the same trend.A tenfold compute increase raises the optimal loop count about 1.7 times and token count about 5.8 times.Source: Fixed model sizeloops versus data · pp. 15–16

Free model sizethe law favors no looping

  • With model size free, the law favors no looping at every compute, while (2, 2) is the best looped setting and comes close to two untied copies in effective size.The fitted looping coefficient g(2, 2) is approximately 2.Source: Free model sizethe law favors no… · p. 17

Two fully trained loops are competitive with fewer parameters

  • At equal compute, the compute-optimal (2, 2) model stays within 0.008 nats of the non-looped model while using 27–32% fewer non-embedding parameters.At the largest budget, the loss gap is smaller than its standard error across seeds.Source: Two fully trained loops are… · p. 17
  • With 8% more compute, the (2, 2) model’s loss falls 0.007 nats below the non-looped model while retaining 27% fewer non-embedding parameters.The extra compute is used to train the (2, 2) model on 8% more tokens.Source: Two fully trained loops are… · p. 17

Conclusion

Truncation helps under the fitted law only when the forward-loop count is prescribed; when both loop counts are chosen, full BPTT is optimal across the analyzed sizes and compute levels. Compute allocation favors more tokens at fixed model size, while two fully trained loops are a strong option when model size is free.

Conclusion

  • When both loop counts are chosen, full BPTT is optimal at every analyzed model size and compute level.With a prescribed forward-loop count, truncation can help by freeing compute for a larger model and more tokens, offsetting discounted loop gains.Source: Conclusion · pp. 17–18
  • At fixed model size, most additional compute should go to data, while fully trained loops grow only as about C^0.24.When model size is free, two fully trained loops nearly match the compute-optimal non-looped model at equal compute with 27–32% fewer non-embedding parameters, and achieve lower loss with 27% fewer parameters at 8% more compute.Source: Conclusion · p. 18
  • The fitted law favors no looping when model size is free, although two fully trained loops come close to the compute-optimal non-looped model.The two-loop setting achieves lower loss with 27% fewer non-embedding parameters at 8% more compute.Source: Conclusion · p. 18

Limitations

  • The study focuses on dense models and does not establish whether its findings extend to looped MoE models, downstream task performance, or reasoning ability.Experiments are capped at T ≤16, so the gain ceiling and later switches to more loops remain extrapolations.Source: Conclusion · p. 18

 Introduction 

The listed appendix sections cover model architecture, BPTT, experimental setup, compute-optimal pairs, an untied model, and learning-rate and batch-size rules.

B.1 Model architecture

  • Appendix B.1 covers model architecture.The contents list gives page 27.Source: §B.1 · p. 25

B.2 Full and truncated backpropagation through the loops

  • Appendix B.2 covers full and truncated backpropagation through the loops.The contents list gives page 27.Source: §B.2 · p. 25

B.3 Experimental setup

  • Appendix B.3 covers the experimental setup.The contents list gives page 28.Source: §B.3 · p. 25

E.1 Compute-optimal pairs

  • Appendix E.1 covers compute-optimal pairs.The contents list gives page 42.Source: §E.1 · p. 25

E.2 The untied model

  • Appendix E.2 covers the untied model.The contents list gives page 43.Source: §E.2 · p. 25

F.1 Learning-rate rule

  • Appendix F.1 covers the learning-rate rule.The contents list gives page 43.Source: §F.1 · p. 25

F.2 Batch-size rule

  • Appendix F.2 covers the batch-size rule.The contents list gives page 46.Source: §F.2 · p. 25

A Extended related work

Prior work uses truncated gradients across looped language, reasoning, and other recurrent models, but few studies measure truncation's costs. This paper varies backward depth independently of forward depth and compares the compute saved with the loop gain lost.

Training through recurrent depth

  • Truncated BPTT is used across looped language models, recurrent-depth models, diffusion refinement, and recursive or hierarchical reasoning models.Examples include training converted pretrained models and using short backward horizons for recursive reasoning (Park et al., 2026; Salvatori et al., 2026).Source: A Extended related work · p. 26
  • Few studies measure truncation costs, and reported effects vary with task or model structure.Some compare recipes at fixed forward depth, whereas this work independently varies backward depth to weigh saved compute against lost loop gain.Source: A Extended related work · p. 26

Scaling laws for recurrence

  • Prior looped-transformer scaling laws either fit separate architecture-specific laws or model loop gains with different forms.Iso-Depth and Parcae are particularly related: each leaves joint forward- and backward-loop optimization unresolved in its scaling study (Schwethelm et al., 2026; Prairie et al., 2026).Source: A Extended related work · p. 26
  • The unified law recovers a truncation benefit when forward loops are prescribed but favors full BPTT when both loop counts are free.For dense models, it ranks two fully trained loops first among looped settings; equal-compute experiments nearly match non-looped models with substantially fewer non-embedding parameters.Source: A Extended related work · p. 27
  • Two-loop choices also appear in prior looped MoE studies, which report benefits over three or four loops or the largest marginal recurrence benefit at two.Those results are from MoE models, while this paper's two-loop finding concerns dense models (Wang et al., 2026b; Gao et al., 2026).Source: A Extended related work · p. 27

B Model architecture and experimental setup

The study uses shared-weight recurrent transformer blocks and varies forward and backward loop counts independently. It compares truncated and full gradients through a fixed-point perspective, then evaluates loop settings across model sizes, token budgets, and training configurations.

Model architecture

  • Each loop reuses the same recurrent-block weights, while the prelude and coda surround the recurrent computation.The block applies input injection, pre-norm transformer layers, and RMSNorm; the architecture holds unique parameters fixed across loop settings.Source: B Model architecture and experimental… · pp. 27–28

Full and truncated backpropagation through the loops

  • A (T, B) setting executes T loops but backpropagates only through the last B, detaching the state entering those loops.The truncated gradient keeps the B contributions with the fewest Jacobian products and drops the others, including corresponding prelude-gradient contributions.Source: B Model architecture and experimental… · p. 27
  • Under a fixed-point assumption, truncation approximates the fixed-point gradient with the first B Neumann-series terms, whereas full BPTT uses the first T.The approximation may fail when few loops do not approach a fixed point; even at a fixed point, dropped terms make the truncated gradient biased.Source: B Model architecture and experimental… · pp. 27–28

Experimental setup

  • The scaling suite spans 1.1–1,291 training tokens per parameter and fifteen loop settings, including forward depths up to 16.Loop-dependent compute grows with T + 2B because each backward loop costs twice a forward loop; token budgets are based on D0 ≈1.07 × 10^9 tokens.Source: B Model architecture and experimental… · p. 28
  • Training uses FineWeb-Edu and reports cross-entropy in nats per token on held-out shards after removing duplicated training documents.The corpus is tokenized into sequences of 2,048 tokens with the Llama 2 tokenizer.Source: B Model architecture and experimental… · p. 29
  • The reported study comprises 592 training runs using MuonH and AdamW, with measured compute including optimizer updates.The runs used 2.7 × 10^22 FLOPs of training compute; optimizer updates add 0.06–5.7% to forward and backward pass compute.Source: B Model architecture and experimental… · p. 29

C The law in detail

Measured runs show that forward loops without gradients add little initially, while trained backward loops improve loss with diminishing returns. The appendix details how effective loop gains are estimated and modeled, and evaluates the bounded law against held-out configurations and a power-law alternative.

What the runs show before any law is fitted

  • At fixed backward-loop counts, doubling forward loops from B to 2B changed loss by less than 0.007 nats in all 24 combinations, while doubling again lowered loss in all 24.At equal compute per token, (2, 2) beat (4, 1) and (4, 4) beat (8, 2) in all 24 combinations; (2, 2) improved loss by 0.023 nats on average over (4, 1). With full BPTT, each doubling from T = B = 1 to 8 lowered loss, by less each time.Source: C The law in detail · p. 29

Empirical looping coefficient

  • The empirical looping coefficient rescales recurrent-block parameters after matching a looped run’s loss to a non-looped model at the same token budget.The reference interpolation uses non-looped measurements, passes through observed points, and preserves monotonicity; estimates are reported only for eligible groups and losses within the observed range, without extrapolation.Source: C The law in detail · p. 30

Summaries across configurations

  • Across 163 eligible configurations, the configuration summaries describe agreement with runs, not independent validation; their interquartile ranges reflect variation across sizes and budgets, not confidence intervals.Some reference configurations are shared, so estimates are dependent; separate holdout evaluations are reported in §C.3.Source: C The law in detail · p. 30

Form and fit

  • The bounded gain rises from g(1, 1) = 1 toward a fitted ceiling of about 3.04, an extrapolation rather than an observed asymptotic limit.The gain increases with both forward and backward loop counts; truncation reduces it relative to full BPTT.Source: C The law in detail · p. 31
  • The bounded-gain law had 8.67 × 10^-3 nats held-out-setting error versus 1.65 × 10^-2 for the power-law gain, which was less accurate in every evaluation.Across 230 configurations, the bounded law’s RMSE was 6.40 × 10^-3 for one-configuration-out evaluation, 7.48 × 10^-3 for one-size-out, 6.42 × 10^-3 for one-token-budget-out, and 8.67 × 10^-3 for one-setting-out evaluation.Source: C The law in detail · p. 32

Runs at sixteen loops

  • Holding out all 32 T = 16 configurations gave RMSE 0.0115 nats for the bounded gain versus 0.0340 for the power-law gain; the T = 1 row was the harder extrapolation.The T = 1 RMSE was 0.0172 for the bounded gain and 0.0207 for the power-law gain; four 480M configurations at 2D0 were outside the fitting set.Source: C The law in detail · pp. 32–33

D Training compute and allocation

This section examines how training compute is allocated between loops and tokens, and how preferred loop counts vary with model-size constraints. It presents compute–loss frontier and fitted-law results, conditional on the adopted loss law and compute approximation.

D Training compute and allocation

  • At fixed model size, compute per token depends on forward and backward loop counts through T + 2B, so (2, 2) and (4, 1) cost the same.Attention costs follow the same loop-count accounting; optimizer updates are independent of T and B and their omission moves fixed-size switch points by less than 0.5%.Source: D Training compute and allocation · p. 33, DFixed model sizeallocate loops… · p. 38
  • Full BPTT settings reach the fixed-size loss frontier almost everywhere; only (4, 3) briefly leads between (2, 2) and (4, 4), by at most 0.002 nats.Two loops become best at about 14–17 tokens per parameter, four at 312–413, and eight at 7,000–9,600; sixteen pay only beyond 147,000–210,000. Runs confirm that (2, 2) beats the non-looped model beyond about 30 tokens per parameter, while (4, 4) overtakes (2, 2) at the smaller sizes.Source: D Training compute and allocation · p. 34

DFixed forward-loop countoptimize sizedataand backward loops

  • With a fixed forward-loop count, the backward-loop counts that minimize loss are independent of compute and are ordered by the effective-to-compute parameter price χ.The result assumes a fixed recurrent parameter fraction and the leading-FLOP approximation; when the recurrent block is absent, all backward counts tie.Source: DFixed forward-loop countoptimize… · pp. 35–36

DJoint forward–backward allocation at equal compute

  • At equal compute, joint continuous optimization favors full BPTT when q ≥ 2p; the adopted fit satisfies q > 2p.Replacing a truncated split with full BPTT preserves model size and token count while strictly improving loss, provided the replacement is admissible; discrete loop menus still require checking available settings.Source: DJoint forward–backward allocation… · p. 38

DFixed model sizeallocate loops and data

  • For fixed model size, the fitted condition q ≥ 2p leads the analysis to set B = T before balancing loop count against the tokens affordable at the compute budget.The resulting interior optimum is characterized by a loop–token balance with D = C/FN(T, T); this subsection’s supplied passages introduce the condition but do not state its solution or large-compute limit.Source: DFixed model sizeallocate loops… · pp. 38–39
  • At fixed model size, added compute goes to tokens until a deeper setting reaches lower predicted loss; asymptotically, T* and D* scale with exponents 0.238 and 0.762.The exponents are consequences of the fitted law along B = T, not a refit to measured optima; extrapolation to unbounded depth is not empirically established.Source: DFixed model sizeallocate loops… · p. 40

DFree model sizecompare joint allocations

  • For T ≥ 2, replacing a fully trained T-loop model with untied copies at the same leading compute increases effective size under the adopted coefficient conditions.The adopted coefficients satisfy the conditions, including g(2, 2) ≃ 1.910 < 2; the result is restricted to T ≥ 2.Source: DFree model sizecompare joint… · pp. 40–41
  • For any loop setting, jointly optimizing model size and tokens is equivalent to a non-looped model at budget C/χ; lower loss requires a compute premium λ > χ.The fitted-law optimum minimizes χ at every compute.Source: DFree model sizecompare joint… · pp. 41–42
  • At 3.1% more compute, the (2,2) setting predicts equal loss with 31.6% fewer parameters than (1,1).Among the fifteen settings, (1,1) minimizes χ and (2,2) is the best looped setting; the passage reports these comparisons alongside Table 5 and Figure 16b.Source: DFree model sizecompare joint… · p. 42
  • The free-size identities hold within the specified empirical loss law and compute approximation, not as a universal property of looped models.Size or data bounds and architectural compute terms can break the factorization, so full-compute optimization requires a fresh numerical envelope.Source: DFree model sizecompare joint… · p. 42

E Compute-optimal pairs and the untied model

Compute-optimal looped models can use substantially fewer non-embedding parameters than non-looped models, with a small loss difference at matched compute. A separate comparison tests looped models against untied counterparts.

E Compute-optimal pairs and the untied model

  • Across three compute budgets, (2, 2) models have 27–32% fewer non-embedding parameters and come within (1–8) × 10^-3 nats of non-looped models.Each model is trained at the size and token count that minimize its loss for the same training compute.Source: E Compute-optimal pairs and the untied… · pp. 42–43
  • At 2.2 × 10^19 FLOPs (+8%), the 344M (2, 2) model’s loss is 2.3590 versus 2.3663 for the 471M (1, 1) model (gap −7.3 ± 1.8 × 10^-3).The table compares compute-optimal pairs; the (2, 2) model uses 8% more tokens and compute.Source: E Compute-optimal pairs and the untied… · p. 43
  • With all untied copies initialized equally, the untied model reproduces the looped model’s loss and logits exactly, while holding about 1.5 times as many parameters for two loops.The comparison holds compute per token constant.Source: E Compute-optimal pairs and the untied… · p. 43
  • The wider looped model uses 15–22% more forward-and-backward compute per token than the untied model and has lower loss in all twelve seed pairs.The untied model gives each loop its own copy of the recurrent block, input injection, and normalization.Source: E Compute-optimal pairs and the untied… · p. 43

F Learning rate and batch size

The appendix tests learning-rate and batch-size transfer rules against measured optima, then evaluates the fitted law across a five-size grid. The learning-rate rule needs no depth term, and one batch rule is acceptable across all tested sweeps.

F Learning rate and batch size

  • With γ = 0.40, the learning-rate rule stays within 0.06% of the optimum in every sweep; the optimal learning rate follows training tokens, not depth.The exponent minimizing the rule’s largest loss above the optimum is γ = 0.406 (95% interval 0.400–0.412).Source: F Learning rate and batch size · pp. 45–46
  • Across 33 batch-size sweeps, the rule stays within 0.17% of the best batch, while one ladder step lower reaches 0.30% and one step higher 0.42%.The rule changes comparisons between loop settings by less than 0.004 nats.Source: F Learning rate and batch size · p. 46

G The law across the five-size grid

  • Across the five-size grid, the law’s RMSE is 0.0031–0.0093 nats on fitted configurations and 0.0200 on the rest, most of which lie outside 5–125 tokens per parameter.The grid contains 343 configurations, with 385 accepted runs.Source: G The law across the five-size grid · p. 47
  • One set of coefficients tracks all four full-BPTT settings across 26 size–budget pairs, with RMSE 0.0053 nats on 68 fitted configurations.RMSE on the other 36 configurations is 0.0148 nats.Source: G The law across the five-size grid · p. 48

H Training curves

The training-curve frontier shifts to larger models as compute grows. The plotted curves are training loss, and each run reaches its final loss only at its token-budget endpoint.

H Training curves

  • As compute grows, larger models take over the frontier of run ends; at each panel’s highest compute, the frontier follows the largest model shown.That model is 1B, or 480M when T = 16.Source: H Training curves · p. 48
  • Each token budget is a separate run with a cosine schedule ending at that budget, so the run reaches its final loss only at the endpoint.The frontier joins run ends rather than fitted curves.Source: H Training curves · p. 48
  • Figure 20 plots training loss against compute for each loop setting; its highlighted frontier connects run ends that finish below every run with less compute.Each curve records every tenth training step, not validation loss on held-out documents.Source: H Training curves · p. 48
Loading 2610.26030v2…