🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

The paper introduces Token-Selective Attention (TSA), a lightweight, end-to-end differentiable gating mechanism that enables transformers to dynamically skip redundant layer computations for individual tokens based on contextual difficulty, achieving significant efficiency gains with minimal quality loss and no need for explicit depth regularization.

Original authors: Ahmed Abdelmuniem Abdalla Mohammed

Published 2026-05-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ahmed Abdelmuniem Abdalla Mohammed

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 factory where every single product that comes down the assembly line gets the exact same treatment. Whether it's a simple, perfect widget or a complex, broken one, every item spends exactly the same amount of time at every single station. This is how standard AI models (Transformers) currently work: they process every word in a sentence through the exact same number of "thinking layers," regardless of how easy or hard that word is to understand.

This paper introduces a new system called Token-Selective Attention (TSA). Think of TSA as installing a smart, automatic gatekeeper at every single station in the factory.

The Problem: The "One-Size-Fits-All" Factory

In a standard AI model, if you ask it to write "To be, or not to be," the model treats the word "To" and the word "question" with the same amount of brainpower.

Easy words (like "the" or "is") are like simple widgets. They don't need much processing.
Hard words (like complex concepts or rare names) are like broken widgets. They need extra time and attention to fix.

Currently, the factory wastes energy processing the easy widgets with the same intensity as the hard ones.

The Solution: The Smart Gatekeeper (TSA)

The authors added a tiny, lightweight "gatekeeper" (a small neural network) between every layer of the AI. This gatekeeper looks at each word (token) individually and asks: "Does this word need to go through the next processing station, or can it skip it?"

The Decision: The gatekeeper doesn't just say "Yes" or "No." It gives a probability score (like a dimmer switch). If a word is easy, the gatekeeper might say, "You can skip the next step," or "You only need to do half the work." If a word is hard, it says, "Go all the way through."
The Magic: The best part is that the gatekeeper teaches itself. It doesn't need a human to tell it which words are hard. It learns purely by trying to minimize mistakes. If skipping a step for a specific word causes an error, the gatekeeper learns to keep that word active next time. If skipping it works fine, it learns to save energy.

How It Works in Practice

The paper tested this on two main things:

Simple Logic Puzzles: When the AI had to copy a list of numbers, the gatekeeper realized, "This is easy," and skipped about 65% of the work. When it had to sort numbers (which is harder), it only skipped about 27% of the work. It naturally figured out that harder tasks need more steps.
Writing Stories: When asked to write like Shakespeare or summarize Wikipedia articles, the AI saved 14% to 23% of its computing power.
- It learned that punctuation marks, spaces, and common words are easy and can be processed quickly.
- It learned that unique content words need the full "factory" treatment.

The Results: Faster and Smarter

No Quality Loss: The AI wrote just as well as the standard model, but it used significantly less energy (computing power).
Better than "Early Exit": Other methods try to stop the whole sentence early if the AI feels "confident." TSA is smarter; it stops individual words from going through unnecessary steps while letting hard words keep going. The paper found TSA produced better results than these older methods when they were set to save the same amount of energy.
Real Speed: When the researchers actually ran the code on a computer, they saw a real speedup (about 2.3% faster) because the computer literally didn't have to do the math for the words that were skipped.

The Bottom Line

This paper presents a way to make AI models "lazy" in the smartest way possible. Instead of forcing every word to do the same amount of work, TSA lets the AI decide, on the fly, which words are easy and which need extra help. It's like having a factory where the easy products zip through the line, while the difficult ones get the full VIP treatment, saving time and energy without sacrificing quality.

Technical Summary: Adaptive Computation Depth via Learned Token Routing in Transformers

Problem Statement

Standard Transformer architectures apply a fixed number of layers to every token in a sequence, regardless of the contextual difficulty or information density of that token. This design prioritizes architectural simplicity over per-token adaptability, resulting in significant computational inefficiency. In practice, common tokens in predictable contexts require far less processing than rare tokens in novel constructions, yet both receive identical compute resources during every forward pass. This inefficiency is particularly costly at inference scale, where the dominant cost is the full forward pass through all layers for all tokens. Existing solutions, such as Adaptive Computation Time (ACT) for RNNs, Universal Transformers, Mixture-of-Depths (MoD), and Early Exit mechanisms, attempt to address this but often rely on hard selection, accumulated budgets, or post-hoc confidence thresholds that may not be fully differentiable or optimal.

Methodology: Token-Selective Attention (TSA)

The paper proposes Token-Selective Attention (TSA), a mechanism that introduces a learned, per-token gate on residual updates between consecutive transformer blocks.

Architecture and Routing

TSA inserts a lightweight router $r_l$ after each block $f_l$ (for $l = 0, \dots, L-2$ ) in a pre-norm decoder-only transformer. The stem block $f_0$ always executes unconditionally.

Router Design: Each router is a two-layer Multi-Layer Perceptron (MLP) with a sigmoid output, producing a continuous halting probability $p_l \in (0, 1)$ for each token. The hidden dimension is $d/4$ (floored at 16), adding approximately 1.7% parameter overhead to the base model.
Gated Update: The residual updates for the subsequent block $f_{l+1}$ $f_{l + 1}$ are scaled by $(1 - p_l)$ $(1 - p_{l})$ .
- If $p_l = 0$ , the update is standard.
- If $p_l = 1$ , the state remains unchanged, effectively skipping the block.
- The interpolation is smooth, ensuring the mechanism is fully differentiable without requiring straight-through estimators, Gumbel sampling, or reinforcement learning.

Training and Regularization

During training, all layers execute fully; the gate scales the residual updates but does not skip computation. The total loss is defined as $L = L_{task} + L_{depth}$ .

Depth Regularization: A term $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ is added to gently encourage early halting.
Emergent Behavior: Notably, the paper demonstrates that the routing mechanism learns to skip operations even when $\lambda = 0$ (no explicit depth pressure), driven solely by the task-loss gradient.

Inference

At inference, the mechanism switches to a "sparse" mode. Tokens with $p_l > 0.5$ are excluded from the Feed-Forward Network (FFN) via gather/scatter operations, while attention remains dense to preserve exact key-value semantics. This yields real FLOPs savings and wall-clock speedups.

Key Contributions

Differentiable Token Routing: A simple, continuous soft gate on residual updates that allows per-token, per-layer routing without architectural changes to the base Transformer.
Gradient-Driven Emergence: Evidence that difficulty-proportional routing emerges naturally from the task-loss gradient alone. Even without depth regularization ( $\lambda=0$ ), the router learns to skip ~20% of token-layer operations.
Efficiency-Quality Trade-off: Validation on character-level language modeling showing 14–23% savings in token-layer operations (TLOps) with less than 0.5% quality loss.
Superiority over Early Exit: At matched efficiency, TSA achieves lower validation loss than Early Exit baselines, suggesting end-to-end learned routing decisions are superior to post-hoc confidence thresholding.
Real-World Speedup: Demonstration of real wall-clock speedup (up to 2.3%) on commodity hardware (Apple M1 Pro) via sparse inference.

Experimental Results

Synthetic Tasks

On copy and sort tasks (length-10 sequences):

Copy Task: The router learned that the task was an identity mapping, skipping 54.9% of TLOps ( $\alpha = 0.341$ ) with no accuracy loss.
Sort Task: The router allocated more compute, skipping 22.5% of TLOps ( $\alpha = 0.730$ ), reflecting the higher difficulty.

Character-Level Language Modeling

Tiny-Shakespeare: TSA saved 22.8% of TLOps ( $\alpha = 0.726$ ) with a validation loss increase of only 0.006 nats (+0.4%).
enwik8: On a more diverse corpus (Wikipedia), TSA saved 13.9% of TLOps ( $\alpha = 0.833$ ) with no quality degradation (marginally better validation loss).
Robustness: The mechanism is robust to the regularization hyperparameter $\lambda$ across two orders of magnitude ( $[0, 0.1]$ ). At $\lambda=0.05$ , it achieved 50.4% TLOps savings with negligible quality cost.

Comparison with Early Exit

At a matched active fraction ( $\alpha \approx 0.726$ ), TSA achieved a 0.71% lower validation loss than an Early Exit baseline. This indicates that learning routing decisions end-to-end produces better decisions than applying a confidence threshold at inference time.

Wall-Clock Throughput

Soft-TSA (Training): Added ~1% overhead regardless of $\alpha$ .
Sparse-TSA (Inference): Achieved a 2.3% speedup on Tiny-Shakespeare ( $\alpha=0.726$ ) and break-even performance on enwik8 ( $\alpha=0.833$ ). Speedups required a batch size of $\ge 64$ to overcome CPU-GPU synchronization overhead.

Significance and Claims

The paper claims that TSA offers a minimal, architecturally simple solution to the inefficiency of fixed-depth Transformers. Its primary significance lies in demonstrating that adaptive computation depth can be learned directly from the task objective without explicit supervision regarding difficulty or complex reinforcement learning signals.

The authors emphasize that the router acts as a "learned noise gate," attenuating updates when residual contributions are redundant. This allows the model to dynamically allocate compute based on token difficulty, achieving significant efficiency gains (up to 55% on synthetic tasks, 14–23% on language modeling) with minimal quality trade-offs. The work suggests that end-to-end differentiable routing is a more effective strategy for adaptive depth than discrete selection or post-hoc thresholding, paving the way for more efficient inference on commodity hardware.

The paper remains modest regarding scale, noting that all experiments were conducted on models with approximately 5–6 million parameters, and scaling behavior to larger models (10M–100M+) remains an open question. Additionally, current speedups are batch-size dependent, and attention sparsity has not yet been implemented.