Axion AI.

Learned draft length speeds up cross-model speculative decoding by up to 4.55x, but self-speculative MTP gains are within a few percent (Georgia Tech preprint)

AMTP picks how many tokens to draft each round, with no retraining. Meta's 7B drafting for Llama-2-13B ran 7.6 to 28.2 tok/s on AIME; Meta's 7B alone lost 15%.

By Axion AI Desk · · Source: arXiv cs.AI

Typographic card

MTP must fix the horizon before drafting, as all heads fire in the same forward pass. We therefore formulate horizon selection as an online control problem, using verification feedback to learn how far the draft can be trusted in each context.

The problem it targets

Multi-token prediction (MTP) heads draft several tokens at once and the target model verifies them in one pass. MTP modules now ship in frontier models, and ordinary next-token models can be converted into MTP drafters. That leaves an inference-time decision most servers answer with a constant: how many tokens to draft before the next verification, called the speculation horizon. A long horizon pays off when the draft is reliable and wastes compute when an early token is rejected and the rest is thrown away. Unlike autoregressive drafting, where you can stop token by token, MTP heads fire together, so the horizon has to be fixed before drafting begins.

The Georgia Tech authors propose Adaptive Multi-Token Prediction (AMTP): a small controller that learns, from the accept and reject outcomes it already sees, how far the draft is likely to survive in the current context, then picks the horizon with the highest expected tokens per second. It needs no retraining and leaves the exact verifier unchanged, so output quality is governed by the same acceptance rule as before. The paper also proves that prediction error in the controller translates into bounded throughput regret.

The numbers, from the paper's tables

The authors test five pairings on AIME 2025 and GSM8K: a self-distilled Qwen3-4B, Meta's multi-token-prediction 7B model, Medusa heads on Llama-3.1-8B-Instruct, and two cross-model pairings with a larger separate target (Meta's 7B drafting for Llama-2-13B, and the 4B drafting for Qwen3-14B). The baseline is the largest fixed window each draft supports (16, 4 or 6 tokens).

In the self-speculative setups the gains are small. Self-distilled Qwen3-4B moved from 34.0 to 34.5 tokens per second on AIME and from 38.6 to 38.3 on GSM8K, which is a tie. Medusa on Llama-3.1-8B gained about 7 percent on AIME (23.6 to 25.3) and lost about 3 percent on GSM8K (20.0 to 19.4). Meta's 7B model regressed 15 percent on AIME (37.7 to 31.9), which the authors attribute to a base checkpoint that does not reliably stop after answering.

The cross-model pairings are where the controller earns its keep. Meta's 7B drafting for Llama-2-13B went from 7.6 to 28.2 tokens per second on AIME against the fixed window, a 4.55x speed-up over plain decoding, and from 4.7 to 11.6 on GSM8K. The 4B-to-14B pairing gained 10 percent on AIME (22.5 to 24.7) and tied on GSM8K. Across all runs AMTP drafted fewer tokens per round yet committed about the same number, raising the acceptance rate, for instance from 18 to 34 percent for the self-distilled 4B on AIME. The chosen window also moves: for the 4B on GSM8K it ranged from 1 to 16 with a mean of 11.2.

What does not line up

Read the abstract with the tables open. It says AMTP decodes up to 4.1x faster than standard decoding, yet Table 1 reports 4.55x for the Llama-2-13B pairing on AIME. It says throughput gains reach up to 10 percent over the fixed window, while the discussion text claims 268 percent on AIME for that same pairing. The large figure comes from a fixed window of 4 that was a poor match for a 13B target the drafter was not built for, so it shows how badly a wrong constant can do rather than how much a tuned system gains. The abstract also says AMTP matches or beats the fixed window in 8 of 10 settings; the three regressions above are within what the tables show. We use the tables throughout.

Should you switch

If you run a self-speculative MTP model on a single GPU and have already tuned the window, expect roughly nothing: the paper's own self-speculative results sit within a few percent either way. If you pair a small drafter with a larger separate target, or serve prompts with highly variable draft quality, a learned horizon is worth testing because a wrong constant is expensive. Nothing here covers production serving: every run decodes one sequence at a time with the KV cache disabled, each condition is a single unseeded run, and hardware differs across models, so comparisons hold only within a model. The authors say batched serving with caching changes the cost of each horizon and leave it to future work. Treat this as a promising control idea with proofs behind it, not a drop-in speed-up for your inference stack.

Comprehensive take

If you run speculative decoding with a fixed draft length, the useful lesson is that the constant matters most when drafter and target are different models. Test an adaptive horizon there. For single-model MTP with a tuned window, the paper shows no reliable gain, and its setup (batch of one, no KV cache, single runs) is far from production serving.

Use cases

  • Teams pairing a small drafter with a larger target

    The largest effects were in cross-model setups: Meta's 7B drafting for Llama-2-13B went from 1.23x to 4.55x over plain decoding on AIME 2025 once the window was adaptive.

  • Inference engineers tuning a fixed window

    Log acceptance per position before adopting anything: the paper's controller is cheap, but self-speculative models like the Qwen3-4B setup moved by under 2 percent on both tasks.

  • Researchers building MTP drafters

    The proofs bound completion time by the controller's prediction error, which gives a concrete target for better acceptance predictors.

Specs

Pairings testedSelf-distilled Qwen3-4B; Meta MTP 7B; Medusa on Llama-3.1-8B; Meta MTP 7B to Llama-2-13B; 4B to Qwen3-14B source
Self-distilled Qwen3-4B, AIME 202534.0 tok/s fixed k=16 vs 34.5 adaptive (plain 18.8) source
Meta MTP 7B to Llama-2-13B, AIME 20257.6 tok/s fixed k=4 vs 28.2 adaptive (plain 6.2); 4.55x over plain source
Meta MTP 7B, AIME 202537.7 tok/s fixed vs 31.9 adaptive (15 percent regression) source
Abstract claim vs tableAbstract: up to 4.1x; Table 1: 4.55x on one pairing source
Test conditionsOne sequence at a time, KV cache disabled, single unseeded runs, hardware differs across models source

Insights

  • Axion reading: the headline 268 percent gain is mostly a measure of how poor a fixed window of 4 is for a 13B target the drafter was not built for; a tuned baseline would shrink it.
  • Axion caution: no KV cache and batch size one means verification cost is overstated relative to production engines, which could either inflate or deflate the value of adaptive horizons.
  • Axion note: the abstract's 4.1x and 10 percent figures do not match the tables; readers should rely on the tables.

Community

Share (0): XRedditLinkedinEmail

Comments

No comments yet.

Log in or register to like or comment.