Technical Summary: LEDE – Experience-Driven Dynamic Exits for LLMs with Reinforcement Learning

1. Problem Statement

Large Language Models (LLMs) suffer from significant latency during autoregressive inference because every generated token must traverse the entire stack of transformer layers. While Self-Speculative Decoding (SSD) has been proposed to accelerate this process by using a model's own early layers as an internal "drafter" to generate candidate tokens in parallel, current SSD methods are fundamentally limited by static configurations.

Existing approaches rely on fixed draft depths and fixed speculation lengths applied uniformly across the entire generation process. This rigidity fails to account for contextual sparsity, where the difficulty of predicting tokens varies significantly depending on the local context. Consequently, static methods cannot optimally balance the trade-off between draft quality (acceptance rate) and speculation speed, leading to suboptimal inference performance.

2. Methodology: The LEDE Framework

The authors propose LEDE (Learning-based Dynamic Exit), a framework that reframes the dynamic control of SSD as a Markov Decision Process (MDP). Instead of using rigid, pre-defined rules, LEDE employs offline reinforcement learning (RL) to train a policy that dynamically selects both the optimal exit layer and the speculation length based on the local context of the generated sequence.

2.1. Markov Decision Process Formulation

The framework models the decision-making process during a single generation step as an MDP with the following components:

• State Space ($S$): The agent observes a state vector $s_l$ at each candidate layer $l$. This vector is derived from the model's internal token distribution and includes:
  • Token Confidence ($C_k^l$): The average negative log-probability of the top-$k$ tokens.
  • Token Entropy ($H^l$): The uncertainty of the next-token prediction distribution.
  • Average Trace of Token Confidence ($C_{trace}$): The average negative log-probability of the top-1 token over a sliding window of previous steps.
• Action Space ($A$): At each candidate layer, the policy (a Q-network) decides between two actions:
  • Continue ($a=0$): Proceed to the next candidate layer.
  • Exit ($a=1$): Designate the current layer as the optimal exit layer ($l^*$) for the draft stage.
• Reward Function ($R$): Designed to optimize inference efficiency.
  • Positive Reward: Given if the draft token is accepted ($x_t^{l^*} = x_t^L$), scaled by the depth saved ($1 + \frac{L-l^*}{L}$).
  • Negative Reward: Given if the draft token is rejected ($-2 + \frac{L-l^*}{L}$).
  • Penalty: A small penalty ($-0.01$) is applied for continuing, encouraging early exits when appropriate.

2.2. Policy Optimization via Offline RL

LEDE utilizes a Deep Q-Network (DQN) trained offline using a replay buffer of past inference experiences:

1. Data Collection: An initial exploration phase collects diverse experience tuples $(s_t, a_t, r_t, s_{t+1})$ where the agent interacts with the environment.
2. Offline Training: The Q-network is trained by minimizing the n-step Temporal Difference (TD) error. The target value is computed using a stable target network, balancing the immediate reward of an early exit against the long-term value of the state.
3. Dual-Control Mechanism: During online inference, the trained agent performs two adaptive controls:
   • Dynamic Exit Layer Selection: It determines the optimal depth $l^*$ to start drafting.
   • Adaptive Drafting Length: After selecting $l^*$, the agent re-evaluates the state after each draft token. An "exit" action continues drafting (indicating high confidence), while a "continue" action (indicating low confidence) terminates drafting immediately for verification.

3. Key Contributions

The paper outlines three primary contributions:

1. MDP Formulation: It is the first work to formulate the dynamic control of draft depth and speculation length in SSD as an MDP problem, applying reinforcement learning to derive a control policy.
2. LEDE Framework: The design and implementation of a framework that leverages offline RL and experience replay to learn a context-aware policy, overcoming the limitations of both static configurations and simple heuristic-based methods.
3. Empirical Validation: Comprehensive experiments demonstrating significant speedups over autoregressive generation and substantial efficiency gains over static baselines.

4. Experimental Results

The framework was evaluated on Llama-2 (7B, 13B) and Llama-3 (1B) models, as well as CodeLLaMA (7B, 34B), across diverse tasks including instruction following, summarization, language modeling, and code generation.

• Speedup Performance:
  • LEDE achieves an average speedup of 2.32× over standard autoregressive decoding.
  • Peak speedups reach 2.72× (e.g., on Llama-2-7B for language modeling).
  • LEDE provides an additional 17% speedup over the best static speculative baseline (LayerSkip).
• Comparison with Baselines:
  • LayerSkip (Static): Performance is highly sensitive to fixed configurations and fails to generalize well across tasks.
  • Rule-based Methods (LITE, Draft & Verify): Constrained by predefined thresholds, leading to sub-optimal trade-offs between depth and acceptance rates.
  • LEDE: By learning a fine-grained control policy, it consistently outperforms all baselines in speedup while maintaining comparable generation quality (measured by Rouge-L).
• Ablation Study: Removing either the dynamic exit layer selection or the adaptive drafting length significantly degrades performance (dropping speedup to ~2.0×), confirming that both mechanisms are synergistic and essential.

5. Significance and Claims

The paper claims that LEDE represents a conceptual shift from rigid, heuristic-driven inference to adaptive, policy-driven inference. By learning to co-optimize computational depth and speculation length based on real-time context, the framework enables more robust and efficient decoding across diverse models and tasks.

The authors position this work as a step toward building more intelligent and computationally efficient LLM systems. They note that future work will focus on scaling these benefits to larger models (70B scale and beyond), but the current contribution establishes the viability of using offline reinforcement learning to dynamically manage the inference pipeline of large language models.