PaperlayerView exact revision on arXiv

Source-linked AI summary

Parallel Loop Transformer for Efficient Test-Time Computation Scaling

Bohong Wu, Mengzhao Chen, Xiang Luo, Shen Yan, Qifan Yu, Fan Xia, Tianqi Zhang, Hongrui Zhan, Zheng Zhong, Xun Zhou, Siyuan Qiao, Xingyan Bin

arXiv:2510.24824v1cs.CL

TL;DR

Inference costs constrain practical LLM deployment, while sequential looped transformers increase latency and KV-cache requirements with loop count. PLT parallelizes loop computation across tokens and combines shared first-loop KV caches with gated sliding-window attention. Experiments report looped-model accuracy with negligible latency and memory overhead relative to standard transformers.

  • Problem

    LLM inference can be computationally costly, and vanilla looped transformers scale latency and KV-cache size linearly with the number of loops.

  • Method

    PLT uses Cross-Loop Parallelism to overlap loops across tokens and Efficient Representation Enhancement with shared first-loop KV caches and gated sliding-window attention.

  • Results

    PLT significantly outperforms vanilla transformer baselines in accuracy with minimal latency, while preserving loop-level accuracy near vanilla efficiency.

  • Takeaways & Limitations

    PLT enables scalable test-time computation by decoupling loop count from wall-clock latency and keeping inference overhead close to vanilla transformers.

Abstract

from arXiv · show

Large Language Models (LLMs) are powerful but often too slow and costly for real-world use during inference. Looped transformers save on parameters by reusing the same weights for multiple computational steps, or "loops." However, this approach has a major flaw: the loops run one after another, causing inference latency and memory requirements to increase with each added loop. This makes them impractical for fast applications. To solve this problem, we introduce the Parallel Loop Transformer (PLT). PLT is a new architecture that delivers the performance benefits of a deep, looped model but with the low latency of a standard, non-looped model. PLT works using two key techniques. First, Cross-Loop Parallelism (CLP) breaks the sequential dependency by computing different loops for different tokens at the same time, all within a single pass. Second, to prevent memory costs from growing, we use an Efficient Representation Enhancement strategy. This method shares the memory (KV cache) from the first loop with all other loops. It then uses a Gated Sliding-Window Attention (G-SWA) to combine this shared global information with local information, maintaining high accuracy. Our experiments show that PLT achieves the high accuracy of a traditional looped model but with almost no extra latency or memory cost compared to a standard transformer.

1 Introduction

PLT addresses the inference bottlenecks of sequential looped transformers by parallelizing different loops across tokens while preserving parameter-efficient deep computation. Its experiments report higher accuracy than vanilla transformers with minimal latency overhead, including a smaller PLT outperforming a larger vanilla model.

  • Vanilla looped transformers reuse weights for greater effective depth but incur O(L) scaling in per-token compute, wall-clock latency, and KV-cache size.
  • PLT breaks sequential loop dependence by computing different loops for different tokens in parallel within a single forward pass.
  • Cross-Loop Parallelism overlaps the l-th loop of the current token with later loops of previous tokens, collapsing L sequential steps into one.
  • PLT combines Cross-Loop Parallelism with KV-cache sharing and gated sliding-window attention to overcome looped models’ latency and memory bottlenecks.
  • PLT significantly outperforms vanilla transformer baselines in accuracy while adding minimal latency across in-house and open-source models.
  • A 1.7B-activated-parameter PLT achieves superior performance and lower latency than a 2.5B-activated-parameter vanilla model.

2 Method

PLT converts sequential loop computation into cross-loop parallel execution while sharing first-loop KV caches and adding gated local attention to retain efficiency and representation quality.

  • Vanilla Loop Transformer: Vanilla loop transformers execute loops sequentially, so per-token compute, latency, and KV-cache size scale as O(L) despite parameter sharing.This makes their advantage primarily relevant under equal-parameter rather than equal-latency comparisons.
  • Cross-Loop Parallelism: Right-shifting previous-loop states and adding the original embeddings removes same-index dependencies between consecutive loops, enabling parallel decoding.PLT repeats this transformation for loops 2 through L while reusing the first loop’s shared KV cache.
  • Cross-Loop Parallelism: PLT overlaps different loops across different tokens, replacing L sequential inference passes with one parallel pass per token.For L=3, decoding uses a displaced micro-batch containing successive loop computations for consecutive tokens.
  • Efficient Representation Enhancement: PLT shares the first loop’s Kshare and Vshare across later loops, reducing storage to one global KV cache while preserving global information.Non-first loops retain private queries and attend globally to the shared first-loop keys and values.
  • Efficient Representation Enhancement: Gated sliding-window attention combines private local attention with shared global attention in non-first loops using a sigmoid gate and fixed window w=64.The window size does not increase with sequence length, keeping additional memory and computation small when w ≪ L.
  • Inference Efficiency: PLT maintains near-baseline decoding latency and avoids the L-dependent KV-cache growth while adding computation intended to improve accuracy.The design targets the memory-bound regime, where extra decoding FLOPs have negligible latency impact.

3 Experiments

The experiments evaluate PLT against vanilla and looped transformers for accuracy, latency, memory, and scalability. PLT preserves loop-level accuracy while keeping latency and KV-cache overhead near vanilla transformers, including at matched accuracy with a shallower model.

  • 3.1 Experiments: +5.0 average-accuracy points from two-loop computation, while CLP reduces latency from 9.4 to 5.9 ms with nearly unchanged accuracy.CLP retains a +4.9-point gain over vanilla at approximately 1.23× latency.
  • 3.1 Experiments: KV-cache sharing cuts cache from 560M to 280M and latency from 5.9 to 4.8 ms, while G-SWA restores accuracy from 36.2 to 39.7 with only 1.4% cache overhead.Adding G-SWA changes latency from 4.8 to 4.9 ms.
  • 3.1 Experiments: 40.8 average accuracy with PLT-3, improving 1.1 points over PLT-2 while increasing latency only from 4.8 to 4.9 ms and KV cache from 284M to 287M.PLT-2 matches the naive loop transformer at 39.7 accuracy with near-vanilla efficiency.
  • 3.1.3 Inference Efficiency: 47% lower PLT-2 latency than naive looping at batch sizes 32 and 64, with residual overhead versus vanilla limited to 1.06× and 1.04×.PLT-2 remains within 4–6% of vanilla for batch sizes at least 32.
  • 3.2 Latency Comparisons under Same Accuracy: A 1.7B/40B PLT-2 MoE reaches 62.6 average accuracy, 0.5 points above the vanilla 2.5B/60B MoE, with about 30% lower latency and roughly two-thirds the KV cache.The PLT model uses two-thirds as many layers as the baseline.

4 Related Works

Related work studies looped and latent chain-of-thought transformers that add computation but remain inefficient because their computation is sequential. PLT instead overlaps loop computation across tokens to improve inference efficiency, addressing limitations in prior parallelization methods.

  • 4.1 Latent Reasoning: Looped Transformers reuse parameters across computational steps, providing effective depth and stronger reasoning without increasing model storage.Prior work reports parameter efficiency, data-fitting capability, and length generalization.
  • 4.1 Latent Reasoning: Both looped Transformers and latent CoT suffer inferior inference efficiency because their additional computation is inherently sequential.PLT overlaps different loops across different tokens.
  • 4.2 Parallelized Computation: PHD shares KV cache and uses chunk-wise sliding-window attention, but its repeated-token strategy is described as inefficient for parallel computation.The supplied passage truncates the final comparison detail.
  • 4.2 Parallelized Computation: ParScale introduces P× KV-cache overhead when P inference streams are active, especially in high-throughput serving.This creates both KV-cache and latency overhead.
  • 4.2 Parallelized Computation: StagFormer provides partial layer-level parallelism, but incomplete attention parallelism and memory access limit its efficiency gains.It splits layers into stacks and uses cross-attention to previous-step hidden states.
  • 4.2 Parallelized Computation: PLT applies cross-loop parallelism to looped Transformers, maintaining their performance while achieving improved inference efficiency.The paper presents this as a higher utilization of parallel computation.

5 Conclusion

PLT combines Cross-Loop Parallelism with gated sliding-window attention to parallelize loop computation and KV-cache access. The paper reports negligible latency overhead versus vanilla Transformers and better efficiency than vanilla looped Transformers at similar or better performance.

  • 5 Conclusion: PLT overlaps later-loop computation and memory access across tokens, while gated sliding-window attention enables parallel KV-cache access without performance degradation.These mechanisms form the paper’s core architecture.
  • 5 Conclusion: PLT reports negligible latency overhead versus vanilla Transformers in memory-access-bottlenecked settings.The conclusion frames this as improved parallel-computation utilization.
  • 5 Conclusion: PLT achieves clearly better inference efficiency than vanilla looped Transformers with the same inference computation budget and similar or better performance.This is the paper’s stated comparative conclusion.

6 Contributions and Acknowledgments

The supplied section lists names in alphabetic order. It contains author-name entries but no substantive contribution or acknowledgment claims.

  • 6 Contributions and Acknowledgments: The section lists names in alphabetic order.The supplied passages provide no additional section-level content.
  • 6 Contributions and Acknowledgments: The supplied entries list Mengzhao Chen, Xiang Luo, Bohong Wu, Shen Yan, Fan Xia, Tianqi Zhang, Hongrui Zhan, Zheng Zhong, Qifan Yu, and Xun Zhou.These are the names present in the provided passages.

A Experiments opensource

The open-source experiments cover dense and MoE model series, with additional ablation studies on dense models.

  • Open-source experiments evaluate both dense and MoE model series with over 1 billion activated parameters.
  • Ablation studies on dense models complement the in-house Seed-MoE experiments.

A.1 Evaluation Datasets

The evaluation uses six open-source datasets spanning language understanding, commonsense reasoning, and related benchmark tasks.

  • The evaluation includes MMLU, HellaSwag, ARC, PIQA, Winogrande, and CommonsenseQA.

A.2.1 Training and Evaluation Settings

The experiments compare PLT with vanilla and looped transformers under matched activated-parameter settings across dense and MoE models. Evaluation measures benchmark performance and serving efficiency across prefilling lengths and batch sizes, with PLT-2 matching looped-transformer performance while improving efficiency.

  • Evaluation setup: The open-source implementation uses OLMo and OLMoE, comparing PLT against vanilla and vanilla looped transformers with the same activated parameters.
  • Training setup: The dense and MoE recipes use 16 transformer layers, hidden dimension 2048, and 400B training tokens, with architecture-specific attention and MLP configurations.Dense models use GQA, while MoE models use SwiGLU experts with an 8 in 64 recipe.
  • Ablation analysis: The component analysis examines two extra PLT components for inference speed and performance, with latency reported in milliseconds.
  • Evaluation setup: Efficiency evaluation measures latency and throughput while varying prefilling length from 1024 to 2048 and serving batch size from 1 to 256.The settings shift from memory-access bottlenecks toward compute bottlenecks, and latency is averaged over five independent decoding runs.
  • Results: PLT-2 achieves very similar performance to vanilla looped transformers and obviously better efficiency than them.These observations hold across the reported dense- and MoE-model evaluations.
Loading 2510.24824v1…