ArXiv: 2404.18911
🎯 Pitch
Kangaroo achieves up to 1.68× speedup over standard decoding with zero loss in output quality—while using 88.7% fewer extra parameters than previous methods—by repurposing a shallow sub-network of the target model as its own draft model. The key insight is that confidence-guided early termination during drafting eliminates wasted computation on hard tokens, making self-drafting viable even when the small model isn't dramatically faster than the large one.
1. Executive Summary
This paper introduces Kangaroo, a lossless self-speculative decoding framework that accelerates LLM inference without external draft models by reusing a fixed shallow sub-network of the target model as a self-draft model, connected through a lightweight adapter (a single multi-head attention block with two normalization layers) — and further introduces a double early-exiting mechanism that halts draft generation when confidence per token drops below a threshold (dynamic drafting steps of up to 6 tokens with a confidence threshold η = 0.6), avoiding wasted computation on harder tokens. On Spec-Bench with Vicuna-7B and 13B, Kangaroo achieves speedups up to 1.68×, outperforming Medusa-1 while using 88.7% fewer additional parameters (67M compared to 591M), establishing that self-drafting with shared shallow-layer features can surpass dedicated head-based methods only when draft token efficiency and acceptance rate are jointly optimized via confidence-guided early termination.
2. Context and Motivation
The Core Problem: Speculative Decoding's Draft Model Costs Are Prohibitive
The central tension this paper addresses is deceptively simple: speculative decoding works well in theory, but obtaining a good draft model is expensive in practice. To understand why this matters, we need to unpack the autoregressive decoding bottleneck first.
LLM inference is memory-bandwidth-bound, not compute-bound. During autoregressive generation, each new token requires reading the full model weights from GPU memory — a multi-gigabyte operation — to perform relatively modest arithmetic computations. The paper cites a concrete example: Vicuna-33B on four NVIDIA V100 GPUs produces only seven new tokens per second (Section 1). The computational hardware sits idle while memory transfers dominate wall-clock time.
Speculative decoding (SD) offers an elegant escape: use a smaller, faster "draft" model to propose multiple future tokens cheaply, then verify them all in parallel through a single forward pass of the large "target" model (Chen et al., 2023; Leviathan et al., 2023). If the draft proposes tokens, the system generates between and new tokens per target model forward pass — a direct multiplier on throughput. Crucially, SD is lossless: through a carefully designed acceptance-rejection sampling mechanism, the output distribution remains identical to what the large model would produce autoregressively.
The problem is not with SD's verification mechanism. The problem is with obtaining the draft model.
Why a Good Draft Model Is Hard to Get
The effectiveness of SD depends on two interacting factors that the paper identifies (Section 1):
1. The gap between draft and target models. The draft model must predict tokens that the target model would itself produce. The closer the draft's output distribution matches the target's, the higher the token acceptance rate — the probability that a proposed token is verified as correct. A low acceptance rate means many proposed tokens are discarded, wasting the draft's computation and limiting speedups.
2. The inference latency of the draft model itself. The draft model's forward pass takes time. If the draft is slow relative to the target, even a high acceptance rate yields disappointing end-to-end acceleration because the drafting phase consumes a meaningful fraction of the total budget.
The standard approach — training a standalone "tiny" draft model from scratch on a large corpus (e.g., LLaMA-68M for LLaMA-7B, as cited from Miao et al., 2023) — addresses factor (1) reasonably well but introduces a substantial training cost that the paper characterizes as "costly, limiting its application in real-world scenarios" (Section 1). This is not merely an incremental expense: training a draft model from scratch requires a separate pretraining pipeline, a large dataset, and significant computational resources — essentially replicating the effort of building a smaller LLM. For organizations deploying models of various sizes (7B, 13B, 33B, etc.), this multiplies across model variants. The paper cites DistillSpec (Zhou et al., 2023) as an example of this external-draft-model paradigm, acknowledging its effectiveness while lamenting its cost.
The practical consequence is that speculative decoding, despite its theoretical appeal, has not been as widely adopted as its performance numbers suggest. The barrier is not the inference-time mechanism — it's the upfront investment required to train a task-aligned draft model.
The Self-Drafting Alternative and Its Shortcomings
In response, several recent methods have explored self-drafting: generating draft tokens from the target model itself, without any external model. The paper positions itself against three representative approaches, each with distinct limitations:
Medusa (Cai et al., 2024) trains multiple independent feed-forward network (FFN) heads attached to the last decoder layer of the target model. Each head predicts a token at a different future position (next token, next-next token, etc.) in a single forward pass, making draft generation extremely fast — essentially zero additional inference cost beyond the single target model forward pass that generates the features. This is the strength that makes Medusa a compelling baseline.
However, Medusa has a weakness that Figure 1(a) quantifies: its token acceptance rate declines steeply with distance. The heads are time-independent — each head predicts its position based solely on the same hidden state, without conditioning on the tokens predicted by earlier heads. This independence limits how well the heads can model the sequential dependencies of language. On the mathematical reasoning subtask of Spec-Bench, Medusa's consistent token acceptance rate at window size (predicting the next-next token correctly given the next token was correct) drops substantially compared to autoregressive self-drafting methods like Kangaroo and Lookahead. The paper characterizes this as "its token acceptance rate is not yet satisfactory" (Section 1).
Lookahead (Fu et al., 2024) takes a different approach using Jacobi iteration: generate multiple tokens in parallel, then iteratively refine them until convergence. This produces an autoregressive-like dependency structure that yields a token acceptance rate competitive with Kangaroo on mathematical reasoning (Figure 1(a), where Lookahead's CTAR decay curve closely tracks Kangaroo's). However, Lookahead pays for this acceptance rate with lower drafting efficiency: the iterative refinement process consumes more wall-clock time per drafted token than Medusa's single-pass head predictions. The paper demonstrates this tradeoff concretely in Figure 1(b): on the mathematical reasoning subtask, Lookahead achieves a higher compression rate than Medusa but a lower wall-time speedup — it generates better draft tokens, but takes too long to generate them. This is the critical insight the paper wants to convey: focusing exclusively on token acceptance rate without considering draft latency leads to suboptimal end-to-end acceleration (Section 1, paragraph 4).
REST (He et al., 2023) generates draft tokens by retrieving relevant text spans from a reference database. This avoids training entirely but introduces a different cost: the retrieval step itself has latency, and the draft quality depends on having a sufficiently rich and relevant reference corpus. The paper includes REST as a comparison point on Spec-Bench (Table 1) where it generally underperforms Kangaroo in both speedup and compression rate across subtasks.
Draft & Verify (Zhang et al., 2023) deserves mention as the most direct precursor to Kangaroo's approach. It also uses early exiting — skipping intermediate layers of the target model to create a faster self-draft model. However, the paper identifies a critical flaw: while Draft & Verify "could achieve a high token acceptance rate, the inference latency of the 'small model' is exceptionally high, which can hinder end-to-end acceleration efficiency" (Section 2). The problem is that skipping too few layers (to maintain high acceptance) leaves the draft model nearly as large and slow as the target, eliminating the speculative advantage. Skipping too many layers (for speed) destroys acceptance rate. Draft & Verify, as the paper presents it, cannot resolve this tension effectively.
The Gap Kangaroo Identifies
The paper's diagnosis of the landscape reveals a missing design point: a self-draft model that simultaneously achieves:
- High token acceptance rate through autoregressive generation (capturing sequential token dependencies, unlike Medusa's independent heads).
- Low draft inference latency through aggressive layer reuse (sharing a shallow sub-network with the target, unlike Draft & Verify's deep early exit) and dynamic early termination (stopping draft generation on hard-to-predict tokens, unlike both Medusa's fixed heads and Lookahead's fixed-step iteration).
- Minimal training cost (a lightweight adapter with fewer parameters than Medusa's heads, trained on a single dataset for only 10 epochs).
The key intellectual move is recognizing that factors (1) and (2) are not independent knobs to trade off — they can be jointly optimized through the double early-exiting mechanism. The first early exit (shallow layer extraction) creates a fast base draft model. The second early exit (confidence-thresholded termination) dynamically adjusts the number of draft tokens per target-model forward pass, avoiding wasted drafting cycles on tokens where the draft model's confidence is low.
The paper frames this through a new evaluation metric it introduces, consistent token acceptance rate (CTAR), defined in Section 3. CTAR measures the probability that consecutive draft tokens are all accepted — not just any individual token. The key observation from Figure 1(a) is that CTAR decays with , meaning distant tokens are increasingly unlikely to be correct. A fixed drafting budget of tokens per verification step (as in Medusa with 3 independently-predicted heads, or Kangaroo's maximum step) wastes computation on tokens at positions 4, 5, and 6 where acceptance probability is low. By dynamically halting draft generation when confidence drops, Kangaroo reallocates that wasted drafting time to the next verification step, improving overall throughput.
The Broader Landscape and Why This Matters Now
The paper situates its contribution against two broader research trends (Section 2):
Speculative decoding maturation. The survey cited (Xia et al., 2024) indicates that speculative decoding has moved from a proof-of-concept to a practical deployment technique, with multiple competing approaches (external draft models, self-drafting, retrieval-based). At this stage, the field needs methods that work in real deployments, not just under idealized assumptions. Training cost, memory overhead, and robustness across different types of text (Spec-Bench's six subtasks span translation, summarization, math reasoning, RAG, question answering, and multi-turn conversation) become the relevant benchmarks — not just peak speedup on a single task.
Early exiting for efficiency. The paper acknowledges a parallel research thread on early exiting for inference acceleration (Schuster et al., 2022; Varshney et al., 2023), where predictions are made from intermediate layers to skip remaining computation. However, the paper notes a critical limitation: "since early exiting accelerates inference by saving subsequent computations, it inevitably incurs the issue of performance degradation" (Section 2). That is, standard early exiting is lossy — the predictions from intermediate layers are worse than those from the final layer, and there is no mechanism to recover the lost quality. Kangaroo's key insight is to embed early exiting within a speculative decoding framework, where the early-exit predictions are merely proposals that the full model verifies and can reject. This makes the system lossless despite using early-exited draft tokens — the final output distribution remains identical to the full model's because verification corrects any errors.
The timing of this work also coincides with the broader adoption of memory-bandwidth-bound LLM inference as a production bottleneck. As models grow (7B, 13B, 33B, 70B) and deployment shifts toward latency-sensitive applications (chatbots, real-time translation, interactive coding assistants), any technique that accelerates decoding without quality degradation and without requiring expensive bespoke training pipelines has immediate practical value. The paper explicitly names this tension in its abstract: "the conventional approach of training a separate draft model... can be costly."
How Kangaroo Positions Itself
Kangaroo's contribution is not a single new technique but rather a specific architectural combination — shallow-layer early exit + lightweight autoregressive adapter + confidence-thresholded dynamic drafting — that the paper argues achieves a previously unrealized sweet spot in the self-drafting design space:
- Against external draft models (DistillSpec, SpecInfer): Kangaroo eliminates the need for a separately trained draft model entirely, reducing the deployment cost to a small adapter network (67M parameters for Vicuna-7B, compared to training a full 68M-parameter LLaMA draft model from scratch).
- Against Medusa: Kangaroo achieves higher token acceptance rates (especially at longer speculative distances) due to autoregressive dependency among draft tokens, while using dramatically fewer additional parameters ( reduction). The trade-off is that Kangaroo's drafting is sequential and therefore slower per token than Medusa's parallel heads — but the paper argues that the acceptance rate advantage and the dynamic early termination compensate.
- Against Lookahead: Kangaroo achieves comparable acceptance rates but with lower draft latency due to the shallower draft model (only layers for Vicuna-7B) and the dynamic early exit that avoids Lookahead's fixed multi-iteration refinement.
- Against Draft & Verify: Kangaroo uses a much shallower early exit (layer 2 rather than skipping only a few intermediate layers), making the draft model genuinely faster, and bridges the resulting quality gap with the dedicated adapter network rather than hoping that the raw intermediate representations suffice.
- Against lossy early exiting: Kangaroo is lossless by construction — verification against the full model ensures the output distribution is preserved, converting the early-exit quality gap from a permanent defect into a recoverable efficiency cost (rejected tokens waste draft computation but do not affect output quality).
The paper's positioning is most concisely captured in Figure 1(b): Kangaroo's bar is consistently the highest across all four Spec-Bench subtasks (math reasoning, RAG, summarization, MT-bench), demonstrating that its design choices translate to actual wall-time improvements rather than just favorable theoretical metrics. The paper explicitly warns against evaluating speculative decoding methods solely by compression rate (Section 3, introducing CTAR as a more informative metric), and its own evaluation prioritizes the metric that matters in deployment: end-to-end speedup.
3. Technical Approach
This is primarily a systems and methods paper whose core idea is that self-speculative decoding can achieve both high token acceptance rates and low draft-model latency by combining a shallow-layer early exit (reusing the target model's first few layers as a self-draft model connected through a lightweight adapter) with a confidence-thresholded dynamic drafting mechanism that stops generating draft tokens as soon as the self-draft model's top-1 probability drops below a threshold, avoiding wasted computation on harder-to-predict tokens.
3.1 Reader Orientation
Kangaroo is a system that accelerates LLM text generation by using a small, fast "draft model" to propose multiple future tokens at once, then verifying all proposals in parallel through the full model — except the draft model is not a separate trained network but rather a reused shallow slice of the target model itself (the first 2–3 transformer layers) connected to a tiny learned adapter, and the number of tokens it proposes is dynamically determined per verification step by monitoring its own confidence. The problem it solves is obtaining a high-quality draft model without the prohibitive training cost of a standalone draft model, and the shape of the solution is a double early-exit mechanism: a fixed architectural early exit (running only the first few layers) creates the draft model, and a dynamic behavioral early exit (stopping draft generation when confidence drops below $ \eta = 0.6 $) prevents wasted drafting cycles.
3.2 Big-Picture Architecture (Diagram in Words)
The Kangaroo system has four major components, arranged in an inference-time pipeline:
-
Shared shallow sub-network
$ \mathcal{M}^{b}[:l] $: The first$ l $transformer layers of the target LLM$ \mathcal{M}^{b} $. These layers are frozen — their weights are unchanged from the pretrained target model — and their output hidden states serve as input to the adapter. For Vicuna-7B,$ l = 2 $; for Vicuna-13B,$ l = 3 $. -
Lightweight adapter network
$ \mathcal{A} $: A small trainable module consisting of exactly one multi-head attention block and two RMS normalization layers. It transforms the shallow-layer hidden states into predictions that approximate what the full model would produce. It shares the target model's LM Head (the final linear projection from hidden states to vocabulary logits), so the adapter's output hidden states are fed through the same unembedding matrix as the target model. -
Confidence-thresholded drafting controller: A decision mechanism that, after each draft token is generated by the adapter + LM Head, checks whether the top-1 probability (the softmax score of the most likely token) exceeds
$ \eta = 0.6 $. If confidence is above the threshold and fewer than$ \gamma = 6 $draft tokens have been generated, drafting continues. If confidence drops below the threshold, drafting halts immediately — the current token is still proposed, but no further tokens are generated in this round. -
Full-model verifier
$ \mathcal{M}^{b}[l:] $: The remaining layers of the target model (from layer$ l+1 $to the final layer$ L $), which process all proposed draft tokens in parallel using the hidden states collected from the draft model's forward passes as input features. The verifier applies the standard speculative decoding acceptance-rejection mechanism to determine which draft tokens are correct, ensuring the output distribution is identical to the target model's autoregressive output.
Information flows as follows: given a prefix $ x^{t} $, the shallow sub-network processes it to produce a hidden state at layer $ l $ → the adapter transforms this hidden state and feeds it through the shared LM Head to produce the first draft token $ x'_{1} $ → the confidence controller checks whether $ \max_n \mathcal{M}^{s}_n(x_{t+1}' ) > \eta $ → if yes, the draft token is fed back (autoregressively) as input to the shallow sub-network + adapter to produce the next draft token → this repeats until confidence drops below $ \eta $ or $ \gamma $ tokens have been drafted → the hidden states from all draft positions are concatenated into a parallel compute unit and fed through the remaining layers $ \mathcal{M}^{b}[l:] $ → the full model's logits at each position determine acceptance/rejection via standard speculative sampling → accepted tokens become part of the output, and the process repeats from the last accepted position.
3.3 Roadmap for the Deep Dive
- First, the formal notation framework that underpins the speculative decoding analysis, including the definition of consistent token acceptance rate (CTAR) and why it matters — these definitions establish the metrics by which all design decisions are evaluated.
- Second, the architecture and training of the self-draft model (the first early exit): how the shallow sub-network is selected, what the adapter architecture looks like, how it is trained, and why cross-entropy loss against the target model's distribution is used.
- Third, the dynamic drafting mechanism (the second early exit): how the confidence threshold
$ \eta $determines when to stop drafting, why this outperforms fixed-step drafting, and the relationship between threshold choice, compression rate, and wall-time speedup. - Fourth, the full decoding algorithm that integrates drafting and verification, including the parallel verification procedure and the speculative sampling acceptance mechanism that guarantees losslessness.
- Fifth, the hyperparameter selection methodology revealed through ablation studies (exit layer
$ l $, confidence threshold$ \eta $, maximum draft steps$ \gamma $, adapter architecture choices).
3.4 Detailed, Sentence-Based Technical Breakdown
Notation and the Consistent Token Acceptance Rate (CTAR)
The paper operates within the standard speculative decoding notation, but introduces a new evaluation metric — the consistent token acceptance rate (CTAR) — that drives the design of the dynamic drafting mechanism. Understanding this metric is essential because it motivates the entire second early exit.
The standard speculative decoding setup uses the following notation:
$ x^{t} $denotes the discrete token sequence$ (x_1, x_2, \dots, x_t) $— the context up to and including position$ t $.$ x^{i:j} $denotes the subsequence$ (x_i, x_{i+1}, \dots, x_j) $.$ \mathcal{V} $is the discrete vocabulary set, with size$ |\mathcal{V}| $.$ \mathcal{M}(\cdot \mid x^{t}) \in \mathbb{R}^{|\mathcal{V}|} $is the probability distribution over the vocabulary produced by a language model$ \mathcal{M} $given prefix$ x^{t} $.$ \mathcal{M}_n(\cdot \mid x^{t}) $is the$ n $-th entry of this distribution — the probability assigned to token$ n $.$ \mathcal{M}^{b} $is the large target language model (the one whose output distribution we want to preserve).$ \mathcal{M}^{s} $is the speculative small draft model (which proposes tokens cheaply).
Why standard evaluation metrics are insufficient. The paper introduces CTAR because the commonly reported compression rate (CR) obscures a crucial pattern:
The compression rate (CR) is defined as:
where $ S = [s_1, s_2, \dots, s_{|S|}] $ is the list of accepted token counts per forward pass of the target model $ \mathcal{M}^{b} $ during a decoding run that generates $ N $ total tokens, so $ \sum_k s_k = N $, and $ |S| $ is the number of target-model forward passes.
What it computes: CR is the average number of tokens accepted per verification step. If CR = 2.3, then on average each time we run the large model, we accept 2.3 draft tokens (plus the one guaranteed from the target model, for a total speedup factor of roughly 2.3× relative to a naive forward pass that produces one token).
Why it is insufficient: CR averages across all positions in the drafted sequence. It does not tell us whether the second drafted token has a 90% acceptance rate or a 30% acceptance rate. The paper observes empirically that acceptance rates decay with speculative distance — tokens further into the future are harder to predict correctly. A fixed drafting budget of $ \gamma $ tokens inevitably wastes computation on distant positions with low acceptance probability. The CR for a fixed-step drafter conflates high-acceptance close tokens with low-acceptance distant tokens, making it impossible to diagnose whether extending the draft length actually helps.
The consistent token acceptance rate (CTAR) is defined as:
where $ w $ is a window size (a speculative distance), $ s_k $ is the number of accepted tokens in the $ k $-th verification step (as in the CR definition), $ \mathbb{I}(\cdot) $ is the indicator function (1 if the condition holds, 0 otherwise), and the sum is over all verification steps.
What it computes: For a given window size $ w $, CTAR($ w $) is the fraction of verification steps in which at least $ w+1 $ tokens are accepted (equivalently, the fraction of steps where the $ w $-th drafted token — the token $ w $ positions ahead — is accepted, given that all earlier tokens were accepted). In greedy decoding, CTAR($ x^t, w $) for a specific prefix $ x^t $ is a binary value: it equals 1 if the top-1 predictions of $ \mathcal{M}^{s} $ and $ \mathcal{M}^{b} $ agree at every position from $ t+1 $ through $ t+w $, and 0 otherwise (there is at least one disagreement, which triggers rejection and discards all subsequent tokens regardless of their correctness).
Why this form: Unlike CR which averages token counts, CTAR directly measures the probability of successfully speculating $ w $ tokens ahead conditioned on success at all previous steps. This is exactly the right quantity for deciding whether to allocate drafting effort to position $ w+1 $: if CTAR($ w $) is low, the tokens at and beyond position $ w $ will rarely survive verification even if individually plausible, because the rejection mechanism discards everything after the first mismatch.
The key empirical finding from CTAR (Figure 1(a)): On the mathematical reasoning subtask of Spec-Bench, CTAR decays rapidly with window size for all methods, but the decay rate differs dramatically across architectures. Medusa's time-independent heads show steep decay (CTAR drops quickly as $ w $ increases from 1 to 6), while autoregressive drafters (Kangaroo, Lookahead) show slower decay because each draft token conditions on previously drafted tokens, capturing sequential dependencies. This empirical pattern is the direct motivation for Kangaroo's dynamic drafting: rather than always drafting $ \gamma = 6 $ tokens and wasting computation on low-CTAR positions, stop drafting when the draft model's own confidence signal indicates that the next position's acceptance probability is likely low.
The Self-Draft Model via Early Exiting (First Early Exit)
The first early exit is architectural: instead of training a separate draft model from scratch or attaching prediction heads to the final layer, Kangaroo simply cuts the target model at layer $ l $ and trains a small adapter to map from that shallow representation to full-model-quality predictions. The draft model is thus:
where $ \mathcal{M}^{b}[:l] $ denotes the first $ l $ transformer layers of the target model $ \mathcal{M}^{b} $ (the shared shallow sub-network), $ \mathcal{A} $ is the adapter network (a small trainable module), $ \circ $ denotes functional composition, and $ \mathcal{M}^{s} $ is the resulting self-draft model. The draft model also shares the LM Head of the target model — the final linear projection from hidden states to vocabulary logits — so $ \mathcal{M}^{s} $ produces predictions through the same unembedding matrix as $ \mathcal{M}^{b} $.
Why only the first $ l $ layers? The depth of the shared sub-network controls a trade-off between draft model quality and draft model speed. Using more layers (deeper early exit, larger $ l $) gives the adapter access to richer features, potentially increasing token acceptance rate. But it also increases the inference cost per draft token, because more transformer layers must execute for each sequentially drafted token. The paper explores this trade-off experimentally (Figure 3(a)) and selects $ l = 2 $ for Vicuna-7B and $ l = 3 $ for Vicuna-13B — extremely shallow exits, using only 2–3 of the 32–40 total transformer layers.
Why this is a first "early exit": In the standard autoregressive usage, the target model $ \mathcal{M}^{b} $ must execute all $ L $ layers for every generated token. The draft model $ \mathcal{M}^{s} $ instead "exits early" after only $ l \ll L $ layers, producing predictions much faster (fewer memory reads, fewer FLOPs). The adapter then corrects for the missing layers — it is explicitly trained to bridge the representation gap between layer $ l $ and the final layer $ L $.
Adapter Network Architecture and Design Rationale
The adapter $ \mathcal{A} $ has a deliberately minimal architecture: one multi-head attention (MHA) block and two RMS Layer Normalization layers (RMSNorm). There is no feed-forward network (FFN) component. The paper describes this as "surprisingly... efficient but powerful" (Section 1), and Table 2 quantifies the trade-offs.
The architecture in detail (Figure 2):
- Input: The hidden state
$ h_l $from layer$ l $of the target model, which is the output of$ \mathcal{M}^{b}[:l] $given the current prefix. - First RMSNorm: Normalizes the input hidden state.
- Multi-head attention: Applies self-attention over the available context (the same key-value cache that
$ \mathcal{M}^{b}[:l] $uses, since the shallow sub-network is frozen). This is the only learned transformation with substantial capacity. - Second RMSNorm: Normalizes the attention output.
- The resulting hidden state is then passed through the shared LM Head (the target model's unembedding matrix) to produce logits, and a softmax yields the draft probability distribution.
What is conspicuously absent: A feed-forward network (FFN). In a standard transformer block, the FFN accounts for roughly 67% of the total parameters (the paper cites this fraction in Section 4.2 when discussing Table 2). By omitting the FFN entirely, the adapter remains extremely lightweight — only the MHA weights plus the two RMSNorm parameter vectors (scale and shift per dimension). For Vicuna-7B, the adapter totals 67M parameters, compared to Medusa-1's 591M parameters for its multiple independent prediction heads — an 88.7% reduction in additional parameters (stated in the abstract and Section 1).
Why this minimal design works: The self-attention mechanism in the adapter serves a critical role that the paper implies but does not belabor: it allows the draft model to be autoregressive. Unlike Medusa's heads, which each predict independently from the same hidden state, the adapter's MHA can attend to the full context including previously drafted tokens. When generating the second draft token $ x'_{2} $, the attention can look back at $ x'_{1} $, enabling sequential dependencies among draft tokens. This autoregressive structure is why Kangaroo's CTAR decays more slowly than Medusa's (Figure 1(a)).
The two RMSNorm layers serve standard purposes: the first stabilizes the input distribution to the attention block (important because the adapter receives hidden states from a frozen sub-network whose statistics it cannot modify), and the second stabilizes the output before the LM Head projection.
Ablation evidence for the design (Table 2): The paper tests three adapter variants for Vicuna-7B on Spec-Bench:
- With FFN + separate LM Head: 173M additional parameters, achieves an average speedup of 1.36×.
- Without FFN + separate LM Head: 136M additional parameters, achieves 1.37× speedup.
- Without FFN + shared LM Head (the Kangaroo design): 67M additional parameters, achieves 1.50× speedup.
The progression is telling. Removing the FFN barely affects speedup (1.36→1.37) despite eliminating ~37M parameters — the FFN was providing negligible acceptance-rate benefit for its cost. But sharing the LM Head instead of training a separate one produces a substantial jump (1.37→1.50), likely because sharing ensures the draft model's vocabulary-level predictions are computed through the same projection as the target model, eliminating a potential source of distributional mismatch. The smaller parameter count also reduces memory overhead during deployment.
Training the Adapter: Objective and Procedure
The adapter $ \mathcal{A} $ is the only component that is trained — the target model's layers (both the shallow sub-network $ \mathcal{M}^{b}[:l] $ and the remaining layers $ \mathcal{M}^{b}[l:] $) remain frozen throughout.
Training loss. The natural objective would be to maximize the token acceptance rate — that is, maximize the probability that the draft model's prediction agrees with the target model's prediction at each position. However, the paper observes that cross-entropy loss converges faster as a training objective while being a close proxy:
where the outer sum is over all token positions $ t $ in the training data, the inner sum is over the vocabulary indices $ n = 1, \dots, |\mathcal{V}| $, $ \mathcal{M}^{b}_{n}(x_{t}) $ is the target probability that the full model $ \mathcal{M}^{b} $ assigns to token $ n $ at position $ t $ (given the true prefix $ x^{t-1} $), and $ \mathcal{M}^{s}_{n}(x_{t}) $ is the draft probability that the self-draft model $ \mathcal{M}^{s} $ assigns to the same token $ n $ at the same position.
What it computes: This is the standard cross-entropy between two categorical distributions — the target model's full output distribution serves as the soft "label" distribution, and the draft model's output distribution is trained to match it. At each training position, the loss penalizes the draft model for putting probability mass on tokens that the target model considers unlikely, weighted by how much probability the target model assigns to each token.
Why this form over directly maximizing acceptance rate: The token acceptance rate in greedy decoding depends only on whether the draft and target agree on the argmax (top-1) prediction. Training to directly maximize this binary agreement would produce a non-smooth, non-differentiable objective. Cross-entropy is smooth and differentiable, and minimizing it naturally pushes the draft model's distribution to match the target model's distribution globally — when the distributions match exactly, the argmaxes will agree with high probability. The paper claims (without extensive empirical justification) that this objective "exhibits faster convergence rate" than alternatives, likely because the soft distributional matching provides richer gradient signal per example than a binary agree/disagree signal.
Why the target model's distribution is the "label" rather than the ground-truth next token: The goal is not to train a model that predicts correct tokens per se — it is to train a model that mimics the target model as closely as possible, including mimicking the target's errors and idiosyncrasies. If the draft model were trained to predict ground-truth next tokens while the target model makes different predictions, the draft-target mismatch would cause low acceptance rates during speculative decoding and therefore poor speedups. Training against the target model's own distribution (effectively, distillation from the target to the draft) ensures that the draft's behavior is aligned with the verifier's behavior, maximizing the probability that they agree.
Training procedure (Section 4.1):
- Dataset: ShareGPT, following the same training data choice as Medusa (Cai et al., 2024) — a dataset of multi-turn conversations between users and AI assistants, chosen because it provides diverse, natural language contexts representative of chatbot deployment scenarios.
- Optimizer: AdamW (Loshchilov & Hutter, 2017), a variant of Adam with decoupled weight decay regularization that is standard for transformer fine-tuning.
- Epochs: 10, a relatively short training schedule consistent with the adapter's small parameter count and the goal of low training cost.
- Hardware: The adapter training runs on NVIDIA V100 GPUs (the same hardware used for inference benchmarking), though the paper does not specify the number of GPUs or training wall-clock time.
The paper emphasizes that this training procedure is "low-cost" compared to training a standalone draft model from scratch — the adapter has only 67M parameters (less than 1% of the 7B parameter target model), trains on a single conversational dataset rather than a massive pretraining corpus, and converges in 10 epochs. This positions Kangaroo as deployable with minimal upfront investment.
Dynamic Drafting via Confidence-Thresholded Early Exiting (Second Early Exit)
The second early exit is behavioral rather than architectural: during the drafting phase, Kangaroo stops generating additional draft tokens when the self-draft model's confidence in its top prediction falls below a threshold $ \eta $. This is an "exit" in the sense that drafting terminates early based on a runtime signal, rather than proceeding to a fixed maximum step count.
The stopping criterion is:
where $ \mathcal{M}^{s}_n(x) $ is the draft model's predicted probability for token $ n $ given the current context $ x $ (which includes all previously generated draft tokens in the current round), $ \max_n $ selects the maximum probability across all vocabulary entries, and $ \eta = 0.6 $ is the predefined confidence threshold.
What this means operationally: After the draft model produces a probability distribution for the next draft token, the system examines the top-1 probability — the softmax score of the most likely token. If this score is 0.6 or lower, drafting halts immediately. The low-confidence token is still proposed for verification (it is not discarded), but no further tokens are drafted in this round. If the score exceeds 0.6, drafting continues to the next position, up to the maximum of $ \gamma = 6 $ tokens.
Why 0.6? The threshold $ \eta $ quantifies draft-model confidence. A threshold of 1.0 would mean drafting continues only as long as the model is absolutely certain (never, in practice, so drafting would be 0 steps — no speculative benefit). A threshold of 0.0 would mean drafting never stops early, equivalent to fixed-step drafting of $ \gamma $ tokens regardless of confidence. The value 0.6 is an empirically chosen sweet spot (from the ablation in Figure 3(b)) that balances two competing effects:
-
Stopping too early (high
$ \eta $): Drafting terminates after only 1–2 tokens even when the draft model might successfully predict 3–4 tokens at useful acceptance rates. This loses the opportunity to verify more tokens per target-model forward pass, reducing compression rate and potential speedup. -
Stopping too late (low
$ \eta $): Drafting continues to low-confidence positions (4–6) where acceptance probability is low. These tokens will likely be rejected during verification, meaning the drafting time spent on them was wasted — the draft model's sequential forward passes consumed GPU time without producing accepted output tokens. This is exactly the problem visible in Figure 1(a), where CTAR(5) and CTAR(6) are low for all methods.
Why this is an "early exit": The standard speculative decoding procedure always drafts exactly $ \gamma $ tokens (or a fixed tree of tokens, in the case of Medusa and SpecInfer). Kangaroo's drafting procedure can terminate after $ d $ tokens where $ 1 \leq d \leq \gamma $, depending on the draft model's confidence trajectory. The word "early" refers to exiting before the maximum draft length $ \gamma $ is reached.
The relationship between dynamic drafting and CTAR: The CTAR metric from Section 3 directly motivates this mechanism. If CTAR($ w $) is low for larger $ w $, then drafting tokens at positions beyond $ w $ is unlikely to produce accepted output — those tokens will be discarded during verification. The draft model's own confidence signal serves as a proxy for CTAR at the current position: when the draft model is uncertain about a token, it is likely that the full target model would produce a different token (either the target model's top-1 is different, or the draft model's uncertainty reflects genuine ambiguity that the target model might resolve differently). By halting on low confidence, Kangaroo avoids incurring the drafting cost at positions where the probability of contributing to accepted output is low.
The paper implicitly argues that confidence is a better signal than any fixed schedule because it is context-dependent: the model's confidence naturally varies with the difficulty of the next prediction. On easy-to-predict tokens (common function words, predictable continuations of formulaic text), confidence remains high and drafting continues to the full $ \gamma $ steps. On harder tokens (unusual vocabulary, creative continuations, math problem steps requiring computation), confidence drops early, and drafting stops to avoid wasted cycles. This adaptivity is what gives Kangaroo its name — it "hops" (speculates) a variable distance depending on the terrain (context difficulty).
Empirical validation of dynamic vs. fixed drafting (Figure 3(b)): The paper compares dynamic drafting at various thresholds $ \eta $ against fixed-step drafting ($ \eta = 0 $, which means the stopping condition is never triggered since $ \max_n \mathcal{M}^{s}_n(x) > 0 $ always holds for a proper probability distribution). The results on Spec-Bench show:
$ \eta = 0 $(fixed-step, maximum$ \gamma $) achieves the highest compression rate — it generates the most draft tokens per verification step, and on average more of those tokens are accepted than are wasted (since CTAR, while decaying, remains above zero for positions 2–6 on many token sequences).- However,
$ \eta = 0 $achieves a sub-optimal wall-time speedup — the extra drafting time consumed generating low-CTAR tokens outweighs the marginal increase in accepted tokens. $ \eta = 0.6 $achieves the optimal wall-time speedup across most maximum draft step settings, confirming that the confidence threshold effectively prunes wasted drafting cycles.- The optimal
$ \eta $is "consistent across different maximum different steps" (Section 4.2), meaning the right threshold does not depend strongly on the choice of$ \gamma $— a practically useful property since it means$ \eta $and$ \gamma $can be tuned independently.
The maximum draft step $ \gamma = 6 $: The paper sets the upper bound on draft tokens to 6 per verification step. This value likely comes from the typical speculative decoding literature (where $ \gamma $ is often set to 3–8 based on empirical performance) and the observation in Figure 1(a) that CTAR beyond position 6 is near zero for all methods on the mathematical reasoning subtask — there is essentially no benefit to drafting more than 6 tokens ahead, regardless of drafting quality. The choice of 6 is thus a conservative upper bound that ensures the dynamic drafting mechanism has room to operate (it can stop anywhere from 1 to 6 tokens) without wasting time exploring impossibly long speculative sequences.
The Full Decoding Algorithm: Drafting and Verification
The Kangaroo decoding algorithm interleaves draft generation (using the shallow sub-network + adapter, with dynamic early termination) and parallel verification (using the full model's remaining layers). The algorithm is lossless — it preserves the sampling distribution of the target model $ \mathcal{M}^{b} $ — through the standard speculative sampling acceptance mechanism.
Step-by-step, one verification round proceeds as follows:
Phase 1: Context preparation for drafting. The system has a prefix $ x^{t} $ (the tokens generated and verified so far, up to position $ t $). The shallow sub-network $ \mathcal{M}^{b}[:l] $ performs a forward pass on this prefix, producing a hidden state $ h_l(x^{t}) $ at layer $ l $ — the representation of the last token after $ l $ transformer layers. This hidden state serves as the starting point for the adapter.
Phase 2: Autoregressive drafting with confidence monitoring. The adapter $ \mathcal{A} $ transforms $ h_l(x^{t}) $ through its MHA + RMSNorm pipeline, producing a hidden state $ \tilde{h}(x^{t}) $. This is projected through the shared LM Head to produce $ \mathcal{M}^{s}(\cdot \mid x^{t}) $, a probability distribution over the vocabulary. The token with maximum probability is selected as the first draft token $ x'_{1} $.
The confidence monitor checks: $ \max_n \mathcal{M}^{s}_n(x'_{1}) > \eta $? If yes, drafting continues. The draft token $ x'_{1} $ is appended to the context, the shallow sub-network processes the extended sequence $ (x^{t}, x'_{1}) $ to produce $ h_l(x^{t}, x'_{1}) $, and the adapter generates $ x'_{2} $. This autoregressive loop continues until either:
$ \max_n \mathcal{M}^{s}_n(x'_d) \leq \eta $(confidence drops below threshold), or$ d = \gamma = 6 $(maximum draft length reached).
At termination, the system has generated $ d $ draft tokens $ x'_{1}, x'_{2}, \dots, x'_{d} $ where $ 1 \leq d \leq 6 $.
Phase 3: Parallel verification. The key efficiency insight of speculative decoding is that verification does not need to proceed sequentially — the full model can evaluate all draft tokens in parallel using a single forward pass. Kangaroo achieves this by collecting the hidden states at the early exit layer for each draft position. Specifically, during the drafting phase, each draft token $ x'_i $ is generated by the shallow sub-network processing the prefix including previous draft tokens, producing a hidden state $ h_l(x^{t}, x'_{1}, \dots, x'_{i-1}) $ at layer $ l $. These hidden states are saved.
For verification, the system concatenates all these early-exit hidden states into a parallel compute unit: $ [h_l(x^{t}), h_l(x^{t}, x'_{1}), \dots, h_l(x^{t}, x'_{1}, \dots, x'_{d-1})] $. This is a sequence of $ d $ hidden state vectors, one for each position that needs verification. The remaining layers of the target model, $ \mathcal{M}^{b}[l:] $ (layers $ l+1 $ through $ L $), process this entire sequence in a single forward pass, producing full-model logits at each of the $ d $ positions.
Why this works: The transformer architecture processes sequences in parallel within each layer — the computation at position $ i $ depends only on positions $ \leq i $ through the causal attention mask. By providing the hidden states that the shallow sub-network would produce for each prefix (including the draft tokens), the remaining layers can "continue" processing from where the draft model left off, as if the full model had been run autoregressively up to those points. This is the standard speculative decoding verification trick: the full model evaluates all draft positions in one forward pass because the hidden states capture the autoregressive state.
Phase 4: Acceptance-rejection via speculative sampling. The verifier produces probability distributions $ \mathcal{M}^{b}(\cdot \mid x^{t}), \mathcal{M}^{b}(\cdot \mid x^{t}, x'_{1}), \dots, \mathcal{M}^{b}(\cdot \mid x^{t}, x'_{1}, \dots, x'_{d-1}) $ — the target model's predictions at each draft position. For each draft token $ x'_i $ (at position $ t+i $), standard speculative sampling (Leviathan et al., 2023; Chen et al., 2023) is applied:
- Compare
$ p = \mathcal{M}^{s}(x'_i \mid \dots) $and$ q = \mathcal{M}^{b}(x'_i \mid \dots) $— the draft and target probabilities assigned to the drafted token. - Accept
$ x'_i $with probability$ \min(1, q/p) $. If accepted, move to position$ i+1 $. - If rejected, the token at position
$ t+i $is instead sampled from the residual distribution:$ \text{norm}(\max(0, \mathcal{M}^{b}(\cdot \mid \dots) - \mathcal{M}^{s}(\cdot \mid \dots))) $— the target model's distribution minus the draft model's distribution, renormalized. All subsequent draft tokens$ x'_{i+1}, \dots, x'_{d} $are discarded regardless of their quality.
This procedure guarantees (mathematically, as proven by Leviathan et al., 2023) that the output distribution is identical to running the target model $ \mathcal{M}^{b} $ autoregressively. The draft model's role is purely as a proposal mechanism; any biases it introduces are corrected by the acceptance-rejection step.
Phase 5: Preparing for the next round. The last accepted token position becomes the new context prefix. Importantly, the paper notes an optimization (Figure 2 caption): if all $ d $ draft tokens are accepted, the next drafting round can start from position $ t+d $ directly, without needing to re-process the accepted tokens through the shallow sub-network to obtain their early-exit hidden states — those hidden states were already computed during the verification phase (since the full model's forward pass necessarily produces hidden states at all layers, including layer $ l $, for all positions it processes). This saves the latency of a redundant forward pass through the shallow sub-network at the start of the next round, providing a small but consistent efficiency gain across rounds where all draft tokens are accepted.
Hyperparameter Selection and Design Trade-offs
The paper's ablation studies (Section 4.2, Figures 3(a) and 3(b), Table 2) reveal the empirical basis for Kangaroo's hyperparameter choices and illuminate the trade-offs that govern its design.
Early exit layer $ l $ (Figure 3(a)): The paper sweeps the exit layer depth and measures both compression rate (CR) and wall-time speedup on Spec-Bench. The optimal depth differs by model scale:
- Vicuna-7B:
$ l = 2 $achieves the best speedup. Deeper exits ($ l = 4, 6 $) achieve higher compression rates (more accepted tokens per verification step, since the draft model has access to richer features) but lower wall-time speedup because the draft model's forward pass becomes slower. The extreme case is$ l = 0 $(no shared layers — the adapter receives raw token embeddings), which is fast but has unacceptably low compression rate. - Vicuna-13B:
$ l = 3 $is optimal. The larger model can afford a slightly deeper exit because the absolute cost of running 3 layers of a 13B model as a draft model is still acceptable relative to the target model's total depth (40 layers for Vicuna-13B), while 2 layers might not capture enough representational capacity to achieve competitive acceptance rates.
The general principle: the optimal exit layer is the shallowest layer that still provides enough representational quality for the adapter to produce useful draft tokens. The paper selects the depth that maximizes speedup (the deployment-relevant metric), not compression rate (the intermediate metric).
Confidence threshold $ \eta $ (Figure 3(b)): Sweeping $ \eta $ from 0 (fixed-step drafting, no early termination) to 0.9 (aggressive early termination) across different maximum draft step settings ($ \gamma = 4, 6, 8 $) reveals a consistent pattern:
$ \eta = 0 $(fixed-step in all configurations): Maximum compression rate, sub-optimal speedup.$ \eta = 0.6 $: Near-optimal speedup across all$ \gamma $configurations tested. This is the chosen value.$ \eta > 0.7 $: Speedup decreases as drafting terminates too aggressively, losing the opportunity to draft and verify moderately-confident tokens that would have been accepted.
The paper notes that the optimal threshold is "consistent across different maximum different steps" (Section 4.2), which is a practically important finding: it means the practitioner can choose $ \eta = 0.6 $ as a default without needing to re-tune it when adjusting $ \gamma $ for different deployment scenarios.
Maximum draft steps $ \gamma = 6 $: The paper selects $ \gamma = 6 $ based on empirical performance in Figure 3(b), where $ \gamma = 6 $ with $ \eta = 0.6 $ achieves the best speedup among the tested configurations ($ \gamma = 4, 6, 8 $). The marginal benefit of $ \gamma = 8 $ over $ \gamma = 6 $ is small, likely because CTAR near position 8 is close to zero for most contexts, meaning the dynamic termination mechanism would almost always stop before reaching 8 tokens anyway — the extra $ \gamma $ headroom simply doesn't get used often enough to affect average performance.
Adapter architecture (Table 2, discussed in detail above): The progression of ablations (with FFN, without FFN, shared vs. separate LM Head) converges on the minimal design: MHA only, shared LM Head. This achieves the best speedup (1.50×) with the fewest parameters (67M). The key insight from Table 2 is that parameter count for the draft model directly impacts deployment efficiency, not just training cost — the adapter's weights must be resident in GPU memory alongside the target model during inference, so reducing adapter size both lowers the memory footprint and reduces the adapter's own inference latency.
Why these hyperparameters matter beyond the specific values: The ablation studies collectively demonstrate a principle that generalizes across all speculative decoding methods: the metric to optimize is wall-time speedup, not compression rate or token acceptance rate in isolation. A method with a higher compression rate (more tokens accepted per verification) can be slower end-to-end if the drafting phase consumes too much time — this is Lookahead's problem in Figure 1(b). A method with extremely fast drafting (like Medusa's parallel heads) can underperform if acceptance rate is too low — this is Medusa's problem. Kangaroo's design navigates this trade-off by:
- Making the draft model fast enough (shallow exit, lightweight adapter) that its sequential drafting cost remains manageable.
- Making the acceptance rate high enough (autoregressive adapter capturing sequential dependencies) that drafting effort converts to accepted tokens at a useful rate.
- Dynamically adjusting the draft length so that the ratio of useful work (accepted tokens) to wasted work (rejected tokens + drafting overhead for low-confidence predictions) stays favorable across diverse contexts.
4. Key Insights and Innovations
Innovation 1: Framing Speculative Draft Generation as an Adaptive-Confidence, Not Fixed-Budget, Problem
The dominant assumption in speculative decoding — both in external-draft-model and self-drafting paradigms — is that the draft model should propose a fixed number of tokens per verification step. Medusa always predicts from exactly $\gamma$ parallel heads. Lookahead iterates for a fixed number of Jacobi steps. Even standard speculative decoding with a standalone draft model generates exactly $\gamma$ tokens autoregressively per round. The field's optimization problem was: given a fixed drafting budget $\gamma$, how do we maximize token acceptance rate while minimizing per-token draft cost?
Kangaroo redefines the problem entirely. The core insight, crystallized in the consistent token acceptance rate (CTAR) metric and validated in Figure 1(a), is that token acceptance probability decays systematically with speculative distance — and that this decay is predictable from the draft model's own confidence signal. The innovation is not the dynamic termination mechanism itself (early exiting based on confidence exists in other inference contexts, e.g., Schuster et al., 2022), but rather recognizing that speculative decoding's sequential rejection dynamics (one rejection discards all subsequent tokens) make fixed-budget drafting inherently wasteful in a way that standard metrics like compression rate fail to capture.
This is a conceptual reframing, not an incremental refinement. Prior work implicitly treated all draft positions as equally valuable — the compression rate metric averages across them, masking the fact that position-6 tokens contribute far less to throughput than position-2 tokens. Kangaroo's introduction of CTAR as a diagnostic tool (Section 3, Definition 1) makes this heterogeneity visible, and the confidence-thresholded termination is the natural policy response: stop drafting when the marginal expected benefit of the next token falls below the marginal cost of generating it.
What makes this significant beyond the empirical speedup is that it converts a static resource-allocation problem into an adaptive, per-context decision problem. The optimal number of draft tokens becomes a function of the input, not a global hyperparameter. On formulaic text continuations, the model drafts to the maximum $\gamma = 6$; on creative or computationally-demanding tokens, it stops after 1–2. This adaptivity is what gives Kangaroo its conceptual name — it "hops" variable distances depending on the terrain — and it implies that future speculative decoding methods should be evaluated not just by their peak speedup on a benchmark, but by how well they dynamically allocate drafting effort across heterogeneous contexts.
The evidence base is Figure 3(b): fixed-step drafting ($\eta = 0$) achieves the maximum compression rate but sub-optimal wall-time speedup, while $\eta = 0.6$ achieves the best speedup despite lower compression rate. This directly proves that compression rate — the field's standard evaluation metric — is an inadequate proxy for the true objective (throughput), and that methods optimizing for compression rate (like fixed-step drafting with large $\gamma$) leave performance on the table. This is a cautionary finding for the speculative decoding literature broadly, extending beyond Kangaroo's specific mechanism.
Innovation 2: Reconciling Self-Drafting Speed and Quality via Extreme Architectural Asymmetry
The self-drafting literature prior to Kangaroo faced an apparent trade-off: high acceptance rate requires a draft model that closely approximates the target, but close approximation implies the draft model is large and slow, undermining the speculative speedup. Draft & Verify (Zhang et al., 2023) demonstrated this tension concretely: using intermediate-layer early exits as the draft model achieved "high token acceptance rate" but "exceptionally high inference latency" for the draft, crippling end-to-end acceleration (Section 2, paragraph on Draft & Verify). The field's default solution was to abandon autoregressive self-drafting entirely in favor of parallel prediction heads (Medusa), which achieve extremely low draft latency at the cost of degraded acceptance rate on longer speculative sequences.
Kangaroo's architectural innovation is demonstrating that this trade-off is not fundamental — it can be broken by combining an aggressive early exit (layer 2 of 32, for Vicuna-7B) with a small learned adapter that bridges the resulting massive representation gap. The key conceptual move is recognizing that the draft model does not need to be a good language model in general — it only needs to produce predictions that the target model will accept at inference time on the deployment distribution. This is a much weaker requirement, and it justifies an extreme design choice (exiting after ~6% of the total layers) that would be absurd for any standalone model.
The depth of the early exit ($l = 2$ for Vicuna-7B) is the critical number here — it is aggressively shallow, far more so than prior early-exit approaches for inference acceleration. Standard early exiting typically exits at 50% or more of total depth to preserve reasonable prediction quality (Schuster et al., 2022; Varshney et al., 2023). Kangaroo exits at ~6% of depth, accepting a dramatic quality degradation that the adapter — and, crucially, the verification mechanism — corrects. The adapter is not asked to make the draft model good; it is asked to make the draft model mimic the target model well enough on average that the speculative yield (accepted tokens per drafting FLOP) exceeds 1.0.
This is a fundamental reframing of what "draft model quality" means in speculative decoding. Prior work implicitly evaluated draft models by their perplexity or standalone accuracy. Kangaroo evaluates them by a single number: the ratio of accepted tokens produced to drafting forward passes consumed. This metric can be high even for a "bad" language model, provided its errors align with the target model's behavior and it is fast enough. The extreme-layer-exit design works because even the shallowest transformer layers capture substantial lexical and syntactic information (word identity, part-of-speech, simple collocations) that the adapter can leverage, while being dramatically faster than the full model.
The evidence is in the ablation studies. The adapter architecture is pared to the absolute minimum — one MHA block, no FFN, shared LM Head — and still delivers speedups superior to all prior self-drafting methods (Figure 1(b), Table 1) while using 88.7% fewer additional parameters than Medusa. The speedup advantage over Draft & Verify (which the paper implies in Section 2 but does not benchmark directly) stems directly from the shallower exit: Draft & Verify's "intermediate redundant layers" are still too deep to achieve net acceleration, while Kangaroo's layer-2 exit is fast enough that even sequential autoregressive drafting — fundamentally slower than Medusa's parallel heads — still yields net throughput gains.
Innovation 3: CTAR as a Diagnostic Tool That Reveals Why Drafting Strategies Succeed or Fail
The paper introduces the consistent token acceptance rate (CTAR) as a new evaluation metric (Definition 1, Section 3), and while this may appear to be a minor methodological contribution, it functions as a diagnostic instrument that explains the performance landscape of self-drafting methods in a way that compression rate cannot. This is a conceptual innovation, not a mechanical one — the metric itself is simple (fraction of verification steps where at least $w$ consecutive tokens are accepted), but the insight it enables is substantial.
Compression rate answers "how many tokens, on average, are accepted per verification step?" CTAR answers "what is the probability that the $w$-th drafted token contributes to throughput, given that tokens 1 through $w-1$ did?" The distinction matters because speculative decoding's rejection rule — discard all tokens after the first mismatch — makes the value of position $w$ conditional on success at all earlier positions. A method might have a compression rate of 2.5, but if CTAR(2) = 0.4, then 60% of the time, position-3 tokens (and beyond) are wasted regardless of their individual quality. CTAR makes this hidden waste visible.
The paper uses CTAR diagnostically in Figure 1(a) to explain why different self-drafting methods achieve different end-to-end speedups despite similar compression rates. Medusa's time-independent heads achieve reasonable CTAR(1) (the first predicted token is often correct) but CTAR decays steeply at positions 2–3 because the heads cannot model sequential token dependencies effectively. Lookahead and Kangaroo, which use autoregressive draft generation, maintain higher CTAR at positions 2–4 because each draft token conditions on previously drafted tokens. This explains the speedup pattern in Figure 1(b): Lookahead's high CTAR is offset by its slow drafting, while Medusa's fast drafting is offset by its low CTAR; Kangaroo achieves the best of both by combining autoregressive drafting (for high CTAR) with shallow-exit speed and dynamic termination (for low drafting latency).
The significance of CTAR extends beyond this paper. It provides the speculative decoding literature with a tool to diagnose failure modes separately: a method with low CTAR(1) has a fundamental draft-target mismatch problem (the draft model doesn't agree with the target on the immediate next token); a method with low CTAR(3) but high CTAR(1) has a long-range dependency problem (the draft captures local structure but not broader coherence). These distinct failure modes call for different fixes — better distillation vs. better sequential modeling vs. dynamic termination — and CTAR separates them in a way that the scalar compression rate cannot. The paper does not elaborate this diagnostic framework explicitly, but it is implicit in the analysis of Figure 1(a), and it constitutes a transferable conceptual contribution to speculative decoding evaluation methodology.
Innovation 4: Making Early Exiting Lossless by Embedding It in a Speculative Verification Framework
Early exiting for inference acceleration — producing predictions from intermediate layers to skip remaining computation — is well-established but has a fundamental limitation: it is lossy. The intermediate-layer predictions are worse than final-layer predictions, and there is no mechanism to recover the lost quality. The paper cites this explicitly: "since early exiting accelerates inference by saving subsequent computations, it inevitably incurs the issue of performance degradation" (Section 2, citing Schuster et al., 2022).
Kangaroo's innovation is to recognize that early exiting and speculative decoding are complementary in a specific, non-obvious way: speculative decoding provides a verification mechanism that makes lossy draft predictions lossless overall, while early exiting provides a way to create a draft model with essentially zero training cost and shared KV cache. This synthesis converts a weakness (early-exit predictions are imperfect) into a manageable cost (rejected draft tokens waste some computation but do not affect output distribution).
This is more than a simple combination of existing techniques. The intellectual leap is understanding that the quality requirements for a speculative draft model are fundamentally different from those for a standalone early-exit model. A standalone early-exit model must be good enough to use directly — its predictions are the output. A speculative draft model only needs to be good enough that the expected benefit of accepted draft tokens exceeds the expected cost of rejected ones. This difference dramatically relaxes the acceptable quality threshold, enabling much more aggressive early exits (layer 2 rather than layer 16) than would be tolerable in a lossy system.
The paper does not make this argument explicitly, but it is the logical foundation for why Kangaroo works when Draft & Verify (which also combines early exiting with verification) struggles. Draft & Verify uses a relatively deep early exit to maintain high acceptance rate, but this makes the draft model too slow for net speedup. Kangaroo uses an extreme early exit (because losslessness is guaranteed by verification, so draft quality can be much lower) and compensates with the adapter for what little quality is needed. The verification guarantee is what enables the aggressive depth reduction — without it, layer-2 predictions would be unusably bad, but with it, they only need to be right often enough to justify their drafting cost.
The evidence is implicit in Kangaroo's design choices: the selection of $l = 2$ (not 4, not 8) as optimal for Vicuna-7B in Figure 3(a) only makes sense under the logic that draft quality requirements are relaxed by verification. If the draft model's predictions were the final output, layer 2 would be catastrophically inadequate. Under the speculative framework, it is optimal because the speed gain from extreme shallowness outweighs the acceptance-rate degradation.
This insight suggests a broader principle for inference acceleration: any lossy acceleration technique can be made lossless by wrapping it in a speculative verification loop, provided the technique produces candidate outputs that the full model can verify in parallel. This principle — which the paper demonstrates but does not articulate as a general design pattern — could extend beyond early exiting to quantization, pruning, or retrieval-based generation, where lossy fast approximations propose candidates that a full-precision model verifies. </output>
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All experiments use Spec-Bench (Xia et al., 2024), a benchmark specifically designed for evaluating speculative decoding methods. Spec-Bench contains six subtasks spanning diverse text generation scenarios: translation, summarization, mathematical reasoning, retrieval-augmented generation (RAG), question answering, and multi-turn conversation. The paper reports per-subtask speedup and compression rate in Table 1, and uses aggregated "average" results for ablation studies. The use of Spec-Bench rather than a standard language modeling benchmark (e.g., WikiText, C4) is deliberate: it tests speculative decoding under realistic deployment conditions where the model generates diverse output types, not just open-ended text continuations. The six subtasks vary substantially in difficulty — mathematical reasoning requires precise symbolic computation, RAG involves integrating external knowledge, and multi-turn conversation requires maintaining dialogue coherence — providing a more thorough stress test than homogeneous perplexity-based benchmarks.
-
Base model(s). Experiments are conducted on Vicuna-7B and Vicuna-13B (Chiang et al., 2023), open-source chatbot models fine-tuned from LLaMA (Touvron et al., 2023) on user-shared conversational data. The choice is pragmatic: Vicuna models are widely used in the speculative decoding literature (Medusa, REST, and Lookahead all evaluate on Vicuna), enabling direct comparison with prior work. The 7B and 13B scales bracket a deployment-relevant range — small enough to fit on single GPUs for interactive applications, large enough to exhibit the memory-bandwidth bottleneck that motivates speculative decoding. The paper does not evaluate on models larger than 13B (e.g., Vicuna-33B or LLaMA-70B), meaning the scaling behavior of Kangaroo's design choices (exit layer depth, adapter capacity, confidence threshold) at larger model sizes remains unexplored. The 7B model has 32 transformer layers (per the LLaMA architecture), making
$l=2$an exit after ~6% of total depth; the 13B model likely has 40 layers, making$l=3$an exit after ~7.5% depth. -
Metrics. Two primary metrics are reported:
- Wall-time speedup ratio: the factor by which Kangaroo reduces end-to-end generation latency compared to vanilla autoregressive decoding with the target model. This is measured directly by clocking the time to generate a fixed number of tokens (or complete a fixed set of prompts) under each configuration, with all experiments conducted on the same hardware (NVIDIA V100 GPUs, Section 4.1). Speedup is the deployment-relevant metric — it captures not just algorithmic efficiency but also the real-world impact of draft model latency, KV cache management, and parallel verification overhead.
- Compression rate (CR): defined in Equation 1 as the average number of tokens accepted per forward pass of the target model
$\mathcal{M}^{b}$. CR is an algorithmic efficiency metric — it measures how effectively the draft model proposes tokens that the target model accepts, independent of hardware-dependent factors like memory bandwidth utilization or kernel launch overhead. CR is reported alongside speedup in Table 1 to enable diagnosis of whether a method's speedup limitations stem from poor acceptance rate (low CR) or high drafting overhead (reasonable CR but low speedup).
The paper does NOT report standalone quality metrics (perplexity, BLEU, ROUGE, task accuracy) because speculative decoding is mathematically guaranteed to be lossless. The acceptance-rejection mechanism preserves the target model's sampling distribution exactly (Leviathan et al., 2023), so output quality is identical to vanilla autoregressive decoding by construction — there is nothing to evaluate on this dimension. This is a critical point for understanding the experimental design: the paper is evaluating efficiency, not quality, and any metric that measures output tokens rather than the runtime performance of the system would be tautologically identical between Kangaroo and the baseline.
-
Baselines. The paper compares against three self-drafting speculative decoding methods:
- Medusa-1 (Cai et al., 2024): The primary baseline. Medusa trains multiple independent FFN heads attached to the target model's last decoder layer, predicting tokens at positions
$t+1, t+2, \dots, t+k$in a single forward pass. The paper compares against Medusa-1 (not Medusa-2, the tree-attention variant) because the single-sequence verification setting (verifying one candidate sequence per forward pass, not a tree of candidates) matches Kangaroo's verification approach. Medusa-1 requires 591M additional parameters for Vicuna-7B. - Lookahead (Fu et al., 2024): Uses Jacobi iteration for draft generation — parallel token prediction followed by iterative refinement until convergence. Lookahead achieves high acceptance rates (visible in Figure 1(a)) but slower draft generation due to the multi-step refinement process.
- REST (He et al., 2023): Retrieval-based speculative decoding that generates draft tokens by retrieving relevant text spans from a reference database, avoiding any draft model training.
The paper also implicitly compares against Draft & Verify (Zhang et al., 2023) in the related work discussion (Section 2), but does not benchmark it directly. The critique — that Draft & Verify's intermediate-layer early exit is too deep for effective acceleration — is analytical rather than empirical.
- Medusa-1 (Cai et al., 2024): The primary baseline. Medusa trains multiple independent FFN heads attached to the target model's last decoder layer, predicting tokens at positions
-
Generation budget / compute accounting. The relevant "generation budget" in speculative decoding is not a fixed hyperparameter but rather the ratio of drafting cost to verification cost. The paper accounts for this implicitly through the wall-time speedup metric: any draft tokens generated during the drafting phase consume GPU time, and the speedup captures whether the accepted-token throughput from verification outweighs this drafting overhead. There is no explicit FLOPs accounting or token-generation budget — the experimental unit is wall-clock time on identical hardware, which naturally accounts for all computational costs (draft model forward passes, adapter inference, parallel verification, KV cache management, attention computation) without requiring analytical FLOP models that might miss constant-factor differences in memory access patterns or kernel efficiency. The maximum draft length is
$\gamma = 6$, and the confidence threshold$\eta = 0.6$dynamically reduces this budget per verification step. -
Cross-validation / statistical protocol. The paper does not report standard cross-validation or statistical significance testing. The primary results (Table 1) report speedup and compression rate per subtask on the standard Spec-Bench evaluation protocol, which presumably uses fixed prompts for each subtask with deterministic evaluation (given the lossless guarantee, there is no output variance from sampling — the acceptance-rejection mechanism produces the identical distribution regardless, but greedy decoding is deterministic). The ablation studies (Figures 3(a), 3(b), Table 2) report "average" speedup and compression rate across all Spec-Bench subtasks, computed as a simple arithmetic mean across the six subtask-level results. The absence of error bars, confidence intervals, or multiple-seed experiments means the reported differences (e.g., 1.50× vs. 1.37× speedup in Table 2) cannot be assessed for statistical reliability — a limitation given that speedup differences of ~0.1× could plausibly arise from hardware measurement noise.
Main Quantitative Results
Overall Speedup Comparison on Spec-Bench
Table 1 reports the headline results: wall-time speedup ratio and compression rate for Kangaroo, Medusa-1, Lookahead, and REST across all six Spec-Bench subtasks, for both Vicuna-7B and Vicuna-13B.
For Vicuna-7B, Kangaroo achieves the following per-subtask speedups (extracted from Table 1):
- Translation: 1.54× speedup (Medusa-1: 1.36×, Lookahead: 1.31×, REST: 1.17×)
- Summarization: 1.53× speedup (Medusa-1: 1.48×, Lookahead: 1.51×, REST: 1.32×)
- Mathematical reasoning: 1.68× speedup (Medusa-1: 1.42×, Lookahead: 1.38×, REST: 1.23×)
- RAG: 1.49× speedup (Medusa-1: 1.37×, Lookahead: 1.36×, REST: 1.21×)
- Question answering: 1.60× speedup (Medusa-1: 1.41×, Lookahead: 1.37×, REST: 1.30×)
- Multi-turn conversation: 1.50× speedup (Medusa-1: 1.42×, Lookahead: 1.25×, REST: 1.23×)
Aggregating across subtasks, Kangaroo achieves an average speedup of approximately 1.56×, outperforming Medusa-1 (average ~1.41×), Lookahead (average ~1.36×), and REST (average ~1.24×). The maximum single-subtask speedup is 1.68× on mathematical reasoning — notably, the subtask where Figure 1(a) showed the largest CTAR gap between autoregressive drafters (Kangaroo, Lookahead) and Medusa, confirming that the autoregressive dependency advantage is most impactful on tasks requiring sequential reasoning over isolated token predictions.
The compression rate results (Table 1, right columns) contextualize these speedups:
- Kangaroo's CR ranges from 1.97 (multi-turn conversation) to 2.28 (translation), averaging approximately 2.1× — meaning the target model accepts roughly 2.1 tokens per verification step on average.
- Medusa-1's CR ranges from 1.62 (math) to 2.20 (translation), averaging approximately 1.95× — lower than Kangaroo on most subtasks, particularly math where the gap is largest (Kangaroo 2.10 vs. Medusa 1.62).
- Lookahead's CR is consistently higher than Medusa's and competitive with Kangaroo's (e.g., 2.15 on math vs. Kangaroo's 2.10), but Lookahead's speedup is lower than Kangaroo's on every subtask despite similar or higher CR — the critical demonstration that CR alone does not predict wall-time performance. Lookahead's drafting phase (iterative Jacobi refinement) consumes more time per drafted token, offsetting its acceptance-rate advantage.
For Vicuna-13B, the speedup pattern is similar but the absolute numbers are slightly higher due to the larger model's greater memory-bandwidth bottleneck (making verification relatively more expensive and draft overhead relatively cheaper):
- Kangaroo achieves speedups ranging from 1.49× (question answering) to 1.67× (translation), averaging approximately 1.57×.
- Medusa-1 ranges from 1.35× to 1.60×, averaging approximately 1.48×.
- The compression rate gap favors Kangaroo across most subtasks (e.g., math: Kangaroo 2.22 vs. Medusa 1.80), consistent with the 7B pattern.
Key observation from the subtask breakdown: The subtasks where Kangaroo shows the largest advantage over Medusa — mathematical reasoning and question answering — are those requiring sequential reasoning and multi-step inference, where autoregressive dependency among draft tokens (capturing the logical flow from one step to the next) provides the greatest benefit over Medusa's time-independent parallel heads. On translation and summarization, where the output is more formulaic and individual token predictions are more locally predictable from the source text, the gap is narrower. This is consistent with the CTAR analysis in Figure 1(a), which shows Medusa's CTAR decay is steepest on mathematical reasoning.
Difficulty-Dependent Performance via CTAR Analysis
Figure 1, which the paper uses in the introduction to motivate Kangaroo's design, contains results that are best understood as experimental evidence rather than just motivation. The two subfigures make distinct empirical claims:
Figure 1(a) plots the consistent token acceptance rate (CTAR) as a function of window size $w$ (from 1 to 6) for Kangaroo, Medusa, and Lookahead on the mathematical reasoning subtask of Spec-Bench. The data show:
- All three methods have CTAR(1) in the range of approximately 0.65–0.75 — the first drafted token is accepted roughly two-thirds to three-quarters of the time.
- As
$w$increases, CTAR decays for all methods, but the decay rate differs dramatically:- Medusa's CTAR drops steeply, reaching approximately 0.25 at
$w = 3$and near zero by$w = 5$. The time-independent heads cannot model sequential token dependencies, so the probability of correctly predicting three consecutive tokens is the product of nearly independent probabilities, producing rapid decay. - Kangaroo and Lookahead show slower decay, with CTAR remaining above 0.3 at
$w = 4$and above 0.15 at$w = 6$. The autoregressive generation captures dependencies — predicting the third token benefits from knowing the actual values (not just distributions) of tokens 1 and 2.
- Medusa's CTAR drops steeply, reaching approximately 0.25 at
- Kangaroo's CTAR curve closely tracks Lookahead's, with slight variations at different window sizes — the two autoregressive methods have similar acceptance-rate characteristics on this subtask.
Figure 1(b) compares end-to-end wall-time speedup for Kangaroo, Medusa, Lookahead, and REST across four Spec-Bench subtasks (mathematical reasoning, RAG, summarization, MT-bench). The bar chart reveals:
- Kangaroo achieves the highest speedup on all four subtasks, with the largest margin on mathematical reasoning (where the CTAR advantage over Medusa is largest) and the smallest margin on summarization (where Medusa's CR is competitive).
- Lookahead achieves speedups consistently below Medusa despite having higher compression rates (Table 1) — the Lookahead bar is shorter than Medusa's bar on three of four subtasks (Figure 1(b)). This is the empirical smoking gun for the paper's central argument: Lookahead's drafting process (iterative Jacobi refinement) is slow enough per draft token that its acceptance-rate advantage over Medusa is consumed by drafting overhead. Optimization must target wall-time speedup, not just CR.
- REST consistently achieves the lowest speedup across all subtasks, reflecting the latency of the retrieval step and the quality limitations of reference-based draft tokens.
The two subfigures together establish the joint requirement that a successful speculative decoding method needs both high CTAR (to produce useful draft tokens at distance) AND low draft latency (to ensure those draft tokens are generated cheaply enough to improve net throughput). Kangaroo satisifies both through autoregressive adapter generation (for CTAR) and shallow exit + dynamic termination (for low draft latency).
Parameter Efficiency
The abstract and Section 1 emphasize that Kangaroo achieves its speedup using 88.7% fewer additional parameters than Medusa-1: 67M for Kangaroo's adapter versus 591M for Medusa's multiple FFN heads (both measured for Vicuna-7B). This comparison is with Medusa-1 specifically, not Medusa-2 (which adds tree attention and additional complexity).
The parameter count is not abstracted — it is a direct deployment consideration. The additional parameters must reside in GPU memory alongside the target model during inference. For a 7B model (requiring ~14GB in FP16), Medusa's 591M heads add ~1.2GB (roughly 8.5% overhead), while Kangaroo's 67M adapter adds ~134MB (roughly 0.95% overhead). The smaller memory footprint matters for multi-model serving scenarios, edge deployment, and GPU memory-constrained environments.
Table 2 provides the ablation evidence that justifies this parameter efficiency. The three adapter variants compared (for Vicuna-7B, averaged across Spec-Bench):
- With FFN, separate LM Head: 173M parameters, speedup 1.36×, CR not reported but implied lower than Kangaroo's.
- Without FFN, separate LM Head: 136M parameters, speedup 1.37× — removing the FFN saves 37M parameters with negligible speedup impact, confirming the FFN provides minimal acceptance-rate benefit for its cost.
- Without FFN, shared LM Head (Kangaroo design): 67M parameters, speedup 1.50× — sharing the LM Head both reduces parameters (by ~69M) AND improves speedup (from 1.37× to 1.50×). This is a non-obvious finding: sharing the unembedding matrix not only saves parameters but actually improves performance, likely because it eliminates distributional mismatch between the draft model's logits and the target model's logits at the vocabulary level, increasing token acceptance rate.
The counterfactual is instructive: if the paper had only reported the 173M variant (matching the architectural richness of Medusa's heads), Kangaroo's speedup advantage over Medusa (1.36× vs. ~1.41× averaged across subtasks from Table 1 estimation) would not have existed. The finding that a smaller, simpler adapter outperforms a larger one is the key experimental discovery enabling Kangaroo's efficiency claims.
Ablation Studies and Robustness Checks
Early exit layer depth $l$ (Figure 3(a)): For Vicuna-7B, the paper sweeps $l$ from 0 (adapter receives raw token embeddings, no shared layers) to 6 and measures both compression rate (CR) and wall-time speedup averaged across all Spec-Bench subtasks. CR increases monotonically with $l$ — deeper exits give the adapter richer features, improving draft quality and acceptance rate. Speedup, however, is non-monotonic: it peaks at $l = 2$ (speedup ~1.50×) and declines for $l \geq 3$ despite higher CR. The decline occurs because the draft model's per-token latency grows with $l$ (each additional shared layer adds a full transformer block's computation to every sequentially drafted token), and beyond $l = 2$, the marginal acceptance-rate improvement from deeper features no longer compensates for the increased drafting time. At $l = 0$, speedup is lowest (~1.20×) because the adapter operating on raw embeddings cannot produce useful draft tokens (low CR). The paper selects $l = 2$ for Vicuna-7B and $l = 3$ for Vicuna-13B (data point stated in text but the 13B sweep is not plotted — the paper asserts the optimal shift to $l = 3$ for the larger model based on analogous experimentation).
Confidence threshold $\eta$ (Figure 3(b)): Sweeping $\eta$ from 0 to 0.9 across three maximum draft step settings ($\gamma = 4, 6, 8$), Figure 3(b) plots both CR and speedup averaged across Spec-Bench subtasks for Vicuna-7B. The key findings:
- CR decreases monotonically with increasing
$\eta$across all$\gamma$settings — more aggressive early termination means fewer draft tokens generated per verification step, directly reducing the average tokens accepted. - Speedup is concave (inverted-U shaped), peaking at
$\eta \approx 0.6$for all three$\gamma$values. At$\eta = 0$(fixed-step drafting, no early termination), speedup is sub-optimal despite maximum CR because the drafting phase wastes time on low-CTAR tokens. At$\eta \geq 0.7$, speedup declines because drafting terminates too aggressively — the system generates too few draft tokens per step, underutilizing the parallel verification capacity. - The optimal
$\eta$is "consistent across different maximum different steps" (Section 4.2): the CR-vs-speedup trade-off curves for$\gamma = 4, 6, 8$all peak near$\eta = 0.6$. This is a robustness result — the threshold does not require re-tuning when$\gamma$is adjusted. - The
$\gamma = 6, \eta = 0.6$configuration achieves the highest absolute speedup among all tested combinations, motivating the final hyperparameter selection.
Maximum draft steps $\gamma$ (Figure 3(b), implicitly): Comparing the three curves ($\gamma = 4, 6, 8$) at their respective optimal $\eta$ values, $\gamma = 6$ achieves the highest speedup. $\gamma = 8$ provides minimal additional benefit even at $\eta = 0.6$ — the dynamic termination mechanism likely stops drafting before reaching 8 tokens in most contexts, so the extra headroom is rarely used. $\gamma = 4$ is unnecessarily restrictive — on easy contexts where the draft model could successfully propose 5–6 tokens, the system truncates prematurely, leaving throughput on the table. The paper selects $\gamma = 6$ as the best operating point.
Adapter architecture components (Table 2): The three-row ablation shows a monotonic improvement in speedup (1.36× → 1.37× → 1.50×) as the adapter is simplified:
- Removing the FFN: Negligible speedup impact (+0.01×) with substantial parameter savings (173M → 136M). The FFN, representing 67% of transformer block parameters, contributes little to the adapter's ability to produce accept-worthy draft tokens — the MHA's self-attention over context is doing the heavy lifting. This is a non-obvious negative result: the natural inclination when designing a draft model bridge network would be to include FFN capacity for non-linear transformation, but the data show it is redundant.
- Sharing the LM Head: Both reduces parameters (136M → 67M) AND improves speedup (1.37× → 1.50×). This is a genuinely surprising result — sharing the output projection is typically assumed to be a parameter-saving compromise that might hurt performance (since the draft and target models operate on different hidden states), but it actually helps. The likely mechanism: the shared LM Head forces the adapter to produce hidden states in the same representational space as the target model's final layer, acting as an implicit regularization that improves the draft-target distributional alignment. A separate LM Head could learn to compensate for adapter-induced distortions in ways that increase per-token accuracy but reduce agreement with the target model's exact probability distribution — counterproductive for speculative decoding where agreement is the sole objective.
Model scale consistency (Table 1, Vicuna-7B vs. 13B): The Kangaroo design transfers from 7B to 13B without architecture changes (only the exit layer $l$ is adjusted from 2 to 3). The speedup pattern across subtasks is qualitatively similar, and Kangaroo maintains its advantage over Medusa-1 at the larger scale. This provides evidence that the design principles (extreme shallow exit, lightweight adapter, confidence-thresholded termination) are not idiosyncratic to a single model size. However, the 7B→13B scaling is modest (less than 2× parameters), and the paper does not test on 33B or 70B models where the memory-bandwidth bottleneck is more severe and the relative cost of the draft model changes — a 3-layer draft model for a 70B target is proportionally cheaper than for a 7B target (since the depth ratio is 3/80 vs. 2/32), potentially allowing deeper exits and higher acceptance rates.
Pretraining data not revisited: The adapter is trained on ShareGPT only, for 10 epochs with AdamW. The paper does not ablate the training data (e.g., comparing ShareGPT to WikiText, C4, or task-specific data), training duration (2 epochs vs. 10 vs. 20), or optimizer. These choices are inherited from Medusa and not interrogated. It is possible that domain-specific adapter training (e.g., on mathematical reasoning data for the math subtask) would further improve token acceptance rate on that subtask, but the generalist ShareGPT training aligns with the goal of a general-purpose deployment-ready system.
Dynamic termination vs. oracle: The confidence threshold $\eta = 0.6$ is a heuristic — it thresholds the draft model's top-1 softmax probability, which is a noisy and potentially miscalibrated signal. The paper does not compare against an "oracle" termination policy that would stop drafting exactly at the position where the target model would first disagree with the draft model. Such an oracle would reveal how much performance is left on the table due to imperfect confidence calibration — if the oracle achieves substantially higher speedup, it would indicate that better termination signals (e.g., a learned rejection predictor, or a small verifier network) could unlock further gains. This ablation is absent.
Critical Assessment
The experiments in this paper are well-designed to support its central engineering claim: Kangaroo achieves superior wall-time speedup compared to existing self-drafting speculative decoding methods on Spec-Bench, while using dramatically fewer additional parameters. The evidence for this claim is strong and multi-faceted: per-subtask breakdowns (Table 1), per-method CTAR analysis revealing why the speedup advantage exists (Figure 1(a)), and ablation studies confirming the necessity of each design component (Figures 3(a), 3(b), Table 2). The combination of speedup data AND diagnostic metrics (CTAR, CR) allows the reader to understand not just that Kangaroo wins, but where and why it wins.
However, the scope of the experimental validation is narrower than the paper's framing sometimes implies. The claim that Kangaroo is a general solution for self-speculative decoding is supported only for a specific model family (Vicuna-7B/13B) on a specific benchmark (Spec-Bench) with a specific hardware configuration (NVIDIA V100 GPUs). Several unexamined dimensions would strengthen or qualify the conclusions:
Single hardware platform. All speedup measurements are on NVIDIA V100 GPUs (Section 4.1). The V100 is an older architecture (2017) with different memory bandwidth characteristics than newer GPUs (A100, H100) or consumer hardware (RTX series). The memory-bandwidth bottleneck that speculative decoding exploits is architecture-dependent — on an H100 with dramatically higher bandwidth, the relative cost of the draft model's sequential forward passes might change, potentially altering the optimal exit layer depth or even the speedup ranking of methods. The paper does not discuss this, and the speedup numbers should not be assumed to transfer directly to other hardware without validation.
Evaluation set size and composition. Spec-Bench's six subtasks cover diverse generation scenarios, but the paper does not report the number of prompts per subtask, the total number of tokens generated during evaluation, or any measure of variance across prompts within a subtask. Speedup measurements on small prompt sets can be noisy due to hardware-level factors (GPU clock speed variation, memory controller contention, operating system scheduling). Without reporting standard deviations or running multiple evaluation seeds, it is impossible to assess whether the reported speedup differences (e.g., 1.50× vs. 1.41× average across subtasks) are statistically distinguishable from measurement noise. The Spec-Bench paper (Xia et al., 2024) may specify standard evaluation protocols, but this paper does not reference them.
No direct Draft & Verify comparison. The paper critiques Draft & Verify (Zhang et al., 2023) in Section 2 for having "exceptionally high" draft model latency, but never benchmarks it. This is a missed opportunity — Draft & Verify is the closest conceptual predecessor (early exiting + verification), and a head-to-head comparison would directly demonstrate that Kangaroo's shallower exit + adapter design solves the problem Draft & Verify identified but couldn't resolve. The absence of this baseline weakens the claim that Kangaroo's specific depth choice (layer 2) is necessarily superior to Draft & Verify's intermediate-layer exits.
Training cost not empirically quantified. The paper claims "low-cost" training for the adapter (Section 1, contributions bullet 2), but reports only the number of epochs (10) and dataset (ShareGPT) — not the wall-clock training time, the number of GPUs used, the total FLOPs, or a comparison to the cost of training a standalone draft model (e.g., the LLaMA-68M cited in Section 1). "88.7% fewer parameters than Medusa" quantifies the deployment overhead, not the training cost. A practitioner deciding between Kangaroo and Medusa needs to weigh not just the inference-time speedup and memory overhead, but also the upfront training investment. The paper does not provide the data for this comparison.
Confidence threshold generalizability. The optimal $\eta = 0.6$ is selected based on Figure 3(b) ablation and appears robust across $\gamma$ values, but this is evaluated only on Vicuna-7B and only aggregated across all Spec-Bench subtasks. The optimal per-subtask threshold might differ — on mathematical reasoning where CTAR decays rapidly (Figure 1(a)), a more aggressive threshold (e.g., 0.7–0.8) might be better, while on translation where token predictions are more formulaic, a less aggressive threshold (e.g., 0.5) might be optimal. The paper's claim that the threshold is "consistent" could be an artifact of averaging across diverse subtasks with opposing optimal values, masking per-subtask variation. Subplot-breakdowns of Figure 3(b) by subtask are not provided.
Adapter training does not use on-policy data. The adapter is trained on ShareGPT using cross-entropy against the target model's distribution (Equation 3), with training data generated by the target model processing real conversational text. However, during speculative decoding inference, the draft model receives its own previously generated draft tokens as context — these are off-policy relative to the training distribution (which used real tokens as context). This distribution shift could cause the draft model to perform worse during autoregressive drafting than during training, because the context it conditions on (a sequence of its own potentially suboptimal predictions) differs from the context it was trained on (a sequence of real tokens). The paper does not investigate this — there is no comparison between adapter performance on in-distribution context (real next-token prediction) vs. autoregressive context (conditioning on its own draft tokens). If this gap is large, it could partially explain why fixed-step drafting underperforms expectations based on the adapter's standalone token prediction accuracy. An on-policy training approach (generating draft sequences and training to maximize acceptance rate on those sequences) might close this gap, but is not explored.
The 13B results are under-reported. The Vicuna-13B experiments appear only in Table 1 (speedup and CR per subtask) and a single data point in the text (exit layer $l = 3$ mentioned but not plotted in Figure 3(a)). The ablation studies (Figures 3(a), 3(b), Table 2) are Vicuna-7B only. This asymmetry makes it impossible to verify whether the hyperparameter optimization that produced the 7B speedup numbers (exit layer, threshold, adapter architecture) transfers to 13B or whether independent tuning would yield different choices. The 13B speedup numbers in Table 1 should be interpreted as applying the 7B-optimized design to the larger model, not as representing a 13B-optimized Kangaroo.
Losslessness is assumed, not verified. The paper states that Kangaroo is lossless (title, abstract, Section 3) because it uses standard speculative sampling verification. This is mathematically correct given the Leviathan et al. (2023) acceptance-rejection mechanism. However, the paper does not empirically verify losslessness — for example, by comparing the output distribution (or any downstream metric like perplexity, BLEU, or task accuracy) of Kangaroo against vanilla autoregressive decoding on a sample of prompts. In practice, numerical issues (floating-point differences in softmax computation between the draft path and verification path) or implementation bugs could introduce subtle distributional discrepancies. A brief empirical confirmation would strengthen the losslessness claim, particularly since the shared LM Head and the adapter introduce a different computation graph than standard two-model speculative decoding.
The "double early exiting" framing is somewhat overstated. The paper's title and framing emphasize the "double early exiting" — the shallow layer exit AND the confidence-thresholded termination. The second "exit" is not an architectural exit (no additional exit layers, no auxiliary prediction heads at intermediate positions) but rather a simple runtime heuristic: stop when softmax confidence < 0.6. This is a useful practical trick, but calling it "early exiting" links it conceptually to the substantial literature on confidence-based early termination in neural networks (Schuster et al., 2022; Varshney et al., 2023), where the exit mechanism involves architectural modifications (prediction heads at intermediate layers) and learned confidence calibration. The paper's dynamic termination is simpler and more heuristic; whether it deserves the "early exiting" label or is better described as "confidence-thresholded drafting" is debatable. The experiments do not compare against alternative dynamic termination strategies (e.g., a learned stopping policy, or using the target model's own confidence on the first draft token to decide how many more to draft) that might more fully realize the "early exiting" analogy.
Despite these limitations, the central experimental finding — that a minimal adapter on an extremely shallow exit can achieve competitive or superior speedup to Medusa's elaborate multi-head architecture, and that dynamic confidence-thresholded termination further improves performance — is well-supported within the tested scope. The experiments are internally consistent, the ablation studies cleanly isolate the effect of each design choice, and the diagnostic use of CTAR to explain the performance landscape is genuinely insightful. The paper identifies a previously overlooked design point in the self-drafting space and validates it empirically; the open questions concern how broadly that design point generalizes, not whether it works in the tested configuration.
6. Limitations and Trade-offs
6.1 Hardware and Model-Scale Generalization Is Unverified
The assumption or constraint. All speedup measurements are conducted on NVIDIA V100 GPUs with a single model family (Vicuna-7B and 13B, derived from LLaMA). The paper does not claim these results should transfer, but also provides no analysis of how the speedup ranking or optimal hyperparameters might shift across hardware generations or model architectures. The V100 (2017, 900 GB/s memory bandwidth) has substantially different memory-bandwidth characteristics from modern datacenter GPUs (A100: 2 TB/s, H100: 3.35 TB/s) or consumer cards — and the entire speculative decoding mechanism fundamentally trades off memory-bandwidth utilization against extra computation. On a higher-bandwidth GPU, the relative cost of the sequential draft model forward passes increases relative to the parallel verification pass, potentially shifting the optimal exit layer depth, the speedup advantage over Medusa, or even the basic decision of whether autoregressive drafting beats parallel heads.
The consequence. A practitioner deploying Kangaroo on an A100 or H100 cluster cannot reliably predict their speedup from these numbers. The 1.68× headline figure might be substantially higher (if the speedup bottleneck is memory bandwidth and the V100 under-represents the target model's cost relative to the draft model) or lower (if the draft model's sequential latency becomes proportionally more expensive when the target model's memory reads are faster). The scaling from 7B to 13B — a less-than-2× increase in parameters — provides minimal evidence about behavior at 33B, 70B, or larger scales where the ratio of target model depth to exit layer depth grows, potentially enabling deeper exits with higher acceptance rates.
What evidence exists in the paper. The Vicuna-7B and 13B results in Table 1 show qualitatively similar speedup patterns, but this is a single model family with identical architecture (LLaMA-based transformer). The 13B experiments apply hyperparameters optimized on 7B (exit layer $l = 3$ versus $l = 2$ is the only adjustment, mentioned in Section 4.2 but not plotted in Figure 3(a), which shows only 7B sweeps). The ablation studies (Figures 3(a), 3(b), Table 2) are exclusively on Vicuna-7B. No other model architecture (Mistral, Qwen, Pangu-π cited in the references) is tested. No other GPU architecture is tested.
Mitigation status. The paper does not address this limitation explicitly. The abstract and conclusion state speedup claims (up to 1.68×) without qualifying their hardware or model specificity. The code is released (GitHub link in abstract), enabling practitioners to benchmark on their own hardware and models, but the paper provides no guidance on how to adapt hyperparameters to different configurations.
6.2 Dynamic Confidence Threshold Is a Heuristic Without Calibration Guarantees
The assumption or constraint. The second early exit mechanism uses a raw softmax top-1 probability thresholded at $\eta = 0.6$ (Equation 4) as a proxy for whether the target model will accept the draft token. This assumes that the draft model's confidence is calibrated to acceptance probability — that low confidence predicts target-model disagreement. However, neural network softmax outputs are known to be poorly calibrated in general (tending toward overconfidence), and the adapter is trained with cross-entropy loss against the target model's distribution (Equation 3), which optimizes for distributional matching but not for calibration of the acceptance-rejection decision boundary.
The consequence. There are two failure modes. If the draft model is overconfident (predicting high softmax scores on tokens the target model will reject), the threshold fails to terminate drafting when it should — the draft model continues generating tokens at positions where acceptance probability is actually low, wasting computation exactly as fixed-step drafting does. If the draft model is underconfident (assigning low softmax scores to tokens the target model would accept), the threshold terminates drafting prematurely, losing the opportunity to draft and verify additional tokens. In either case, the speedup is suboptimal relative to a properly calibrated termination signal. The paper's ablation in Figure 3(b) selects $\eta = 0.6$ empirically, but this is a global threshold applied uniformly to all contexts and all tokens — it cannot adapt to the draft model's varying calibration quality across different types of text (mathematical reasoning tokens may have fundamentally different confidence distributions than conversational tokens).
What evidence exists in the paper. Figure 3(b) shows that speedup is sensitive to $\eta$ — varying the threshold from 0 to 0.9 changes speedup by approximately 0.15× to 0.20× across the tested configurations. The optimal $\eta = 0.6$ produces the peak, but the paper provides no analysis of whether this threshold is optimal per-subtask or simply optimal on average. The CTAR analysis in Figure 1(a) shows that acceptance probability decays with distance, but the relationship between the draft model's softmax confidence at position $w$ and the actual CTAR($w$) is never plotted — the paper assumes confidence predicts acceptance but never validates this correlation. There is no comparison against an oracle termination policy (stop at the position where the target model first disagrees) that would reveal how much speedup is lost due to imperfect confidence calibration.
Mitigation status. The paper does not address calibration directly. The consistency of the optimal $\eta \approx 0.6$ across $\gamma = 4, 6, 8$ (noted in Section 4.2) provides some evidence that the threshold is not hypersensitive, but this is tested only on Vicuna-7B aggregated across all Spec-Bench subtasks. The paper does not suggest calibration improvements (temperature scaling, learned stopping policies, or a dedicated acceptance predictor) as future work.
6.3 Training Cost Claims Are Asserted but Not Measured
The assumption or constraint. The paper repeatedly characterizes Kangaroo's adapter training as "low-cost" (Section 1 contributions, Section 3.2 para 2, and abstract: "a low-cost approach to train a lightweight small model") and uses the 88.7% parameter reduction versus Medusa (67M versus 591M) as the quantitative evidence. However, parameter count measures deployment overhead (GPU memory), not training cost. Training cost is a function of dataset size, number of training steps, optimizer state memory, and total FLOPs — none of which are reported.
The consequence. A practitioner deciding between Kangaroo and an external draft model approach (e.g., training a LLaMA-68M from scratch or using DistillSpec-style distillation) needs to compare the total resource investment required to reach deployment readiness, not just the inference-time memory overhead. Kangaroo trains an adapter for 10 epochs on ShareGPT with AdamW — but how many GPU-hours is this? How does it compare to the cost of training Medusa's heads (also 10 epochs on ShareGPT, but with 591M parameters versus 67M)? If the per-epoch training time for the adapter is proportional to its parameter count, Kangaroo's training would be roughly 8.8× cheaper than Medusa's — but the adapter training involves forward passes through the frozen target model's shallow layers (to produce the hidden states at layer $l$ that the adapter receives as input), which could dominate the training cost and reduce or eliminate the parameter-count advantage. The paper provides no data to resolve this.
What evidence exists in the paper. The only training-related metrics reported are: dataset (ShareGPT), optimizer (AdamW), epochs (10), and the loss function (Equation 3). There is no wall-clock training time, no total FLOPs estimate, no GPU count, and no comparison to Medusa's training cost (which would be the natural baseline for the "low-cost" claim). The adapter's 67M parameter count is compared to Medusa's 591M, but this is a deployment metric, not a training metric.
Mitigation status. The paper does not acknowledge this gap. The "low-cost" framing in the abstract and contributions is presented as an established finding rather than an asserted but unmeasured property. The GitHub code release may enable practitioners to measure training cost themselves, but the paper provides no reference numbers.
6.4 Off-Policy Drafting Gap Is Not Diagnosed
The assumption or constraint. The adapter is trained using cross-entropy loss against the target model's output distribution on real text tokens (Equation 3, trained on ShareGPT). The training context for each position is the ground-truth prefix — real conversational tokens. However, during inference, the draft model generates tokens autoregressively — it conditions on its own previously generated draft tokens, which may differ from the ground-truth tokens that would appear at those positions. This is an off-policy setting: the draft model is trained to predict the next token given correct context, but at inference it must predict given potentially incorrect context (its own earlier draft tokens).
The consequence. The draft model's predictions during autoregressive inference may be systematically worse than its predictions on the training distribution. If an early draft token is suboptimal (even if accepted by the target model — speculative decoding accepts tokens probabilistically, so non-argmax tokens can be accepted), all subsequent draft tokens are conditioned on this suboptimal context, potentially degrading their quality below what the adapter's standalone accuracy would suggest. This effect compounds with draft length: the third draft token is conditioned on two potentially noisy tokens, the fourth on three, etc. This off-policy degradation could partially explain why CTAR decays with distance (Figure 1(a)) beyond what would be expected from mere task difficulty — the draft model might maintain higher acceptance rates if it were always conditioned on ground-truth prefixes. It also means that the adapter's standalone token-prediction accuracy (measurable during training) may overestimate its practical utility during speculative decoding.
What evidence exists in the paper. None directly. The CTAR measurements in Figure 1(a) are generated during actual speculative decoding (on-policy evaluation), but these are never compared to the adapter's accuracy when conditioned on ground-truth prefixes (off-policy evaluation). The paper does not discuss the off-policy issue or acknowledge it as a potential limitation. The cross-entropy training objective (Equation 3) is purely off-policy — it minimizes KL divergence between the draft and target distributions given correct context — with no on-policy fine-tuning or data augmentation using draft-model-generated prefixes.
Mitigation status. Not addressed. On-policy training (generating draft sequences with the current adapter, scoring acceptance, and fine-tuning to maximize acceptance rate on those sequences) would directly address this gap but is not explored. The paper does not mention this as future work. The negative result with the ReST-style revision model in the companion paper (Appendix K, Figure 16) — where on-policy training degraded rather than improved performance — suggests that naive on-policy approaches can backfire, but the paper draws no connection to this potential limitation.
6.5 No Empirical Verification of Losslessness
The assumption or constraint. The paper's title, abstract, and Section 3 all state that Kangaroo is lossless — it preserves the target model's output distribution exactly. This claim rests entirely on the mathematical guarantee of the speculative sampling acceptance-rejection mechanism (Leviathan et al., 2023; Chen et al., 2023). The paper does not empirically verify losslessness by comparing outputs from Kangaroo against vanilla autoregressive decoding.
The consequence. In practice, losslessness can fail for reasons beyond the mathematical proof. The adapter introduces a different computation graph than standard two-model speculative decoding: the draft model shares layers and the LM Head with the target model, and the verification phase processes hidden states produced by the draft model's forward pass through the remaining target model layers. If there are any numerical inconsistencies — floating-point nondeterminism, subtle differences in how attention masks are applied during parallel verification versus sequential drafting, KV cache management bugs — the output distribution could deviate from the target model's. These are implementation concerns, not theoretical ones, but the paper provides no evidence that the implementation is correct.
Additionally, the dynamic termination mechanism introduces a decision (when to stop drafting) that affects which tokens get verified but should not affect the final output distribution — the acceptance-rejection mechanism guarantees losslessness regardless of how many draft tokens are proposed. However, if the termination decision inadvertently influences the verification process (e.g., through subtle interactions with KV cache state or attention masking), losslessness could be compromised.
What evidence exists in the paper. None. There is no perplexity comparison, no output-distribution KL divergence measurement, no BLEU/ROUGE/accuracy comparison between Kangaroo and vanilla decoding on any benchmark. The paper relies entirely on the theoretical guarantee.
Mitigation status. The paper does not acknowledge the absence of empirical losslessness verification. The standard for speculative decoding papers varies — some include brief empirical confirmation (e.g., measuring perplexity of generated text), others rely on the proof. Given that Kangaroo's architecture differs from standard two-model speculative decoding (shared layers, shared LM Head, adapter in the draft path), a brief empirical check would substantively strengthen the losslessness claim.
6.6 Dynamic Termination Is Evaluated Only on Aggregated Metrics, Masking Per-Subtask Variation
The assumption or constraint. The confidence threshold $\eta = 0.6$ and the dynamic termination mechanism are evaluated only through aggregated Spec-Bench metrics — average speedup and compression rate across all six subtasks (Figure 3(b), Table 2). The individual subtask results in Table 1 report speedup and CR for the final Kangaroo configuration but do not show how the optimal $\eta$ varies per subtask.
The consequence. The paper's claim that the optimal threshold is "consistent across different maximum different steps" (Section 4.2) refers to consistency across $\gamma$ values within the aggregated evaluation, not consistency across subtasks. The six Spec-Bench subtasks span fundamentally different text generation challenges: mathematical reasoning requires precise symbolic manipulation, translation involves cross-lingual alignment, summarization requires content compression, RAG involves factual integration. The draft model's confidence calibration — and therefore the optimal termination threshold — likely differs across these regimes. On mathematical reasoning, where CTAR decays steeply (Figure 1(a)), an aggressive threshold (higher $\eta$) might prevent wasted drafting on low-CTAR distant positions. On translation, where token predictions may be more formulaic, a lower threshold might capture more speculative benefit. The paper's aggregated evaluation cannot detect whether the single $\eta = 0.6$ is genuinely near-optimal for all subtasks or merely the best compromise on average, leaving per-subtask performance on the table.
What evidence exists in the paper. Table 1 shows that Kangaroo's speedup advantage over Medusa varies substantially by subtask: 1.68× vs. 1.42× on math (0.26× gap) versus 1.53× vs. 1.48× on summarization (0.05× gap). This variation could arise from differences in CTAR, draft model latency, or the effectiveness of the confidence threshold — but the paper does not decompose these factors. Figure 3(b) would need to be replicated per subtask to assess whether $\eta = 0.6$ is uniformly optimal or whether per-subtask tuning could improve speedup further.
Mitigation status. The paper does not discuss per-subtask threshold variation or provide the per-subtask ablation data. The Spec-Bench evaluation protocol (six diverse subtasks) is well-chosen for stress-testing speculative decoding, but the ablation analysis does not exploit this diversity — it treats all subtasks as interchangeable by averaging their metrics. Future work on adaptive, context-dependent thresholds (which the confidence signal itself could inform) is not suggested.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper makes two contributions that shift how the speculative decoding field thinks about draft model design: an architectural reframing (extreme shallow-exit self-drafting with a minimal adapter breaks the assumed trade-off between draft quality and draft speed) and a methodological reframing (consistent token acceptance rate reveals that fixed-budget drafting is inherently wasteful in ways compression rate masks). Neither is a paradigm shift — speculative decoding's verification mechanism and losslessness guarantee remain unchanged — but together they open a design space that prior work had prematurely foreclosed.
The architectural reframing. Prior to Kangaroo, the self-drafting landscape was polarized between two unsatisfactory extremes: parallel head-based methods (Medusa) that achieve extremely low draft latency but poor long-range token acceptance because heads are time-independent, and early-exit methods (Draft & Verify) that achieve high acceptance rate but crippling draft latency because the exit is too deep. The field's implicit assumption was that autoregressive self-drafting — which captures sequential token dependencies and thus achieves higher CTAR at speculative distances 2–6 — was fundamentally incompatible with low draft latency because autoregressive generation requires many sequential forward passes through the draft model, and a draft model good enough to produce useful tokens must be large enough to be slow.
Kangaroo falsifies this assumption through a specific architectural combination: exit after layer 2 of 32 (for Vicuna-7B), connect a single MHA block, share the LM Head. The paper's ablation (Table 2) demonstrates that the adapter can be stripped to its absolute minimum — no FFN, 67M parameters — and still produce draft tokens that the target model accepts at rates competitive with Lookahead's iterative Jacobi refinement. The finding that removing the FFN and sharing the LM Head improves speedup (1.37× → 1.50×) is the empirical linchpin: it proves that draft model quality, in the speculative decoding sense, is not about representational capacity but about distributional alignment with the target model, and that a tiny learned correction on top of the shallowest possible features achieves sufficient alignment.
This finding redirects research attention away from what was becoming the field's default trajectory: ever-more-elaborate prediction head architectures on the target model's final layer (Hydra's sequentially-dependent heads, Eagle's feature-uncertainty modeling, Recurrent Drafter's recurrence). Kangaroo demonstrates that these heads are solving the wrong problem — they are compensating for the fundamental limitation of time-independent prediction from a single hidden state, when a simpler approach (autoregressive generation from a shallower-but-still-autoregressive draft model) achieves higher acceptance rates with fewer parameters. The parameter advantage over Medusa (67M vs. 591M for the heads alone, not counting Medusa's tree attention overhead) is not an incremental improvement — it represents a qualitatively different scaling regime for deployment overhead.
The methodological reframing. The introduction of consistent token acceptance rate (CTAR) as an evaluation metric is, in one sense, a minor technical contribution — it is a simple modification of the compression rate formula. But its diagnostic power transforms how speculative decoding methods should be compared. Prior work reported compression rate (average tokens accepted per verification step) as the primary algorithmic efficiency metric, implicitly treating all draft positions as equally valuable. CTAR decomposes this single number into a curve that reveals where a method's acceptance rate degrades: a steep CTAR decay indicates that distant draft tokens contribute little to throughput regardless of the average compression rate, while a flat CTAR curve indicates room to extend the draft horizon.
The paper uses CTAR diagnostically in Figure 1(a) to explain why Medusa's compression rate advantage over Lookahead on some subtasks doesn't translate to speedup advantage — Medusa's CTAR decays rapidly at positions 2–3, meaning its higher CR is driven by the first one or two tokens where all methods perform well, while its distant predictions are rarely accepted. This diagnostic tells practitioners where to invest engineering effort: if CTAR(1) is low, improve draft-target distributional matching; if CTAR(3–4) is low, improve sequential dependency modeling or implement dynamic termination.
The broader implication is that speculative decoding evaluation should move beyond scalar metrics to curve-based diagnostics. The field's current practice of reporting a single speedup number and a single compression rate number obscures the performance heterogeneity that Kangaroo's dynamic termination exploits. A method that achieves 1.5× speedup with CTAR decaying gracefully to position 6 is fundamentally different from one that achieves 1.5× speedup with CTAR near zero at position 4 — the former can be improved by extending the draft horizon, while the latter is already at its ceiling unless draft quality improves. CTAR gives researchers the tool to make this distinction.
Reconciling prior contradictions. The paper resolves an apparent tension in the self-drafting literature between methods that optimize for acceptance rate (Lookahead, Draft & Verify) and methods that optimize for draft speed (Medusa). The resolution is not that one approach is correct — it's that optimizing either metric in isolation produces suboptimal end-to-end performance, and that the field's reliance on compression rate as a proxy for algorithmic quality systematically favored acceptance-rate-oriented methods while masking their drafting overhead. Figure 1(b) captures this concisely: Lookahead has higher CR than Medusa on mathematical reasoning but lower speedup. Kangaroo's explicit optimization for wall-time speedup rather than CR — embodied in the choice of exit layer $l$ (Figure 3(a), where speedup peaks at $l=2$ despite CR continuing to increase with $l$) and the confidence threshold $\eta$ (Figure 3(b), where speedup peaks at $\eta=0.6$ despite CR decreasing monotonically with $\eta$) — provides a template for how the field should evaluate speculative decoding methods going forward.
Research directions that become more attractive. The paper's demonstration that extreme shallow exits work for self-drafting opens a design space that prior work had not explored: the relationship between exit depth, adapter capacity, and token acceptance rate on different types of text. The finding that a single MHA block suffices as the adapter for a 7B model raises the question of how adapter capacity should scale with target model size — does a 70B model need a proportionally larger adapter, or does the representational quality of the shallow layers improve with scale such that the same minimal adapter works? The paper provides no scaling evidence beyond 7B→13B, and the 13B experiments lack adapter architecture ablations.
Research directions that become less attractive. The paper's results weaken the case for investing in elaborate parallel prediction head architectures on the final layer. Medusa's 591M additional parameters — a substantial engineering effort spanning multiple FFN heads, tree attention, and specialized training procedures — are outperformed by a 67M adapter on a layer-2 exit with simple autoregressive generation. This does not mean parallel heads are obsolete (they have latency advantages that may dominate in certain deployment scenarios, particularly where draft model memory overhead is the binding constraint), but it does suggest that the marginal return on architectural complexity in head design is low compared to the return on better sequential dependency modeling through autoregressive drafting.
The finding that fixed-step drafting with $\gamma=8$ provides negligible benefit over $\gamma=6$ with dynamic termination (Figure 3(b)) also weakens the case for methods that invest in extending the draft horizon through more elaborate prediction mechanisms. The bottleneck is not the number of draft tokens that can be generated, but the number that should be generated given the decaying CTAR. Dynamic termination — even with a simple heuristic threshold — captures most of the benefit that a perfectly chosen fixed $\gamma$ would provide, while being robust to context variation.
Follow-Up Research This Work Enables
1. Scaling Kangaroo to 33B–70B models and characterizing how optimal exit depth scales with target model size. The paper establishes that $l=2$ works for Vicuna-7B and $l=3$ for Vicuna-13B, but provides no theoretical or empirical guidance for larger models. The key open question is whether the optimal exit layer (measured as a fraction of total depth) increases, decreases, or stays constant as the target model grows. Intuition cuts both ways: larger models may have richer shallow-layer representations (enabling even shallower exits), or they may require deeper exits because the representational gap between layer $l$ and layer $L$ grows with model depth (requiring more capacity to bridge). A strong follow-up would measure the speedup-vs-$l$ curve (analogous to Figure 3(a)) for LLaMA-33B, LLaMA-70B, and ideally a single model family at multiple scales to extract a scaling law. The prediction: if shallow-layer representational quality improves with scale, optimal $l/L$ should decrease, making Kangaroo more efficient at larger scales. If the gap grows, adapter capacity might need to increase with model size, partially offsetting the parameter-efficiency advantage over Medusa.
2. Training a dedicated acceptance predictor to replace the heuristic confidence threshold. The paper's dynamic termination uses the draft model's raw softmax top-1 probability thresholded at $\eta=0.6$ (Equation 4). This is a heuristic that conflates two distinct signals: the draft model's uncertainty about what the correct token is, and the draft model's uncertainty about what the target model would predict. These can diverge — the target model might confidently produce a token that the draft model considers a low-probability alternative, or the draft and target might agree on a low-confidence token (both are uncertain in the same way). A strong follow-up would train a small binary classifier (acceptance predictor) that takes as input the draft model's full probability distribution at a position, the adapter's hidden state, and optionally the target model's hidden state at the exit layer, and predicts whether the target model will accept the draft token. This classifier could be trained on the same ShareGPT data used for the adapter, by comparing draft model predictions against target model predictions at each position and labeling agreement/disagreement. The evaluation would compare speedup with the learned predictor against the heuristic $\eta=0.6$ threshold, and against an oracle that knows the true acceptance outcome. The hypothesis: a learned predictor can capture calibration patterns that a scalar threshold misses (e.g., the relationship between the shape of the entire probability distribution and acceptance probability), closing some of the gap to oracle performance.
3. On-policy adapter training to close the off-policy drafting gap. The limitation identified in Section 6.4 — that the adapter is trained to predict tokens given correct context but at inference must predict given its own potentially incorrect draft tokens — is addressable through on-policy fine-tuning. A concrete experiment: after training the adapter with the standard cross-entropy objective (Equation 3), run speculative decoding on a held-out portion of the training data to generate draft sequences and record which draft tokens are accepted by the target model. Then fine-tune the adapter (or a small correction head) to maximize the probability of accepted tokens given the draft-model-generated prefix (the on-policy context). The evaluation metric is the change in CTAR at positions 2–6, comparing the on-policy fine-tuned adapter against the off-policy-only baseline. The ReST negative result cited in the paper (from a different domain — revision model training — but conceptually related) suggests caution: naive on-policy training can amplify spurious correlations. A careful ablation comparing different on-policy training objectives (maximizing acceptance rate directly, distillation from the target model on draft-generated prefixes, or a hybrid) would clarify whether the off-policy gap is practically significant and whether it can be closed without destabilizing the adapter.
4. Per-subtask adaptive thresholding using the draft model's confidence trajectory, not a static $\eta$. The paper selects $\eta=0.6$ as a global threshold aggregated across all Spec-Bench subtasks, but the six subtasks differ substantially in the difficulty of token prediction. On mathematical reasoning, CTAR decays steeply (Figure 1(a)), suggesting an aggressive threshold (higher $\eta$) might be optimal to avoid wasted drafting on low-CTAR distant tokens. On translation, where the output is more formulaic and CTAR likely decays more slowly, a lower threshold might capture more speculative benefit. A follow-up study would measure the optimal per-subtask $\eta$ for each of the six Spec-Bench subtasks independently, and then design an adaptive mechanism that selects the threshold based on features of the input prompt or the draft model's early confidence signals. A simple approach: on the first draft token of each round, use the draft model's confidence to select among a small set of per-subtask-optimized thresholds — high first-token confidence suggests an "easy" context suitable for aggressive drafting (low $\eta$), while low first-token confidence suggests a "hard" context where early termination (high $\eta$) is prudent. The evaluation compares the adaptive threshold against the static $\eta=0.6$ and against per-subtask oracle thresholds.
5. Cross-architecture validation on non-LLaMA model families to test whether Kangaroo's design principles transfer. All experiments use Vicuna models derived from LLaMA. The LLaMA architecture has specific properties — pre-normalization with RMSNorm, SwiGLU activation in FFN layers, rotary position embeddings — that may affect the quality of shallow-layer representations and thus the viability of extreme early exits. A necessary stress test is to replicate the core finding (that a layer-2 exit with a minimal MHA adapter outperforms Medusa-style heads) on architecturally distinct model families: Mistral (sliding window attention, grouped query attention), Qwen, or Gemma. The key measurement is whether the optimal exit depth changes with architecture, and whether the MHA-only adapter design (no FFN, shared LM Head) remains optimal or whether architectural differences require adapter modifications. A negative result — e.g., Kangaroo underperforms Medusa on a non-LLaMA architecture — would bound the generality of the paper's claims and identify architectural features that enable or prevent extreme shallow-exit self-drafting.
6. Combining Kangaroo's autoregressive adapter with tree-structured verification (SpecInfer-style) to push speedup beyond the 1.68× single-sequence ceiling. Kangaroo uses single-sequence verification — one candidate sequence of draft tokens per round. Medusa-2 and SpecInfer demonstrate that verifying a tree of candidate sequences (multiple alternative tokens at each position, with the verifier processing the tree in parallel) can increase the expected number of accepted tokens per forward pass beyond what a single sequence achieves. The paper explicitly uses Medusa-1 (single-sequence) as its baseline and does not compare against tree-based methods. A natural extension is to generate multiple candidate tokens at each drafting position using the adapter's autoregressive generation (e.g., top-2 or top-3 tokens from the adapter's distribution) and construct a verification tree, then benchmark against Medusa-2 and Eagle (which already use tree verification). The hypothesis: Kangaroo's higher CTAR at positions 2–4 (due to autoregressive dependency) combined with tree verification could yield multiplicative gains — more accepted tokens per position AND more positions verified per forward pass. The risk is that the adapter's sequential generation makes producing a tree expensive (multiple forward passes per position for alternative tokens), potentially offsetting the verification benefit.
Practical Applications and Downstream Use Cases
1. Cost-efficient serving of chatbot models in production with minimal infrastructure changes. The most immediate deployment scenario is accelerating open-weight chatbot models (Vicuna, LLaMA-chat, Mistral-instruct) in production serving environments. Kangaroo's key practical advantage over external draft model approaches is deployment simplicity: the adapter is a single 67M-parameter module (for 7B target models) that shares the target model's shallow layers and LM Head, requiring minimal additional GPU memory (~134MB in FP16, approximately 0.95% of the 7B model's ~14GB footprint) and no separate model server, KV cache management, or inter-model communication. For an organization already serving Vicuna-7B with vanilla autoregressive decoding at, say, 10 tokens per second on V100 GPUs, deploying Kangaroo requires swapping the inference code (the paper's implementation is open-sourced), loading the adapter weights, and immediately achieving approximately 15–16 tokens per second (1.5–1.6× speedup from Table 1 averages) with identical output quality and no additional GPU provisioning. The 1.68× peak on mathematical reasoning is particularly relevant for educational or coding-assistant applications where users submit math or logic queries. The speedup translates directly to reduced latency per request (users wait 40% less time for responses) or increased throughput per GPU (serve 50–60% more concurrent users on the same hardware).
2. On-device or edge deployment of medium-sized LLMs where draft model memory overhead is the binding constraint. For LLM deployment on consumer hardware (laptops, phones, edge servers), GPU memory is often the primary bottleneck — a 7B model in 4-bit quantization requires roughly 4–5GB, leaving limited headroom for additional components. Medusa's 591M additional parameters (even quantized, ~300MB in 4-bit) can push a deployment over the memory threshold of a consumer GPU (e.g., 6GB or 8GB cards). Kangaroo's 67M adapter (even in FP16, ~134MB; in 4-bit, ~30MB) is negligible, making speculative decoding viable in memory-constrained environments where Medusa would not fit. The speedup benefit (1.5–1.6× on average across Spec-Bench subtasks) is meaningful for interactive applications (code completion, writing assistance, local chatbots) where perceived latency directly affects user experience. The dynamic termination mechanism also provides an implicit quality-of-service benefit: on easy continuations (common in boilerplate code or formulaic text), the system drafts to the full 6 tokens and achieves high speedup; on harder tokens requiring more computation, it gracefully degrades to shorter draft sequences without the abrupt performance cliffs that fixed-$\gamma$ methods exhibit.
3. Batch inference pipelines for synthetic data generation or evaluation. For organizations running large-scale batch inference — generating training data via LLM distillation, evaluating models on benchmark suites, or processing document collections — throughput (tokens per second per GPU) is the dominant cost driver, and latency variance across individual prompts is tolerable. Kangaroo's speedup applies uniformly regardless of batch size (since speculative decoding accelerates the autoregressive decoding bottleneck, which persists at any batch size on memory-bandwidth-bound hardware), and the dynamic termination adapts per-prompt without manual tuning. A batch inference job that would take 100 GPU-hours with vanilla decoding on Vicuna-7B would take approximately 63–67 GPU-hours with Kangaroo (based on the average 1.5–1.6× speedup from Table 1). The cost savings are directly proportional and do not require changes to the model, the prompts, or the evaluation pipeline — only a swap of the inference backend. The losslessness guarantee (output distribution is identical) means that downstream metrics (perplexity, accuracy, BLEU) calculated on the generated text are unaffected, eliminating the need to revalidate model quality after switching to the speculative decoding backend.
4. Rapid deployment of self-drafting for fine-tuned or domain-adapted models without retraining the draft component. A common production pattern is to fine-tune a base LLM (e.g., LLaMA-7B) on domain-specific data (medical, legal, code). External draft model approaches require training a separate draft model for each fine-tuned variant, multiplying the training cost. Medusa requires retraining the prediction heads on the fine-tuned model's output distribution. Kangaroo's adapter training is presumably lightweight enough (10 epochs on ShareGPT) that retraining on domain-specific data would be fast, but the paper provides no evidence on this. More interestingly, the paper's finding that the adapter is trained by distilling the target model's distribution (Equation 3) — not by training on ground-truth tokens — suggests a potential transfer property: if the fine-tuned model's shallow-layer representations are similar to the base model's (which they likely are, since fine-tuning primarily affects later layers), the adapter trained on the base model might transfer partially to the fine-tuned model without retraining. A practical experiment would measure the drop in acceptance rate when using a base-model-trained adapter with a fine-tuned target model, versus training a new adapter from scratch. If the acceptance rate remains above, say, 70% of the from-scratch rate, organizations could deploy Kangaroo on fine-tuned models immediately, with adapter retraining as an optional optimization step rather than a prerequisite.
When to Prefer This Method
The paper explicitly positions Kangaroo against three self-drafting alternatives (Medusa, Lookahead, REST) and implicitly against external draft model approaches. The choice criteria emerge from the experimental results and design analysis:
-
Prefer Kangaroo over Medusa when: (a) deployment GPU memory is constrained and 591M additional parameters is prohibitive, or (b) the text generation task involves sequential reasoning (math, logic, multi-step QA) where Medusa's time-independent heads show steep CTAR decay (Figure 1(a)), or (c) the deployment environment cannot easily manage separate training pipelines for prediction heads on each model variant. The speedup advantage over Medusa is largest on mathematical reasoning (1.68× vs. 1.42×) and question answering, and smallest on summarization (1.53× vs. 1.48×) — the benefit is task-dependent.
-
Prefer Kangaroo over Lookahead when: wall-time speedup is the objective, full stop. Lookahead achieves competitive or higher compression rates (Table 1) but lower speedup on every Spec-Bench subtask (Figure 1(b)) due to its iterative refinement drafting cost. Kangaroo's single-pass autoregressive drafting is faster per draft token.
-
Prefer Kangaroo over external draft model approaches (DistillSpec, SpecInfer, standalone draft models) when: training cost for a separate draft model is unacceptable and the target model's shallow layers are accessible. Kangaroo's 10-epoch adapter training on ShareGPT and 67M-parameter footprint are likely cheaper than training a LLaMA-68M draft model from scratch, though the paper does not quantify this.
-
Prefer Medusa over Kangaroo when: draft generation latency must be absolutely minimal (Medusa's parallel heads produce all draft tokens in one forward pass through lightweight FFNs) and the slight acceptance rate disadvantage is acceptable, or when the deployment infrastructure already supports Medusa's tree attention and multi-head verification.
-
The paper does not establish a clear preference between Kangaroo and tree-based methods (Medusa-2, SpecInfer, Eagle), since the comparison is against Medusa-1 only. If tree verification provides substantial additional speedup (as the Medusa-2 paper claims), the Kangaroo + tree combination suggested in follow-up direction 6 above may be necessary to remain competitive with state-of-the-art tree-based methods.