Experience-Driven Dynamic Exits for LLMs with Reinforcement Learning

2026-06-01

AI Summary

This work addresses the inefficiency of autoregressive inference in large language models by introducing offline reinforcement learning into speculative decoding—a first in the field. It formulates the dynamic selection of exit layers and speculation lengths as a Markov decision process, enabling real-time optimization of the trade-off between computational cost and draft quality based on local context. Evaluated on Llama-2 and Llama-3, the proposed method achieves up to 2.7× speedup over standard autoregressive decoding and improves inference throughput by 17% compared to static speculative baselines, thereby overcoming the performance limitations inherent in fixed-configuration approaches.

Abstract

Large Language Models suffer from slow autoregressive inference. While self-speculative decoding accelerates this process, its efficiency is hampered by static configurations like fixed exit layers and speculation lengths. We reframe this optimization as a \textbf{Markov Decision Process} and propose \textbf{LEDE}, a framework that uses offline reinforcement learning. LEDE learns a policy to dynamically select the optimal exit layer and speculation length based on the local context of the generated sequence at each step, balancing computational cost and draft quality. Comprehensive evaluations on Llama-2 and Llama-3 models show LEDE achieves up to a $2.0\times$$\sim$$2.7\times$ speedup over autoregressive decoding and and provides an additional 17\% speedup over the static speculative baselines.

Problem

• Large Language Models
• autoregressive inference
• speculative decoding
• static configuration
• inference efficiency

Innovation

• dynamic exit
• speculative decoding
• offline reinforcement learning
• Markov Decision Process
• LLM inference acceleration