LiteStage: Latenzbewusstes Layer-Skipping für mehrstufiges Reasoning

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

October 16, 2025

papers.authors: Beomseok Kang, Jiwon Song, Jae-Joon Kim

cs.AI

papers.abstract

Mehrstufiges Denken hat sich als effektive Strategie erwiesen, um die Denkfähigkeit kleiner Sprachmodelle zu verbessern, indem komplexe Probleme in sequenzielle Teilphasen zerlegt werden. Dies geht jedoch mit erhöhter Latenz einher. Wir beobachten, dass bestehende adaptive Beschleunigungstechniken wie das Überspringen von Schichten Schwierigkeiten haben, in diesem Kontext Effizienz und Genauigkeit auszugleichen, und zwar aufgrund von zwei zentralen Herausforderungen: (1) phasenweise Variation in der Empfindlichkeit des Überspringens und (2) die Erzeugung redundanter Ausgabetokens. Um diese Probleme zu adressieren, schlagen wir LiteStage vor, ein latenzbewusstes Framework zum Überspringen von Schichten für mehrstufiges Denken. LiteStage kombiniert eine phasenweise Offline-Suche, die optimale Schichtbudgets zuweist, mit einem Online-Frühausstieg basierend auf Konfidenz, um unnötige Dekodierung zu unterdrücken. Experimente auf drei Benchmarks, z. B. OBQA, CSQA und StrategyQA, zeigen, dass LiteStage eine Beschleunigung von bis zu 1,70x mit einem Genauigkeitsverlust von weniger als 4,0 % erreicht und damit bisherige trainingsfreie Methoden zum Überspringen von Schichten übertrifft.

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.