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

AI summary

Overview

Research area: Machine learning — specifically large language model pretraining, parameter-efficient architectures (looped/weight-tied transformers), and empirical scaling laws.

Technical level: Advanced. The core question is intuitive, but the paper's central artifact is a nine-parameter nonlinear scaling law fitted to hundreds of training runs; understanding the validation protocol and the "effective parameter" decomposition requires comfort with scaling-law methodology.

Scope: A systematic scaling study across 592 training runs that jointly varies model size, training tokens, forward loop count, and backward (gradient) loop count to determine whether truncated backpropagation-through-time (BPTT) is worth using when training looped transformers.

What This Paper Is About

Looped transformers reuse one shared block of layers many times, gaining computational depth without gaining parameters — but training them requires backpropagating through repeated applications of that block. To save time and memory, practitioners usually truncate: run all the loops forward, but only backpropagate through the last few. That truncation saves compute but biases the gradient.

The paper asks a single, practical question: under a fixed training compute budget, should you truncate, and how should you split that budget among forward loops, backward loops, model size, and training tokens? Prior scaling studies fit a separate law per backpropagation recipe (e.g., full vs. half-length), so they can't answer this — they never let the backward-loop count be an independent variable. This paper builds a unified law that does.

Key Contributions

1. A unified scaling law over effective parameters. The authors extend the coupled size–data scaling law of Videau et al. (2026) by replacing the parameter count with an effective count $N_{\text{eff}} = N_{\text{once}} + g(T,B)\cdot N_{\text{rec}}$, where the looping coefficient $g(T,B) = \exp[c(1 - T^{-p}B^{-q})]$ depends on the forward loop count $T$ and backward loop count $B$ independently. The prelude and coda count once; the shared recurrent block is multiplied by $g$.

2. An explanation of when truncation actually pays. The law shows the answer flips depending on whether the forward-loop count is prescribed. At a fixed $T$, truncation frees backward compute for a larger model and more tokens, reproducing the benefit reported by prior work. When both $T$ and $B$ are free, the law favors full BPTT.

3. Compute-allocation rules for fixed and free model size. At fixed model size, the law characterizes how the optimal loop count and token budget grow with compute. When model size is also free, it ranks two fully trained loops as the best looped configuration.

4. A large, controlled measurement suite. 592 training runs, including a scaling suite of 347 distinct configurations spanning 58M–919M non-embedding parameters, 1.1–1,291 tokens per parameter, and fifteen $(T,B)$ settings with $T \in {1,2,4,8,16}$.

Main Findings

• Full BPTT wins when you can choose both loop counts. Reducing the forward-loop count frees compute for additional backward loops until every remaining loop is trained, and under the fitted law this reallocation yields more effective parameters at the same compute than the original truncated setting. Matched-compute experiments support this: full BPTT through fewer loops consistently beats truncated settings of equal compute.

• Truncation does pay when the forward-loop count is fixed. At a prescribed $T$, the saved backward compute supports a larger model and more training tokens, which more than offsets the discounted loop gain — reproducing the result of Schwethelm et al. (2026) within a single unified law.

• Loop gains saturate; truncation discounts them. Figure 1(b) shows that at $8D_0$ tokens loops add effective parameters while truncation discounts them. Panel (a) of Figure 8 shows the full-BPTT gain levelling off at 240M parameters and $2D_0$ tokens, whereas the power-law gain assumed in prior work keeps growing.

• Training more loops beats running more loops. At fixed $T$, increasing $B$ lowers the loss substantially; at fixed $B$, adding forward loops helps little (Figure 8b,c). The fitted exponent on $B$ ($q = 1.02181$) is larger than the fitted exponent on $T$ ($p = 0.23677$).

• At fixed model size, loop count grows slowly with compute. The optimal loop count and training-token budget both grow with compute, but the law allocates most additional compute to tokens and increases loop count only gradually (Figure 1c shows this at 120M parameters). This extends Parcae's fixed-size analysis beyond a prescribed truncation rule.

• Two fully trained loops is the best looped setting when model size is free. A $(2,2)$ model with 344M non-embedding parameters achieves a lower validation loss than the compute-optimal non-looped model — 2.36 vs. 2.37 nats — at $2.1 \times 10^{19}$ FLOPs, with about 27% fewer non-embedding parameters and only 8% more training compute. Figure 2 places every truncated setting below the full-BPTT frontier.

• Truncating the same fraction costs less at larger $T$. The empirical looping coefficient's log–log slope against $r = B/T$ falls as $T$ grows, which is why the coefficient couples $T$ and $B$ rather than treating them separably.

• The law fits and predicts tightly. In-sample RMSE is $6.14 \times 10^{-3}$ nats/token and MAE $4.91 \times 10^{-3}$, with $R^2 = 0.9992$ over 230 configurations. Leave-one-model-size-out gives RMSE $7.48 \times 10^{-3}$; leave-one-token-budget-out $6.42 \times 10^{-3}$; leave-one-configuration-out $6.40 \times 10^{-3}$; and leave-one-$(T,B)$-setting-out is hardest at $8.67 \times 10^{-3}$.

Methodology in Plain English

The team takes a standard recipe — a prelude–recurrent–coda architecture following Schwethelm et al. (2026) — and holds the architecture fixed so that only the training of the loops varies.

They train five model sizes (57.8M, 115.4M, 239.2M, 470.8M, and 918.7M non-embedding parameters) on the 100B-token FineWeb-Edu corpus, using MuonH for transformer-block and input-injection matrices and AdamW for embedding, output, and normalization parameters. Evaluation is validation cross-entropy in nats per token on deduplicated held-out shards.

For each size they sweep token budgets and fifteen loop settings. A setting $(T,B)$ means: run $T$ loops forward, but detach the state after the first $T-B$ loops so gradients flow through only the last $B$. $(1,1)$ is just a standard transformer, $(T,T)$ is full BPTT.

The modeling move is to ask: how big a non-looped model would achieve this loss at this token budget? That gives an "effective parameter count," and the ratio of that to the actual recurrent-block size gives the looping coefficient $g$. The authors compute this coefficient empirically from the runs (interpolating non-looped measurements with PCHIP at matched token budgets), then fit a three-parameter surface to it.

The chosen functional form is constrained to satisfy four properties: $g \geq 0$ with $g(1,1)=1$ (a single loop recovers the non-looped model), monotonicity in both $T$ and $B$, $g(T,0)=0$ (no trained loop adds no effective size), and a finite ceiling (finite weights store finite information). Nine coefficients in total are fitted — six for the size–data law and three for the loop surface — by minimizing squared loss with equal weight per configuration, refitting 64 times from random starts and keeping the best solution. Reference scales are fixed at $N_0 = 2.4\times10^8$ parameters and $D_0 = 2^{30}$ tokens.

Of the configurations, 230 are used for fitting and 117 held out for validation. Holdouts are structured: entire model sizes, entire token budgets, single configurations, and entire $(T,B)$ settings are removed in turn.

Why This Matters

Impact on research. The paper reframes the truncation debate. Prior work compared fixed recipes (full vs. half-length BPTT) and concluded truncation helps at a given depth; this work shows that conclusion is an artifact of holding $T$ fixed. Making $B$ an independent axis turns truncation from a compute-saving trick into a compute-reallocation decision — and the reallocation favors training every loop you run. The effective-parameter formulation also gives the community a common currency for comparing architectures with different parameter-sharing schemes.

Real-world applications.

• Memory-constrained pretraining. Teams whose GPUs cannot hold activations for many backpropagated loops can use the matched-compute result to decide whether a smaller model with full BPTT beats a larger model with truncation, rather than defaulting to truncation.
• On-device and edge deployment. The ~27% non-embedding parameter saving at near-equal loss and only 8% more compute matters where the parameter budget, not the FLOP budget, is the binding constraint.
• Compute-budget planning. The allocation rules let a fixed FLOP budget be divided among loop count, tokens, and model size without a separate hyperparameter sweep per scale.
• Architecture selection for recurrent-depth models. The finding that two fully trained loops captures most of the benefit sets a concrete default for dense looped models, echoing the two-loop convention already common in looped MoE work.

Industry relevance. Pretraining budgets are dominated by decisions made before any run starts. A law that predicts loss as a function of $(N, D, T, B)$ — including held-out model sizes and token budgets — is directly usable for planning, and the specific claim that truncation is dominated when $T$ is free could change default training configurations for any organization training looped or recurrent-depth models.

Future Directions

• Joint optimization beyond the fitted grid. The law is fitted up to $T=16$ and five model sizes below 1B parameters; whether the preference for full BPTT and the two-loop optimum survive at larger scale, or under the much longer token budgets used in production pretraining (Ouro, for comparison, scaled looped pretraining to 7.7T tokens), is untested here.

• Architecture generalization. All runs use a dense prelude–recurrent–coda model. The two-fully-trained-loop recommendation is derived for dense models; how the looping coefficient changes under mixture-of-experts routing, hyper-connections, or other recurrence variants is open.

• Closing the loop-setting prediction gap. Held-out $(T,B)$ settings are the hardest extrapolation (RMSE $8.67\times10^{-3}$ vs. $6.14\times10^{-3}$ in-sample), and the authors note that extrapolating back to the non-looped baseline is the weakest loop-count check. A functional form that predicts unseen loop settings more accurately would make the law more reliable for design choices outside the fitted grid.

• Reconciling the law with training dynamics. The law is a loss predictor, not a mechanism. Why backward loops matter more than forward loops ($q = 1.02181$ vs. $p = 0.23677$) — whether this reflects gradient bias, optimization difficulty, or representational capacity — is not explained, and connecting the coefficient to the truncated-Neumann-series view of the fixed-point gradient would be a natural next step.

Target Audience

Researchers and engineers working on parameter-efficient or recurrent-depth language model architectures, and anyone doing compute-budget planning for pretraining. It is most valuable to readers already familiar with scaling laws (Kaplan/Hoffmann-style) and with the looped-transformer literature (Universal Transformers, Huginn, Ouro, Parcae, Iso-Depth), since much of the contribution is defined relative to those baselines. Practitioners deciding between truncation and full BPTT for a fixed GPU memory budget will find the matched-compute results directly actionable; readers without scaling-law background will find the qualitative conclusions accessible but the fitting and validation sections demanding.

Authors’ 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.