LiteStage: Salto di Livelli con Consapevolezza della Latenza per il Ragionamento Multi-stadio

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

October 16, 2025

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

cs.AI

Abstract

Il ragionamento multi-stadio è emerso come una strategia efficace per potenziare la capacità di ragionamento dei piccoli modelli linguistici, scomponendo problemi complessi in sotto-stadi sequenziali. Tuttavia, ciò avviene a scapito di una maggiore latenza. Osserviamo che le tecniche esistenti di accelerazione adattiva, come il salto di strati, faticano a bilanciare efficienza e accuratezza in questo contesto a causa di due sfide principali: (1) la variazione nella sensibilità al salto tra gli stadi, e (2) la generazione di token di output ridondanti. Per affrontare questi problemi, proponiamo LiteStage, un framework di salto di strati consapevole della latenza per il ragionamento multi-stadio. LiteStage combina una ricerca offline per stadio che assegna budget ottimali di strati con un'uscita anticipata basata sulla confidenza durante la generazione online, per sopprimere la decodifica non necessaria. Esperimenti su tre benchmark, come OBQA, CSQA e StrategyQA, dimostrano che LiteStage raggiunge un'accelerazione fino a 1.70x con una perdita di accuratezza inferiore al 4.0%, superando i precedenti metodi di salto di strati senza addestramento.

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.