LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning

TL;DR

LiteStage addresses the latency of multi-stage reasoning in small LLMs by combining a stage-aware offline layer-budget search with online confidence-based generation early exit. It estimates sub-layer importance via cosine similarity, then greedily allocates skipped layers from the slowest to the fastest stage, ensuring accuracy remains within a threshold. An online module terminates decoding early when token confidence decays, reducing extra generation tokens. Across OBQA, CSQA, and StrategyQA, LiteStage delivers up to 1.70x speedup with minimal accuracy loss, outperforming prior training-free methods and demonstrating the value of stage-aware efficiency in reasoning tasks.

Abstract

Multi-stage reasoning has emerged as an effective strategy for enhancing the reasoning capability of small language models by decomposing complex problems into sequential sub-stages. However, this comes at the cost of increased latency. We observe that existing adaptive acceleration techniques, such as layer skipping, struggle to balance efficiency and accuracy in this setting due to two key challenges: (1) stage-wise variation in skip sensitivity, and (2) the generation of redundant output tokens. To address these, we propose LiteStage, a latency-aware layer skipping framework for multi-stage reasoning. LiteStage combines a stage-wise offline search that allocates optimal layer budgets with an online confidence-based generation early exit to suppress unnecessary decoding. Experiments on three benchmarks, e.g., OBQA, CSQA, and StrategyQA, show that LiteStage achieves up to 1.70x speedup with less than 4.0% accuracy loss, outperforming prior training-free layer skipping methods.

Figures (10)

• Figure 1: Multi-stage Reasoning. (a) Example of multi-stage reasoning introduced in TinyThinker piao2024tinythinker. The process consists of three stages, where each inference plays a distinct yet complementary role, jointly leading to a stronger conclusion. Stage 1 (Recall) generates an initial idea for solving the problem, Stage 2 (Analysis) evaluates each option with reasoning, and Stage 3 (Summary) delivers the final conclusion. (b) Entropy of token IDs generated across the three stages. Higher entropy indicates greater diversity in the generated tokens. (c) Accuracy-speedup trade-off in recent training-free layer-skipping techniques. Speedup is normalized to the latency of the full-layer model.
• Figure 2: Motivational Study. (a) Accuracy and latency when the same number of layers are skipped across all reasoning stages. (b)-(d) Accuracy and latency when layer skipping is independently applied to a single stage, from Stage 1 through Stage 3, respectively. Notably, the accuracy trend in (d) closely aligns with that in (a). Furthermore, greater layer skipping does not always guarantee lower latency. Results are obtained on the TinyThinker OBQA dataset using AdaSkip-based layer importance in TinyLlama-1.1B.
• Figure 3: Overview of LiteStage. The proposed method consists of an offline configuration (Steps 1–2) and an online adjustment (Step 3). In the offline phase, (1) layer importance is estimated at the sub-layer level (MHSA and FFN) and accumulated across stages, followed by (2) a search for the optimal layer budget that minimizes latency within a target accuracy. In the online phase, (3) generation early exit dynamically terminates decoding when the average confidence of recent tokens falls below a threshold, preventing excessive generation length and ensuring consistent efficiency gains.
• Figure 4: Confidence of Output Tokens. (a) Average token confidence over decoding steps with varying layer skipping levels. (b) Number of remaining samples during decoding, where more samples are associated with generation early exit in larger skip budgets due to confidence decay. Results include skipping of 10 (yellow), 15 (orange), and 20 (red) sub-layers in Stage 1 of OBQA validation using TinyLlama-1.1B.
• Figure 5: Comparison with Baselines. (a)–(c) present the accuracy–speedup trade-offs on the OBQA, CSQA, and StrategyQA datasets, respectively. The performance of the full-layer model is marked in the upper-left region of each plot (i.e., 64.0% for OBQA, 54.8% for CSQA, and 62.4% for StrategyQA), and the speedup is normalized with respect to its latency. The highest achievable speedup by our method, along with its corresponding accuracy, is indicated in the lower-right region. All results are obtained using TinyLlama-1.1B models.