PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopSpec: Pipelined Self-Speculative Decoding for Looped Transformers

SangLyul Cho, Langqing Cui, Sehoon Kim, Dongsu Han, Insu Han

arXiv:2609.17184v1cs.LGcs.CL

TL;DR

Repeated access to shared weights makes Looped Transformer decoding slow, despite their compact parameter counts. LOOPSPEC uses early recurrent states to draft tokens in a lossless speculative pipeline, achieving speedups of up to 6.83× across reasoning, math, and coding benchmarks.

  • Problem

    Looped Transformers reduce parameter counts through recurrently shared blocks, but repeated weight access at each recurrent depth increases decoding latency.

    Source: §1 · p. 1

  • Method

    LOOPSPEC is a training-free self-speculative pipeline that drafts from early recurrent states and uses a gated residual proposal from a deeper state as a fallback.

    Source: §1 · p. 2

  • Results

    Up to 6.83× speedup across the Ouro and Raven families on general reasoning, math, and coding benchmarks.

    Source: §7 · p. 10

  • Takeaways & Limitations

    LOOPSPEC converts recurrent depth into token-position parallelism while retaining lossless decoding through cascade verification.

    Source: §7 · p. 10

    The optimal-depth analysis assumes a non-increasing power-law envelope that upper-bounds expected total variation distance between intermediate and target distributions.

    Source: §5 · p. 7

Abstract

from arXiv · show

Looped Transformers achieve strong performance with compact parameter sizes by repeatedly applying a shared stack of Transformer blocks across recurrent depths. However, they incur higher decoding latency than standard Transformer models of comparable parameter size because shared weights are accessed at every recurrent depth. To improve decoding efficiency, self-speculative decoding is particularly well suited to Looped Transformers, as their intermediate recurrent states can directly provide draft predictions without an auxiliary draft model. We therefore propose LoopSpec, a training-free self-speculative decoding framework tailored for Looped Transformers. LoopSpec extracts draft tokens from early recurrent states and operates in a pipelined manner, overlapping draft generation of future tokens with target verification of the current token. To improve draft accuracy without excessive compute overhead, we introduce a selective second proposal from deeper recurrent depth while ensuring lossless decoding under both greedy and sampling regimes. Furthermore, we derive the optimal proposal depths in closed form and show the prediction matches measurement. Across reasoning and coding benchmarks, LoopSpec achieves up to 6.83$\times$ inference speedup across diverse Looped Transformers.

1 INTRODUCTION

Looped Transformers reduce parameter counts through recurrently shared blocks, but repeated weight access increases decoding latency. LoopSpec overlaps speculative token generation with verification and reports up to 6.83× speedup across Looped Transformer families.

1 INTRODUCTION

  • Repeatedly accessing shared Transformer blocks across recurrent depths makes Looped Transformer decoding memory-bandwidth bound and increases latency.Four blocks applied eight times provide the effective depth of a 32-layer model while storing only four layers of parameters.Source: §1 · p. 1
  • Intermediate recurrent states make Looped Transformers suitable for self-speculation, which avoids a separate draft model and its training costs, distribution mismatch, and extra memory.Each recurrent depth provides an increasingly refined proposal toward the final prediction.Source: §1 · p. 1
  • LOOPSPEC drafts from an early recurrent state and starts computing the next token before the current token reaches final-depth verification.If the draft is rejected, the speculative continuation is pruned and computation restarts from the verified prefix.Source: §1 · p. 2
  • A gated residual proposal from a deeper recurrent state provides a fallback when the first draft is rejected, creating fewer pipeline stalls.The gate creates the fallback continuation only when the deeper state disagrees with the first draft.Source: §1 · p. 2, Figure 1 · p. 2
  • The analysis shows restart probability depends on the second proposal depth, not the first, and derives closed-form optimal depths that match measured optima.The depth-selection analysis uses the empirically observed power-law decay of intermediate-to-target total variation distance.Source: §1 · p. 2
  • 6.83× lossless speedup across seven checkpoints from two Looped Transformer families; LOOPSPEC also outperforms training-based DFlash by up to 1.7×.The implementation uses branch-wise KV management and CUDA Graph optimization.Source: §1 · p. 2

2 RELATED WORK

The section reviews speculative decoding, recurrently shared Transformer architectures, and approaches to reducing their inference cost, including their respective trade-offs.

Speculative Decoding

  • Speculative decoding drafts candidate tokens for parallel target verification; self-speculative methods instead draft from the target model without an external drafter (Chen et al., 2023).Draft & Verify skips intermediate layers, while LayerSkip uses early exits.Source: §2 · p. 3

Looped Transformers

  • Looped Transformers repeatedly apply shared blocks to increase effective depth without proportionally increasing parameters, but recurrent computation raises inference memory traffic and limits throughput (Geiping et al., 2025; Zhu et al., 2025; McLeish et al., 2026; Park et al., 2026; Nanbeige Lab et al., 2026).Source: §2 · p. 3
  • LT2 changes attention and distills hybrid checkpoints, while Think-at-Hard adjusts recurrent passes with routing; SPEED pipelines decoding but requires custom training and supports only greedy decoding (Deng et al., 2026).Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries describe recurrent computation in Looped Transformers and the rejection-sampling rule used to make speculative decoding lossless.

3.1 LOOPED TRANSFORMERS

  • A Looped Transformer applies pre-layers, recurrent blocks, and post-layers to produce a next-token distribution; Ouro places all Transformer blocks in the recurrence, while Raven places some in pre- and post-layers.The post-layers include the LM head and output-distribution processing.Source: §3.1 · pp. 3–4
  • Intermediate readouts are distributions at recurrent depth r, and the target distribution pR is the readout at the model's final depth R.Source: §3.1 · pp. 3–4

3.2 REJECTION SAMPLING AND SPECULATIVE DECODING

  • Speculative decoding samples a candidate from proposal q and accepts it according to its probability under target p relative to q; after rejection, it samples from the residual distribution.Source: §3.2 · p. 4
  • Rejection sampling yields a token distributed exactly as a target-model sample for any valid proposal, making the procedure lossless.Source: §3.2 · p. 4

4 LOOPSPEC: PIPELINED SELF-SPECULATIVE DECODING FOR LOOPED TRANSFORMERS

LOOPSPEC uses intermediate recurrent states as drafts and overlaps their computation with final-depth verification. A second, gated proposal and shared-block batching extend the pipeline beyond a single draft depth.

4.1 OBSERVATION: EARLY RECURRENT STATES ARE EFFECTIVE DRAFTERS

  • On Raven-Llama-3.2, depth-1 greedy agreement on GSM8K exceeds 86% and surpasses 99% by depth 8; intermediate-to-target TV distance also decreases sharply with depth.Figure 3 compares greedy top-1 agreement and sampling TV distance across recurrent depths on GSM8K and MATH-500.Source: §4.1 · p. 4, Figure 3 · p. 5
  • These observations motivate using early recurrent states to draft future tokens while recurrence continues toward verification, with deeper states available as fallback proposals.Source: §4.1 · p. 4

4.2 PIPELINED SELF-SPECULATIVE DECODING WITH A SINGLE PROPOSAL

  • In the single-proposal pipeline, a branch drafts at depth d1 < R, starts a child branch, and continues toward final-depth verification; active branches advance together through batched shared-block calls.The draft computation can start the next token before the parent branch is verified.Source: §4.2 · p. 5

3. Verification and Continuation. When the branch reaches R, it verifies ex(1)

LoopSpec verifies an early draft at full recurrent depth, but a rejection discards speculative progress and restarts the next-token branch. A deeper residual proposal and selective gating provide a fallback path that can preserve work after rejection.

3. Verification and Continuation. When the branch reaches R, it verifies ex(1)

  • At depth R, the target distribution verifies the first draft; rejection commits a replacement and restarts the next-token branch at depth 0.A rejection discards speculative descendants, motivating a more accurate fallback draft.Source: §3 · p. 5

4.3 RESIDUAL SECOND PROPOSAL WITH GATING MECHANISM

  • The residual second proposal consistently improves decoding speedup over the non-residual second proposal across benchmarks.The fallback proposal is constructed from the residual target distribution.Source: §4.3 · p. 6
  • The gate creates a fallback branch only when the deeper state disagrees with the primary candidate, limiting the growth that unconditional branching would cause.Under greedy decoding, the gate opens when the deeper readout’s argmax differs from the first draft.Source: §4.3 · p. 6
  • Cascade verification checks the primary candidate before the fallback and samples a residual replacement if neither is accepted.When the gate is closed and the first proposal is rejected, the replacement is sampled from ρ1.Source: §4.3 · p. 6
  • At d2, a branch drafts an alternative from an estimate of pR and starts a fallback child; accepting it preserves completed recurrent steps, reducing the next commitment to d2 more steps.The second proposal estimates the residual target needed after rejecting the first proposal.Source: §4.3 · p. 6

5 THEORETICAL ANALYSIS

The analysis bounds rejection using only the second proposal depth’s approximation error, then derives proposal depths under a power-law error envelope. The predicted depths closely match the empirically best configuration for Raven-Llama-3.2.

5 THEORETICAL ANALYSIS

  • With residual proposals, the probability both candidates are rejected is at most ε2, independent of d1; gating raises the bound to at most 2ε2.The rejection bound depends on the total variation distance between pd2 and pR.Source: §5 · p. 7
  • Assuming a non-increasing power-law envelope for intermediate-to-target total variation distance, the analysis derives proposal depths minimizing expected recurrent steps per committed token.The power-law envelope is ϕ(r) = βr^-α.Source: §5 · p. 7
  • For Raven-Llama-3.2 with R = 32, predicted depths (1.26, 8.85) closely match the empirically best depths (2, 8).The prediction uses β = 0.175 and α = 1.03 from the empirical upper bound.Source: §5 · p. 7

6 EXPERIMENTS

Experiments evaluate LOOPSPEC on Ouro and Raven models across reasoning, math, and coding benchmarks, with greedy and sampling settings for reasoning models. The reported results include speedups over autoregressive decoding and DFlash, plus an ablation of the second proposal and gating mechanism.

6 EXPERIMENTS

  • The evaluation covers Ouro with R = 4 and Raven with R = 32 across mathematical reasoning, general reasoning, and code generation; reasoning models use both greedy and sampling settings.Base models use greedy decoding; evaluations use a single-batch setting.Source: §6 · p. 7

6.1 BASE MODELS

  • Up to 6.83× speedup is reported for Raven base models, with mean accepted length γ exceeding 14.5; the second proposal compensates for lower acceptance at d1 = 2.For Ouro base models, speedups reach 3.36× with γ exceeding 3.5 out of R/d1 = 4.Source: Table 1 · p. 8, §6.1 · p. 8

6 EXPERIMENTS

  • Mean accepted length γ counts generated tokens per full R-step model execution, but excludes pre- and post-layer computation, which can make higher γ coexist with lower wall-clock speedup.Wall-clock speedup is measured during decoding, after prompt prefilling.Source: §6 · p. 8, §6.1 · p. 9

6.2 REASONING MODELS

  • Up to 2.95× speedup is reported on long-context reasoning benchmarks, while stochastic sampling reaches 2.83×.On Ouro-2.6B-Thinking under GSM8K, the residual second proposal raises γ from 3.11 to 3.54.Source: Table 2 · p. 9, §6.2 · p. 9

6.3 ABLATION STUDY

  • Residual second proposal and gating improve speedups from 3.24×–4.59× to 3.74×–5.00× on Raven-Llama-3.2 under sampling.The ablation masks the first proposal token for the non-residual and residual variants to ensure the second proposal differs from the first.Source: Table 3 · p. 9, §6.3 · pp. 9–10
  • LOOPSPEC consistently outperforms DFlash, reaching up to 1.7× decoding-time speedup on Ouro thinking models under greedy and stochastic sampling.Figure 4 compares decoding speedup over standard autoregressive decoding for both methods.Source: Figure 4 · p. 10, §6.3 · p. 10

7 CONCLUSION

LOOPSPEC converts recurrent depth into token-position parallelism through early-state drafting and verification, with a gated residual second proposal and cascade verification. Across Ouro and Raven families, it achieves speedups up to 6.83× on reasoning, math, and coding benchmarks.

7 CONCLUSION

  • LOOPSPEC converts recurrent depth into token-position parallelism using early-state drafts and a gated residual second proposal with cascade verification.The framework is training-free and is evaluated across Ouro (R = 4) and Raven (R = 32), reaching up to 6.83× speedup.Source: §7 · p. 10

A EXPERIMENTAL SETUP DETAILS

The appendix records evaluation configurations and implementation details for pipelined decoding, including branch-specific KV-cache handling and execution graphs for changing batch sizes.

A EXPERIMENTAL SETUP DETAILS

  • Table 4 maps evaluated checkpoints to the model names used in the paper.Source: Table 4 · p. 14
  • Benchmark specifications, prompting protocols, and maximum generation budgets are provided for base and reasoning models.Reasoning evaluations enable thinking mode, while code-generation prompts are zero-shot and request self-contained Python scripts.Source: Table 5 · p. 14, A EXPERIMENTAL SETUP DETAILS · p. 14

B IMPLEMENTATION DETAILS OF PIPELINED DECODING

  • Each branch keeps separate KV caches by recurrent depth, and new branches reuse cached prefixes by copying KV indices rather than KV states.After verification, the selected branch and descendants are retained while unneeded KV states are released.Source: B IMPLEMENTATION DETAILS OF PIPELINED… · p. 14
  • CUDA Graphs are prepared for supported batch sizes, while depth choices require d2 and R to be multiples of d1 with d1 < d2 < R.This aligns proposal and verification stages and allows pre- and post-layer computations to be batched.Source: B IMPLEMENTATION DETAILS OF PIPELINED… · pp. 14–15

C DFLASH IMPLEMENTATION DETAILS

For comparison with LOOPSPEC, the authors train DFlash drafters on Ouro thinking models and use matching inference hyperparameters and SGLang configurations.

C DFLASH IMPLEMENTATION DETAILS

D EMPIRICAL RESULTS ON GATING STRATEGY AND EFFECTIVE BATCH SIZE

Gating reduces effective batch size and selectively limits the creation of additional speculative branches. Gate-opening rates vary across models and benchmarks.

D EMPIRICAL RESULTS ON GATING STRATEGY AND EFFECTIVE BATCH SIZE

E TOKEN COMMITMENT SOURCES ACROSS PROPOSAL DEPTHS

First proposals commit most tokens, and second proposals recover more tokens than full-depth fallback across the evaluated models and benchmarks. A power-law fit for Raven-OLMo-2-0425 predicts proposal depths close to measured optima.

E TOKEN COMMITMENT SOURCES ACROSS PROPOSAL DEPTHS

F EMPIRICAL POWER-LAW DECAY AND OPTIMAL PROPOSAL DEPTHS FOR RAVEN-OLMO-2-0425

G PROOFS

  • Total variation distance is used to evaluate speculative decoding and measures a probability of rejection.Source: G PROOFS · p. 18

G.1 PROOF OF THEOREM 1

  • Theorem 1’s proof bounds joint rejection using residual-distribution TV errors, with the gate-closed case yielding a 2ε2 bound.Source: §G.1 · p. 19

G.2 PROOF OF THEOREM 2

  • Theorem 2’s proof minimizes expected recurrent-step cost, assigning d1, d2, or R according to first-proposal, second-proposal, or fallback acceptance.Source: §G.2 · p. 19

H LOOPSPEC ALGORITHM PSEUDOCODE

Related papers

Loading 2609.17184v1…