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

TL;DR

GateSkip introduces a differentiable residual gating mechanism for decoder-only transformers to enable token-wise layer skipping, addressing the inefficiency of uniform compute. The method places a sigmoid-linear gate on each Attention/MLP branch and uses per-token budgets to rank and skip low-importance tokens, with hard skips realized at inference from smooth, trainable gates. Empirical results across Llama-3.2 and Gemma-2 demonstrate up to 15% compute savings while preserving over 90% of baseline accuracy on long-form reasoning, and gains on instruction-tuned models with near-baseline quality at roughly 50% savings; larger models show even more pronounced trade-offs. Gate activations reveal interpretable information flow, e.g., BOS tokens and punctuation acting as anchors, and GateSkip remains compatible with 4-bit quantization, speculative decoding, and pruning, making it a practical, stable adaptive-depth approach for real-world deployment.

Abstract

We introduce GateSkip, a simple residual-stream gating mechanism that enables token-wise layer skipping in decoder-only LMs. Each Attention/MLP branch is equipped with a sigmoid-linear gate that condenses the branch's output before it re-enters the residual stream. During inference we rank tokens by the gate values and skip low-importance ones using a per-layer budget. While early-exit or router-based Mixture-of-Depths models are known to be unstable and need extensive retraining, our smooth, differentiable gates fine-tune stably on top of pretrained models. On long-form reasoning, we save up to 15% compute while retaining over 90% of baseline accuracy. For increasingly larger models, this tradeoff improves drastically. On instruction-tuned models we see accuracy gains at full compute and match baseline quality near 50% savings. The learned gates give insight into transformer information flow (e.g., BOS tokens act as anchors), and the method combines easily with quantization, pruning, and self-speculative decoding.

Figures (24)

• Figure 1: We introduce gating mechanisms that regulate the flow of information into the residual stream and can be used to skip layers altogether. Our mechanism enhances downstream accuracy of instruction-tuned models even when skipping $\sim$25% of the model.
• Figure 2: Mean gate value for each token in a sample sequence (Llama-3.2-1b, vector gate, shared across layers). The first Attention and MLP layer, as well as BOS tokens and punctuation, receive elevated importance. This hints at the model using BOS tokens to "dilute" attention, avoiding over-mixing. Another hypothesis is that punctuation and BOS tokens are used as critical reference points for establishing contextual boundaries.
• Figure 3: The top-10 tokens with the highest mean gate values across the entire test sets of GSM8K, CommonsenseQA, and validation set of PIQA, when evaluating Llama-1b with shared vector gates at varying token budgets. The mean gate activation for the BOS token is considerably higher than any other activation. While there is a noticeable jump between the BOS token's activation and the next activation, the remaining activations are rather uniform. This pattern persists along varying skipping ratios.
• Figure 4: Mean gate value for each token in another sample sequence that exhibits a "forbidden request" (Llama-3.2-1b, vector gate, shared across layers). The model evidently puts much more importance on the tokens "chemical weapon" and "please", showing that the model needs more capacity from its layers to process this part of the prompt.
• Figure 5: Distribution of gate values across the PIQA validation dataset.