Research map / Large language models / Reasoning and chain of thought
Mixture of Attentions For Speculative Decoding
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 a Mixture of Attentions drafter for speculative decoding (Layer Self-Attention, Self/Cross-Attention, Target Layer Inference) that improves EAGLE-2 speedups on a single device and enables a client-server setting that keeps generating during disconnection.
Problem
Self-drafting speculative decoding models that reuse LLM activations suffer from partial observability and lack of on-policyness during training, which degrade draft quality; also, serving LLMs from small client devices needs fast decoding and useful output when the server connection drops.
Approach
A drafter that compresses the LLM's per-layer KV cache via Layer Self-Attention (addressing partial observability), uses a causal Self-Attention plus a masked Cross-Attention over LLM activations so it can draft multiple tokens in one pass during training (more on-policy), and a Target Layer Inference knob that reuses the last N frozen LLM decoder layers to trade speed for accuracy.
Key results
- Ours (TLI=1) and TLI=0 achieve SOTA SpecBench speedups, about 9.5% faster than EAGLE-2 on average in tokens-per-second, with acceptance length up to 25% higher (single device).
- In a simulated 5G client-server setting on HumanEval, Ours (TLI=0) reaches 34.1 tokens/s versus 24.3 for EAGLE-2, with 24.1 LLM calls versus 36.4.
- On HumanEval with 1.3B drafter under single device, the method reaches acceptance length 4.30 versus 2.81 for EAGLE-2 (strict stopping setup).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Mixture of Attentions with TLI=1 instead of EAGLE-2 | Llama3-8B-Instruct drafter · SpecBench mean | speedup: 1.78 -> 1.95 | Helped |
| Use Mixture of Attentions with TLI=3 instead of EAGLE-2 | Llama3-8B-Instruct drafter · SpecBench mean | acceptance length: 3.60 -> 4.79 | Helped |
| Use Mixture of Attentions (TLI=0) instead of EAGLE-2 | Llama3-8B-Instruct drafter · HumanEval, 5G | tokens per second: 24.3 -> 34.1 | Helped |
| Use Mixture of Attentions (TLI=0) instead of EAGLE-2 | Llama3-8B-Instruct drafter · HumanEval, 4G | tokens per second: 13.6 -> 25.1 | Helped |
| Use Mixture of Attentions (TLI=3) instead of EAGLE-2 | Llama3-8B-Instruct drafter · HumanEval, disconnection after 10 tokens | pass@1: 8.07 % -> 11.18 % | Helped |
| Add LSA to the TLI=1 drafter (ablation) | Mixture of Attentions · HumanEval, chain drafting | tokens per second: 21 -> 39 | Helped |
| Add LSA to the TLI=1 drafter (ablation) | Mixture of Attentions · HumanEval, chain drafting | acceptance length: 1.28 -> 2.25 | Helped |
| Add the CA-based on-policy and T-step-bounded components versus EAGLE (ablation) | Mixture of Attentions without LSA, TLI=0 · HumanEval, chain drafting | tokens per second: 30 -> 38 | Helped |
| Add the CA-based on-policy and T-step-bounded components versus EAGLE (ablation) | Mixture of Attentions without LSA, TLI=0 · HumanEval, chain drafting | acceptance length: 1.45 -> 1.93 | Helped |
| Increase TLI from 0 to 3 (ablation) | Mixture of Attentions · HumanEval, chain drafting | acceptance length: 2.14 -> 2.54 | Helped |
| Increase TLI from 0 to 3 (ablation) | Mixture of Attentions · HumanEval, chain drafting | tokens per second: 40 -> 39 | Mixed |
| Use TLI=1 instead of vanilla decoding | Llama3-8B-Instruct · HumanEval | pass@1: 62.5% -> 62.5% | No clear effect |
Surprising findings
- Ours (TLI=0) in client-server setting is faster than independent small models that exchange only token ids despite sending activation tensors.
- Removing the LSA layer cut tokens per second from 39 to 21 on HumanEval ablation, indicating larger drops than expected from partial observability alone is not what the text reports; the paper says LSA's effect is about 6% in tokens per second.
Limitations
- Client-server results rely on simulated network conditions via Linux traffic control rather than real deployment
- Higher TLI improves acceptance but increases drafting compute, reducing tokens/s in some settings
- Not on-policy in full: query inputs for future positions are not generated by the LM during training
Experimental setup
Datasets: UltraChat SpecBench MT-Bench HumanEval GSM8K Alpaca CNN/Daily Mail Natural Questions
Benchmarks: SpecBench HumanEval GSM8K CNN/DM MT-Bench Alpaca Natural Questions
Metrics: speedup ratio tokens per second acceptance length number of LLM calls pass@1 accuracy ROUGE-L
Baselines: EAGLE-2 EAGLE-2 off. Independent 1.3B Independent 1.7B Glide Ouroboros Vanilla decoding
Builds on: EAGLE-2 EAGLE Medusa Tandem Transformers
Closest papers by meaning
- SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.93)
- MoESD: Unveil Speculative Decoding's Potential for Accelerating Sparse MoE (NeurIPS 2025 · similarity 0.92)
- MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding (ICLR 2025 · similarity 0.91)
- A Theoretical Perspective for Speculative Decoding Algorithm (NeurIPS 2024 · similarity 0.91)
- Online Speculative Decoding (ICML 2024 · similarity 0.91)
- LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding (ICML 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.