💬 NLP

What Layers When: Learning to Skip Compute in LLMs with Residual Gates

GateSkip is a stable, fine-tuneable gating mechanism that enables token-wise layer skipping in decoder-only language models, allowing for significant computational savings during inference with minimal loss in accuracy.

Original authors: Filipe Laitenberger, Dawid Kopiczko, Cees G. M. Snoek, Yuki M. Asano

Published 2026-02-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Filipe Laitenberger, Dawid Kopiczko, Cees G. M. Snoek, Yuki M. Asano

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 reading a massive, 1,000-page legal contract.

As you read, your brain naturally works in two ways: when you hit a sentence like "The party of the first part shall...", you might skim it because it’s standard filler. But when you hit a sentence like "The total penalty is $50 million," your brain snaps into high gear, focusing every ounce of your attention to make sure you don't miss a single detail.

Currently, Artificial Intelligence (LLMs) doesn't work this way. Whether it's reading a boring "the" or a critical "million dollars," the AI uses the exact same amount of "brainpower" (computation) for every single word. This is incredibly wasteful.

GateSkip is a new method designed to give AI a "focus filter," allowing it to skim the boring parts and spend its energy only on the important stuff.

The Analogy: The Smart Toll Booth

Imagine a highway (the Residual Stream) where cars (the Information/Tokens) are traveling. Along this highway, there are several massive, expensive toll plazas (the Attention and MLP Layers) where every car must stop, pay, and undergo a full inspection.

In a standard AI model, every single car—even a tiny bicycle or a single passenger—must stop at every single toll plaza. This makes the journey very slow and expensive.

GateSkip introduces a "Smart Gate" at each toll plaza:

The Inspection (The Gate): Before a car enters the expensive inspection zone, a quick, lightweight sensor (the Gate) looks at the car. It asks: "Is there anything important in this vehicle, or is it just empty space?"
The Decision (The Skip):
- If the sensor sees a heavy truck full of gold (a high-importance word like "million" or a complex math problem), it opens the gate and lets the car through the full inspection.
- If the sensor sees a single piece of lint (a low-importance word like "the" or "of"), it tells the car: "You're not important enough for this plaza. Just stay on the highway and skip this stop."
The Budget (The Quantile): To make sure the highway doesn't get too crowded, the system has a "budget." It says, "We only have enough money to inspect 80% of the cars at this plaza." The system then ranks all the cars and only lets the top 80% through, letting the bottom 20% skip ahead.

Why is this a big deal?

Usually, when scientists try to make AI faster by skipping parts, the AI "gets dumber" very quickly—it's like trying to learn to read by skipping every third word; eventually, you lose the plot.

GateSkip is special because:

It’s Smooth, Not Choppy: Instead of making "hard" decisions that confuse the AI during training, it uses a "soft" mathematical approach that lets the AI learn how to skip without breaking its brain.
It Can Actually Make AI Smarter: Surprisingly, the researchers found that on certain tasks, the "focusing" effect actually helped the AI perform better than the original version, because it learned to prioritize the right information.
It Plays Well with Others: It’s like a new type of fuel additive; it doesn't break the engine. You can use it alongside other efficiency tricks (like "quantization," which is like making the cars smaller) to make the AI even faster.

The Result

In short, GateSkip turns the AI from a worker who stares intensely at every single grain of sand into a smart reader who skims the fluff and focuses on the meaning. This allows the AI to run much faster and use much less electricity, while still being just as smart (or even smarter) than before.

Technical Summary: GateSkip: Learning to Skip Compute in LLMs with Residual Gates

Published at: ICLR 2026

1. Problem Statement

Large Language Models (LLMs) currently suffer from computational inefficiency due to uniform resource allocation. Standard transformer architectures apply the same amount of computation (every layer, every attention head, every MLP unit) to every token, regardless of whether the token is a simple punctuation mark or a complex semantic concept.

Existing attempts to solve this—such as Mixture-of-Depths (MoD) (which uses discrete, unstable routers) or Early-Exit methods (which require auxiliary heads and alter pretrained hidden states)—often suffer from training instability, require training from scratch, or fail to maintain performance during long-form generative reasoning.

2. Methodology: The GateSkip Mechanism

The authors propose GateSkip, a lightweight, differentiable residual-stream gating mechanism designed for decoder-only transformers.

Residual Gating: Instead of a discrete router, GateSkip adds a small, trainable linear gate with a sigmoid activation at the exit point of every Attention and MLP module. The output of a module ( $o_\ell$ ) is element-wise multiplied by a gate value ( $g_\ell$ ) before being added back to the residual stream:
$h_{\ell+1} = h_\ell + (o_\ell \odot \sigma(W_G h_\ell + b))$
Training Objective: The model is trained using a joint loss function:
$\mathcal{L} = \mathcal{L}_{CE} (\text{Cross-Entropy}) + \lambda_S \mathcal{L}_S (\text{Sparsity Penalty})$
The sparsity penalty ( $\mathcal{L}_S$ ) is an $L_2$ distance that encourages gate activations to stay close to zero, effectively teaching the model to "compress" or skip unimportant information.
Token-Wise Skipping (Inference): During inference, the model calculates importance scores for all tokens in a layer. It uses a quantile-based threshold to select only the top tokens (based on a user-defined budget $b_t$ ) to process. Skipped tokens simply copy their hidden states and KV-cache entries from the previous layer to the next, saving significant FLOPs.
Stability: Because the gates are smooth and differentiable, GateSkip can be fine-tuned on top of existing pretrained models without the instability associated with hard-routing methods.

3. Key Contributions

Stable Adaptive Depth: Introduced a method that enables token-wise layer skipping that is stable to train and compatible with pretrained backbones.
Superior Trade-offs: Demonstrated state-of-the-art compute-vs-accuracy trade-offs, particularly in generative reasoning tasks where prior methods (like MoD) often collapse.
Orthogonal Compatibility: Proved that GateSkip works seamlessly with other efficiency techniques, including 4-bit quantization, structured pruning, and self-speculative decoding.
Interpretability Insights: Provided an analysis of gate activations, revealing that the model uses BOS (Beginning-of-Sequence) tokens and punctuation as "structural anchors" to manage information flow.

4. Key Results

The authors evaluated GateSkip on Llama 3.1 (up to 8B) and Gemma 2 (2B) models.

Generative Reasoning (e.g., GSM8K, CSQA): On long-form reasoning tasks, GateSkip saves up to 15% compute while retaining >90% of baseline accuracy.
Instruction-Tuned Models: On models like Llama-3b-Instruct, GateSkip actually improves accuracy at full compute and can match baseline quality even with ~50% compute savings.
Log-Likelihood Benchmarks: It matches or outperforms prior state-of-the-art methods (MoD, CALM, FREE) across multiple benchmarks (MMLU, HellaSwag, etc.).
Real-World Efficiency: In optimized environments (vLLM), the theoretical FLOP savings translate directly to measurable increases in throughput and decreases in latency.
Robustness: Unlike "random skipping" which causes catastrophic accuracy collapse, GateSkip degrades gracefully as the compute budget decreases.

5. Significance

GateSkip represents a significant step toward dynamic, efficient inference. By turning the residual stream into a control mechanism, it allows LLMs to "think harder" about difficult tokens and "skim" through easy ones. Its ability to be applied to pretrained models without massive retraining and its compatibility with quantization makes it a highly practical solution for deploying large-scale models in latency-sensitive or resource-constrained environments. Furthermore, the discovery that gates can act as safety/interpretability tools (highlighting "policy-relevant" tokens like "chemical weapon") opens new avenues for AI safety research.