LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning
Authors: name: Beomseok Kang, name: Jiwon Song, name: Jae-Joon Kim
Organization: Seoul National University VLSI Lab

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.