PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Beomseok Kang, Jiwon Song, Jae-Joon Kim

arXiv:2510.14211v2cs.CLcs.AI

TL;DR

Sequential multi-stage reasoning adds latency, and layer-skipping trade-offs vary across stages while approximate decoding can produce redundant tokens. LiteStage combines offline stage-wise layer-budget search with confidence-based early exit; across OBQA, CSQA, and StrategyQA, it surpasses prior training-free methods with a large speedup margin.

  • Problem

    Sequential multi-stage reasoning adds latency, while stage-dependent skip sensitivity and extra generated tokens complicate efficient layer skipping.

    Source: §1 · pp. 1–2

  • Method

    LiteStage searches stage-wise layer budgets offline under an accuracy threshold and applies online confidence-based generation early exit.

    Source: §1 · p. 2

  • Results

    Across OBQA, CSQA, and StrategyQA, LiteStage surpasses prior training-free methods with a large speedup margin; on OBQA it achieves 1.49× speedup with 60.6% accuracy, while AdaSkip collapses to nearly 0% beyond 1.10×.

    Source: §4.1 · p. 6, §5 · p. 8

  • Takeaways & Limitations

    The results highlight stage-aware optimization and adaptive decoding as important to efficient multi-stage reasoning.

    Source: §5 · p. 8

    The analysis primarily focuses on short multi-stage reasoning, and extending LiteStage to long-horizon generation may require additional investigation.

    Source: §6 · p. 9

Abstract

from arXiv · show

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 outperforms prior training-free layer skipping methods.

1 Introduction

Short multi-stage reasoning suits practical tasks and small language models, but sequential stages add latency and differ in sensitivity to layer skipping. LiteStage responds with stage-wise layer-budget search and confidence-based generation early exit.

1 Introduction

  • Short multi-stage reasoning decomposes practical question-answering and decision tasks into structured steps such as recall, local inference, and decision aggregation.Small language models may be unable to reason reliably in a single step because of limited capacity.Source: §1 · p. 1
  • Sequential multi-stage reasoning adds non-trivial latency, and stages differ in decoding length, token diversity, and tolerance for acceleration.The resulting non-uniform information density means some stages tolerate aggressive acceleration while others are highly sensitive.Source: §1 · pp. 1–2
  • Layer skipping has asymmetric stage trade-offs: Stage 1 is notably robust to skipping, while latency benefits are greatest in the longer-generating Stage 2.Some approximate decoding also induces extra tokens, potentially making end-to-end latency slower than full-layer inference.Source: §1 · p. 2
  • LiteStage searches stage-wise layer budgets offline under an accuracy threshold, then uses token confidence online to stop generation early and avoid redundant decoding.The offline search proceeds from the longest stage to the shortest.Source: §1 · p. 2

2 Related Works

Prior work studies multi-stage reasoning, layer skipping, and generation early exit, but efficiency in multi-stage settings remains comparatively under-addressed. Training-free approaches are practical when curated reasoning paths for training are rarely accessible, while uniform skipping can yield suboptimal trade-offs.

2 Related Works

  • Multi-stage reasoning works largely prioritize reasoning quality over computational efficiency.Examples span staged reasoning, cascading chain-of-thought distillation, dynamic structure discovery, and multimodal extensions.Source: §2 · p. 2
  • Layer-skipping methods include intermediate-layer exits and router-based dynamic skipping, with router approaches requiring model and router training or router-only fine-tuning.The passage contrasts these training-based approaches with training-free methods discussed next.Source: §2 · p. 2
  • Training-free layer skipping is more practical when curated reasoning paths for training are rarely accessible, but common uniform policies can produce suboptimal efficiency–accuracy trade-offs.Examples include gradually skipping deeper layers, periodic skipping, and cosine-similarity-based block selection.Source: §2 · pp. 2–3
  • Generation early-exit research often targets verbose long-reasoning outputs, using trained probing heads or training-free repetition and logit-based heuristics.The cited work gives less attention to prolonged generation induced by model compression.Source: §2 · p. 3

3 Proposed Methods

LiteStage searches stage-specific layer-skipping budgets to minimize latency within an accuracy threshold, then applies confidence-based early exit to curb excessive decoding.

3 Proposed Methods

  • LiteStage searches stage-wise layer budgets L to minimize inference latency while keeping accuracy within a specified threshold.The objective compares a model using the skip budget with the full-layer model on a test dataset.Source: §3 · p. 3

3.1 Overview of LiteStage

  • LiteStage pairs offline configuration of skipped sub-layers with online early exit when recent-token confidence falls below a threshold.The offline phase estimates importance and searches a budget; the online phase terminates generation to prevent excessive output length.Source: §3.1 · p. 3, Figure 2 · p. 4

3.2 Step 1: Estimate Layer Importance

  • LiteStage ranks MHSA and FFN sub-layers by input-output cosine similarity, treating greater similarity as greater redundancy and lower importance.Importance is estimated during prefilling, averaged over validation samples and stages, and accumulated using each stage's prompt; the process runs offline once per dataset.Source: §3.2 · pp. 3–4

3.3 Step 2: Search Layer Budget

  • Greedy profiling searches budgets from the longest stage, jointly reflecting non-monotonic latency, cross-stage interactions, and differing stage sensitivity.In the CSQA example, a 1.0% accuracy threshold makes 50.7% the target from a 51.7% baseline; skipping four layers yields latency 0.2986.Source: §3.3 · pp. 4–5

3.4 Step 3: Generation Early Exit

  • Confidence-based exit targets extra low-confidence tokens, which become more evident as layer skipping increases and may contribute little to the final reasoning outcome.The figure compares confidence trajectories and remaining samples for different skip levels on OBQA validation data using TinyLlama-1.1B.Source: §3.4 · p. 5, Figure 3 · p. 5
  • The early-exit rule replaces a token with EOS when mean confidence over the recent-token cache falls below 0.5, with n=5 tokens used to stabilize decisions.Token confidence is the maximum logit value; confidence-based exit is also included in search so the accuracy-latency profile reflects evaluation conditions.Source: §3.4 · pp. 5–6

4 Experimental Results

The section evaluates LiteStage against training-free layer-skipping baselines on three question-answering benchmarks, then examines stage-wise budgets and generation early exit. It also diagnoses behavior on deep-reasoning tasks.

4.1 Comparison with Baselines

  • 1.49× speedup with 60.6% accuracy on OBQA, while AdaSkip falls to nearly 0% beyond 1.10×; LiteStage extends baseline performance boundaries across three datasets.AdaSkip maintains comparable accuracy only up to 1.10× speedup on OBQA.Source: Figure 4 · p. 6, §4.1 · p. 6

4.2 Benefits of Non-uniform Layer Budget

  • LiteStage caps skipped layers in sensitive Stage 3 at five for TinyLlama and three for Qwen, shifting budgets toward Stage 1 and then Stage 2 as greater speedup is sought.On OBQA, Stage 1 budgets increase from 5→21 for TinyLlama and 6→15 for Qwen; more aggressive settings shift budget toward Stage 2.Source: §4.2 · p. 7

4.3 Benefits of Generation Early Exit

  • Early exit counters end-to-end speedup saturation: LiteStage often exceeds per-token speedup, unlike baselines whose end-to-end speedup saturates or falls.Figure 6 compares per-token and end-to-end speedup against a y=x reference.Source: §4.3 · p. 7, Figure 6 · p. 8
  • As skipping increases, early exit shortens generated sequences and reduces latency proportionally; it can also improve accuracy when low-confidence redundant tokens hinder correct outputs.With 12 skipped TinyLlama sub-layers and 2 Qwen sub-layers, differences in accuracy, latency, and decoding steps are marginal.Source: §4.3 · pp. 7–8, Figure 7 · p. 8

4.4 Diagnostic Study on Deep Reasoning

  • On deep-reasoning tasks, combining layer skipping with early exit consistently raises speedups but causes substantial accuracy degradation even under moderate approximation, indicating a different efficiency–accuracy trade-off.The diagnostic study examines layer skipping with and without generation early exit.Source: §4.4 · p. 8

5 Conclusion

LiteStage combines stage-wise layer-budget optimization with confidence-based generation early exit to balance accuracy and latency in multi-stage reasoning. On OBQA, CSQA, and StrategyQA, it surpasses prior training-free methods with a large speedup margin.

5 Conclusion

  • LiteStage jointly optimizes stage-wise layer budgets and applies confidence-based generation early exit to balance accuracy and latency.Experiments on OBQA, CSQA, and StrategyQA show it surpasses prior training-free methods with a large speedup margin.Source: §5 · p. 8

6 Limitations

The paper identifies profiling overhead, a simple confidence-based exit criterion, and limited focus on short multi-stage reasoning as limitations. Its appendix outlines additional implementation, ablation, and deep-reasoning analyses.

6 Limitations

  • LiteStage uses one-time profiling that may add overhead versus fixed heuristics, a simple confidence exit criterion, and focuses mainly on short reasoning; extending it to long-horizon tasks needs investigation.Profiling is performed once per model and setting and is amortized over repeated inference.Source: §6 · p. 9

A Overview

  • The appendix outlines supplementary stage-wise decoding statistics, implementation and training details, sublayer-importance analyses, ablations, and deep-reasoning experiments.Source: A Overview · p. 11

B.1 Data Statistics

Full-layer models show an imbalance in decoding length across the three reasoning stages.

B.1 Data Statistics

  • Across OBQA, CSQA, and StrategyQA, Analyze has the most decoding steps, Summarize the fewest, and Recall falls between them.Summarize typically follows a fixed answer template; the statistics use full-layer models without layer skipping or early exit.Source: §B.1 · p. 11

B.2 Training Setups

The training setup fine-tunes small models on structured reasoning paths, and reports averaged decoding-step counts for each reasoning stage across three datasets.

B.2 Training Setups

  • TinyLlama-1.1B and Qwen2.5-0.5B are fine-tuned on OBQA, CSQA, and StrategyQA using reasoning paths generated by larger models under the TinyThinker protocol.The reasoning paths provide a form of knowledge distillation from GPT to the smaller models.Source: §B.2 · p. 11
  • Table 4 estimates and averages test-set decoding steps separately for Recall, Analyze, and Summarize across OBQA, CSQA, and StrategyQA for both models.The supplied passage describes the table scope but does not provide its cell values.Source: Table 4 · p. 11
  • Fine-tuning updates all model parameters with AdamW on a single NVIDIA A6000 GPU, using weight decay 0.0, β1=0.9, β2=0.999, and ϵ=10−8.The authors use hyperparameters adapted for decoder-only architectures rather than TinyThinker’s original T5 setup.Source: §B.2 · p. 11

B.3 Training Dynamics

The appendix reports test accuracy and training behavior for two fine-tuned models, illustrates generation early exit on a CSQA example, and analyzes sub-layer importance across datasets and architectures.

B.3 Training Dynamics

  • T5-Large reports 65.4% CSQA, 68.8% OBQA, and 69.0% StrategyQA, compared with TinyLlama-1.1B’s 54.8%, 64.0%, and 62.4%, respectively.T5-Large is reported as the largest TinyThinker model, with 770M parameters.Source: §B.3 · pp. 11–12
  • Qwen2.5-0.5B scores 64.4% CSQA, 70.2% OBQA, and 65.5% StrategyQA, exceeding TinyLlama-1.1B on all three datasets.TinyLlama-1.1B scores 54.8% CSQA, 64.0% OBQA, and 62.4% StrategyQA.Source: §B.3 · p. 12
  • Across datasets, training loss decreases smoothly while validation accuracy saturates early; Figure 8 plots both measures and marks selected evaluation checkpoints.Validation accuracy is lower than test accuracy because test results use ten-generation self-consistency while validation uses single-pass decoding.Source: §B.3 · p. 12, Figure 8 · p. 12

C.1 Example of Generation Early Exit

  • In the CSQA example, early exit truncates the low-confidence suffix after “visual representation,” while both models retain the same final answer, rejecting C and selecting D.Tokens remain identical to the baseline up to the exit point.Source: §C.1 · pp. 12–13

C.2 Layer Importance Estimation

  • Across datasets, sub-layer importance ordering is largely stable, but the balance between MHSA and FFN importance shifts with dataset and model family.Figure 9 reports cosine-similarity-based importance estimates separately for MHSA and FFN sub-layers.Source: §C.2 · p. 13, Figure 9 · p. 13
  • For TinyLlama-1.1B on StrategyQA, greater FFN importance leads the skipping policy to prioritize MHSA while preserving FFN computation.This differs from Qwen2.5-0.5B, where MHSA is more important than FFN across all three datasets, so skipping is concentrated on MHSA.Source: §C.2 · p. 13
  • The authors attribute sub-layer importance trends primarily to model architecture and use them to motivate separate MHSA and FFN skip decisions.Dataset-level shifts in the balance between sub-layer types remain relevant to the computation budget.Source: §C.2 · p. 13

C.3 Ablation Studies

The ablations show that practical acceleration depends on profiling stage-specific accuracy–latency behavior and pairing layer skipping with generation early exit.

C.3 Ablation Studies

  • The experiments vary skipped sub-layers in one reasoning stage at a time, measure validation accuracy and normalized end-to-end latency, and compare configurations with and without early exit.The ablations cover TinyLlama-1.1B and Qwen2.5-0.5B on OBQA, CSQA, and StrategyQA.Source: §C.3 · pp. 13–14
  • Across datasets, Stage 1 is generally most robust to skipping, Stage 2 next, and Stage 3 most sensitive, but sensitivity varies substantially by model and dataset.On OBQA Stage 2, Qwen2.5-0.5B is markedly more accuracy-robust than TinyLlama-1.1B; static cosine similarity alone does not reliably reveal such differences.Source: §C.3 · p. 14
  • Accuracy and latency do not change monotonically with skipped sub-layer count, so LiteStage searches the observed accuracy–latency profile for a layer budget.The ablations include cases where skipping more layers raises validation accuracy or skipping fewer layers yields higher speedup.Source: §C.3 · p. 14
  • Without early exit, aggressive skipping can lengthen generation enough to negate or reverse latency gains, especially in Stage 2, which typically has the longest generation.Across the ablations, early exit keeps normalized latency stable or lower even under aggressive skipping.Source: §C.3 · p. 14
  • With generation early exit, normalized latency remains stable or decreases under aggressive skipping across both models, all three datasets, and all reasoning stages.The ablations find that early exit suppresses latency growth across these settings.Source: §C.3 · p. 14

D Diagnostic Study on Deep Reasoning

The diagnostic study finds that combining layer skipping, periodic full decoding, and generation early exit can improve deep-reasoning speed, but approximation still carries substantial accuracy costs. It also records the work’s ethics and language-model-use statements.

D Diagnostic Study on Deep Reasoning

  • Without validation splits for the deep-reasoning benchmarks, the study used uniform skipping for a single-stage process and raised the generation early-exit confidence threshold to 0.8.Layer importance was estimated using OBQA, CSQA, and StrategyQA.Source: D Diagnostic Study on Deep Reasoning · pp. 14–15
  • 1.40× on AIME 25, 1.84× on GPQA-D, and 1.46× on LCB-v5 are the reported speedups for LS+PF+GE; generation early exit improves efficiency–accuracy trade-offs.The comparison uses Qwen3-1.7B and reports end-to-end speedup normalized against full-layer decoding.Source: Table 5 · p. 15, D Diagnostic Study on Deep Reasoning · p. 16, Figure 10 · p. 16
  • Skipping two layers reduced accuracy from 30.0% to 13.3% on AIME 25 and 34.8% to 11.8% on LCB-v5; one skipped layer reduced GPQA-D from 46.0% to 32.8%.Skipping applies during decoding, except in the first and last four layers; prefill remains full-depth.Source: D Diagnostic Study on Deep Reasoning · p. 15
  • Long-horizon intermediate states make deep reasoning vulnerable to accumulating layer-skipping errors, while multi-stage prefill reprocesses prior-stage output at full depth.The authors contrast deep reasoning’s extended decoding with relatively independent stages in short multi-stage reasoning.Source: D Diagnostic Study on Deep Reasoning · p. 15
  • At # Skip: 3, periodic full decoding improved accuracy from 0.0% to 20.0% on AIME 25, 20.2% to 32.8% on GPQA-D, and 0.0% to 25.5% on LCB-v5.The heuristic alternates 1,000 full-layer tokens with 1,000 skipped-layer tokens throughout generation.Source: D Diagnostic Study on Deep Reasoning · p. 15
  • Layer skipping with periodic full decoding still reached only 1.15× speedup at best, on GPQA-D at # Skip: 3, as longer generations limited end-to-end gains.All evaluated configurations showed noticeable accuracy degradation, even under moderate approximation.Source: D Diagnostic Study on Deep Reasoning · pp. 15–16
  • The study characterizes stage-aware skipping as suited to short multi-stage reasoning, while deep reasoning’s long-horizon consistency demands remain a mismatch beyond this work’s scope.The authors suggest periodic full decoding as a complementary mechanism for deep reasoning.Source: D Diagnostic Study on Deep Reasoning · p. 16

E Ethics Statement

  • The authors report no anticipated additional risks beyond those associated with base models and say language models were used only for editorial refinement.They state that LLMs did not contribute to research ideas, methods, analyses, or scientific content.Source: E Ethics Statement · p. 16
Loading 2510.14211v2…