Reinforcement learning improves speed of language model speculative decoding
RCSD: Reinforcement and Continual Learning for Domain-Specific Speculative Decoding
Computation and Language
Summary
Generating text with large language models can be slow, so people use a smaller model to guess next words and check them with the big model to speed things up. The usual training forces the small model to copy the big one too closely, causing delays due to uncertain guesses. The authors created RCSD, a new training method that rewards the small model for better guessing sequences that the big model will accept. This approach makes the process faster by about 35% without making the output worse.
What this means in practice
- •For ai infrastructure teams: Speed up large language model inference by training draft models to improve acceptance rates and reduce rejections.
- •For real-time chatbot developers: Achieve faster response times in chatbots by using a trained draft model that better aligns with target language models during decoding.
Authors
Mingkang Liu, Weiming Meng, Peipei Jiao, Jiawei Zhang, Tianyuan Liu, Junyan Li, Lichen Pang
Abstract
Speculative decoding accelerates large language model inference by employing a compact draft model to propose candidate tokens that are verified in parallel by a target model. Conventional training relies on soft-label distillation, which forces the draft model to mimic the target's full output distribution, including high-entropy regions that degrade acceptance rates and increase early rejection risk. We propose RCSD (Reinforcement and Continual Learning for Speculative Decoding), a training framework that aligns draft model optimization with the system-level objective of maximizing acceptance length. RCSD leverages naturally occurring accepted and rejected sequences during speculative inference to construct implicit preference pairs, with masked loss to avoid penalizing target-aligned tokens. Evaluated across multiple benchmarks, RCSD improves average acceptance length by 35% over HASS and achieves 1.35x faster end-to-end inference than HASS on practical workloads, without compromising generation quality.