Speculative Sampling with Reinforcement Learning

2026-01-18

AI Summary

This work addresses the high inference latency of large language models and the limitations of existing speculative sampling methods, which rely on static hyperparameters and struggle to adapt across diverse contexts and domains. The paper proposes Re-SpS, the first reinforcement learning–based dynamic speculative sampling framework, which achieves an optimal trade-off between generation aggressiveness and computational overhead by dynamically adjusting draft tree hyperparameters in real time. Re-SpS introduces a context-aware reinforcement learning mechanism and a multi-step action persistence policy, significantly enhancing inference efficiency. Evaluated on five diverse benchmarks, Re-SpS accelerates inference by up to 5.45× over the base large language model and outperforms the current state-of-the-art method, EAGLE-3, by 1.12×, all while preserving output quality without degradation.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Inference time latency has remained an open challenge for real world applications of large language models (LLMs). State-of-the-art (SOTA) speculative sampling (SpS) methods for LLMs, like EAGLE-3, use tree-based drafting to explore multiple candidate continuations in parallel. However, the hyperparameters controlling the tree structure are static, which limits flexibility and efficiency across diverse contexts and domains. We introduce Reinforcement learning for Speculative Sampling (Re-SpS), the first reinforcement learning (RL)-based framework for draft tree hyperparameter optimization. Re-SpS dynamically adjusts draft tree hyperparameters in real-time, learning context-aware policies that maximize generation speed by balancing speculative aggression with computational overhead. It leverages efficient state representations from target model hidden states and introduces multi-step action persistence for better context modeling. Evaluation results across five diverse benchmarks demonstrate consistent improvements over the SOTA method EAGLE-3, achieving up to 5.45$\times$ speedup over the backbone LLM and up to 1.12$\times$ speedup compared to EAGLE-3 across five diverse benchmarks, with no loss in output fidelity.

Problem

• speculative sampling
• large language models
• inference latency
• hyperparameter optimization
• draft tree

Innovation

• Reinforcement Learning
• Speculative Sampling
• Draft Tree Optimization
• Inference Acceleration
• Large Language Models