PaperlayerView exact revision on arXiv

Source-linked AI summary

Shallow Queries, Mature Values: Depth-Asynchronous Self-Speculation for Looped Transformers

Guanghao Li, Zihan Su, Hao Yu, Jinyang Jiang, Tao Ren, Zehao Li, Feng Lu, Ming Tang, Chun Yuan

arXiv:2609.34538v1cs.LGcs.AI

TL;DR

Looped Transformers make autoregressive decoding costly through sequential recurrent passes, while self-speculative drafts typically read prefix representations from a similarly shallow depth. DAS separates draft depth from verified-prefix representation depth, and DAS-Wave achieves 4.00–6.96× mean throughput speedup over paired full-depth autoregressive decoding across four checkpoints and mathematics and code workloads.

  • Problem

    Autoregressive decoding in Looped Transformers is expensive because each generated token requires many sequential recurrent passes, while self-speculative drafts typically read prefix representations from the same or a similarly shallow recurrent depth.

    Source: §1 · p. 1

  • Method

    DAS decouples draft computation depth from verified-prefix representation depth, and DAS-Wave combines depth-asynchronous reads with carried parallel refinement and independent full-depth verification.

    Source: §1 · p. 2

  • Results

    4.00–6.96× mean throughput speedup over paired full-depth autoregressive decoding across four recurrent-model checkpoints and mathematics and code workloads.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results identify prefix-information depth as a distinct design axis for recurrent self-speculation.

    Source: §5 · p. 10

Abstract

from arXiv · show

Looped Transformers reuse a shared block across recurrent depths, making autoregressive decoding expensive because every generated token requires many sequential recurrent passes. Self-speculative decoders reduce this cost by drafting at an early depth and verifying at full depth, but typically bind draft computation to prefix representations from the same recurrent depth. We find that queries and keys approach their final-depth representations earlier than values, and controlled prefix-channel interventions show that mature values substantially improve shallow draft predictions. Motivated by this asymmetry, we introduce Depth-Asynchronous Self-Speculation (DAS), which decouples the depth of draft computation from the depth of verified-prefix representations it reads. Its Mature-V primitive lets shallow queries retrieve full-depth prefix values without additional recurrent computation. We further develop DAS-Wave, which combines depth-asynchronous prefix reads with carried parallel refinement, progressive block growth, and an independent full-depth verifier. Across four recurrent-model checkpoints and mathematics and code workloads, DAS-Wave achieves 4.00--6.96$\times$ mean throughput speedup over paired full-depth autoregressive decoding in the same inference stack. These results identify prefix-information depth as an effective design axis for recurrent self-speculation.

1 INTRODUCTION

DAS separates the depth of draft computation from the depth of verified-prefix information it reads. DAS-Wave combines this approach with parallel refinement and full-depth verification, achieving 4.00–6.96× mean throughput speedup over paired full-depth autoregressive decoding.

1 INTRODUCTION

  • Looped Transformers require many sequential recurrent passes before each generated token, making autoregressive generation expensive as recurrent depth grows.Self-speculation proposes future tokens with a cheaper drafter and verifies them with the target; shallower exits are cheaper but typically less reliable.Source: §1 · p. 1
  • Queries and keys approach final-depth representations earlier than values, and replacing shallow-draft prefix values with full-depth values substantially improves agreement with the full-depth target.The intervention does not increase draft recurrent depth or require additional draft computation passes.Source: §1 · p. 2
  • DAS decouples draft computation depth from the depth of verified-prefix representations, and Mature-V lets shallow queries retrieve full-depth prefix values without additional draft computation passes.Mature-V can be applied to sequential self-speculation without changing its verification procedure.Source: §1 · p. 2
  • DAS-Wave refines candidate positions in parallel, carries recurrent states across rounds, progressively expands the speculative block, and uses a separate full-depth target pass to verify proposals.The target model and standard speculative verification rule remain unchanged.Source: §1 · p. 2
  • 4.00–6.96× mean throughput speedup over paired full-depth autoregressive decoding across four recurrent-model checkpoints and mathematics and code workloads.DAS-Wave combines depth-asynchronous prefix reads with carried parallel refinement and independent full-depth verification.Source: §1 · p. 2

2 BACKGROUND AND RELATED WORK

Looped language models repeatedly apply shared computation across recurrent depth, exposing intermediate predictions for self-speculation. DAS studies which depth of verified-prefix representations drafts read, while DAS-Wave combines those reads with carried parallel refinement and independent full-depth verification.

2 BACKGROUND AND RELATED WORK

  • Looped language models repeatedly apply a shared computation across recurrent depth, growing effective depth and inference-time compute without proportionally increasing parameter count.The model notation defines P as the prelude, F as the shared recurrent core, H as the coda, r as recurrent depth, and R as full target depth.Source: §2 · p. 2
  • Figure 1 depicts queries attending to full-depth prefix K/V, parallel refinement with carried states and window growth, and independent full-depth verification of candidates.In the illustrated example, the first two candidates are accepted and the third is replaced by the target prediction.Source: Figure 1 · p. 3
  • Self-speculative methods obtain cheaper proposals from the target itself, while DAS focuses on which recurrent-depth representation of the already verified prefix the draft reads.Prior approaches vary in proposal construction or refinement; looped Transformers expose intermediate recurrent predictions.Source: §2 · p. 3
  • Parallel Samplers refine candidate positions concurrently while carrying latent states across rounds; DAS-Wave adopts that refinement and progressive window growth while changing verified-prefix read depth.DAS-Wave uses an independent full-depth target verifier; carry and window growth are inherited from prior parallel sampling.Source: §2 · p. 3
  • Figure 2 compares Q/K/V similarity to final-depth representations and draft–target agreement under prefix-channel replacements; its text reports earlier Q/K maturation than V.The figure caption specifies R = 32 and reports 12 GSM8K prompts and 756 evaluated next-token prediction positions per model.Source: Figure 2 · p. 4

3 DEPTH-ASYNCHRONOUS SELF-SPECULATION

DAS separates draft-computation depth from verified-prefix representation depth, then applies mature prefix reads in a carried parallel-refinement decoder with independent full-depth verification.

3.1 DESIGN MOTIVATION: SHALLOW QUERIES, MATURE VALUES

  • Queries and keys mature earlier than values; at draft depth r = 1, replacing prefix V raises agreement from 47.22% to 66.93% on Huginn and from 83.99% to 87.96% on Raven-Llama.Replacing K alone gives 46.83% on Huginn.Source: §3.1 · p. 4

3.2 DEPTH-ASYNCHRONOUS PREFIX READS

  • DAS varies verified-prefix K/V read depths while draft-side K/V stay current; Mature-V pairs shallow K with full-depth V, already cached by target verification.Deep V and Mature-V name the same read-depth assignment in sequential drafting.Source: §3.2 · p. 5

3.3 DAS-WAVE: CARRIED PARALLEL REFINEMENT

  • DAS-Wave applies k recurrent updates in parallel to carried states, expands active width by one each round up to γ, and discards provisional-position K/V between rounds.Existing positions resume from saved latent states; newly activated positions begin from fresh states.Source: §3.3 · pp. 5–6

3.4 INDEPENDENT FULL-DEPTH VERIFICATION

  • An independent target runs all R recurrent steps on the anchor and proposals, accepts greedy matches in order, and emits the target prediction at the first mismatch.The unchanged target computation and verification rule preserve greedy target output in exact arithmetic under the standard speculative-decoding argument.Source: §3.4 · p. 6

4 EXPERIMENTS

Across four recurrent-model checkpoints and four mathematics and code workloads, DAS-Wave leads paired full-depth autoregressive decoding in throughput. The experiments show how mature prefix reads, carried refinement, and proposal-cost choices shape that performance.

4.1 SETUP

  • The evaluation covers four recurrent-model checkpoints and four mathematics and code workloads, using batch-size-1 latency measurements on a single NVIDIA L20X GPU.The reported speedup is mean throughput relative to the corresponding paired full-depth AR baseline, including prefill.Source: §4.1 · pp. 6–7

4.2 DECODING PERFORMANCE

  • 4.00–6.96× mean throughput speedup over paired full-depth AR, with the largest measured speedup in all sixteen settings, is reported for DAS-Wave.It also produces the longest accepted blocks in every setting; on GSM8K it accepts 15.6–18.1 draft tokens per verification, versus 5.2–12.2 for tuned sequential self-speculative baselines.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 PREFIX READS AND LATENT CARRY

  • Mature-Spec improves over conventional sequential self-speculation in all sixteen settings, while DAS-Wave combines Deep K/V reads with carried parallel refinement for substantially larger speedups.In sequential drafting, Deep K/V improves accepted length and speed in all four shown model–task settings; with carry, combining Deep K/V yields much larger gains than either intervention alone.Source: §4.2 · p. 8, §4.3 · p. 8, Table 2 · p. 9

4.4 REFINEMENT COST–ACCEPTANCE TRADEOFF

  • On Huginn, increasing recurrent updates per refinement round from 2 to 8 raises accepted length A from 16.49 to 19.56 but reduces speedup from 4.33× to 2.37×.Increasing verifier capacity from 8 to 32 raises speedup from 1.98× to 4.33× on Huginn and from 2.36× to 5.47× on Raven-Llama; performance depends on balancing accepted work against proposal cost.Source: §4.4 · p. 8, Figure 3 · p. 9

4.5 APPLICABILITY OF DAS AND DAS-WAVE

  • DAS applies when intermediate predictions and depth-indexed verified-prefix K/V are available, while DAS-Wave combines depth-asynchronous reads with carried recurrent refinement.Carried latent states can incorporate revised token guesses across rounds, allowing proposal computation to accumulate as the candidate sequence is refined.Source: §4.5 · pp. 8–9

5 CONCLUSION

The conclusion presents DAS as separating draft-computation depth from verified-prefix representation depth, with Mature-V using full-depth prefix values for shallow drafts. DAS-Wave combines that reuse with carried parallel refinement and achieves reported speedups across four checkpoints and mathematics and code workloads.

5 CONCLUSION

  • DAS decouples draft depth from the depth of the verified-prefix representations it reads, and Mature-V lets shallow drafts benefit from full-depth values already materialized for the verified prefix.DAS-Wave combines depth-asynchronous prefix reuse with carried parallel refinement.Source: §5 · pp. 9–10
  • Across four recurrent-model checkpoints and mathematics and code workloads, DAS-Wave achieves 4.00–6.96× mean throughput speedup over paired full-depth autoregressive decoding.The conclusion identifies prefix-information depth as a distinct design axis for recurrent self-speculation.Source: §5 · p. 10

A EXPERIMENTAL PROTOCOL

The appendix provides protocol details for reproducing the main experiments and describes the accompanying diagnostic materials.

A EXPERIMENTAL PROTOCOL

  • The appendix provides protocol details needed to reproduce the main experiments.It also presents algorithmic details, mechanism controls, refinement-budget studies, and performance diagnostics.Source: A EXPERIMENTAL PROTOCOL · p. 12
  • Each diagnostic uses its stated prompt set and paired AR reference.Estimates from separate timing sessions are not pooled.Source: A EXPERIMENTAL PROTOCOL · p. 12
  • The appendix groups reproducibility protocol details with algorithmic details, mechanism controls, refinement-budget studies, and performance diagnostics.Source: A EXPERIMENTAL PROTOCOL · p. 12

A.1 MODELS AND TASKS

The experiments use recurrent-depth checkpoints and fixed task prompts, then compare frozen neural-draft configurations and external proposal baselines. Draft depth and block length are selected on a development grid before evaluation.

A.1 MODELS AND TASKS

  • Main experiments use R = 32 recurrent steps, and Table 3 summarizes the evaluated checkpoints.A core layer is one layer inside the repeatedly applied shared block; its index differs from recurrent depth.Source: §A.1 · p. 12, Table 3 · p. 12
  • The main evaluation uses 40 fixed prompts per model–task setting, shared across methods and separate from configuration-selection prompts.GSM8K uses eight worked examples, MATH-500 uses four demonstrations, and HumanEval and MBPP use code-completion prompts.Source: §A.1 · p. 12, Table 4 · p. 12
  • Main latency measurements use batch size 1 and bf16 on a single NVIDIA L20X GPU.Source: §A.1 · p. 12

A.2 DRAFT CONFIGURATIONS AND BASELINES

  • The self-speculative drafter’s development grid varies recurrent draft depth d ∈{4, 8} and block length b ∈{8, 16, 32}, then freezes the selected settings.Component experiments fix d = 4 to compare prefix-read policies at the same draft depth.Source: §A.2 · p. 12
  • Mature-Spec uses the fixed Deep K/V policy, conventional self-speculation uses aligned prefix reads, and DAS-Wave uses the one-position growth schedule.Table 5 summarizes the neural drafting configurations used in the main experiments.Source: §A.2 · p. 12, Table 5 · p. 12
  • External baselines use their corresponding proposal mechanisms in the same reference stack as paired AR, with construction, lookup, and update costs included in timing.The mechanisms include Jacobi refinement with an n-gram pool, prompt lookup, successor reuse, and dynamic suffix retrieval.Source: §A.2 · p. 13

A.3 TIMING, ACCEPTANCE, AND UNCERTAINTY

The appendix defines paired timing and acceptance measures, then describes DAS-Wave’s block lifecycle and its target-preservation condition. Under matching verifier conditionals, cache retention, and stopping rules, verification preserves the target’s greedy sequence in exact arithmetic.

A.3 TIMING, ACCEPTANCE, AND UNCERTAINTY

  • CUDA-event timing includes prefill and decoding, and each accelerated method is paired with a full-depth AR reference from the same session and stack.Source: §A.3 · p. 13
  • Mean accepted draft length is the average number of accepted draft tokens per verification, not an acceptance probability; prefill time does not count as a verification event.Source: §A.3 · p. 13
  • Unless otherwise stated, 95% confidence intervals use 10,000 paired-prompt bootstrap resamples; persistent candidate stores use contiguous five-request blocks.Direct method comparisons use paired measurements from the same prompts and timing session.Source: §A.3 · p. 13

B.1 STATE AND CACHE LIFETIMES

  • DAS-Wave carries draft latent states between refinement rounds while discarding provisional draft K/V, and only target-materialized verified positions supply mature prefix reads.A final emitted token may anchor the next block before its full-depth K/V are materialized by the subsequent target pass.Source: §B.1 · p. 13

B.2 ONE DAS-WAVE BLOCK

  • Each DAS-Wave block initializes provisional guesses, refines active positions with Deep K/V prefix reads, carries latent states, verifies independently at full depth, then commits proposals through the first mismatch.The target prediction at the mismatch is emitted and the speculative suffix is discarded.Source: §B.2 · p. 13

B.3 GREEDY TARGET PRESERVATION

  • Speculative verification emits the same greedy sequence as target AR in exact arithmetic when verifier conditionals, accepted-prefix cache, and stopping rule match target AR.The guarantee does not require the carried draft trajectory to match an ordinary target recurrence.Source: §B.3 · p. 14

C.1 DEPTH-DEPENDENT CHANNEL REPLACEMENT

The component studies test depth-dependent prefix reads, latent carry, and refinement-budget choices. Their findings support joint mature K/V reads with carry, while the initial-width sweep finds no consistent advantage from a wider starting window.

C.1 DEPTH-DEPENDENT CHANNEL REPLACEMENT

  • The Huginn depth grid fixes draft depth at r = 2 and varies verified-prefix K and V depths independently; stronger variation along V supports the V-centric motivation, while remaining K dependence leaves K depth as a separate choice.This diagnostic uses a separate prompt cohort from Table 6.Source: §C.1 · pp. 14–15, Table 6 · p. 14, Figure 4 · p. 15

C.2 SEQUENTIAL DEEP V AND DEEP K/V READS

  • Across all four tasks on three Raven checkpoints, joint Deep K/V yields larger speedup and accepted length than Deep V alone.Table 7 reports the sequential mature-prefix policy comparisons.Source: §C.2 · p. 15, Table 7 · p. 15

C.3 FIXED-SETTING READS AND LATENT CARRY

  • The fixed-setting component comparisons have uncertainty estimates in Table 8, and their paired timing sessions are separate from the main-system study.Within each table, comparisons retain their corresponding AR reference.Source: §C.3 · p. 15, Table 8 · p. 16

C.4 COMPLETE CARRY/K/V FACTORIAL

  • The carry/K/V interaction is positive across evaluated model–task settings; for Huginn and Raven-Llama it is 2.85 and 1.20 tokens on GSM8K, and 2.76 and 0.83 tokens on MATH-500, respectively.The Raven-OLMo and Raven-TinyLlama controls show the same qualitative pattern; Table 9 reports the complete factorial results.Source: §C.4 · p. 16, Table 9 · p. 17

C.5 RETUNED CARRIED CONTROLS

  • Retuning improves aligned and Deep V carried controls, but their best measured GSM8K configurations remain below full Deep K/V DAS-Wave on both models.Table 10 reports selected configurations and confirmation measurements.Source: §C.5 · p. 16, Table 10 · p. 17

D.1 RECURRENT UPDATES AND VERIFIER CAPACITY

  • The refinement-budget sweeps vary recurrent updates k and verifier capacity γ; increasing k adds computation to every active position, while γ permits a wider verifier and, until the schedule limit, active window.Table 11 reports mean speedup with 95% confidence intervals and accepted length A.Source: §D.1 · p. 17, Table 11 · p. 17

D.2 INITIAL ACTIVE WIDTH

The initial active-width sweep finds only modest variation and no consistent benefit from widening the starting window, so the paper retains Winit = 2 as a fixed default rather than a tuned optimum.

D.2 INITIAL ACTIVE WIDTH

  • Changing Winit also changes total active-position work, making this a schedule-sensitivity study rather than an equal-compute comparison.Across tested widths, results vary modestly with no consistent advantage from increasing the initial window; Winit = 2 remains the fixed default, not a tuned optimum.Source: §D.2 · pp. 17–18, Table 12 · p. 18

E.1 MAIN-RESULT UNCERTAINTY

The appendix reports uncertainty intervals for AR-relative speedups and a separate diagnostic comparing sequential proposal costs. The diagnostic shows that depth-4 Mature-Spec reduces proposal time relative to depth-8 self-speculation at similar accepted length.

E.1 MAIN-RESULT UNCERTAINTY

  • Table 13 reports 95% confidence intervals for each method’s AR-relative speedup, not direct pairwise tests between accelerated methods.Token Recycling uses moving-block bootstrap over its fixed request stream; other intervals use paired prompt bootstrap.Source: §E.1 · p. 18, Table 13 · p. 19

E.2 SEQUENTIAL PROPOSAL COST

  • At similar accepted length, depth-4 Mature-Spec substantially reduces proposal time relative to depth-8 self-speculation in the separate Huginn diagnostic.Draft time measures cumulative CUDA-event time for proposal construction per request, excluding prefill and target verification.Source: §E.2 · p. 18, Table 14 · p. 19

Related papers

Loading 2609.34538v1…