PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopFold: Parallelizing Recurrent Depth in Looped Transformers

Feichen Song, Boyi Zeng, Yiqin Hao, Shixiang Song, He Li, Yizhou Yang, Zitong Wang, Yifan Liu, Ziwei He, Zhouhan Lin

arXiv:2610.35556v1cs.CL

TL;DR

Looped Transformers increase recurrent computation but sequential dependencies make deeper decoding slower, while existing parallel approaches have limited marginal accuracy gains at larger depths. LoopFold predicts later-loop inputs and corrects their parallel outputs as earlier results arrive. In 1.4B models, it improves quality across increasing depth while reducing sequential decoding stages, with additional parallel computation.

  • Problem

    Looped Transformers can improve language modeling and downstream performance at fixed backbone parameter count, but sequential loop dependencies make greater depth increase per-token decoding time, while existing parallel methods face diminishing accuracy gains.

    Source: §1 · p. 1

  • Method

    LoopFold batches states from multiple recurrent depths for the same token prefix, predicting later-loop inputs and correcting their outputs as earlier-loop results become available.

    Source: §1 · p. 1

  • Results

    At D=9, LoopFold has 2.20 percentage points higher Avg-8 accuracy and 46% lower decoding latency than the vanilla looped model, using additional parallel backbone computation.

    Source: §4.4 · p. 8

  • Takeaways & Limitations

    The results suggest recurrent depth and sequential decoding depth need not coincide, because additional recurrent computation can be shifted to parallel evaluation when provisional states are corrected.

    Source: §5 · p. 11

    LoopFold uses higher FLOPs and increases KV-cache requirements; KV sharing mitigates the overhead but brings a modest quality and/or latency trade-off.

    Source: §5 · p. 11

Abstract

from arXiv · show

Looped Transformers increase computation per token by repeatedly applying a shared backbone, offering a way to improve language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop, so increasing recurrent depth also increases decoding latency. Existing approaches include Parallel Loop Transformers, which reduce latency by overlapping loop computation across token positions, but yield rapidly diminishing marginal accuracy gains at larger depths. To address this challenge, we introduce LoopFold, which folds 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. A corrector then refines the outputs of these later loops using newly available results from earlier loops. All components are trained jointly using only the final next-token prediction loss. In experiments with 1.4B models, deeper LoopFold configurations achieve progressively higher average zero-shot accuracy across eight tasks, rising from 53.41% to 56.38%. Vanilla looped models and PLT show limited gains over the same depth range. Compared with a nine-loop vanilla model, LoopFold uses five sequential stages and additional parallel computation to reduce decoding latency by 46%.

1. Introduction

Looped Transformers add recurrent computation at fixed backbone parameter count, but their sequential dependencies increase decoding time and existing parallel approaches lose marginal accuracy gains at greater depths. LoopFold predicts later-loop inputs and corrects them as earlier-loop results arrive, improving quality while reducing sequential stages.

1. Introduction

  • Greater recurrent depth increases per-token computation at fixed backbone parameter count but also extends decoding time because each loop depends on the preceding loop’s output.Looping can improve language modeling and downstream task performance.Source: §1 · p. 1
  • PLT overlaps computation across token positions, but positional mismatch persists across loops and marginal accuracy gains quickly diminish at larger depths.Parallel recurrent-depth samplers also trade accuracy for faster decoding when they reduce refinement steps between draft updates.Source: §1 · p. 1
  • LoopFold batches hidden states from different recurrent depths for the same token prefix, using a lightweight predictor to start later loops early and a corrector to refine their outputs.Corrected states feed the next stage, and the deepest corrected state predicts the next token.Source: §1 · pp. 1–2
  • At target depth index D=9, LoopFold reduces latency by 46% versus a nine-loop vanilla model and raises Avg-8 accuracy by 2.20 percentage points, with additional parallel computation.It uses five sequential stages rather than nine; across eight tasks, accuracy improves.Source: §1 · p. 2

2. Related Work

Prior work shortens recurrent decoding through token-position pipelining, parallel token-draft refinement, or parallel correction. LoopFold instead predicts later-loop inputs and corrects parallel outputs, optimizing fixed-stage computation for next-token prediction rather than reproducing the sequential trajectory.

Scaling latent computation

  • Other recurrent-depth approaches include PonderLM’s probability-weighted embedding inputs, sampled recurrent depths, Ouro’s input-dependent exit gate, and Hidden Decoding’s multiple embedding streams.Looping can approach the accuracy of deeper untied models on arithmetic and question-answering tasks.Source: §2 · p. 2

Parallel recurrent computation

  • PLT pipelines recurrent depths with states shifted across token positions, while recurrent-depth samplers refine multiple token drafts concurrently using per-loop input injection and decodable intermediate states.LoopFold instead runs multiple loops in parallel for the same token prefix using predicted inputs and output correction.Source: §2 · p. 2

Parallel correction methods

  • LoopFold learns a fixed-stage parallel computation optimized for next-token prediction, rather than iteratively recovering the sequential fine solution as classical Parareal does.Its predictor initializes later-loop inputs once, while its corrector incorporates newly available earlier-loop states.Source: §2 · p. 2

3. Method

LoopFold batches predicted states from multiple recurrent depths into parallel backbone evaluations, then corrects provisional outputs as earlier-loop results become available. Its predictor and corrector are jointly trained with the backbone for next-token prediction.

3. Method

  • The schedule 1 + 3 + 3 uses three sequential stages and width 3 to reach target depth index D=5 with seven backbone evaluations.It first computes depth 1, then predicts provisional states at indices 2 and 3; later stages produce corrected states through indices 4 and 5.Source: §3 · pp. 2–3, Figure 1 · p. 3

3.3. Training Objective and Computational Cost

  • For parallel width n and S stages, schedule 1 + n + ··· + n reaches target depth index D=n+S−1 using 1+n(S−1) backbone evaluations across S sequential stages.All components are trained jointly using only next-token language-modeling loss; the two maps add 2d^2 parameters, or 8.39M at 1.4B.Source: §3 · p. 3, §3.3 · p. 4

3.1. Predicting Inputs for Later Loops

  • A shared token-wise predictor transforms the first loop’s state change and accumulates increments to initialize inputs for later loops.At identity initialization, this is linear extrapolation; prediction runs once to supply the n−1 states for the next batch.Source: §3.1 · p. 3

3.2. Parallel Evaluation and Correction

  • Each later stage batches one sequential state with provisional states, evaluates the n lanes in parallel, and corrects stale outputs using changes in their estimated inputs.The shared corrector translates input changes into output adjustments, and corrected states propagate newly available information in depth order.Source: §3.2 · p. 4

4. Experiments

The experiments compare LoopFold with baselines under matched training-compute, comparable-latency, and matched-target-depth regimes, keeping those comparisons distinct. Models are evaluated on language-modeling perplexity and zero-shot accuracy, alongside measured decoding latency.

4. Experiments

  • The study uses matched-training-compute, comparable-latency, and matched-target-depth comparisons because each holds a different quantity fixed.The regimes are reported separately rather than interpreted interchangeably.Source: §4 · p. 5

4.1. Experimental Setup

  • The main models are 1.4B parameters, trained from scratch on the Pile for a common 30B-token budget; ablations use 410M models trained on 25B tokens.Schedules are trained separately with the same tokenizer, data order, and training recipe.Source: §4.1 · p. 5
  • Baselines include sequential vanilla looped models, PLT without KV sharing, and Hidden Decoding with 5, 7, or 9 embedding streams.Recurrent-depth parallel samplers are omitted because they require per-loop input injection that this backbone does not use.Source: §4.1 · p. 5
  • Evaluation reports Pile validation PPL and zero-shot accuracy on eight tasks, while decoding latency is measured in milliseconds per generated token with a 1,024-token prompt and continuation.Main latency comparisons use batch size 1; decoding runs on the same GPU in BF16 with tensor parallelism disabled.Source: §4.1 · p. 5

4.2. Parallel Width at Fixed Stage Count

The section tests whether wider parallel batches improve Pile perplexity when the number of sequential stages is fixed. Wider configurations improve PPL, but each added lane brings diminishing gains.

4.2. Parallel Width at Fixed Stage Count

  • Increasing parallel width at fixed stage count consistently improves PPL, but the gain per added lane decreases as width grows.The comparison fixes sequential stages at S = 3, 4, or 5 while varying width.Source: §4.2 · p. 5
  • Figure 2 varies width n along curves with fixed sequential stages S to compare Pile perplexity across LoopFold schedules.The caption describes schedules of the form 1 + n + · · · + n.Source: Figure 2 · p. 6
  • Figure 3 compares Pile perplexity against training compute and decoding latency for LoopFold and vanilla looped models.Panel (a) varies LoopFold stage count at width 2; panel (b) shows LoopFold schedules at target depths 5, 7, and 9.Source: Figure 3 · p. 6

4.3. Quality versus Training Compute and Decoding Latency

At matched training backbone compute and at comparable decoding latency, LoopFold achieves lower PPL and higher mean zero-shot accuracy than vanilla looped models. It also exceeds PLT and Hidden Decoding in mean accuracy at matched training budgets.

4.3. Quality versus Training Compute and Decoding Latency

  • At matched training backbone compute, LoopFold achieves lower PPL and higher mean zero-shot accuracy than vanilla looped models at 1.68 and 2.16 ZFLOPs.The vanilla looped model leads at 1.20 ZFLOPs, while LoopFold exceeds PLT and Hidden Decoding in mean accuracy across all three budgets.Source: §4.3 · p. 6, Table 2 · p. 7
  • At comparable decoding latency, LoopFold has lower PPL and higher mean accuracy than paired vanilla looped models at target depths 5, 7, and 9.The comparisons use the same number of sequential stages; the two deeper LoopFold schedules gain on all eight tasks.Source: §4.3 · p. 6, Table 3 · p. 7

4.4. Quality Across Target Depth Indices

Across target depths 5, 7, and 9, LoopFold configurations increase parallel width, sequential stages, and total backbone computation together. Mean zero-shot accuracy rises more across these configurations than for vanilla looped models and PLT.

4.4. Quality Across Target Depth Indices

  • Figure 4 reports mean zero-shot accuracy at target depths D = 5, 7, and 9 for LoopFold configurations.The corresponding schedules are 1 + 3 + 3, 1 + 4 + 4 + 4, and 1 + 5 + 5 + 5 + 5; width, stage count, and total computation rise together.Source: §4.4 · p. 7, Figure 4 · p. 8
  • 56.38% at D = 9 versus 53.41% at D = 5 is LoopFold’s mean zero-shot accuracy, while vanilla looped and PLT results vary less.At D = 9, LoopFold also has 2.20 percentage points higher Avg-8 accuracy, lower PPL, and 46% lower decoding latency than the vanilla looped model, using additional parallel computation.Source: §4.4 · p. 8

4.5. Throughput with Larger Decoding Batches

LoopFold’s throughput advantage is measured against a five-loop vanilla model across decoding batch sizes. At batch size 1, it achieves 1.65× the baseline throughput, with per-sequence throughput remaining stable as batch size increases.

4.5. Throughput with Larger Decoding Batches

  • 1.65× is LoopFold’s throughput relative to the five-loop vanilla looped model at B = 1.LoopFold uses the 1 + 3 + 3 schedule, with seven backbone evaluations in three stages.Source: §4.5 · p. 8, Table 4 · p. 8
  • LoopFold throughput grows approximately linearly with batch size, maintaining stable per-sequence decoding throughput across the evaluated batch sizes.Table 4 reports throughput over batch sizes from 1 to 16.Source: §4.5 · p. 8

4.6. KV Sharing Across Recurrent Stages

LoopFold reduces KV-cache storage through reuse policies, with a modest quality or latency trade-off. Hard overwrite-64 refreshes recent cached states to recover most of the PPL loss from First-KV reuse.

4.6. KV Sharing Across Recurrent Stages

  • First-KV reuse retains the first evaluation’s cache at each loop depth; Hard overwrite-64 also refreshes the latest 64 positions using newly computed sequential states.Each variant uses the same cache policy in training and inference.Source: §4.6 · p. 9
  • The comparison uses a 1.4B model with a 1 + 2 + 2 + 2 schedule at target depth D=5.Timing uses batch 1, 1,024-token prompts and continuations, greedy decoding, BF16, and no tensor parallelism.Source: Table 5 · p. 9
  • 7× to 5× KV storage with First-KV reuse, at slightly higher PPL and decoding latency.Hard overwrite-64 retains this memory saving and recovers most of the PPL loss by refreshing the latest 64 cached positions, with a small further latency overhead.Source: §4.6 · p. 9, Table 5 · p. 9

4.7. Contributions of Prediction and Correction

The component analyses find that prediction and correction both improve PPL, with correction contributing more, and that learning both maps further improves on identity maps. State measurements show that correction is associated with stronger alignment to sequential references, while the complete model achieves better final PPL than repeated-H1 initialization.

4.7. Contributions of Prediction and Correction

  • The 410M component analyses use a 1 + 4 + 4 + 4 schedule to examine component removal, identity maps, and provisional-state evolution.The variants are trained separately on 25B Pile tokens.Source: §4.7 · p. 9, Table 6 · p. 9, Table 7 · p. 10

4.7.1. Component Removal

  • Removing either component worsens PPL under the same backbone budget, with the larger penalty from removing correction.Without correction, raw backbone outputs pass to the next stage rather than incorporating newly available earlier-loop results.Source: §4.7.1 · p. 9, Table 6 · p. 9

4.7.2. Learning the Prediction and Correction Maps

  • Learning both maps further lowers PPL beyond identity variants, while linear extrapolation and direct discrepancy addition already outperform removing their corresponding components.Each identity variant fixes one map at identity throughout training while learning the backbone and the other map.Source: §4.7.2 · pp. 9–10, Table 7 · p. 10

4.7.3. Evolution of Learned States

  • At depth index j=4, the complete model’s state moves closer to its sequential reference across stages, while removing correction leaves weaker alignment.The table compares provisional states with each variant’s own sequential backbone states on the same input.Source: §4.7.3 · p. 10, Table 8 · p. 10
  • Repeated-H1 initialization starts closer to its own sequential states at depth indices 2 and 3, yet the complete model achieves better final PPL.The authors suggest that minimizing instantaneous distance to the sequential trajectory may not be necessary for good parallel refinement.Source: §4.7.3 · p. 10

5. Conclusion

LoopFold combines a backbone, predictor, and corrector to shift recurrent computation into parallel stages. The conclusion reports improved language modeling and zero-shot performance with fewer sequential stages, while noting increased computation and KV-cache costs.

5. Conclusion

  • LoopFold combines a jointly trained backbone, predictor, and corrector to fold recurrent computation into parallel stages.The conclusion says recurrent depth and sequential decoding depth need not coincide when provisional states are subsequently corrected.Source: §5 · p. 11
  • Experiments show improved language modeling and zero-shot downstream performance with fewer sequential backbone stages, and the component analyses highlight correction’s importance.The conclusion frames this as shifting additional recurrent computation onto the parallel dimension when provisional states are corrected.Source: §5 · p. 11
  • Reducing sequential depth costs additional backbone computation and higher FLOPs, and increases KV-cache requirements despite mitigation through sharing.The reported KV-sharing experiments show a modest trade-off in quality and/or latency.Source: §5 · p. 11

Related papers

Loading 2610.35556v1…