# LoopFold: Parallelizing Recurrent Depth in Looped Transformers

> LoopFold decouples recurrent depth from sequential decoding depth by predicting and correcting parallel loop states, cutting latency 46% while improving zero-shot accuracy by 2.20 points.

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

## Summary

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

| Model | Target depth index $D$ | Backbone evaluations | Sequential stages $S$ | Width $n$ |
|-------|----------------------|---------------------|----------------------|-----------|
| Vanilla looped model | $D$ | $D$ | $D$ | 1 |
| LoopFold | $n + S - 1$ | $1 + n(S-1)$ | $S$ | $n$ |

---

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

| Model | Training (ZFLOPs) | Avg-8 Accuracy |
|-------|-------------------|----------------|
| Vanilla Looped | 1.68 | 54.20% |
| **LoopFold** ($n=2, S=4$) | 1.68 | **54.71%** |
| Vanilla Looped | 2.16 | 54.18% |
| **LoopFold** ($n=2, S=5$) | 2.16 | **55.13%** |

#### 2. Quality vs. Decoding Latency (Comparable Latency)

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

| Model | Latency (ms/token) | Avg-8 Accuracy |
|-------|-------------------|----------------|
| Vanilla Looped $D=3$ | 73.11 | 51.98% |
| **LoopFold** $D=5$ | 74.41 | **53.41%** |
| Vanilla Looped $D=4$ | 98.40 | 53.31% |
| **LoopFold** $D=7$ | 98.66 | **55.23%** |
| Vanilla Looped $D=5$ | 122.88 | 53.69% |
| **LoopFold** $D=9$ | 121.68 | **56.38%** |

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

| Variant | Pile PPL ↓ | ΔPPL |
|---------|-----------|------|
| Full LoopFold | 9.13 | – |
| No prediction | 9.29 | +0.16 |
| No correction | 9.47 | +0.34 |

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

---

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