LiteStage: Пропуск слоев с учетом задержек для многоэтапного рассуждения

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

October 16, 2025

Авторы: Beomseok Kang, Jiwon Song, Jae-Joon Kim

cs.AI

Аннотация

Многоэтапное рассуждение стало эффективной стратегией для повышения способности к рассуждению у небольших языковых моделей за счет декомпозиции сложных задач на последовательные подэтапы. Однако это достигается ценой увеличения задержки. Мы отмечаем, что существующие адаптивные методы ускорения, такие как пропуск слоев, сталкиваются с трудностями в балансировке эффективности и точности в данном контексте из-за двух ключевых проблем: (1) вариативности чувствительности к пропуску на разных этапах и (2) генерации избыточных выходных токенов. Для решения этих проблем мы предлагаем LiteStage — фреймворк для пропуска слоев с учетом задержки в многоэтапных рассуждениях. LiteStage сочетает в себе поэтапный оффлайн-поиск, который распределяет оптимальные бюджеты слоев, и онлайн-ранний выход на основе уверенности для подавления ненужного декодирования. Эксперименты на трех тестовых наборах данных, таких как OBQA, CSQA и StrategyQA, показывают, что LiteStage достигает ускорения до 1,70x с потерей точности менее 4,0%, превосходя предыдущие методы пропуска слоев, не требующие обучения.

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.