Summary (Overview)

• LoopFold is a novel method for parallelizing recurrent depth in looped Transformers, folding computation at multiple recurrent depths into parallel batches for the same token prefix.
• A lightweight predictor estimates inputs to later loops so they can start before earlier loops finish, while a corrector refines outputs using newly available earlier-loop results.
• At target depth index $D = 9$, LoopFold achieves 46% lower decoding latency than a nine-loop vanilla model, requiring only five sequential stages instead of nine.
• Deeper LoopFold configurations achieve progressively higher mean zero-shot accuracy, rising from 53.41% to 56.38% across eight tasks, while vanilla looped models and PLT show limited gains.
• All components are trained jointly using only the final next-token prediction loss, with no auxiliary objectives.

---

Introduction and Theoretical Foundation

Background and Motivation

Looped Transformers increase computation per token by repeatedly applying a shared backbone, improving language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop's output, creating a fundamental trade-off:

• Greater recurrent depth → better performance but higher decoding latency
• Parallel approaches (e.g., PLT) → reduced latency but diminishing marginal accuracy gains at larger depths

The core challenge: How can later loops start before their inputs are available while still incorporating the results of earlier loops?

Theoretical Foundation

LoopFold draws inspiration from:

1. Universal Transformers (Dehghani et al., 2019) – recurrent-depth architectures with shared backbone
2. Parallel Loop Transformers (PLT) (Wu et al., 2025) – pipelining loop depths across token positions
3. Parareal (Lions et al., 2001) – classical parallel-in-time methods using fine/coarse propagators
4. Recurrent-depth parallel samplers (Geiping et al., 2025b) – parallel draft refinement

Unlike classical Parareal which iteratively seeks to recover the sequential fine solution, LoopFold learns a fixed-stage parallel computation optimized directly for next-token prediction. Its objective is not to match the sequential recurrent trajectory, but to produce predictive states under a constrained number of sequential stages.

---

Methodology

Architecture Overview

Let $\mathbf{H}_0 = \text{Embed}(\mathbf{x}_{1:L}) \in \mathbb{R}^{d \times L}$ be the input embeddings and $\mathbf{F}$ be the shared causal Transformer backbone. Sequential states satisfy:

$$\mathbf{H}_j = \mathbf{F}(\mathbf{H}_{j-1})$$

A schedule $1 + n + \cdots + n$ uses one initial backbone evaluation followed by $S - 1$ stages of $n$ parallel evaluations, where:

• $n$ = parallel width
• $S$ = total stage count
• $D = n + S - 1$ = target depth index

1. Predicting Inputs for Later Loops

After computing $\mathbf{H}_1 = \mathbf{F}(\mathbf{H}_0)$, the state change $\mathbf{V}_1 = \mathbf{H}_1 - \mathbf{H}_0$ is used as an initial increment. A shared token-wise linear map $\mathbf{J}_{\text{pred}}$ transforms this increment:

$$\mathbf{Z}^{(1)}_{1+\ell} = \mathbf{H}_1 + \sum_{m=1}^{\ell} \mathbf{J}^m_{\text{pred}} \mathbf{V}_1, \quad \ell = 1, \ldots, n-1$$

For $\mathbf{J}_{\text{pred}} = \mathbf{I}$, this reduces to linear extrapolation: $\mathbf{Z}^{(1)}_{1+\ell} = \mathbf{H}_1 + \ell\mathbf{V}_1$.

2. Parallel Evaluation and Correction

At each stage $s = 2, \ldots, S$, states are batched and evaluated in parallel:

$$\left(\mathbf{H}_s, \mathbf{Y}^{(s)}_{s+1}, \ldots, \mathbf{Y}^{(s)}_{s+n-1}\right) = \text{Split}\left[\mathbf{F}\left(\text{Concat}_{\text{batch}}\left(\mathbf{H}_{s-1}, \mathbf{Z}^{(s-1)}_{s}, \ldots, \mathbf{Z}^{(s-1)}_{s+n-2}\right)\right)\right]$$

The corrector then refines outputs in depth order using a shared token-wise linear map $\mathbf{J}_{\text{corr}}$:

$$\mathbf{Z}^{(s)}_{j+1} = \mathbf{Y}^{(s)}_{j+1} + \mathbf{J}_{\text{corr}}\left(\mathbf{Z}^{(s)}_{j} - \mathbf{Z}^{(s-1)}_{j}\right), \quad j = s, \ldots, s+n-2$$

3. Training Objective

The deepest state $\hat{\mathbf{H}}_D = \mathbf{Z}^{(S)}_{S+n-1}$ produces next-token logits:

$$\ell_t = \mathbf{W}_{\text{LM}} \text{RMSNorm}(\hat{\mathbf{H}}_D)_t, \quad \mathcal{L}_{\text{LM}} = -\sum_{t=1}^{L-1} \log \text{softmax}(\ell_t)_{x_{t+1}}$$

Both maps are parameterized as $\mathbf{J} = \mathbf{I} + \Delta\mathbf{J}$ with $\Delta\mathbf{J} \in \mathbb{R}^{d \times d}$ initialized to zero, adding only $2d^2$ parameters (8.39M at 1.4B).

Computational Cost

---

Empirical Validation / Results

Experimental Setup

• Models: 1.4B parameters trained from scratch on the Pile (30B tokens)
• Ablations: 410M models trained on 25B tokens
• Baselines: Vanilla looped models, PLT, Hidden Decoding
• Metrics: Pile validation perplexity, zero-shot accuracy on 8 tasks

Key Results

1. Quality vs. Training Compute (Matched Budgets)

At training compute of 1.68 and 2.16 ZFLOPs, LoopFold achieves lower PPL and higher mean zero-shot accuracy than vanilla looped models:

2. Quality vs. Decoding Latency (Comparable Latency)

At matched decoding latency, LoopFold outperforms vanilla looped models in all three pairs:

3. Throughput with Larger Batches

At $D=5$, LoopFold achieves 1.61–1.65× throughput compared to vanilla looped models across batch sizes 1–16.

4. Ablation Studies (410M models)

Both components improve PPL; removing correction has the larger penalty, highlighting its importance.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Decoupling recurrent depth from sequential depth: LoopFold demonstrates that recurrent depth and sequential decoding depth need not coincide—additional recurrent computation can be shifted onto the parallel dimension with proper correction mechanisms.

2. Learned parallel refinement: Unlike classical Parareal which seeks to recover the sequential solution, LoopFold learns a fixed-stage parallel computation optimized directly for next-token prediction, suggesting that minimizing instantaneous distance to the sequential trajectory is not necessary for good parallel refinement.

3. Component analysis: The predictor and corrector play complementary roles—the predictor provides reasonable initializations while the corrector propagates newly available information across parallel states.

Practical Implications

1. Latency reduction: 46% lower decoding latency at $D=9$ while improving accuracy by 2.20 percentage points.
2. Scalability: Throughput gains remain stable (1.61–1.65×) across batch sizes, making LoopFold practical for production deployment.
3. Memory trade-offs: KV sharing can reduce cache storage from 7× to 5× with modest quality/latency trade-offs (Hard overwrite-64: PPL 7.63 vs. 7.62 baseline).

---

Conclusion

LoopFold introduces a principled approach to parallelizing recurrent depth in looped Transformers through:

• Predicted inputs to enable early start of later loops
• Corrected outputs to incorporate newly available earlier-loop results
• Joint training with only the final next-token prediction loss

The method achieves lower PPL and higher zero-shot accuracy at matched training compute and comparable decoding latency, with the key insight that recurrent depth and sequential decoding depth can be decoupled.

Future Directions

1. Reducing computational overhead (additional FLOPs from parallel evaluations)
2. Further optimizing KV-cache memory costs
3. Preserving the benefits of parallel recurrent evaluation while minimizing resource requirements