LiteStage: Latencybewuste Laagoverslaan voor Multi-staps Redeneren

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

October 16, 2025

Auteurs: Beomseok Kang, Jiwon Song, Jae-Joon Kim

cs.AI

Samenvatting

Multi-stage reasoning is naar voren gekomen als een effectieve strategie om het redeneervermogen van kleine taalmodellen te verbeteren door complexe problemen op te delen in opeenvolgende substappen. Dit gaat echter ten koste van een verhoogde latentie. We observeren dat bestaande adaptieve versnellingsmethoden, zoals het overslaan van lagen, moeite hebben om efficiëntie en nauwkeurigheid in deze context in balans te brengen vanwege twee belangrijke uitdagingen: (1) variatie in gevoeligheid voor het overslaan van lagen per stap, en (2) de generatie van overbodige uitvoertokens. Om deze problemen aan te pakken, stellen we LiteStage voor, een latentiebewust raamwerk voor het overslaan van lagen bij multi-stage reasoning. LiteStage combineert een stapgewijze offline zoektocht die optimale laagbudgetten toewijst met een online, op vertrouwen gebaseerde vroege beëindiging van de generatie om onnodige decodering te onderdrukken. Experimenten op drie benchmarks, zoals OBQA, CSQA en StrategyQA, tonen aan dat LiteStage een versnelling tot 1,70x bereikt met minder dan 4,0% nauwkeurigheidsverlies, wat beter presteert dan eerdere trainingsvrije methoden voor het overslaan van lagen.

English

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.