PaperlayerView exact revision on arXiv

Source-linked AI summary

APEX: Speculate smarter, not deeper

Manvi Jha, Zach Zhang, Zhichao Xu, Linbo Liu, Sai Muralidhar Jayanthi, Vinayak Arannil

arXiv:2610.07780v1cs.CLcs.LG

TL;DR

Speculative decoding's speed and token utilization depend on proposal mechanism, draft depth, and decoding state, while fixed configurations cannot adapt to changing conditions. APEX selects a speculative expert per request and adapts depth at each verification block using learned acceptance, runtime, and utility signals. Across workloads, APEX-S provides the highest aggregate speedup, while APEX-B offers lower waste at reduced acceleration.

  • Problem

    Speculative decoding must balance speed and draft-token utilization as proposal costs, acceptance, and decoding conditions vary.

    Source: §1 · pp. 1–2

  • Method

    APEX combines request-level selection among speculative experts with block-level depth adaptation using causal decoding state and recent verifier feedback.

    Source: §1 · p. 2

  • Results

    4.27× aggregate speedup with 55.30% waste for APEX-S compares with 4.22× speedup and 79.90% waste for fixed n-gram speculation at k = 16.

    Source: §5.2 · p. 13

  • Takeaways & Limitations

    APEX-S provides the highest aggregate speedup, while APEX-B offers lower aggregate waste at reduced acceleration, with preferred profiles varying by workload and desired operating point.

    Source: §5.6 · p. 18, §6 · p. 19

    On code generation, n-gram at k = 16 exceeds all three APEX profiles on speedup and WAS, and fixed configurations also dominate some direct speed–waste comparisons.

    Source: §5.3 · pp. 16–17

Abstract

from arXiv · show

Speculative decoding reduces large language model inference latency by drafting multiple tokens before target-model verification, but its effectiveness depends on both the proposal mechanism and draft depth. Fixed configurations cannot respond to changes in predictability, repetition, and acceptance during generation, so deeper drafting can increase wasted computation without proportional speedup. We introduce APEX, a learned controller that balances decoding speed and draft-token waste through request-level expert selection and block-level depth adaptation. APEX-Router selects among EAGLE-3, n-gram, and draft-model speculation for each request, while APEX-Depth adjusts draft length at each verification block using causal decoding signals and recent verifier feedback. APEX models accepted draft length as censored survival feedback, learning position-wise rejection hazards, block execution costs, and an action utility that balances throughput, accepted progress, and wasted tokens. This allows the controller to adapt speculation while retaining the target model's verification procedure. We integrate APEX into vLLM and evaluate it with Qwen3-8B across six workloads, achieving up to 5.24X speedup over autoregressive decoding. Across the aggregate evaluation, APEX-S achieves 4.27X speedup, while APEX-B achieves 3.27X speedup with a 41.0% relative reduction in wasted-token percentage compared with fixed n-gram speculation at k=16, providing distinct operating points for balancing acceleration and draft-token utilization.

1 INTRODUCTION

Speculative decoding trades verified progress against draft-token waste, and the best proposal mechanism and depth depend on decoding conditions. APEX pairs request-level expert selection with block-level depth adaptation and models accepted draft length as censored survival feedback.

1 INTRODUCTION

  • Speculation must balance wall-clock latency and draft-token utilization: deeper drafts can enable more progress but expose more proposed work to rejection.High speedup can coexist with substantial waste, while low waste can coexist with slow execution.Source: §1 · p. 1
  • Proposal mechanisms differ in draft quality and execution cost, so the same acceptance rate or waste percentage can have different latency consequences.EAGLE-3 may help when upcoming tokens are predictable, n-gram reuse when structure repeats, and a draft model across varied inputs.Source: §1 · p. 2
  • APEX-Router selects a proposal mechanism and initial depth per request, while APEX-Depth adapts draft length at each verification block using causal prefix features and recent verifier outcomes.The selected mechanism stays fixed for the request; depth can change within that path.Source: §1 · p. 2
  • APEX models accepted draft length as censored survival feedback, using position-wise rejection hazards to distinguish observed rejection from full acceptance that only establishes survival through the chosen depth.Runtime prediction complements acceptance modeling because accepting more tokens need not yield greater speedup when drafting or verification costs rise.Source: §1 · p. 2

1.1 CONTRIBUTIONS

  • 4.27× aggregate speedup over autoregressive decoding is reported for APEX-S; APEX-B reaches 3.27× and reduces wasted-token percentage by 41.0% versus fixed n-gram speculation at k = 16.The reported mean wasted-token percentage falls from 79.90% to 47.17%.Source: §1.1 · pp. 2–3

2 BACKGROUND AND RELATED WORK

Speculative decoding reduces sequential target-model calls by drafting token blocks for verification, but its efficiency depends on accepted progress relative to proposal and verification costs. APEX treats heterogeneous speculators as experts and adapts mechanism and depth to decoding state.

2.1 AUTOREGRESSIVE DECODING

  • A speculative block advances generation by its accepted draft length, not its proposed depth; early rejection discards the remaining draft tokens.Larger depth raises possible progress per verifier call but can also increase wasted drafting and verification work.Source: §2.1 · p. 3

2.2 SPECULATIVE MECHANISMS

  • Fixed mechanism choices can retain a costly or poorly matched speculative path, since proposal mechanisms differ in costs and acceptance behavior.Retrieval-based methods can be fast in repetitive regimes but brittle on novel continuations, whereas neural speculators can better model semantic continuation.Source: §2.2 · pp. 3–4
  • APEX treats EAGLE-3, n-gram, and draft-model speculation as candidate experts, learning which mechanism to use and how deeply to speculate while retaining their existing algorithms.This complements adaptive-depth approaches by combining mechanism selection with depth control.Source: §2.2 · p. 4

2.3 ADAPTIVE SPECULATIVE DECODING

  • Adaptive-speculation work and serving schedules address changing conditions, motivating joint selection of proposal mechanism and depth to balance verified progress, runtime, and draft-token waste.The related-work discussion names dynamic lookahead, AdaEAGLE, BanditSpec, DSDE, and manually specified vLLM schedules.Source: §2.3 · p. 4

3 EMPIRICAL MOTIVATION: FIXED SPECULATION IS NOT ENOUGH

Speculative performance varies with workload, proposal mechanism, depth, and decoding signals. The characterization motivates selecting depth and mechanism using acceptance, execution cost, entropy, repetition, and verifier feedback.

3 EMPIRICAL MOTIVATION: FIXED SPECULATION IS NOT ENOUGH

  • The characterization compares autoregressive decoding with EAGLE-3, n-gram, and draft-model speculation across six workloads using Qwen3-8B in vLLM.The workloads include code generation, mathematical reasoning, long-context completion, long-chain reasoning, conversational generation, and SWE-style tasks.Source: §3 · p. 4
  • Figure 2a shows draft-model speculation near or below autoregressive performance even at relatively high acceptance rates, while other mechanisms achieve substantial speedup.Acceptance alone therefore does not determine acceleration; execution cost also matters.Source: §3 · p. 5, Figure 2 · p. 5
  • Increasing n-gram depth from k = 8 to k = 16 raises mean speedup from 3.59× to 4.22× but increases mean waste from 63.56% to 79.90%.For EAGLE-3, increasing depth from k = 4 to k = 16 reduces mean speedup from 2.35× to 1.37× and raises mean waste from 66.25% to 97.67%.Source: §3 · p. 5
  • Figure 2b relates median speedup to token entropy and repetition density, with performance patterns varying across those conditions and between mechanisms.These associations motivate using prefix entropy and repetition alongside recent verifier feedback as controller inputs.Source: §3 · p. 5, Figure 2 · p. 5
  • A first rejection at position r leaves r − 1 accepted tokens and k − r + 1 discarded tokens, whereas full acceptance only establishes survival through depth k.The censored feedback structure motivates modeling position-wise rejection hazards for adaptive control.Source: §3 · p. 5

4 APEX: BREATHING ADAPTIVITY INTO STATIC SPECULATION

APEX combines request-level expert selection with block-level depth adaptation to balance speed and draft-token waste. Its controller uses decoding state and verifier feedback to score candidate depths while retaining the target model's verification procedure.

4 APEX: BREATHING ADAPTIVITY INTO STATIC SPECULATION

  • APEX-Router selects a speculative expert and initial depth for each request, while APEX-Depth adjusts depth at each verification block.The router uses prompt features; after a verifier block, the depth controller uses updated decoding state and prior verifier feedback, with observed acceptance and runtime updating controller state.Source: §4 · p. 5, Figure 3 · p. 6
  • The two decision timescales reflect systems constraints: expert selection affects model placement, batching, and worker routing, while depth can change more frequently within the selected path..Source: §4 · p. 6
  • APEX formulates accepted draft length as censored survival feedback and uses the resulting acceptance model to predict acceptance, runtime, and utility from a shared representation.The utility balances accepted progress, runtime, and speculative waste.Source: §4 · p. 6, §4.1 · p. 6

4.1 PROBLEM FORMULATION AND STATE REPRESENTATION

  • Request-level state includes prompt embedding and length, decoding configuration, and prompt repetition and entropy features.The request-level state is the information available before generation begins.Source: §4.1 · p. 6
  • Block-level state combines generation progress, prior verifier feedback, and statistics from the current prefix, with the active mechanism and candidate depth supplied separately for state-action features.Verifier history includes recent acceptance and rejection summaries; prefix features include entropy and repetition, and the current block outcome is excluded from prediction.Source: §4.1 · p. 7

4.2 APEX-ROUTER

APEX-Router selects a speculative mechanism and initial depth for each request, while APEX models accepted draft length as censored survival feedback. This feedback distinguishes observed rejection positions from full acceptance, which only establishes survival through the executed depth.

4.2 APEX-ROUTER

  • APEX-Router scores admissible mechanism–initial-depth pairs from the pre-generation request state to select a request’s speculative path.The selected mechanism remains fixed for the request, and the first block uses the selected initial depth.Source: §4.2 · p. 7
  • The router is trained on completed request-level traces, using observed request–method–initial-depth outcomes as empirical utility labels rather than workload labels.Acceptance, latency, and waste are execution outcomes, not inputs available when routing a new request.Source: §4.2 · p. 7

4.3 ACCEPTANCE AS CENSORED SURVIVAL FEEDBACK

  • A partially accepted block identifies its first rejected position, whereas full acceptance only establishes survival through the executed depth.The full-acceptance observation is right-censored and does not impose rejection at the next position.Source: §4.3 · pp. 7–8
  • APEX models position-wise rejection hazards to derive survival probabilities and expected accepted length from verifier feedback.The model combines acceptance prediction with runtime and utility supervision because accepted length alone does not determine candidate-depth efficiency.Source: §4.3 · p. 8

4.4 APEX-DEPTH

APEX-Depth evaluates candidate depths before each subsequent block and selects the highest predicted utility. Its shared representation supports predictions of survival, execution cost, throughput, and utility, with observed outcomes updating the state for later decisions.

4.4 APEX-DEPTH

  • APEX-Depth selects the candidate depth with the highest predicted utility before each subsequent block, then updates its state from observed execution outcomes.This allows depth decisions to respond to changes in acceptance behavior and execution cost during generation.Source: §4.4 · pp. 8–9
  • The shared encoder evaluates each candidate depth in the current generation state, and four prediction heads estimate survival, cost, throughput, and utility.The utility head predicts the action score used for depth selection; the other heads predict execution outcomes.Source: §4.4 · p. 8
  • The survival head predicts position-wise rejection hazards, trained with censored likelihood so fully accepted and partially rejected blocks provide appropriate supervision.The predicted hazards determine the survival curve and expected accepted length.Source: §4.4 · p. 9
  • The cost head predicts candidate-block wall-clock latency, including drafting and verification, while the throughput head predicts output progress rate.Runtime prediction accounts for cases where proposal or verification overhead makes a candidate depth slow despite accepting many tokens.Source: §4.4 · p. 9

4.5 TRACE-DERIVED TRAINING TARGETS

APEX’s depth-utility targets are derived from execution traces and combine throughput, accepted length, waste, and regret. The trace dataset also supplies comparisons across speculative configurations for constructing regret terms.

4.5 TRACE-DERIVED TRAINING TARGETS

  • Each depth-controller example pairs information available before block execution with post-verification outcomes used to supervise the controller’s predictions.Accepted length supervises survival; latency and throughput supervise runtime heads; the utility target supervises the utility head.Source: §4.5 · p. 9
  • The trace dataset supplies request-level comparisons across observed method-depth actions, including global and fixed-mechanism best-throughput references.The resulting normalized regret penalties compare block throughput with those request-level references and are used only to construct training targets, not at inference.Source: §4.5 · pp. 9–10
  • The block-level utility label combines log throughput and accepted length with penalties for draft-token waste and offline regret.The accepted-length term remains finite at zero accepted tokens through ϵ = 0.1.Source: §4.5 · p. 10

4.6 TRAINING OBJECTIVE FOR APEX-DEPTH

APEX-Depth uses multi-task losses to train its survival, runtime, and utility predictions, with candidate ranking also supervising depth selection. The ranking groups compare candidate depths for the same prompt, method, and block.

4.6 TRAINING OBJECTIVE FOR APEX-DEPTH

  • The depth controller combines survival, cost, throughput, utility-regression, and candidate-ranking objectives to train predictions and the score used for depth selection.Utility regression directly supervises the predicted score, while ranking trains candidate ordering.Source: §4.6 · pp. 10–11
  • Cost and throughput heads are trained against observed action outcomes using SmoothL1 regression losses.The training objective includes separate cost and throughput losses.Source: §4.6 · p. 10
  • Candidate-ranking groups compare observed depths within the same prompt, speculative method, and block index, encouraging the highest-utility depth to receive the highest predicted score.The groups use pre-block state and observed utility labels for the candidate depths.Source: §4.6 · pp. 10–11

4.7 APEX VARIANTS

APEX variants share the same architecture, features, candidate depths, data, and deployment procedure; they differ in utility weights to target distinct speed–waste operating points.

4.7 APEX VARIANTS

  • Only the utility weights in Equation 25 differ across the three trained variants.The variants otherwise use identical architectures, features, candidate-depth interfaces, data, and deployment procedures.Source: §4.7 · p. 11
  • APEX-S emphasizes throughput competitiveness, APEX-B balances throughput, accepted progress, and waste, and APEX-E applies a larger waste penalty.Their respective utility weights are (0.20, 0.20, 0.20, 0.15), (0.30, 0.35, 0.15, 0.20), and (0.40, 0.50, 0.10, 0.25).Source: §4.7 · p. 11
  • The stated training preferences do not guarantee Pareto efficiency or that APEX-E will have the lowest waste.Source: §4.7 · p. 11

4.8 ONLINE DEPLOYMENT

During generation, APEX updates its causal state after each verifier block and selects a depth for the next block while keeping network parameters fixed.

4.8 ONLINE DEPLOYMENT

  • After each verifier block, APEX scores every candidate depth using its updated causal state and selects the next block’s depth.Source: §4.8 · p. 11
  • The controller adapts its actions during generation, but its network parameters remain fixed.Source: §4.8 · p. 11
  • SpeculateAndVerify drafts and verifies a block, returns emitted output and execution observations, and uses those observations to update history for the next decision.The output extends the prefix.Source: §4.8 · p. 11

5 RESULTS AND DISCUSSION

Across six workloads, the evaluation compares APEX profiles with fixed speculative configurations on latency and draft-token waste. Results show distinct operating points and workload-dependent trade-offs, while ablations and depth-set comparisons examine the controller’s components and action choices.

5.2 OVERALL RESULTS

  • 4.27× speedup with 55.30% waste for APEX-S slightly exceeds n-gram k=16’s 4.22× while cutting waste 30.8%; APEX-B yields 3.27× with 41.0% less waste.Table 1 reports aggregate performance; APEX-B’s waste is 47.17%.Source: Table 1 · p. 12, §5.2 · p. 13
  • N-gram k=8 to k=16 raises speedup from 3.59× to 4.22× but waste from 63.56% to 79.90%; EAGLE-3 k=4 to k=16 lowers speedup from 2.35× to 1.37× and raises waste from 66.25% to 97.67%.Draft-SD k=2 has 27.82% waste but only 0.96× speedup, so low waste alone does not ensure lower latency.Source: §5.2 · pp. 13–14

5.3 PERFORMANCE ACROSS WORKLOADS

  • On mathematical reasoning, APEX-S reaches 5.24× speedup with 47.00% waste, while conversational APEX-S reaches 4.88× with 51.50% waste versus n-gram k=8 at 4.63× and 54.63%.On mathematical reasoning, APEX-B’s 5.16× speedup with 44.00% waste yields the highest WAS of 11.72; conversational APEX-B reaches 4.15× with 46.56% waste.Source: Table 3 · p. 14, §5.3 · pp. 14–15, Table 6 · p. 16
  • Long-context n-gram k=16 is fastest at 6.40× with 66.25% waste, while APEX-S yields higher WAS; in long-chain reasoning, APEX-B has the highest WAS at 7.56 with 3.07× speedup and 40.59% waste.For long-context completion, APEX-S records WAS 10.18 versus 9.66 for n-gram k=16; for long-chain reasoning, n-gram k=16 reaches 4.57× with 77.69% waste.Source: Table 4 · p. 15, §5.3 · pp. 15–16, Table 5 · p. 15
  • On software-engineering tasks, APEX-S has the highest WAS at 6.55, while code generation favors fixed configurations: n-gram k=16 reaches 4.25× speedup and WAS 5.33, exceeding all APEX profiles on those metrics.For software engineering, n-gram k=16 is faster at 4.15× but wastes 80.31%, versus APEX-S at 3.15× and 48.12% waste.Source: Table 2 · p. 14, Table 7 · p. 16, §5.3 · pp. 16–17

5.5 SENSITIVITY TO CANDIDATE DEPTHS

  • For APEX-B, excluding shallow actions raises speedup from 3.25× to 3.34× but raises waste from 42.76% to 48.55% and lowers WAS from 7.60 to 6.89.The sparse set gives APEX-S 3.39× speedup and 47.95% waste, with WAS 7.06 versus 7.07 for the dense set.Source: Table 9 · p. 17, §5.5 · p. 18

5.4 ABLATION STUDY

  • Removing survival, cost/TPS, entropy/repetition, or verifier-history supervision worsens measured speed–waste balance, with verifier-history removal sharply reducing speedup.Removing survival loss raises waste to 67.62%–69.38%; removing recent verifier feedback lowers speedup to 1.14×–1.31×.Source: §5.4 · p. 18

5.6 DISCUSSION

  • The results support adaptive control but establish workload-dependent limits; profile choice depends on the workload and desired operating point, and evaluation should consider latency and draft-token utilization together.The ablations support combining survival supervision, runtime-cost supervision, and recent verifier feedback.Source: §5.6 · p. 18

6 CONCLUSION AND FUTURE WORK

APEX combines request-level speculator selection with block-level depth adaptation to balance decoding speed and draft-token waste while retaining target-model verification. The reported gains vary across workloads, and future work targets online updates and broader evaluation settings.

6 CONCLUSION AND FUTURE WORK

  • APEX combines request-level speculator selection with block-level depth adaptation to balance decoding speed and draft-token waste while preserving target-model verification.The controller coordinates heterogeneous proposal mechanisms without replacing their existing proposal and verification algorithms.Source: §6 · p. 18
  • APEX demonstrates substantial acceleration with lower waste than aggressive fixed-depth configurations, with trade-offs that vary across workloads.Source: §6 · p. 19
  • Future work will update the controller online from decoding traces and evaluate multi-turn conversations, reinforcement-learning rollouts, tool-augmented generation, and additional speculators.The planned settings would test adaptation to changing context and generation patterns.Source: §6 · p. 19
Loading 2610.07780v1…