💬 NLP

Two-dimensional early exit optimisation of LLM inference

This paper introduces a model-agnostic two-dimensional early exit strategy that coordinates layer-wise and sentence-wise processing to achieve multiplicative computational savings and significant speed-ups in LLM inference for classification tasks, particularly when semantic information accumulates predictably across input structures.

Original authors: Jan Hůla, David Adamczyk, Tomáš Filip, Martin Pavlíček, Petr Sosík

Published 2026-04-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jan Hůla, David Adamczyk, Tomáš Filip, Martin Pavlíček, Petr Sosík

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a very smart, but incredibly slow, librarian. You have been hired to read thousands of customer reviews and decide if they are "Happy," "Neutral," or "Angry."

The Problem:
Normally, this librarian reads every single word of every review, from the first page to the very last page, and then reads the whole thing again using their deepest, most complex thinking skills before making a decision. Even if the reviewer says "I hate this product!" in the very first sentence, the librarian still reads the entire 50-page document and thinks about it for an hour before saying, "Yes, that's anger." This wastes a massive amount of time and energy.

The Old Solution (Layer-wise Early Exit):
Researchers previously figured out a way to speed this up. They told the librarian: "If you are 90% sure after reading just the first half of the book, you can stop reading and give your answer." This helps, but it's still a bit rigid. The librarian still has to read the entire text before they can even start thinking about stopping.

The New Solution: The "2D Early Exit" Strategy
This paper introduces a brilliant new way to organize the librarian's work. Instead of just reading the whole book and then deciding when to stop, the librarian now works in two dimensions at once: Time (Sentences) and Depth (Thinking Power).

Here is how it works, using a simple analogy:

The "Growing Team" Analogy

Imagine the librarian doesn't work alone. They have a team of assistants, ranging from a Junior Intern (who thinks fast but simply) to a Senior Professor (who thinks slowly but deeply).

In the old way, the whole team would read the whole book together, and then the Senior Professor would decide if they were done.

In this 2D Strategy, the process is dynamic:

Sentence 1 arrives: The Junior Intern reads just the first sentence. If the sentence is "This is the worst movie ever!", the Intern is already 95% sure it's a negative review. They shout, "I'm done!" and the process stops immediately. Result: Super fast.
Sentence 1 & 2 arrive: If the first sentence was vague ("I bought a new phone..."), the Intern isn't sure. So, the Senior Professor joins in to read the first two sentences together. If the second sentence is "It broke after one day!", the Professor is now 95% sure. They shout, "Stop!" Result: Fast, but slightly slower than the first case.
Sentence 1, 2, & 3 arrive: If the text is confusing, the team keeps growing. The Intern reads the first sentence, the Professor reads the first two, and a Middle Manager reads the first three. They combine their thoughts. As soon as the accumulated confidence of the whole team hits a high enough level, they stop.

The "Two-Dimensional" Magic

The paper calls this "Two-Dimensional" because the librarian is optimizing two things simultaneously:

Dimension 1 (Horizontal): They stop reading new sentences as soon as they have enough info. They don't waste time reading the rest of the review.
Dimension 2 (Vertical): They don't use their most expensive, slowest "Senior Professor" brain for the whole text. They only "wake up" the deep thinkers when the simpler thinkers aren't enough.

The Result:
For simple reviews, the librarian stops almost immediately, saving 2 to 3 times the time and energy compared to the old methods. It's like realizing you don't need to drive to the store to buy milk because you saw the store was closed from your front porch.

What Happens with "Fine-Tuning"?

The researchers also tried "training" the librarian to be a perfect expert (Fine-Tuning). Surprisingly, this made the 2D strategy less effective.

Why? When you train a model to be perfect, it tends to become "boringly consistent." It stops having those big, obvious "aha!" moments in the early sentences. It waits until the very end to be sure, just like the old slow librarian.
The Lesson: Sometimes, being a "good enough" generalist who makes quick guesses is actually faster and more efficient for this specific type of task than being a perfectionist who waits for all the data.

Why Should You Care?

Large Language Models (like the ones powering chatbots) are currently very expensive to run. They use huge amounts of electricity and computer power.

This method allows companies to run these models much faster and cheaper for tasks like checking if a tweet is positive or negative.
It doesn't require building new, weird hardware. It just changes how the software reads the text.
It works best on tasks where the answer becomes obvious quickly (like sentiment analysis), but might be less useful for complex math problems where you really do need to read every step.

In a nutshell: This paper teaches AI how to stop reading when it already knows the answer, and how to use a "lightweight" brain for easy parts of the text, saving massive amounts of time and energy.

1. Problem Statement

Large Language Models (LLMs) offer state-of-the-art performance but incur significant computational costs in terms of processing power, memory, and energy, particularly during the inference phase. While standard inference processes every input through all layers of the model, many classification tasks (e.g., sentiment analysis) allow the model to reach a confident prediction before processing the entire text or traversing all layers.

Existing Early Exit (EE) strategies typically optimize only one dimension:

Layer-wise EE: Stops computation at an intermediate layer once confidence is high.
Token/Sequence-wise EE: Stops processing specific tokens or sentences early.

The paper argues that optimizing these dimensions independently yields only additive savings. The authors propose that a Two-Dimensional (2D) Early Exit strategy, which coordinates layer-wise and sentence-wise exiting simultaneously, can achieve multiplicative computational savings, especially when semantic information accumulates predictably across the input structure.

2. Methodology

The proposed 2D Early Exit strategy treats the input as a sequence of semantic units (sentences) and processes them incrementally while progressively deepening the network depth.

Core Mechanism

Input Splitting: Input text is split into sentences.
2D Progression: The algorithm processes sentences one by one. For the $s$ $s$ -th sentence, it activates the first $L_s$ $L_{s}$ layers of the model, where $L_s$ $L_{s}$ increases with the sentence index.
- Formula: $L_s = \min((s + 1) \times \Delta, L)$ , where $\Delta$ is a step size based on total layers and sentence count.
Block Processing: As a new sentence is introduced, it is processed through all active layers ($0$ to $L_s$ ). Previous sentences are re-processed only through the newly activated layers ( $L_{s-1}$ to $L_s$ ). This creates a "rectangular" processing block that grows in both dimensions (layers vs. sentences).
Classification Adapters: Lightweight two-layer fully connected networks (Linear $\to$ ReLU $\to$ Linear $\to$ Softmax) are attached to each transformer layer to generate class probabilities.

Training Strategies

The paper evaluates two training modes:

Adapter-only Training: The base LLM weights are frozen; only the classification adapters are trained. The loss function encourages accurate predictions at intermediate layers based on the full input (aggregated sentence embeddings).
Fine-tuning: The entire model (including adapters) is fine-tuned using a weighted sum of losses across all layers, with higher weights assigned to the final layer.

Inference Logic (Algorithm 1)

Confidence Metric: Confidence is calculated as the difference between the top two softmax probabilities ( $max - second\_max$ ).
Thresholds:
- $\tau_{ignore}$ : If confidence exceeds this, it is added to an accumulator for the predicted class.
- $\tau_{acc}$ : If the accumulated confidence for any class exceeds this threshold, inference terminates immediately, returning the prediction.
Fallback: If no threshold is met after processing all sentences and layers, the final layer's prediction for the last sentence is returned.

3. Key Contributions

Methodological Innovation: The integration of layer-wise early stopping with sentence-wise input trimming into a single synergistic framework. This achieves multiplicative rather than additive computational savings.
Comprehensive Benchmarking: Evaluation across four state-of-the-art LLMs (Llama 3.1/3.2, Gemma, Qwen) ranging from 3B to 8B parameters on three sentiment datasets (Steam, MMS, Amazon-5) with varying complexity (2, 3, and 5 classes).
Empirical Insights:
- Demonstration that 2D EE excels when semantic signals accumulate predictably across sentences.
- Identification of a paradox where fine-tuning, while improving overall accuracy, reduces the 2D advantage by flattening accuracy variations across early layers and partial inputs.
Model Agnosticism: The approach requires only lightweight adapters and is orthogonal to other efficiency methods like quantization, pruning, and Mixture of Experts (MoE).

4. Experimental Results

The study compared 2D EE against optimal layer-wise early exit and the LayerSkip method.

Performance on Vanilla Models (Adapter-only):
- Simple Tasks (2-3 classes): 2D EE achieved 1.4× to 2.3× additional speed-up compared to optimal layer-wise early exit.
- Complex Tasks (5 classes): On the Amazon-5 dataset, the advantage was smaller (1.1×–1.5×) and sometimes negative for Llama models, indicating that 2D strategies struggle when semantic convergence is less predictable.
Performance on Fine-tuned Models:
- Fine-tuning reduced the 2D advantage significantly (e.g., dropping from ~2.3× to ~1.7× on simple tasks).
- Ablation Study: Fine-tuning flattened the accuracy landscape across layers and sentences. While it improved final accuracy, it removed the "steep" accuracy gradients in early layers that allow early exits to be effective.
Comparison with LayerSkip:
- 2D EE outperformed LayerSkip on simpler datasets (Steam, MMS) but was comparable or slightly inferior on the most complex dataset (Amazon-5) with strict accuracy constraints.

5. Significance and Implications

Efficiency: The 2D strategy offers a new paradigm for inference optimization, particularly for high-throughput classification tasks where inputs can be segmented into semantic units.
Trade-off Analysis: The paper highlights a critical trade-off: Fine-tuning for maximum accuracy may inadvertently penalize early-exit efficiency by making the model's internal representations more uniform across layers. This suggests a need for specialized training objectives that preserve "early-exit-friendly" gradients.
Future Directions: The authors suggest that future work should focus on:
- Developing fine-tuning strategies that maintain accuracy gradients across dimensions.
- Extending the approach to other sequence-processing tasks (e.g., intent detection, QA).
- Combining 2D EE with orthogonal techniques (quantization, pruning) for compounded efficiency gains.

In conclusion, the paper demonstrates that coordinating computation across both the depth of the network and the breadth of the input sequence can yield substantial efficiency gains, provided the task allows for predictable semantic accumulation and the model is not over-optimized for final-layer accuracy at the expense of intermediate confidence.