Research map / Large language models / Language, safety and interpretability

Optimized Multi-Token Joint Decoding With Auxiliary Model for LLM Inference

ICLR 2025 · Large language models › Language, safety and interpretability · approach: Parallel and early-exit decoding · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes multi-token joint decoding (MTJD) that samples multiple tokens from their joint distribution, and MTAD/tree-based variant that uses a small auxiliary model via beam sampling and joint-likelihood verification to accelerate it, with energy analysis.

Problem

Speculative decoding speeds up LLM inference but still generates each token from its single-token distribution, so output quality (perplexity) is not improved; sampling from the joint distribution of multiple tokens is exponentially expensive.

Approach

Generate multiple tokens from their joint distribution, approximated by beam sampling from a small draft model and verified by the target model's joint likelihood ratio with a threshold tau that accepts the longest passing prefix; a tree variant verifies all intermediate beams with tree attention and normalizes against the final beam's draft likelihood.

Key results

• Multi-token joint sampling improves Llama-3 Spider score from 22.0 to 52.5 and HumanEval from 15.9 to 36.6 versus single-token sampling.
• Tree-based MTAD (the proposed tree variant) achieves 43.4 tokens/s vs 31.1 for SpD on Spider with Llama-3 and 6.0 vs 8.9 J/token for Llama-3 HumanEval-level comparisons, with the paper reporting 21.1% faster and 23.4% less energy than MCSS on average.
• On average the tree method surpasses lossless baselines by 42.7% in downstream performance.

Measured effects reported in the paper

Surprising findings

• Speculative decoding and its variants reduce energy per token despite increasing FLOPs, because memory access dominates energy use.
• Lower perplexity correlates with higher downstream accuracy even for GPT-3.5-turbo at high temperature.

Limitations

• Approximation of joint distribution relies on beam sampling width and assumed small error epsilon; bounds depend on assumptions like bounded log-ratio U.
• Energy measured via nvidia-smi, which has known measurement error.
• Gains depend on draft-target model alignment; theory bound is loose.

Experimental setup

Datasets: Spider MT-Bench HumanEval ChatGPT-Prompts ChatAlpaca CNN/DailyMail

Benchmarks: Spider MT-Bench HumanEval ChatGPT-Prompts ChatAlpaca CNN/DailyMail

Metrics: tokens/s J/token pass@1 execution accuracy MT-Bench score perplexity Rouge-L

Baselines: SpD Spectr SpecInfer MCSS BiLD typical decoding multinomial sampling

Builds on: speculative decoding SpecInfer Spectr BiLD

Closest papers by meaning

• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (ICML 2024 · similarity 0.91)
• Dynamic-Width Speculative Beam Decoding for LLM Inference (AAAI 2025 · similarity 0.91)
• Reward-Guided Speculative Decoding for Efficient LLM Reasoning (ICML 2025 · similarity 0.90)
• Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment (ICLR 2025 · similarity 0.90)
• AdaDecode: Accelerating LLM Decoding with Adaptive Layer Parallelism (ICML 2025 · similarity 0.90)
• Multi-Branch Self-Drafting for LLM Inference Acceleration (AAAI 2025 · similarity 0.90)