Deeper Backpropagation or More Updates? Training Looped Language Models Across Budgets

Jiarun Shen; Xiao Li

cs.CLOct 28, 2026

Deeper Backpropagation or More Updates? Training Looped Language Models Across Budgets

Authors: Jiarun Shen, Xiao Li

Organizations: The Chinese University of Hong Kong, Shenzhen

Abstract

Repeating transformer layers can improve language model performance without adding parameters. However, the added computational depth increases both forward and backward computation costs during training. Training such models requires choosing a backward window that controls how many final layer repetitions retain gradients during backpropagation. Longer windows improve per-token learning efficiency but increase per-update cost, while shorter windows reduce per-update cost but sacrifice per-token efficiency. Understanding which window performs better requires examining different training budgets. Prior work compares backward windows across computational budgets. When training budgets increase, it remains unclear whether the same window continues to perform best. We measure how the relative performance of two windows shifts as training budgets grow. Our method isolates per-token efficiency by training both windows on identical data, and compares complete training recipes at approximately equal training time at each budget. Results show that shorter windows lose much of their advantage as budgets grow. Controlled experiments show that both benefits decrease at larger budgets, with a greater decline in the benefit of extended training. Backward window comparisons depend on the available budget, and the preferred window among the tested candidates can change as the budget grows.

Figures & tables

Figure 1: Four- and eight-iteration backward windows with the same forward computation. Both rows illustrate 16 forward iterations to the shared two-block Core. The Prelude produces input representation

e

, and

x_{0}

denotes the initial recurrent state. Each Core iteration updates the state using an injection of

e

. The Coda and prediction head map the final state to next-token logits. K4 retains the gradient graph for iterations 13–16, and K8 for iterations 9–16. Gray iterations execute in the forward pass without retaining a gradient graph. Red arrows show backpropagation through the retained iterations. Stop-gradient detaches the state before the retained window, while injections into retained iterations preserve gradients to the Prelude. Gradients from retained iterations accumulate into the shared Core parameters. Training samples forward depth

T_{i}

, so a window

k

retains

\min(k,T_{i})

iterations. One optimizer update follows gradient accumulation across microbatches.

Figure 2: The matched-time advantage of shorter backpropagation diminishes between the tested budgets, while the better of the two windows at the larger budget varies. Negative values favor extended four-iteration training (K4+); positive values favor eight-iteration training (K8). The original 144M comparison retains all ten points, including seed 1103 at only the small budget and the medium budget for seed 1104. In the left panel, dashed lines and open markers denote exploratory seeds; solid lines and filled markers denote replication seeds. After dev-based learning-rate selection, both 144M confirmation seeds favor K8 at

4\times

. Both 388M seeds continue to favor K4+ despite diminishing advantage. Lines join separately trained endpoints from the same seed, not checkpoints from one run. The 144M K4+ horizons are 4,986/9,972/19,944; 388M uses 5,024/20,095. Relative budgets do not imply equal absolute GPU time across configurations. All points use T16. Error bars in the middle and right panels are marginal 95% paired test-block intervals conditional on checkpoints. The original panel preserves its point estimates; intervals are tabulated separately.

Configuration / LR	Seed	$C_{1\times}$	$C_{4\times}$	$\Delta C$	95% interval for $\Delta C$
144M, selected	1201	$-0.020240$	$+0.003377$	$+0.023617$	$[+0.022892,+0.024340]$
144M, selected	1202	$-0.015754$	$+0.002236$	$+0.017990$	$[+0.017253,+0.018735]$
388M, 0.004	1204	$-0.029446$	$-0.005610$	$+0.023837$	$[+0.023152,+0.024522]$
388M, 0.004	1205	$-0.033168$	$-0.006250$	$+0.026918$	$[+0.026247,+0.027601]$

Table 2: Budget changes after learning-rate selection and in the second configuration.

C=L_{4+}-L_{8}

is the test NLL difference in nats/token, and

\Delta C=C_{4\times}-C_{1\times}

. Positive

\Delta C

denotes contraction, not necessarily reversal. The complete common-rate grid remains in Table 14 . Intervals are marginal 95% paired test-block intervals conditional on checkpoints. The 144M rates are selected separately for each window and budget; 388M uses the common coarsely calibrated rate 0.004.

Figure 3: Extended training gains decrease more than per-token efficiency gains from deeper backpropagation. Test comparisons use three seeds at common learning rate 0.008 and T16.

L_{4}

L_{8}

, and

L_{4+}

denote test NLL for four-iteration backpropagation (K4), eight-iteration backpropagation (K8), and extended four-iteration training (K4+) respectively. Blue bars show

G=L_{4}-L_{8}

; orange bars show

R=L_{4}-L_{4+}

. At

1\times

, extending four-iteration training more than offsets the per-token efficiency gap. At

4\times

, both gains are smaller, and

R

falls below

G

. The larger decline in

R

raises

C=G-R

; the smaller value of

G

partly offsets that change. Error bars show marginal 95% paired test-block intervals conditional on checkpoints. The controls do not isolate extra tokens from the full optimization schedule and do not decompose the selected-rate or 388M comparisons. Full dev/test values and budget-change intervals appear in Appendix F.1 .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: The budget definition changes the NLL preference in all three original seeds. Equal-token K4 and K8 train for 4,096 updates; matched-time K4+ uses 4,986. Positive test differences favor K8, negative differences favor the four-loop recipe. Lines connect separate terminal models.

Figure 5: Raw single-quarter gains divided by the constant K4–K8 dev gap. Sums are 1.0387 and 1.0028. The frozen 0.4 trigger is not met in both seeds by any quarter.

Figure 6: Word-NLL contributions from all four correctness groups in all six exploratory pairs. Each sum is divided by all 5,153 examples; signed bars reconstruct the total (black mark). Negative favors K4+. Group membership is checkpoint-dependent, and these are descriptive output decompositions, not causal mechanisms.

Figure 7: Budget comparisons retain individual exploratory and replication seeds. Positive favors K8: test NLL is K4+ minus K8; task scores use K8 minus K4+. Lines connect separately trained endpoints at the measured budgets within the same seed, not a fitted crossover. Downstream bars are marginal 95% source-cluster intervals. Exploratory NLL bars use paired test blocks; replication NLL points retain the historical display. Updated paired-block intervals appear in Appendix F.1 . Dashed/open: exploratory; solid/filled: replication.

Figure 8: All six exploratory LAMBADA pairs and approximate 95% per-pair joint rectangles (two 97.5% marginal source-cluster intervals). The upper-left quadrant is the proposed loss-K4+/accuracy-K8 opposition; no rectangle is wholly contained there. Only 4x/1102 is jointly resolved, favoring K8 on both axes. These are not simultaneous six-pair bounds.

Figure 9: Supplementary loss comparisons: the token/time ranking reversal (left) and finite window grid (right). Lines connect separately trained endpoints. All values are locked-test NLL.

Figure 10: Dev and test budget comparisons, extended K4 minus K8. The intermediate budget exists only for seed 1104. Negative favors extended K4.

Figure 11: Inference-depth coverage for the original 61 final models. Left: each line is one final model (blue: native

\mu=16

; orange: native

\mu=8

). Right: raw marginal dev-NLL gain from recurrence 16 to 32. The narrow terminal range does not establish a universal ceiling.

Figure 12: First-target-token teacher-forced margin ECDFs for the six exploratory comparisons. Each ECDF uses all 5,153 LAMBADA examples and is shown over the plotted margin range. These distributions do not establish calibration or compare complete word sequences.

Figure 13: Minimum teacher-forced target-token margin ECDFs across the same six pairs. Group-conditioned summaries and ties are retained in the numerical supplement; correctness-based separation alone is not independent mechanism evidence.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 353 papers

Tapered Language Models
Jun 22, 2026Reza Bayat, Ali Behrouz, Aaron CourvilleAutoregressive Language Modeling
1Mila · 2Cornell University · 3Université de Montréal +1
Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation
Oct 19, 2026Xuyuan Xiong, Siyu Chen, Zihan Shao +5Truncated Backpropagation Through Time Efficient Language Model Training
MMLab, The University of Hong Kong · Yale University
How Much Gradient Does Truncated Backprop Lose in Looped Transformers?
Oct 20, 2026Aur Shalev MerinTruncated Backpropagation Through Time Backpropagation
Independent researcher

Jun 22, 2026J/K move · Enter open · S save

Tapered Language Models

Reza Bayat, Ali Behrouz, Aaron Courville

1Mila · 2Cornell University · 3Université de Montréal+1

Autoregressive Language Modeling

Open paper AskPDF arXiv

Modern language models, including transformer, recurrent, and memory-based variants, share a common chassis: a stack of identical layers in which parameters are allocated uniformly across depth. This is a default inherited from the original transformer and largely unchanged since, yet a growing body of evidence suggests that layers contribute non-uniformly to the final output, with later layers refining the residual stream rather than transforming it. We ask whether parameter capacity should reflect this asymmetry. Our controlled experiment shows that, under a fixed budget, allocating more capacity to earlier layers and less to later layers improves perplexity over a uniform-width baseline, while the reverse allocation hurts. Building on this result, we introduce Tapered Language Models (TLMs), an architectural principle in which a parameter-bearing component is monotonically tapered across depth under a fixed total budget. MLPs are the natural site for this instantiation: they dominate parameter count across all modern LM families and expose width as a single, clean axis of variation. Across three model scales and four architectures (Transformer, Gated Attention, Hope-attention, and Titans), tapering MLP width via a smooth cosine schedule consistently improves perplexity and downstream benchmark performance over uniform baselines, at no additional parameter or compute cost. These findings establish depth-aware capacity allocation as a simple, architecture-agnostic axis of language model design, a free lever hidden in plain sight.

Abstract

LLM Training

Explore similar work

Tapered Language Models

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

How Much Gradient Does Truncated Backprop Lose in Looped Transformers?