Research map / Large language models / Reasoning and chain of thought

Learning Harmonized Representations for Speculative Sampling

ICLR 2025 · Large language models › Reasoning and chain of thought · approach: Speculative 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 HASS, a speculative sampling draft-model training method that combines a Top-K distillation loss (ranking-distillation inspired) with multi-step context-aligned training, built on EAGLE-2, yielding higher acceptance lengths and speedups on LLaMA2/3 models.

Problem

Speculative sampling draft models suffer from inconsistent context between training and decoding (feature-level exposure bias from using target hidden states unavailable at decode time) and from a mismatch between a full-vocabulary training objective and decoding that only cares about high-probabili

Approach

Add a Top-K distillation loss restricted to the target LLM's most probable tokens (inspired by ranking distillation from recommender systems), and train the draft model over multiple steps so that it conditions on its own previously generated features, with an attention mask that keeps training parallel and adds no inference overhead.

Key results

• On LLaMA2-Chat 7B (T=0), HASS reaches 3.24x mean speedup and 5.15 mean acceptance length vs. 2.81x and 4.61 for EAGLE-2.
• On LLaMA3-Instruct 70B, HASS achieves 4.05x mean speedup (T=0) vs. 3.37x for EAGLE-2.
• Across four LLaMA models, HASS reports 2.81x-4.05x speedups, 8%-20% over EAGLE-2, on NVIDIA H800 GPUs.

Measured effects reported in the paper

Surprising findings

• Feature-only context alignment outperforms adding draft-generated token alignment; adding token alignment degrades acceptance length.
• Using a fixed ShareGPT dataset, HASS trained on 1/4 of the data matches or beats EAGLE-2 trained on the full data.

Limitations

• Training cost grows with the number of alignment steps (extra n-1 draft forward passes), and 5 steps reduces acceptance length.
• Acceptance rate at the first speculation step can drop relative to EAGLE-2 on larger models, requiring loss reweighting.
• Top-K loss is hand-designed; the authors state a better loss may exist.

Experimental setup

Datasets: ShareGPT MT-bench HumanEval GSM8K translation tasks (Multilingual-SpecBench)

Benchmarks: MT-bench HumanEval GSM8K Multilingual-SpecBench translation tasks

Metrics: speedup ratio acceptance length

Baselines: PLD Lookahead Vanilla Speculative Sampling Medusa EAGLE EAGLE-2

Builds on: EAGLE EAGLE-2 ranking distillation speculative sampling

Closest papers by meaning

• Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment (ICLR 2025 · similarity 0.90)
• A Theoretical Perspective for Speculative Decoding Algorithm (NeurIPS 2024 · similarity 0.90)
• SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.90)
• Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies (ICML 2025 · similarity 0.90)
• GRIFFIN: Effective Token Alignment for Faster Speculative Decoding (NeurIPS 2025 · similarity 0.89)
• Flatter Tokens are More Valuable for Speculative Draft Model Training (ICLR 2026 · similarity 0.89)