The decoding bottleneck and multi-token speculation
Language model serving at scale requires every forward pass to count. Autoregressive decoding produces one token per model evaluation, which becomes the dominant cost when a model runs continuously for user queries. Multi-token prediction (MTP) addresses this by maintaining auxiliary draft heads alongside the primary model. During each forward pass, the draft heads produce several candidate tokens, and the target model evaluates all of them in a single pass. If the draft is verified — the target model agrees with each candidate — the system commits those tokens without additional passes. This amortizes the cost across multiple output tokens. The approach can accelerate decoding two to five times depending on draft reliability, but the speed gain is highly sensitive to the speculation horizon — the number of tokens the draft heads produce before verification.
Why fixed horizons fail across varied contexts
In standard MTP, the horizon is fixed before the forward pass begins. All draft heads fire simultaneously, so the system cannot inspect intermediate results and decide to stop early. The horizon must be chosen before seeing the context. A short horizon — say, two draft tokens — amortizes cheaply but leaves most of the work to the autoregressive baseline, yielding modest speedup. A long horizon — eight or more draft tokens — offers greater amortization when the drafts are reliable, but when a draft is rejected, the system discards the entire suffix and the compute spent on those tokens is lost. The optimal horizon therefore depends on how trustworthy the drafts are for a given prompt. Prompts that are easy for the model — those with strong, predictable patterns — tolerate longer drafts. Prompts that are ambiguous or out-of-distribution see higher rejection rates at long horizons, making the long draft wasteful. A horizon chosen for an easy prompt wastes compute on a hard prompt, and vice versa. Prior work addressed this by evaluating many fixed horizons and reporting the best, but that requires knowing the distribution of prompt difficulties beforehand and cannot adapt at serving time.
Formulating horizon selection as online control
AMTP treats horizon selection as an online control problem. At each decoding round, the controller observes verification feedback from the previous round and updates its prediction of how far the next draft is likely to survive. The system does not retrain the controller; it updates weights online using simple gradient estimates. The verifier — the target model's agreement check — remains unchanged. This design means AMTP can be dropped into an existing MTP deployment without modifying the target model or the draft head architecture. The controller's state is lightweight: a running average of acceptance outcomes per horizon, updated after every token. The prediction maps the current context to an expected number of committed tokens per horizon, and the controller selects the horizon maximizing expected tokens per second. Because the update uses verification feedback that is already computed for correctness, no extra model passes are needed for control learning.
Acceptance survival and the AMTP controller
The key quantity AMTP estimates is acceptance survival — the probability that a draft prefix of a given length is fully accepted by the target model. If the draft length is L and the first k tokens are all verified, the acceptance survival at horizon k is the probability that tokens k+1 through L are also verified. AMTP characterizes the expected number of committed tokens at each horizon as a function of these survival probabilities. Before each decoding round, the controller predicts survival probabilities for all feasible horizons using the running averages, computes expected committed tokens per horizon, and picks the horizon with the highest ratio of expected committed tokens to the forward-pass cost. The controller then issues that horizon to the MTP engine. After the round completes, AMTP updates the survival averages: if the draft of length H was fully accepted, the survival counts for all h ≤ H are incremented; if rejection occurred at position r < H, the counts for horizons longer than r are adjusted downward. This update uses only the verification outcome that the system already computes, so the controller learns from every token without additional compute.
Regret bounds and sublinear convergence
Theoretically, the paper shows that any online learner's prediction error for the best horizon translates into throughput regret relative to the horizon that would have been chosen with perfect information about the context. This regret, in turn, bounds the wall-clock completion time difference between the adaptive system and the oracle that always picks the best fixed horizon. With online gradient descent on the survival probability estimates, the regret after T decoding rounds is sublinear in T. Sublinear regret means the average per-round regret vanishes as more tokens are decoded, so the adaptive policy eventually matches the best fixed horizon in hindsight and then surpasses it by exploiting context-dependent variation. The proof assumes the survival probabilities lie in a bounded interval and that the gradient estimates have bounded variance, conditions that the experiments verify empirically.
Experimental setup: draft-target pairings and benchmarks
The authors evaluate AMTP across five draft-target pairings. Three use Llama-3.1-8B-Instruct as the target model, each paired with a different drafting mechanism: Medusa heads with 3, 5, and 8 draft tokens respectively; Meta's MTP models checkpoints trained with column-wise parallelism; and Qwen3-4B-Inst-2507-MTP, a compact model with its own MTP heads. The remaining two pairings are external MTP implementations — Meta's larger MTP checkpoints and Qwen3's 2507-MTP variant. All models are evaluated on two benchmark suites: AIME 2025, a set of challenging mathematical reasoning problems, and GSM8K, grade-school level math word problems. Throughput is measured in tokens per second per GPU, and the baseline for comparison is standard autoregressive decoding with no drafting. The authors also report the oracle fixed horizon that would be selected if the acceptance survival probability were known exactly for each prompt, providing an upper bound on AMTP's achievable gain.
Results: speedups and horizon selection accuracy
AMTP decodes up to 4.1 times faster than standard autoregressive decoding across the evaluated pairings. In 8 of the 10 draft-target settings, AMTP matches or outperforms the best fixed horizon, and in the remaining 2 settings the loss over the oracle fixed horizon is less than 3 percent. Throughput gains of up to 10 percent over the best fixed window are reported, with the largest gains observed on the Medusa-8 pairing with AIME 2025 prompts, where the draft reliability is high enough for long horizons but variable enough that adaptive selection matters. On GSM8K, where prompts are simpler and draft rejection rates are lower, AMTP still achieves 2.8 to 3.5 times speedup over baseline, with gains of 4 to 7 percent over the fixed horizon. The controller consistently selects shorter horizons on hard prompts and longer horizons on easy prompts, confirming the core intuition that horizon trustworthiness varies with context. The oracle fixed horizon, which knows the survival probability in advance, achieves up to 4.8 times speedup, showing that AMTP captures most of the available gain.
Where adaptive selection still struggles
The paper acknowledges several limitations. When the draft heads and target model share most of their parameters — as in some of the smaller Medusa pairings — the verification feedback is less informative about the true reliability horizon, and the controller's estimates carry higher variance. In those settings AMTP's gains narrow to 1 or 2 percent over the fixed horizon. The sublinear regret bound assumes the survival probabilities remain in a bounded interval; extreme distribution shifts — such as prompts from a domain completely unlike the training data — can cause the controller to misestimate for the first several dozen rounds until enough verification data accumulates. The paper also notes that AMTP's per-round overhead, while negligible compared to the model forward pass, is non-zero; on GPUs with extremely low token rates the control-update cost could represent a larger share of total time. Finally, the approach assumes the draft heads are trained independently of the controller; if the draft training objective changes, the controller's survival estimates would need recalibration, though the online update rule would continue to function.
What this means for serving engineers
For LLM serving infrastructure, AMTP offers a practical path to closer-to-optimal drafting speed without requiring domain-specific horizon tuning. The controller can be initialized with uniform survival estimates and learns on the first few dozen decoded tokens, adapting as the conversation or task distribution evolves. The per-GPU memory overhead is minimal — a few floats per horizon — and the compute overhead is limited to running the survival average update after each token, which the authors measure at less than 0.5 percent of total decode time. Serving systems that already use MTP can enable AMTP by adding the controller loop and pointing it at the existing verification output; no draft head retraining or target model modification is required. The authors release the controller implementation alongside the paper, configured for common MTP architectures, and report that integration into a ROS 2-based serving stack takes under an hour for engineers familiar with the model's parallel serving primitives. Energy per token drops proportionally with the throughput gain, which for data-center operators is a secondary benefit alongside the latency reduction.