🤖 machine learning

TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference

TIDE is a post-training system that enables per-token early exit in large language models by attaching lightweight learned routers to checkpoint layers, allowing tokens to skip redundant computations once their hidden states converge without requiring model retraining.

Original authors: Jaber Jaber, Osama Jaber

Published 2026-03-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jaber Jaber, Osama Jaber

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 running a massive, high-end factory that produces custom stories, one word at a time. This factory has 32 (or even 80) different workstations, or "layers," that every single word must pass through.

In a standard Large Language Model (LLM), every word gets the full VIP treatment.

The word "the" (a boring, common word) goes through all 32 stations, getting polished, analyzed, and processed just like a complex math equation or a profound philosophical insight.
It's like sending a simple "Hello" through a 32-step security check, a 32-step quality control inspection, and a 32-step gourmet cooking process. It's safe, but it's incredibly wasteful. You are spending a fortune in electricity and time on words that don't need it.

Enter TIDE: The "Smart Bouncer"

TIDE (Token-Informed Depth Execution) is a new system that acts like a super-smart bouncer for this factory. It doesn't stop the words from entering the factory, but it decides exactly when they can leave early.

Here is how it works, using simple analogies:

1. The "Practice Run" (Calibration)

Before the factory opens for business, TIDE does a quick "rehearsal." It feeds the model 2,000 simple sentences (like a Wikipedia article) and watches how the words change as they move through the layers.

It notices that for the word "the," the meaning stops changing after Layer 8. By Layer 9, it's just the same word again.
For a complex math word, the meaning keeps evolving all the way to Layer 31.
TIDE trains tiny, lightweight "routers" (like little decision-making chips) to recognize this pattern. It takes only 3 minutes to learn this, and the resulting "rulebook" is tiny (only 4 MB).

2. The "Post-Game Review" (Inference)

This is the clever part. TIDE doesn't stop the word from going through all 32 layers while the word is being processed. Why? Because stopping early breaks the factory's memory (the "KV Cache"), which causes errors.

Instead, TIDE lets the word go through all 32 layers just like normal, but it does something special in the background:

As the word passes each checkpoint, the tiny router whispers, "Hey, this word is done! It's fully formed!"
Once the word finishes the full 32 layers, the system looks back at the router's notes.
If the router said, "This word was ready at Layer 8," the system ignores the work done from Layer 9 to 32 for that specific word. It grabs the result from Layer 8 and moves on.

Think of it like a student taking a 32-question exam:

Old Way: The student answers all 32 questions, even if they knew the answer to question 1 immediately.
TIDE Way: The student answers all 32 questions (to keep the test flow smooth), but the teacher grades them based on the earliest question where the student got it right. If they got question 1 right, the teacher ignores the rest of the paper for that specific answer, saving time on grading.

3. The Results: Speed and Savings

Because TIDE realizes that most words (like "the," "is," "and") are simple and don't need deep thinking, it cuts out the unnecessary work.

For simple words: They might exit at Layer 11, skipping 21 layers of work.
For hard words: They go all the way to the end.
The Math: In a test with a math problem, 99% of the words were simple enough to exit early, yet the model still solved the complex problem perfectly.

Why is this a big deal?

No Re-training: You don't need to rebuild the factory or retrain the workers. TIDE just adds a small "add-on" to existing models (like those from Hugging Face).
Works on Any Model: Whether it's a Llama, a Qwen, or a GPT, TIDE can figure out how to plug in.
Faster & Cheaper: By skipping the extra work for simple words, the factory runs faster (lower latency) and can handle more customers at once (higher throughput). On powerful GPUs, this means generating text up to 8% faster with less energy.

The Catch (Limitations)

TIDE is currently a "post-game review" system. It still runs the word through all the layers physically, just to be safe with the memory. It doesn't physically stop the word from entering the later rooms yet.

Analogy: It's like a bouncer who lets everyone into the club, but then checks their ID at the door. If they are on the list, he tells the DJ to stop playing music for them. It saves energy on the music, but they still walked through the door.
Future: The next step is to make the bouncer stop them before they even enter the later rooms, which would save even more time.

Summary

TIDE is a smart efficiency tool. It realizes that not every word needs a PhD-level analysis. By attaching tiny, learned "smart sensors" to existing AI models, it lets simple words leave the factory early while letting complex words stay for the full treatment. The result? Faster AI, lower costs, and the same high-quality answers.

1. Problem Statement

Current Large Language Models (LLMs) operate under a uniform compute allocation strategy: every token, regardless of its semantic complexity or difficulty, passes through every layer of the network.

Inefficiency: Simple tokens (e.g., function words like "the") receive the same 32-layer processing as complex reasoning steps, leading to significant computational waste.
Limitations of Existing Solutions:
- Encoder-only methods (e.g., DeeBERT): Designed for classification, they do not generalize to autoregressive generation with KV caches.
- Retraining-based methods (e.g., LayerSkip, EE-LLM): Require access to the training pipeline and hundreds of GPU-hours, making them impractical for users of pre-trained checkpoints.
- Confidence-based heuristics: Relying on softmax entropy is unreliable for generation tasks where entropy is naturally high.
Goal: Develop a system that enables per-token early exit in autoregressive LLMs without retraining the base model, preserving KV cache integrity, and reducing latency/energy consumption.

2. Methodology: TIDE

TIDE is a post-training system that attaches lightweight learned routers to periodic checkpoint layers of a frozen model. It operates in two stages: Offline Calibration and Online Inference.

A. Offline Calibration (One-time)

Data Collection: The frozen model processes 2,000 samples from WikiText-103 with output_hidden_states=True.
Convergence Labeling: For each token $i$ at checkpoint layer $k$ , the system computes the cosine similarity between the hidden state $h_k$ and the final layer hidden state $h_L$ :
$s_k^{(i)} = \frac{h_k^{(i)} \cdot h_L^{(i)}}{\|h_k^{(i)}\| \|h_L^{(i)}\|}$
A binary label is assigned ( $y=1$ ) if similarity exceeds a threshold $\tau$ (default 0.98), indicating the token has "converged."
Router Training: A lightweight binary classifier (router) $\phi_k$ $ϕ_{k}$ is trained for each checkpoint layer.
- Architecture: A 2-layer MLP with a bottleneck dimension of 128 ( $d \times 128 + 128$ parameters).
- Training: 100 epochs of binary cross-entropy loss on a single GPU (completes in <3 minutes).
- Output: A tiny checkpoint file (~4 MB).

B. Online Inference (Per Request)

Full Forward Pass: The model runs a complete forward pass through all $L$ layers for every token, ensuring the KV cache remains fully populated and intact.
Post-Hoc Evaluation: After the forward pass, the trained routers evaluate the hidden states at checkpoint layers.
Exit Selection: The system selects the earliest layer $k$ where the router score exceeds a threshold $\theta$ .
Logit Computation: The final logits are computed using the hidden state from the selected exit layer (via LMHead), skipping the remaining layers for that specific token.

C. Technical Implementation Details

Universal Adapter: A UniversalAdapter probes 17 attribute paths to automatically detect transformer structures (LLaMA, GPT-2, Qwen, Phi, etc.) without manual code adjustments.
Fused CUDA Kernels: To minimize overhead, TIDE uses fused kernels for:
- RMSNorm + Router evaluation (single launch).
- Batch compacting and scatter operations for exiting tokens.
- Support for float32, float16, and bfloat16.
Compatibility: Works with any HuggingFace causal LM and is compatible with transformers v5.3+ (handling decorators that intercept exceptions).

3. Key Contributions

Post-Training Early Exit: The first system to enable per-token early exit for autoregressive LLMs without modifying the base model weights or requiring retraining.
Universal Architecture Support: An adapter that supports 17+ model architectures (LLaMA, Mistral, Qwen, etc.) out of the box.
KV Cache Integrity: Unlike "skip" methods that break cache continuity, TIDE runs all layers but selects the output layer post-hoc, ensuring compatibility with standard inference pipelines.
High-Performance Kernels: Custom fused CUDA kernels for router evaluation and data movement, supporting GPU architectures from V100 to Blackwell.
Open Source: A complete, tested package (tide-inference) with 74 passing tests, available on PyPI and GitHub.

4. Experimental Results

Evaluated on an NVIDIA A100 (bf16) using DeepSeek R1 Distill 8B (32 layers) and Qwen3 8B (36 layers).

Exit Rates:
- Prefill: 100% of tokens find an exit point. On DeepSeek R1, 5% exit at layer 11, while the rest exit at layer 31.
- Decoding: 98–99% of tokens exit early (mostly at the penultimate layer).
Performance Gains:
- Latency: Reduced prefill latency by 5.5% – 7.2%.
- Throughput:
  - DeepSeek R1: +6.6% at batch size 1.
  - Qwen3: +8.1% at batch size 8.
- Note: At high batch sizes (BS=8) for DeepSeek, throughput dropped (-16.3%) due to the output_hidden_states overhead growing superlinearly, highlighting a current limitation.
Quality Preservation:
- On a multi-step math problem, TIDE maintained 100% correctness with 95 unique output tokens (vs. 99 for baseline) while 99.6% of tokens exited early.
Calibration Efficiency: Takes <3 minutes on 2,000 samples to produce a ~4 MB router file.

5. Significance and Future Work

Significance:
TIDE bridges the gap between theoretical efficiency and practical deployment. It proves that token-level convergence is a learnable signal that can be exploited post-training to reduce inference costs without sacrificing model quality or breaking existing inference infrastructure (KV caches). It offers a "drop-in" solution for users of pre-trained models.

Limitations & Future Directions:

Wall-Clock Skipping: Currently, TIDE runs all layers physically to preserve the cache. Future work aims to implement "true" skipping (physically bypassing layers) by managing cache discontinuities.
Threshold Tuning: The current convergence threshold ( $\tau=0.98$ ) is conservative, causing most exits to cluster at the penultimate layer. Dynamic or per-layer thresholds could enable deeper early exits.
Batch Scaling: The overhead of collecting hidden states for all layers becomes a bottleneck at large batch sizes; a hybrid approach collecting states only at checkpoints is proposed.
Deep Models: While tested on 8B models, the potential for greater savings in 70B+ models (with 80+ layers) remains to be empirically validated.