Two-dimensional early exit optimisation of LLM inference
Authors: Jan Hůla, David Adamczyk, Tomáš Filip, Martin Pavlíček, Petr Sosík
Organizations: Institute for Research and Applications of Fuzzy Modelling, University of Ostrava, 70200 Ostrava, Czech Republic · Institute of Computer Science, Faculty of Philosophy and Science, Silesian University in Opava, 74601 Opava, Czech Republic
We introduce a two-dimensional (2D) early exit strategy that coordinates layer-wise and sentence-wise exiting for classification tasks in large language models. By processing input incrementally sentence-by-sentence while progressively activating deeper layers, our method achieves multiplicative computational savings that exceed those from optimizing either dimension independently. Experimental evaluation across four state-of-the-art LLMs (Llama 3.1, Llama 3.2, Gemma, Qwen; 3B-8B parameters) on three sentiment classification datasets demonstrates additional speed-ups of 1.4--2.3× over optimal layer-wise early exit for simpler tasks with vanilla models, with graceful degradation on complex multi-class problems. Fine-tuning reduces but does not eliminate this advantage. The approach is model-agnostic, requires only lightweight classification adapters, and is orthogonal to complementary efficiency methods such as quantization and pruning. Our findings indicate that 2D early exit strategies excel when semantic information accumulates predictably across input structure, suggesting possible applicability to sequence-processing tasks beyond sentiment classification.
Figures & tables
Figure 1: Layer-wise and sentence-wise accuracy of correct classification of Gemma-3n-E4B vanilla (top) and fine-tuned (bottom) on the MMS dataset. Heatmaps based on 7100 samples of 10 sentences each.
Figure 2: Visualization of the 2D early exit strategy for step size Δ=2, see eq. ( 4 ). Arrows Red frames enclose inference progression blocks corresponding to individual sentences.
Input:
example: input text
τignore,τacc: hyperparameters – ignore and acceptance threshold
Output:
predicted_label: prediction for the input example
operations_used: consumed abstract operations to measure speed-up
Algorithm 1 The 2D early exit inference for one input example.
Figure 3: Layer-wise accuracy of LLama-3.1-8B, vanilla model (above) and fine-tuned model (bellow), averaged on the MMS test dataset. Full circles mark the layers with the best accuracy. Stars mark the first layers with the accuracy 2% below maximum.
Model
ModernBERT
Llama-3.1
Llama-3.2
Gemma-3n
Qwen2.5
Dataset
Large
8B
3B
E4B
7B
Steam
0.87
0.97
0.96
0.96
0.96
MMS
0.78
0.80
0.79
0.78
0.79
Amazon-5
0.62
0.65
0.63
0.63
0.64
Table 1: Accuracy of the fine-tuned ModernBERT Large versus other models tested.
Accuracy loss 2%
Accuracy loss 4%
Dataset
Model
Acc
layer-wise EE
2D EE
layer-wise EE
2D EE
Layer
Spd-up
Spd-up
Layer
Spd-up
Spd-up
Gemma-3n-E4B
0.94
14
2.3
4.2
11
2.9
6.4
Llama-3.1-8B
0.95
9
3.2
5.8
7
4.0
9.2
Steam
Llama-3.2-3B
0.94
9
2.8
5.3
7
3.5
7.5
Qwen2.5-7B
0.95
12
2.2
3.1
10
2.5
5.1
Table 2: Results of the 2D early exit strategy in vanilla models. The cases when 2D EE did not reach the 2% accuracy loss threshold are marked by *.
Accuracy loss 2%
Accuracy loss 4%
Dataset
Model
Acc
layer-wise EE
2D EE
layer-wise EE
2D EE
Layer
Spd-up
Spd-up
Layer
Spd-up
Spd-up
Gemma-3n-E4B
0.96
13
2.5
3.2
12
2.7
4.2
Llama-3.1-8B
0.97
9
3.2
4.3
7
4.0
6.9
Steam
Llama-3.2-3B
0.96
8
3.1
4.1
7
3.5
5.8
Qwen2.5-7B
0.96
10
2.5
3.2
9
2.8
4.5
Table 3: Results of the 2D early exit strategy in fine-tuned models.
Accuracy loss 2%
Accuracy loss 4%
Dataset
Model
Acc
LayerSkip
2D EE
LayerSkip
2D EE
Layer
Spd-up
Spd-up
Layer
Spd-up
Spd-up
Steam
Llama-3.1-8B
0.97
11
2.9
4.3
10
3.2
6.9
Llama-3.2-3B
0.96
11
2.5
4.1
10
2.8
5.8
MMS
Llama-3.1-8B
0.80
15
2.1
2.1
15
2.1
3.2
Llama-3.2-3B
0.79
15
1.9
2.1
15
1.9
2.9
Table 4: Results of LayerSkip compared to the 2D EE technique.
Configuration
Max. accuracy
Speed-up
(i) Adapter + Layer-wise EE
0.74
2.3 ×
(ii) Adapter + 2D EE
0.74
4.6 ×
(iii) Fine-tuned + Layer-wise EE
0.78
1.8 ×
(iii) Fine-tuned + 2D EE
0.78
1.9 ×
Table 5: Ablation study results on Gemma-3n-E4B with MMS dataset (accuracy loss of 2%).
Figure 4: Average block accuracies of 2D early exit for Gemma-3n-E4B on the MMS sub-dataset of samples with 10 sentences.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Cosine distances of embedding vectors between consecutive layers of Llama-3.1-8B, averaged over the MMS testing split.
Figure 6: Accuracy and speedup heatmap of the 2D early exit strategy for Llama-3.1-8B on the MMS dataset, no fine-tuning.
LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.
In Large Language Model (LLM) inference, early-exit refers to stopping computation at an intermediate layer once the prediction is sufficiently confident, thereby reducing latency and cost. However, recent LLMs adopt improved pretraining recipes and architectures that reduce layer redundancy, potentially limiting early-exit opportunities. We re-evaluate layer-wise early-exit in modern LLMs and analyze how intermediate representations evolve during training. We introduce a metric to quantify a model's intrinsic suitability for early-exit and propose a benchmark for researchers to explore the potential early-exit benefits on different models and workloads. Our results show a diminishing trend in early-exit effectiveness across newer model generations. We further find that dense transformers generally offer greater early-exit potential than Mixture-of-Experts and State Space Models. In addition, larger models, particularly those with more than 20 billion parameters, and base pretrained models without specialized tuning tend to exhibit higher early-exit potential.
Rui Wei, Rui Du, Hanfei Yu +4
Stevens Institute of Technology · Northeastern University · Stony Brook University
Large Language Models (LLMs) have demonstrated exceptional performance across diverse domains but are increasingly constrained by high inference latency. Early Exit has emerged as a promising solution to accelerate inference by dynamically bypassing redundant layers. However, in decoder-only architectures, the efficiency of Early Exit is severely bottlenecked by the KV Cache Absence problem, where skipped layers fail to provide the necessary historical states for subsequent tokens. Existing solutions, such as recomputation or masking, either introduce significant latency overhead or incur severe precision loss, failing to bridge the gap between theoretical layer reduction and practical wall-clock speedup. In this paper, we propose River-LLM, a training-free framework that enables seamless token-level Early Exit. River-LLM introduces a lightweight KV-Shared Exit River that allows the backbone's missing KV cache to be naturally generated and preserved during the exit process, eliminating the need for costly recovery operations. Furthermore, we utilize state transition similarity within decoder blocks to predict cumulative KV errors and guide precise exit decisions. Extensive experiments on mathematical reasoning and code generation tasks demonstrate that River-LLM achieves 1.53 to 2.16 times of practical speedup while maintaining high generation quality.
Large language model (LLM) inference is often constrained by both computation and memory, especially in offloading-based deployments where model weights are transferred across memory hierarchies during autoregressive decoding. In this setting, reducing the number of executed layers can lower per-token latency while also avoiding costly weight movement. Motivated by this observation, we present FlexEE, an early exiting framework for resource-constrained and offloading-based LLM inference. FlexEE makes early exiting practical for LLM decoding through layer-wise exit supervision for reliable intermediate-layer prediction, self-speculative decoding over a Top-K local vocabulary for low-cost exit decisions, and dynamic hidden state management for KV-cache-correct and memory-aware execution. Across generative and downstream tasks, FlexEE enables efficient early exit with minimal accuracy degradation, delivering up to 1.27×/3.16× and 1.25×/2.83× end-to-end speedups on Llama2-7B and Llama3-8B under 0%/50% weight offloading, respectively.
Qihu Xie, Ziwei Li, Yi Kang
University of Science and Technology of China, Hefei, China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center, Hefei, China