LiteStage: Pular Camadas com Consciência de Latência para Raciocínio em Múltiplos Estágios

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

October 16, 2025

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

cs.AI

Resumo

O raciocínio em múltiplos estágios surgiu como uma estratégia eficaz para aprimorar a capacidade de raciocínio de modelos de linguagem menores, decompondo problemas complexos em subestágios sequenciais. No entanto, isso ocorre ao custo de um aumento na latência. Observamos que as técnicas de aceleração adaptativa existentes, como o salto de camadas, lutam para equilibrar eficiência e precisão nesse cenário devido a dois desafios principais: (1) variação na sensibilidade ao salto entre os estágios e (2) a geração de tokens de saída redundantes. Para abordar esses problemas, propomos o LiteStage, uma estrutura de salto de camadas com consciência de latência para raciocínio em múltiplos estágios. O LiteStage combina uma busca offline por estágio que aloca orçamentos de camadas ideais com uma saída antecipada de geração baseada em confiança online para suprimir a decodificação desnecessária. Experimentos em três benchmarks, como OBQA, CSQA e StrategyQA, mostram que o LiteStage alcança uma aceleração de até 1,70x com uma perda de precisão inferior a 4,0%, superando métodos anteriores de salto de camadas sem treinamento.

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.