LiteStage: Omisión de Capas Consciente de la Latencia para el Razonamiento Multietapa

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

October 16, 2025

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

cs.AI

Resumen

El razonamiento en múltiples etapas ha surgido como una estrategia efectiva para mejorar la capacidad de razonamiento de los modelos de lenguaje pequeños al descomponer problemas complejos en subetapas secuenciales. Sin embargo, esto conlleva un aumento en la latencia. Observamos que las técnicas existentes de aceleración adaptativa, como el salto de capas, tienen dificultades para equilibrar la eficiencia y la precisión en este contexto debido a dos desafíos clave: (1) la variación en la sensibilidad al salto por etapas, y (2) la generación de tokens de salida redundantes. Para abordar estos problemas, proponemos LiteStage, un marco de salto de capas consciente de la latencia para el razonamiento en múltiples etapas. LiteStage combina una búsqueda fuera de línea por etapas que asigna presupuestos óptimos de capas con una salida temprana basada en la confianza durante la generación en línea para suprimir la decodificación innecesaria. Los experimentos en tres puntos de referencia, como OBQA, CSQA y StrategyQA, muestran que LiteStage logra una aceleración de hasta 1.70x con una pérdida de precisión inferior al 4.0%, superando a los métodos anteriores de salto de capas sin entrenamiento.

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.