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

> Full backpropagation through time is optimal for looped transformers when both forward and backward loop counts are free, despite truncation helping when forward loops are fixed.

- **Source:** [arXiv](https://arxiv.org/abs/2610.26030)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/Jnbt7u
- **Whiteboard:** https://picx.dev/p/Jnbt7u/image

## Summary

## Summary (Overview)

- **Core research question**: Should looped transformers use truncated backpropagation through time (BPTT), and how should compute be allocated among forward loops, backward loops, model size, and training tokens?
- **Key contribution**: A unified scaling law with an *effective parameter count* that jointly models model size, training tokens, and independent forward/backward loop counts, fitted on nearly 600 training runs (58M–919M parameters).
- **Main finding on truncation**: When the forward-loop count is prescribed, truncation can be optimal (saving compute for larger models/more tokens); but when both loop counts are free, **full BPTT is optimal** at every model size and compute budget.
- **Compute allocation**: At fixed model size, optimal loop count grows as $C^{0.24}$ and token count as $C^{0.76}$; when model size is free, no looping is optimal, but **two fully trained loops** nearly match the compute-optimal non-looped model with 27–32% fewer non-embedding parameters.
- **Practical result**: A (2,2) looped model achieves lower validation loss than the compute-optimal non-looped model with ~27% fewer parameters and only 8% more training compute.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Looped transformers reuse a shared block of layers to increase computational depth without increasing trainable parameters (Dehghani et al., 2019; Geiping et al., 2025). Training typically uses BPTT, which is often truncated to the last few loops to save compute and activation memory. However, truncation yields biased gradients that may weaken the benefits of looping.

The central question: **How do model size, training tokens, and forward/backward loops jointly affect loss, and should we truncate backpropagation under a fixed training compute budget?**

### Key Gaps in Prior Work

Existing scaling studies have limitations:
- **Schwethelm et al. (2026)** compares full BPTT with truncation to half the loops, fitting separate scaling laws per recipe.
- **Parcae (Prairie et al., 2026)** varies backward-loop counts in small experiments but fixes them to half the forward-loop count in its scaling study.
- No prior work provides a unified law for choosing *both* loop counts across model sizes and token budgets.

### Architecture: Prelude–Recurrent–Coda

The model computes hidden states via a prelude $\mathcal{P}$, a shared recurrent block $\mathcal{R}$ applied over $T$ loops, and a coda $\mathcal{C}$:

$$
s_0 = \mathcal{P}(x), \qquad s_t = \Phi_\theta(s_{t-1}; s_0) = \text{RMSNorm}(\mathcal{R}(W_{\text{inj}}[s_0; s_{t-1}])), \qquad p = \mathcal{C}(s_T). \tag{3.1}
$$

Parameters split as $N = N_{\text{once}} + N_{\text{rec}}$, where $N_{\text{once}}$ counts prelude/coda (run once) and $N_{\text{rec}}$ counts the recurrent block (reused every loop).

### Truncated BPTT

A loop setting $(T, B)$ with $1 \leq B \leq T$ runs $T$ loops forward but backpropagates only through the last $B$ loops. The state $s_{T-B}$ is detached. Full BPTT corresponds to $B = T$; a standard transformer is $(1,1)$. A backward loop costs twice the FLOPs of a forward loop.

### Baseline Scaling Law

The coupled scaling law of Videau et al. (2026) for non-looped models:

$$
L(N, D) = E + \left[ A \cdot \left(\frac{N}{N_0}\right)^{-\alpha} + H \cdot \left(\frac{D}{D_0}\right)^{-\beta} \right]^{\kappa}. \tag{3.2}
$$

---

## Methodology

### Experimental Setup

- **Architecture**: Prelude–recurrent–coda, following Schwethelm et al. (2026).
- **Model sizes**: Five sizes from 58M to 919M non-embedding parameters (Table 2).
- **Corpus**: 100B-token FineWeb-Edu (Penedo et al., 2024).
- **Optimizer**: MuonH for transformer-block and input-injection matrices; AdamW for embeddings, output, and normalization parameters.
- **Runs**: 592 total, including a scaling suite of 347 configurations spanning 1.1–1,291 tokens per parameter and fifteen $(T, B)$ combinations with $T \in \{1, 2, 4, 8, 16\}$.
- **Fitting**: 230 configurations for fitting; 117 reserved for holdout tests.

### Model Sizes

| Size | d | Layers | N | $N_{\text{rec}}/N$ |
|------|---|--------|---|-------------------|
| 60M | 768 | 2/4/2 | 57.8M | 0.510 |
| 120M | 1024 | 2/5/2 | 115.4M | 0.564 |
| 240M | 1280 | 3/6/3 | 239.2M | 0.507 |
| 480M | 1664 | 3/7/4 | 470.8M | 0.506 |
| 1B | 2176 | 4/8/4 | 918.7M | 0.505 |

### The Effective-Parameter Scaling Law

The law replaces $N$ with an effective parameter count:

$$
L = E + \left[ A \cdot \left(\frac{N_{\text{eff}}}{N_0}\right)^{-\alpha} + H \cdot \left(\frac{D}{D_0}\right)^{-\beta} \right]^{\kappa},
$$
$$
N_{\text{eff}} = N_{\text{once}} + g(T, B) \cdot N_{\text{rec}}, \qquad g(T, B) = \exp\big[c \cdot (1 - T^{-p} \cdot B^{-q})\big]. \tag{4.1}
$$

### Design Requirements for $g(T, B)$

The looping coefficient satisfies four requirements:
1. **Normalization**: $g(1,1) = 1$ (single loop recovers non-looped model)
2. **Monotonicity**: $\partial_T g > 0$, $\partial_B g > 0$
3. **Zero-backward limit**: $g(T, 0) = 0$
4. **Saturation**: $\sup_{1 \leq B \leq T} g(T, B) < \infty$ (finite ceiling $\exp(c)$)

### Fitted Parameters

| E | A | H | α | β | κ | c | p | q |
|---|---|---|---|---|---|---|---|---|
| 0.29865 | 6.48695 | 7.63343 | 0.30568 | 0.39281 | 0.32510 | 1.11198 | 0.23677 | 1.02181 |

### Fit and Prediction Errors ($\times 10^{-3}$)

| Evaluation | Fits | RMSE | MAE |
|---|---|---|---|
| In-sample | 1 | 6.14 | 4.91 |
| One configuration out | 230 | 6.40 | 5.11 |
| One model size N out | 7 | 7.48 | 5.89 |
| One token budget D out | 11 | 6.42 | 5.12 |
| One (T, B) setting out | 15 | 8.67 | 6.74 |

---

## Empirical Validation / Results

### Key Finding 1: Loops Add Effective Parameters; Truncation Discounts Them

- Each fully trained loop adds effective parameters with diminishing returns toward a finite ceiling $e^c$.
- Truncation discounts this gain.
- From the same setting, **doubling $B$ adds more effective parameters than doubling $T$** (since $q > p$: $g(T, 2B) > g(2T, B)$ for $2B \le T$).

### Key Finding 2: Truncation Helps at Prescribed Forward-Loop Count

When $T$ is prescribed, truncation can be optimal. The compute-equivalent size is:

$$
C \simeq 6D \cdot N_{\text{comp}}, \quad N_{\text{comp}} = N_{\text{once}} + \frac{T + 2B}{3} \cdot N_{\text{rec}}. \tag{5.1}
$$

The optimization reduces to minimizing $\chi(T, B) = N_{\text{comp}}/N_{\text{eff}}$:

| 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 |

At $T = 8$, the law selects $B^* = 2$ and predicts a 34% compute saving at equal optimized loss.

### Key Finding 3: Full BPTT Is Optimal When Both Loop Counts Are Free

When both $T$ and $B$ can be chosen, the problem becomes:

$$
\underset{T, B \in \mathbb{R}}{\text{maximize}} \quad g(T, B) \qquad \text{subject to} \quad T + 2B = S, \quad 1 \leq B \leq T. \tag{5.2}
$$

The optimal split uses full BPTT whenever $q > 2p$. With fitted values $p \simeq 0.237$ and $q \simeq 1.022$, we have $q > 2p \simeq 0.474$ by a wide margin. **Full BPTT is optimal at every compute.**

Matched-compute experiments confirm: at $D = 16D_0$, (4,4) achieves lower loss than (8,2) at every size:

| Size | $C$ ($10^{19}$) | $L_{8,2}$ | $L_{4,4}$ | $\Delta L$ ($10^{-2}$) |
|---|---|---|---|---|
| 60M | 2.41 | 2.5211 | 2.4933 | 2.78 |
| 120M | 4.59 | 2.3843 | 2.3696 | 1.47 |
| 240M | 8.28 | 2.2778 | 2.2642 | 1.36 |
| 480M | 15.29 | 2.1845 | 2.1738 | 1.07 |
| 1B | 28.38 | 2.1144 | 2.1057 | 0.87 |

### Key Finding 4: Fixed Model Size — Loops Grow Slowly

At fixed model size with full BPTT ($B = T$):

$$
T^{*} \propto C^{\frac{\beta}{\beta + p + q}}, \qquad D^{*} \propto C^{\frac{p + q}{\beta + p + q}} \quad \text{at fixed } N. \tag{6.3}
$$

With fitted values ($p + q \simeq 1.259$, $\beta \simeq 0.393$):

$$
T^{*} \propto C^{0.24}, \qquad D^{*} \propto C^{0.76}. \tag{6.4}
$$

A tenfold increase in compute raises the optimal loop count by ~1.7× and the token count by ~5.8×.

### Key Finding 5: Free Model Size — No Looping Is Optimal

When model size is free, the law favors the non-looped model. The optimal allocation:

$$
N^{*} \propto C^{0.56}, \qquad D^{*} \propto C^{0.44}, \qquad T^{*} = B^{*} = 1. \tag{6.5}
$$

The reason: replacing a shared recurrent block with $T$ untied copies gives greater effective size since $g(T, T) < T$ for $T \geq 2$. However, **$g(2,2) \approx 2$**, so two shared loops provide nearly the effective size of two untied copies.

### Key Finding 6: Two Fully Trained Loops Are Competitive

- At equal compute, (2,2) models stay within 0.008 nats of the non-looped model with 27–32% fewer non-embedding parameters.
- With 8% more compute, the (2,2) model's loss falls **0.007 nats below** the compute-optimal non-looped model (Figure 15), while still having 27% fewer parameters.

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Unified scaling law**: First law jointly modeling forward and backward loop counts independently, using a bounded effective-parameter gain with just three loop coefficients ($c, p, q$).
2. **Explanation of truncation's value**: Resolves the apparent contradiction between prior results — truncation helps only when the forward-loop count is fixed; full BPTT wins when both counts are tunable.
3. **Compute allocation rules**: Extends Parcae's fixed-size analysis beyond prescribed truncation, providing asymptotic exponents for loop/token allocation.

### Practical Implications

- **Parameter efficiency**: Two fully trained loops offer a practical sweet spot — nearly matching non-looped compute-optimal models with 27–32% fewer parameters.
- **Memory-constrained training**: At fixed model size, practitioners should allocate most additional compute to tokens, growing the loop count only slowly ($C^{0.24}$).
- **Training recipe guidance**: When both loop counts can be chosen, full BPTT is preferable to truncation despite its higher per-step cost.

### Comparison with Prior Scaling Laws

| Study | Loop gain | BPTT | Regime |
|---|---|---|---|
| Schwethelm et al. (2026) | power law ($T^{\phi}$) | $T; \lceil T/2 \rceil$ | fixed depth |
| Parcae (Prairie et al., 2026) | linear ($T$) | $\lceil T/2 \rceil$ | fixed N |
| Chen et al. (2026b) | bounded exp. ($T$) | full ($T$) | free N, mem. |
| **Ours** | **bounded ($T, B$)** | **$1 \leq B \leq T$** | **fixed & free N** |

---

## Conclusion

### Main Takeaways

1. **Truncation's value is conditional**: It helps only when the forward-loop count is prescribed; with both loop counts free, full BPTT is optimal at every model size and compute budget.
2. **Effective parameters framework**: Loops add effective parameters with diminishing returns toward a finite ceiling; truncation discounts this gain.
3. **Compute allocation**: At fixed size, $T^* \propto C^{0.24}$ and $D^* \propto C^{0.76}$; with free size, no looping is optimal but (2,2) is the best looped setting.

### Future Directions

- **MoE architectures**: Whether conclusions extend to looped mixture-of-experts models remains open.
- **Downstream tasks**: Only validation loss was measured; reasoning and task performance need study.
- **Deeper loops**: Experiments capped at $T \leq 16$; the gain ceiling and later switches remain extrapolations.
- **Parameter allocation**: The ratio of prelude/coda to recurrent parameters was fixed; varying it could affect loop gains and optimal loop counts.

---

_Markdown view of https://picx.dev/p/Jnbt7u, served by PicX — AI-generated visual whiteboard summaries of research papers._