Parallel Loop Transformer for Efficient Test-Time Computation Scaling

• Vanilla looped transformers reuse weights to gain effective depth without increasing storage, but sequential loops make compute, latency, and KV-cache size scale as O(L).This scaling makes them impractical for latency-sensitive applications.Source: §1 · pp. 1–2
• Cross-Loop Parallelism computes different loops for different tokens concurrently, collapsing L sequential steps into one forward pass.During decoding, loop l for the current token overlaps with loop l + 1 for the previous token; Figure 1b depicts loops within a dashed box computed in parallel.Source: §1 · p. 2, Figure 1 · p. 2
• Efficient Representation Enhancement shares the first loop’s KV cache across subsequent loops and uses gated sliding-window attention to add local context.This strategy reduces KV-cache memory to the level of a single loop while combining shared global and local information.Source: §1 · p. 2
• PLT significantly outperforms vanilla transformer baselines in accuracy while adding minimal latency.Source: §1 · p. 2
• A 1.7B-parameter PLT achieves superior performance and lower latency than a 2.5B-parameter vanilla model.Source: §1 · p. 2

2.1 Preliminaries: Vanilla Loop Transformer

• Vanilla loop transformers reuse weights for greater effective depth, but serial loops make per-token compute, latency, and KV-cache size scale as O(L).The source notes that this parameter efficiency does not reduce latency and can be disadvantageous under equal-latency budgets.Source: §2.1 · p. 3

2.2 Parallel Loop Transformer

• PLT combines cross-loop parallelism with efficient representation enhancement to shorten inference while retaining parameter savings.The design parallelizes L loops and uses shared KV plus gated sliding-window attention to limit time and cache overhead.Source: §2.2 · p. 3

2.2.1 Cross-Loop Parallelism

• Cross-loop parallelism overlaps later-loop computation on earlier tokens, letting PLT process L tokens in one forward pass rather than L sequential passes per token.At each later loop, PLT right-shifts the preceding loop’s states, adds the original embeddings, and reuses first-loop K and V; this removes same-index dependencies across loops.Source: Figure 2 · p. 3, §2.2.1 · pp. 4–5

2.2.2 Efficient Representation Enhancement

• First-loop KV sharing gives later loops global information without separate global caches, while gated sliding-window attention fuses that shared context with private local attention.The local window is w=64 in the experiments; later loops cache only the most recent w KV entries, and the head-wise scalar gate adds negligible parameters and computation.Source: §2.2.2 · p. 5

2.3 Inference Efficiency Analysis

• In the memory-bound analysis, sequential loops raise latency from t to Lt, whereas concurrent loops retain latency ≈t with compute LC.Adding sliding-window attention keeps extra cache and compute overhead small when w ≪ L, and PLT avoids L-fold KV-cache growth.Source: §2.3 · pp. 5–6, Table 1 · p. 5

3.1.1 Settings

• The component and scalability study trained 680M/13B MoE models on 150B tokens, varying PLT loops between 2 and 3 and setting the G-SWA window to 64.It added CLP, KV-cache sharing, and G-SWA sequentially to a vanilla looped Transformer; efficiency tests used FP8 self-attention, W4A8 linear-layer quantization, and batch sizes 4–64.Source: §3.1.1 · p. 6

3.1.2 Results and Analysis

• CLP kept looped-model accuracy nearly unchanged (39.7→39.6) while reducing latency from 9.4 ms to 5.9 ms (−37%).The two-loop model without CLP improved average accuracy over vanilla by 5.0 points, but latency rose from 4.8 ms to 9.4 ms.Source: §3.1.2 · p. 7
• KV-cache sharing cut cache from 560M to 280M (−50%), while adding G-SWA restored accuracy from 36.2 to 39.7 with 1.4% cache overhead.Latency changed from 5.9 to 4.8 ms with sharing, then from 4.8 to 4.9 ms with G-SWA.Source: §3.1.2 · p. 7
• PLT-3 raised average accuracy to 40.8 (+1.1 over PLT-2), with latency at 4.9 ms and KV cache at 287M.PLT-2 matched the naive loop transformer's 39.7 accuracy with latency of 4.9 ms and KV cache of 284M.Source: §3.1.2 · p. 7

3.1.3 Inference Efficiency

• At batch sizes 32 and 64, PLT-2 cut per-token latency by 47% versus naive looping, remaining within 4–6% of vanilla latency.At batch sizes 4 and 8, PLT-2 stayed near vanilla at 1.02×; adding KV sharing to CLP reduced latency by 19% at both sizes.Source: §3.1.3 · p. 7, Table 3 · p. 7

3.2 Latency Comparisons under Same Accuracy

• 62.6 average accuracy—0.5 points above the vanilla 2.5B/60B MoE—came from PLT-2 1.7B/40B 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.Source: §3.2 · pp. 7–8, Table 4 · p. 8

4.1 Latent Reasoning

• Looped Transformers use shared weights across layers or allocate extra depth to complex tokens, supporting reasoning, parameter efficiency, and length generalization [25] [32] [13] [43] [7] [14].Prior work established recurrent layer adoption, demonstrated reasoning and datafitting capabilities, and used token-level looping to increase inference-time reasoning without increasing parameter count.Source: §4.1 · p. 8
• Latent CoT adds discrete non-semantic tokens or compresses reasoning into continuous latent thoughts to improve reasoning capability [46] [16] [35].These methods are described as improving performance on reasoning benchmarks, including GSM8K, NaturalQuestions, and CommonsenseQA.Source: §4.1 · p. 8

4.2 Parallelized Computation

• PHD parallelizes repeated-token forwards and shares KV cache with chunk-wise sliding-window attention, but repeated decoding can reduce throughput in high-throughput serving [42].The paper argues that repeated tokens use parallel computation inefficiently because their hidden representations in earlier layers are similar.Source: §4.2 · p. 9
• ParScale’s sequence repetition with prefix tuning can improve performance, but activating P inference streams incurs P× KV-cache footprint and added latency [6].The paper identifies these costs as especially consequential in high-throughput serving scenarios.Source: §4.2 · p. 9
• StagFormer’s time-staggered layer parallelism remains incomplete, while its weight-sharing variant adds KV-cache and cross-attention overhead; PLT reduces KV-cache footprint by over 50% with no performance degradation [10].The separate-weight variant doubles hardware usage for less than 50% throughput improvement; the paper attributes PLT’s efficiency to parallelism and sharing across the looped Transformer.Source: §4.2 · p. 9
• PLT overlaps different loops across different tokens, aiming to retain looped Transformers’ performance while improving inference efficiency.The paper describes looped Transformers and latent CoT as inherently sequential, whether computation proceeds across loops or tokens.Source: §4.1 · p. 8, §4.2 · p. 9

5 Conclusion

• PLT overlaps computation and memory access across latter loops in previous tokens and previous loops in latter tokens through Cross-Loop Parallelism.This is the architecture’s parallelism strategy.Source: §5 · p. 9
• Gated sliding-window attention parallelizes KV-cache access with no performance degradation.The conclusion identifies gated-SWA as the complementary attention mechanism.Source: §5 · p. 9
• PLT delivers impressive performance improvement with negligible latency overhead versus vanilla Transformers under memory-access bottlenecks.Compared with vanilla looped transformers at the same inference computation budget, it provides better inference efficiency with similar or better performance.Source: §5 · p. 9

A Experiments opensource

• Open-source experiments cover both dense and MoE model series in Section A.2.The models have over 1 billion activated parameters.Source: A Experiments opensource · p. 14
• Section A.3 presents a dense-model ablation study complementary to the in-house Seed-MoE experiments.Source: A Experiments opensource · p. 14

A.1 Evaluation Datasets

• The evaluation included MMLU [17], HellaSwag [47], and ARC [8].Source: §A.1 · p. 14
• The evaluation also included PIQA [2] and Winogrande.Source: §A.1 · p. 14
• CommonsenseQA [36] was another dataset used in the evaluation.Source: §A.1 · p. 14

A.2.1 Training and Evaluation Settings

• The OLMo- and OLMoE-based implementation compares PLT with vanilla and vanilla looped transformers using the same activated parameters.Source: §A.2.1 · p. 14
• Dense models use 16 layers, hidden dimension 2048, MLP hidden dimension 16384, FFN output norm, GQA with 32 query and 8 key/value heads, and tied embedding/LM-head weights.They use query/key/value layer normalization and no attention output norm; norm settings were slightly modified from OLMo-1B for training stability [50].Source: §A.2.1 · pp. 14–15
• Table 6 presents a PLT component analysis against the vanilla looped transformer and defines Lat. as latency in milliseconds.Source: Table 6 · p. 15
• The model is trained for 400B tokens total, with global training batch size 1024, using a cosine learning-rate schedule, a peak learning rate of 3e-4, and 2000 warmup steps.Source: §A.2.1 · p. 15
• MoE models use 16 layers, hidden dimension 2048, SwiGLU experts with the 8-in-64 recipe, 16-head MHA, and tied embedding/LM-head weights.Source: §A.2.1 · p. 15
• The model is trained for 400B tokens with global training batch size 1024, a cosine learning-rate schedule, peak learning rate 4e-4, and 2500 warmup steps.Source: §A.2.1 · p. 15
• Efficiency tests use VLLM [23] on one GPU, varying prefilling lengths of 1024 and 2048 and serving batch sizes from 1 to 256 with FP8 quantization.Latency averages five independent runs decoding 256 tokens; the varied batch sizes span memory-access- to compute-bottleneck regimes, and both latency and throughput are analyzed.Source: §A.2.1 · p. 15

A.2.2 Analysis

• PLT-2 achieves very similar performance to vanilla looped transformers on Dense and MoE models, while showing obviously better efficiency.Table 5 presents the variant performance, and Figure 4 presents the efficiency analysis.Source: §A.2.2 · p. 15

A.3.1 Training/Evaluation Settings

• The Dense recipe changes GQA to 16 query heads and 4 key-value heads, with hidden dimension 1536 and MLP hidden size 8192.Other training settings mainly follow Section A.2.1.Source: §A.3.1 · p. 15
• Dense-model training uses 100B tokens, a 3e-4 peak learning rate, cosine scheduling, and 2000 warm-up steps.Source: §A.3.1 · p. 15
• The MoE recipe changes the layer count to 12 and hidden dimension to 1536.Other training settings mainly follow Section A.2.1.Source: §A.3.1 · p. 15
• MoE training dynamics use 100B tokens, a 3e-4 peak learning rate, cosine scheduling, and 2000 warm-up steps.Source: §A.3.1 · p. 16
• Evaluation follows Section A.2.1 settings except that it uses one H800 GPU.Source: §A.3.1 · p. 16

A.3.2 Analysis

• The OLMo2 ablation study reports observations similar to those in the in-house experiments.The results are presented in Table A.3.1.Source: §A.3.2 · p. 16