LiteStage: マルチステージ推論におけるレイテンシを考慮したレイヤースキップ

要旨

多段階推論は、複雑な問題を逐次的なサブステージに分解することで、小規模言語モデルの推論能力を向上させる効果的な戦略として注目されている。しかし、これにはレイテンシの増加という代償が伴う。既存の適応的加速技術、例えばレイヤースキップは、この設定において効率と精度のバランスを取ることに苦労していることが観察される。これは主に2つの課題によるものである：(1) ステージごとのスキップ感度の変動、(2) 冗長な出力トークンの生成。これらの課題に対処するため、我々は多段階推論のためのレイテンシを意識したレイヤースキップフレームワークであるLiteStageを提案する。LiteStageは、最適なレイヤーバジェットを割り当てるステージごとのオフライン探索と、不要なデコードを抑制するためのオンラインの信頼度ベースの生成早期終了を組み合わせている。OBQA、CSQA、StrategyQAの3つのベンチマークでの実験により、LiteStageは4.0%以下の精度損失で最大1.70倍の高速化を達成し、従来のトレーニング不要なレイヤースキップ手法を上回ることを示した。

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.

LiteStage: マルチステージ推論におけるレイテンシを考慮したレイヤースキップ

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

要旨

Support