LiteStage : Saut de couche conscient de la latence pour le raisonnement multi-étapes

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

Le raisonnement multi-étapes s'est imposé comme une stratégie efficace pour améliorer les capacités de raisonnement des petits modèles de langage en décomposant des problèmes complexes en sous-étapes séquentielles. Cependant, cela se fait au prix d'une latence accrue. Nous observons que les techniques d'accélération adaptatives existantes, telles que le saut de couches, peinent à équilibrer efficacité et précision dans ce contexte en raison de deux défis majeurs : (1) la variation de sensibilité au saut selon les étapes, et (2) la génération de tokens de sortie redondants. Pour y remédier, nous proposons LiteStage, un cadre de saut de couches conscient de la latence pour le raisonnement multi-étapes. LiteStage combine une recherche hors ligne par étape qui alloue des budgets de couches optimaux avec une sortie anticipée de génération basée sur la confiance en ligne pour supprimer le décodage inutile. Les expériences sur trois benchmarks, tels que OBQA, CSQA et StrategyQA, montrent que LiteStage atteint une accélération jusqu'à 1,70x avec une perte de précision inférieure à 4,0 %, surpassant les méthodes de saut de couches antérieures sans entraînement.

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.