LEAP: Zone-Aware MCTS for LLM Self-Speculative Decoding

Self-speculative decoding accelerates LLM inference by using a lightweight draft model for generation and a target model for verification, where the draft model is constructed by a subset of the target model’s layers, and the key challenge lies in layer configuration strategies. To address this challenge, we propose LEAP, a plug-and-play approach that formulates and optimizes the draft model construction problem as a sequential decision-making process by Monte Carlo Tree Search (MCTS). To navigate the prohibitive search space of deep LLMs, we leverage two empirical observations: (i) the prefilling-derived redundancy information remains informative during decoding, and (ii) the layer redundancy exhibits zone-wise characteristics. These observations enable a structured search space through zone partitioning and layer grouping, which serves as an inductive bias to facilitate efficiency of MCTS. Experimental results show that LEAP achieves a speedup of $1.7\times\sim2.0\times$ for LLM inference.

Large language models are powerful, but they can be slow and expensive to use because they usually generate responses token by token. A common way to speed them up is to let a smaller “draft” version of the model guess several upcoming tokens first, and then let the full model quickly verfify them. However, it is difficult to decide which parts of the full model should be used to build this draft version, and a fixed choice may not work well for different tasks. We find that Monte Carlo Tree Search is well suited for this problem because it can repeatedly try different draft versions and adjust its search strategy based on what it has learned. However, directly searching over all possible versions is too costly. To make the search practical, we group similar model layers into zones and groups, so the method can search over groups instead of individual layers. Based on these ideas, we propose LEAP, a method that automatically searches for an efficient draft version during generation. LEAP uses these layer groups to guide its search and decide which parts of the model to keep or skip. It does not require extra training or extra modules.