What Layers When: Learning to Skip Compute in LLMs with Residual Gates
Authors: Filipe Laitenberger
Dawid Jan Kopiczko
Cees G. M. Snoek
Yuki M Asano
Venue: ICLR 2026 Poster (2026)
Keywords: decoder-only language models
large language models
layer skipping
adaptive compute
efficient inference
LLM
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We add learnable gates to the exit points of Attention and MLP modules in GPT-style models, compressing their output so that we are then able to use them to assess token importance and thus as a skipping mechanism.

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 compresses the branch’s output before it re-enters the residual stream. During inference we rank tokens by the gate 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 >90% of baseline accuracy. 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.

Published as a conference paper at ICLR 2026
WHAT LAYERS WHEN: LEARNING TO SKIP COMPUTE IN LLMS WITH RESIDUAL GATES
Filipe Laitenberger Dawid Kopiczko", Cees G.M. Snoek?, Yuki M. Asano 3
Humbblldt University Berlin =Qualcomm-Uva Lab, University of Amsterdam FunAl Lab, University of Technology Nuremberg
ABSTRACT
We introduce GateSkip, simple residual-stream gating mechanism that enables token-wise layer skipping in decoder-only LMs. Each Attention/MLP branch is equipped with sigmoid- linear gate that condense 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 per-layer budget. While early-exit or router- based cexture-of-Deptlee 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 >90% baseline accuracy. On instruction-tun tunion 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.
Gate
+ Delement-wise product
computed
skipped
GateSkip- 3b-lngtruct
Ulmma 3b-Instruct skipeing)
Astention
Inear aigmai
Inear #igmad
The capital of France
1s
Parts
Speed Increase (%)
(a) Add lcarnable gates attention and MLP output
(b) Importance scores can be
(c) Improves accuracy while saving compute
used to skip layers
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-tunel models even when skipping ~55 % of the model.
INTRODUCTION
Large language models have transformed natural language processing, yet their rapid growth has created major challenges for efficient deployment. Current models allocate the same amount of computation to every token at every layer, regardless of difficulty. This uniform allocation is wasteful and makes it hard to deploy models in latency-sensitive or resource-limited environments. Adaptive compute aims to address this by using more resources where they matter and less where they do not.
Most prior approaches fall into two categories. Router-based methods, such Mfxture-of-Depths (Raposo et al. (024), introduce specialized routing components that decide which layers to run. These rely hard, discrete decisions that are often unstable and require careful balancing losses (Zoph et al. 2022). Early-exit methods attach auxiliary language modeling heads at intermediate layers and stop once confidence threshold is reached (Schuster et al. 2022). These approaches alter pretrained hidden states, complicate training, and often fail to calibrate well (Bajpai & Hanawal 2024a). Both approaches usually require implementation during pre-training.

---

Published as a conference paper at ICLR 2026
We introduce GateSkip, a lightweight residual stream gating mechanism for decoder-only transform- ers. Each attention and MLP branch is equipped with a small linear gate and sigmoid activation that squashes the branch output before it is added back to the residual stream. During training, gates optimized to remain sparse while preserving language modeling accuracy. At inference, token-level importance scores derived from the gates allow quantile threshold, while skipped tokens copy their hidden states and key-value cache entries upward.
retain only the top tokens per layer using
This design has several advantages. Because the gates are smooth and differentiable, GateSkip can be trained directly on top of pretrained models without destabilizing optimization. Since it operates entirely within the residual stream, minimally perturbs existing representations. Moreover, the mechanism provides fine-graines control at both the token and module level, enabling nuanced allocation compute. Finally, GateSkip fully compatible with orthogonal efficiency techniques such as quantization, pruning, and self-speculative decoding.
We evaluate GateSkip on Llama 3.1 up to 8B parameters and Gemma 2B models across generative reasoning and log-likelihook benchmarks. On long-form reasoning tasks, GateSkip reduces compu- tation by up to 15% while retaining more than 90% of the original accuracy. On instruction-tuner models, it improves accuracy at full compute and sustains this improvement under reduced budgets, matching baseline quality even with roughly 50% compute savings. Analysis of the learned gate values further reveals consistent patterns: carly layers allocate more computation to the beginning- of-sequence token and punctuation, while deeper layers become increasingly selective and focus on content-bearing words.
Our contributions can be summarized as follows:
1. We propose GateSkip, a residual gating mechanism that enables token-wise layer skipping with smooth training and deterministic hard decisions at inference.
2. We demonstrate state-of-the-art where prior adaptive compute methods often collapse.
trade-off's on generative reasoning tasks,
3. We show that Gate composes seamlessly with quantization, pruning, and self-speculative decoding.
4. We provide an analysis of gate activations that sheds light on information flow within transformers.
GateSkip turns the residual stream into a simple yet effective control mechanism for adaptive depth, delivering efficiency gains without sacrificing stability or performance.
2 RELATED WORK
growing body of work has addressed the inefficiency ever-larger decoder-only Transformers by dynamically adapting computation a per-token or per-sequence basis. Early efforts in layer skipping repurpose sparse routing from Mixture-of-Ei -x-ertssts (MoE) to decide whether to execute each layer all, yielding substantial compute savings without restarting training from scratch. Mixture-of-Depttee (MoD) injects a router at every transformer layer to skip unimportant layers for each token (Raposo et al. 2024), while follow-up methods introduce soft token budgets (Zeng ct al. 2023). sequence-level skipping (Wang et al. (023), and frozen-bace backbone router fine-tuning (He al. [0044). However, discrete routers can be unstable (Zoph et al. 2022 Fedus al. 2022 Puigcerver et al. 2024 Panda et al. 2025), and most require training from scratch. In contrast to hard top-k routers with waad-bllancing losses, GateSkip uses fully differentiable residual gating, avoiding discrete routing during training while still yielding hard skips test time.
In contrast, early exiting methods terminate inference once intermediate representations are deemed confident enough. Pioneered in encoder-only BERT models via entropy or agreement thresholds (Xin ct al., 2020 Zhang et al.) 2022 Zhou et al. 2020 Liu et al. 2020), this concept has been extended to encoder-decoder and decoder-only settings by supervising every layer with an auxiliary language
modeling objective and exiting based on heuristic confidence measures (Tang et al.] 2023) Schuster
et al. 2022, Elbayad et al. 2020 Liu et al. 2021 Bae et al.) 2023 Elhoushi al. 2024 Del Corro
ct al. 2023). Yet these approaches fundamentally alter pretrained hidden representations through
connector modules or self-distillation (Bajpai & Hanawal
2024a b Ji al. 2023). Where early exit
2

---

Published as a conference paper at ICLR 2026
supervises intermediate LM heads and exits confidence thresholds, GateSkip avoids auxiliary heads and maintains pretrained hidden spaces by training gates to compress post-module outputs.
Other efficiency techniques include layer pruning (c.g.. ShortGPT removes redundant transformer
blocks (Men et al., 2025)). KV cache and token pruning (e.g., query-driven pruning (Xu et al. 2025). token-precision trade-offs (Zhang et al.) 2024)), and quantization (e.g.. AWQ (Lin et al.
2023). SpQR (Dettmers al. 2024), BiLLM (Huang al. (024)). These methods operate on different axes of efficiency and are thus orthogonal to our depth-adaptive approach. We further show that GateSkip compatible with pruning as well as quantization, demonstrating its ability combine
with such methods.
3 GATESKIP
We propose adding gating mechanism the residual stream of decoder-only Transformer models and training it with an additional sparsity loss so that the gates learn to assess the importance of certain Attention or MLP module given its preceding hidden state, as shown in Figure
3.1 RESIDUAL GATING MECHANISM
The residual stream layer E in transformer model can be described as the output or E B
of an Attention or MLP layer added to the hidden states hig E BB resulting in the layer output he41, with B being the batch size, the sequence length, and H the hidden dimension:
h he + or
(1)
We propose supplementing the language model with trainable gate g which a sigmoid-activated linear projection of the hidden states he:
he+1 he + of G ge()), ge(be) = a (Wahh + b)
where WG E " and b R" refers to the sigmoid function, and ge refers to the gate at layer E which could theoretically be a shared gate across layers (cf. ablation experiments in $4.4) or separate gates for each layer. The gate is placed the exit point of the module the residual stream, after the output projection, making perfectly compatible with multi-head attention or any variant thereof.
3.2 TRAINING OBJECTIVE
Training minimises a standard language-modee model loss (cross-entropy for next-token prediction)
1
LCR
log po(y) x:)
(3)
B
plus an explicit gate-sparsity penalty (L2 distance on gate activations)
Nr.
Ls
(he))))
(4)
NIH
so that the overall loss becomes
L Lcr + Ag Ls.
(5)
Here Nz is the number of layers, H the hidden dimension, and Às balances accuracy and efficiency. Term (2) encourages each sigmoid gate ge(he) to stay close zero, effectively compressing the module output before it is re-added to the residual stream. Backbone parameters and gate parameters updated jointly with AdamW, with all weights being trainable.
3

---

Published as a conference paper at ICLR 2026
3.3 TOKEN SELECTION
At step ! we allot a frfcctional budget b, E (0, 1], the share of the L tokens that may be processed in the current layer. For every token we collapse the current batch's gate vectors scalar importance scores De,s = Σx ge(he)i,k and form their empirical cumulative distribution function (CDF). We then compute the threshold,
= Quantile ( 1
(6)
using linear interpolation between adjacent order-statistics (see Algorithm 3) thus the expected fraction of scores below 7; equals the desired skip ratio - b;. Tokens with Øe, S Te are skipped, copying the hidden state upwards fig+1,i he, and the rest is processed normally.
During training (see Algorithm 1 the budget decays linearly, as b, b1 (b1 b2) Time that the model learns to tolerate skipped hidden states. During inference (see Algorithm 2 we fix single budget b once, re-use the same post-module gate scores for ranking, and apply the Top-k only to tokens that have not yet emitted the end-of-sequence symbol. Additionally, when a token skips layer, upwards copy the KV cache items from the layer below in order to facilitate K KV-cache reuse.
3.4 IMPLEMENTATION DETAILS
We initialize the gates to ensure the model initially closely resembles the original pre-trained model. Specifically, we initialize the weights of the linear matrix Wa around O using a Gaussian distribution with low standard deviation a = 0.01, and set the biases to so that the module's output remains approximately unchanged (as a (5) xe 1).
key design decision to place the gate at the module input for skipping decisions, but to train it on the module's output, i.e. we multiply the gate element-wise with the module output so that receives its gradient signal downstream of the module. While the input to the gate is the same in both cases (the hidden states he), the learning signal differs compared to if the gate placed at the entry point to the module. The gate is trained to incorporate minimal information from the module's output into the residual stream while maintaining language modeling performance, rather than determining which information should enter module. We empirically found that training the gate after the module leads to better downstream performance (see $4).
Our method is numerically stable compared to other techniques based on hard binary routing (such Mixture-of- DDeettts) (Zoph et 2022 Fedus et al. 2022 Puigcerver al. 2024]] Panda al.) 2025), providing effective control of information flow without introducing training instabilities or convergence issues (see for experimental results).
The added gates introduce negligible parameter overhead to the model, c.g. 0.004% for separate gates with scalar output, and 4% for separate gates with hidden-state sized output on Llama-3.2-11 Table 13 in Appendix shows parameter overhead for each of the variants of GateSkip.
4 EXPERIMENTS
We evaluate Skte on Llama-3 (Meta-AI, 2024) models of varying size as well as on Gemma 2 (Gemma- Team 2024). We then perform ablation studies to isolate the impact of each component, compare against state-of-the of-the-a-r- layer-skipping and early-exit methods, and demonstrate compatibility with 4-bit quantization, speculaece decoding, and structured pruning.
4.1 EXPERIMENTAL SETUP
Models and Training. We primarily evaluate our method Llama-3.2-1b while also experi- menting with 3.2ma-3,2-b,, Llama-3.1-80 1-b,, and Gemma-2-25 to assess scalability and architecture independence. For all experiments, we fine-tune the pretrained backbone train the gates while simultaneously adapting the model to task templates for easier downstream answer extraction. We set the sparsity loss weight À 0.1 and decay the token budget from 100% to 80% during training.
4

---

Published as a conference paper at ICLR 2026
Table 1: Averaged results for loglikelihood- based and longer generation benchmarks for a random skipping baseline, prior adaptive compute methods and GateSkip on Llama-3.2-1b.
Oescratve Beachmnaks
Log: Benchmaals
0%
5%
10%
15%
20%
25%
15%
32%
45%
60%
Liass-10 ib
80.07...4
40.12
(randoms skppiing
10.32
2,20
1.67
1,28
0:67....
28, 23.02...
23.29
28.07,..
3.45....
CALM (softina))
am
3.35
30.73am
31177.x
30
30.77..>
FREE (hddee state samoration)
11:57
11.57
11,57
11:57
11,57
11.07
36.02...
36102..
36.00...
36.12...
36.10...
FREE (solmnaa)
10:20
10.20
10,00cm
10.20
10.20mm
10.0
36.02...
36.00
36.00
36.00
35.00 35.00...
LayeeSkop
10.5
10.05.m
106
10.02mm
106
10
28:
38.35
38.25
38 25
25
0.00..m
0.00
0.00
0.00
51.30...
Stipl-pp
0.70
0.00
0.00
0.04
30.00
78/0
13.00
15.07
15.00
GaseSkop (cers)
20.58
23.05
22.07..
22.14...
20.37...
17.07
47.35
38 38.86,,
31.74 27.00 20.00...
Table 2: Accuracy at 15% saved compute for log- seggl-kelihood-beeee based and generative benchmarks for random skipping baseline, prior adaptive compute methods and GateSkip on Llama-3.2-16
CSQA (Gia.)
GSMSK (Gien.)
MMLU Stem
HellaSwag
CSQA
PIQA
Open- BookQA
Wino- Grande
Llama- Ib (cundom skipping)
3.27
2007
30.87
20-27
53.67
15.74_.
51.13
(hidden state saturation)
5.40
1.47 14720m
21.85
24.00
19.00
61.60
16.20
CALM (softmax)
4.70
1 1.60
21.8 21.85
34.00
19:00
61.00
16.20
50.60
FREE (hidden state saturation)
17.07
6.07
21.26
41.335...
19:40...e
70.93
21.87
50. 20
HREE (sollmaa)
6.00...as
41.30
19.40.c.e
...00
21.90...e
60.99l....
LayerSkip
10g
1.90g
21.31ge-0
10.0....
19.00
72.33
20.50
58.33
MoD (oouee--tmed)
0.00
0.00..as
.0.00
28.87
19.92
53.60
13.81
49.08.1.a
MoD
6.
L.77g0
26.00
38.21g1
20.86
59.88 -
18.58_1
52. 52.1520m
Skipl.ayer
0.10
0.00
21.82
27.76
19.06
50.51 px.0
16.40
28.72
GateSkip (ours)
35.2 25
9.03
30.86
39.05
36.16
70.45
22.69
52.00
Training employs the AdamW W optimizer (Loshchilov & Hutter 2017 A full list of hyperparameters can be found in Appendix K All libraries and their respective versions used for our experiments are listed in Appendix L
Instructions for code access can be found in Appendix !
Generative benchmarks. and GSM8K (Cobbe et al. (021) questions with chain-of ofain-of-thought traces generated by Nemotron 70B (Anonymous) 2024] Reasoning 2024] Wang al., 2024), masking the loss on the question portion so that the model learns both reasoning and answer extraction. For evaluation we measure zero-shot accuracy on the GSMSK and CommonsenseQA test sets using the same prompt template. We sweep the inference budget b over {1.00, 0.95, 0.90, 0.85, 0,80, 0.75). corresponding to compute-savings of (0%, 5%, 10%, 15%, 20%, 25%}; 3 since the realized savings sometimes fall between these targets, we linearly interpolate the measured accuracies to report performance at the exact percentages listed above.
We fine-tune the train sets of CommonsenseQA Talmor al. 2019)
Log-likelihood benchmarks. We fine-tune on Fib data (Penedo et al. 2024 with the same hyperparameters as above. For evaluation we measure five-shot log-likelihowe accuracy on MMLU (Hendrycks et al. 202T). HellaSwag (Zellers et al. 2019). CommonsenseQA (Tal mor et al. 2219)) PIQA (Bisk et al.) (211)). OpenBookQA (Mihaylov et al. 2018). and Wino-
Grande (Sakaguchi et al. 2021) using LM Evaluation Harness (Gao et al. 2024 We sweep over {1.00, 0.85, 0.70, 0.55, 0.40 for compute-savings of (0%, 15%, 30%, 45%, 60%). again using linear interpolation to report exact savings levels.
Variance estimate. We perform five seeds per configuration, measuring mean and standard devia- tion on each metric. All baselines (random skipping, MoD router-tuning, CALM variants) are trained and evaluated with the identical data splits, hyperparameters, inference budgets, and interpolation procedure described above, ensuring a fair comparison.
4.2 COMPARISON TO BASELINE
We begin by evaluating GateSkip against straightforward token-level heuristic: random skipping. At each layer, a fixed fraction of tokens is selected uniformly at random to be omitted from further computation. All experiments use the 1.2 backbone.
5

---

Published as a conference paper at ICLR 2026
Table GateSkip on Llama Instruct.
Generative Benchmarks
Log-Likelihood Benchmarks
saved compute
0%
20%
30%
45%
0%
15%
30%
60%
Llama-3b-Insiten
36.5
46.3
Llama-3b-Insine random skipping
0.5
0.1
34.7
30.4
30.4
GateSkip Allama-3bb-Intreee
49.0
49.0
45.6
35.0
36.7
38.8
32.9
31.0
Table 4: GateSkip out-of- generative tasks.
MMIJU-Gen
PIQA-Gie
saved compute
0%
10%
15%
20%
30%
45%
0%
10%
20%
25%
30%
Llama-Ib
22.8
Llama-Ih (random skipping)
7.5
2.0
1.0
0.3
0.1
7.0
0.8
1.2
0.3
GateSkip
14.0
15.6
18.6
15.9
12.8
5.1
14.3
16,9
21.8
29.8
29.5
Table I presents averaged accuracies across multiple compute-savings levels, while Table reports performance exactly 15% saved compute. Random skipping collapses generative accuracy under 10% even at modest budgets (5-10% savings), while GateSkip retains the majority of performance. Additional results LAMBADA are provided in Appendix A showing that GateSkip maintains stable perplexity and accuracy under compute constraints, while random skipping collapses sharply. Moreover, results on translation are presented in Appendix B showing that GateSkip exhibits analogous performance improvements compared to baseline performance.
4.3 COMPARISON TO PRIOR STATE-OF-THE-ART
Having established the superiority of GateSkip over naive heuristics, we now compare it against prior adaptive-depth methods under identical fine-tuning and evaluation settings: (1) Mixture-of-Depths (MoD) with Router Tuning, (2) CALM in its hidden state saturation and softmax variants, (3) FREE in both variants, and (4) static skipping approaches such as LayerSkip and SkipLayer. Results averaged across different benchmarks well task-level accuracy exactly 15% saved compute can be seen in Tables and anddrreseeee respectively.
On generative benchmarks, GateSkip achieves 36.7% accuracy CSQA and 9.7% GSM8K four times higher than MoD and orders of magnitude above CALM. FREE maintains strong log- likelihood scores across all budgets but remains flat on generation (11.9), while LayerSkip and SkipLayer collapse quickly on longer reasoning, with accuracies of only 13.1 and <2.2 respectively. In contrast, GateSkip sustains high generative accuracy while remaining competitive log-likelihood evaluations.
Regarding fog-likelihood benchmarks, across all metries GateSkip either matches or outperforms MoD, CALM, and the static baselines, while FREE achieves similar log-likelihood scores but without corresponding generative performance. These results confirm that GateSkip consistently delivers superior compute-accuracy trade-off across both reasoning and multiple-choice tasks. On instruction- tuned models (Table GateSkip improves generative accuracy over the Llama-3b-Instin baseline even under aggressive budgets (e.g., +12.5 points at 0-20% saved compute) while also matching or slightly improving log-likelihood accuracy at 15600 savings.
We tested log-likc seg-liillihood-bee based benchmarks following prior literature. However, real-world scenarios would demand robustness to longer generation which is why we performed such experiments well. Notably, there is visible discrepancy between log- ligeeiilldood and generative tasks for prior methods, whereas GateSkip retains accuracy significantly better over longer generation.
Beyond the standard suites. out-of-domain generative evaluations (Table 4h show GateSkip retains competitive performance on MMLU-Gen at reduced compute and, notably, exceeds the unadapted baseline PIQA-Gen at 20-30% saved compute (29.8-29.5 22.9). This suggests that targeted

---

Published as a conference paper at ICLR 2026
Table Ablation of GateSkip's design choices Llama-3.2-11
Generative Benchmarks
Log-L Likelihood Benchmarks
Compute saved
0%
5%
10%
15%
20%
25%
0%
15%
30%
45%
60%
Gate output shape
Scalar gates
23.6
22.7
21.7
20.4
15.9
14.2
42.8
36.8
33.7
30.9
31.8
Vector gates (default)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Gate parameter sharing
Shared
21.8
21.4
21.0
20.7
19.5
15.5
44.3
38.4
35.4
31.8
31.7
Separate (default)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Skipping strategy
Only attention layers
26.8
23.1
19.2
14.9
10.8
6.8
44.3
37.5
30.8
30.6
30.3
Only MLP layers
26.8
24.1
18.6
7.8
1.2
0.4
44.3
32.0
30.4
32.1
32.0
Skip entire layer (attn gate)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Every-second layer
26.8
25.5
22.8
15.0
10.1
9.5
44.3
35.5
33.3
31.6
31.9
Skip all layers (default)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31,2
Gate architecture
MLP-based gate
24.3
24.3
24.4
18.5
17.0
15.6
44.0
33.9
31.8
30.3
30.7
Linear-sigmoid gate (default)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Gate placement
Gate before module (entry)
21.1
5.5
1.3
1.0
0.9
0.8
40.0
35.7
33.2
31.6
30.8
Gate after module (default)
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Gate Loss
L1
23.0
22.6
22.2
21.4
18.7
17.0
47.5
37.6
29.2
27.1
26.3
KL-div.
36.0
18.1
11.9
9.8
7.9
5.3
45.6
26.6
24.5
23.0
23.0
L2 (default)
26.8
25,5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
Frozen Backhone
Frozen B.
14.9
14.1
13.3
12.5
12.6
45.6
37.5
26.7
24.0
23.9
Unfrozen (default)
24.2
23.2
23.6
44.3
37.8
32.7
30.8
31.2
token-level allocation can translate into quality gains on certain OOD generative tasks, not merely lossless efficiency.
4.4 COMPONENT ABLATIONS
To understand the contribution of each design choice in GateSkip. we perform a series of controlled ablations Llama-3.2-1b. Table 5 summarizes the impact of varying the gate parameterization, skipping granularity, gate architecture, and gate placement on both our generative and log-likelihood benchmarks at multiple compute-savings levels.
General Oberservations. Since we condense the output of modules, downstream performance changes even no skipping. Hence, different modifications of method will have differing effects on performance at 0% skipping.
Gate Output Shape. We compare forms of gating output: (1) Vector-gates, which produce an H H-dimensional output residual branch, and (2) Scalar-gates, which produce a single gating value per branch. At 15% compute-savings on our generative benchmarks, vector-gatt -aaes achieve 23. 23.2% accuracy, compared to only 20.4% for scalar-gates, confirming that a full-dimensional gate yields more precise control.
Gate Parameter Sharing. Focusing on the vector-gate design, we then compare (1) Per-layer vector-gates: distinct gate for each Attention and MI.P module, versus (2) Shared vector-gates: one gate shared across all Attention modules and across all MLP modules. Per-layer vector-gates gatr- again lead, with 23.2% at 15% savings, while shared vector-gates lag at 20.7%, showing the value of layer-specific parameters.
Skipping Granularity. Ablating which sub-modules can be skipped reveals that attention and MLP layers are both essential. When only attention layers skipped, accuracy falls to 14.9% 15% compute reduction; skipping only MLP layers drops performance even further, to 7.8% under the same budget. Applying single gate the entire layer performs on par with default per-module

---

Published as a conference paper at ICLR 2026
Table 6: GateSkip on different model sizes and architectures.
Generative Benchmarks
Log-L Log-Likelihood Benchmarks
saved compute
0%
5%
10%
15%
20%
25%
0%
15%
30%
45%
60%
Llama-3.2-1B
26.8
25.5
24.2
23.2
23.6
19.8
44.3
37.8
32.7
30.8
31.2
1lama-3.2--B
45.0
44.4
43.9
43.3
42.7
42.1
55.9
35.6
31.4
29.3
30.4
Llama-3.1-80
57.3
56.5
55.8
55.0
54.3
53.6
62.8
44.2
34.3
31.1
29.1
Gemma-2-2B
38.0
37.4
36.7
36.1
35.4
34.8
52.9
44.1
35.6
31.7
30,9
approach, but skipping every second layer leads to a steep decline, from 23.2% down to 15.0% at the 15% savings level. These results underscore the importance of fine-grained, per- per-module control.
Gate Architecture. We also tested small MLP in place of our lincar-sigmoid gate, but despite the extra parameters underperforms. At 15% compute savings the MLP-based gate achieves only 18.5% on generative tasks compared to 23.2% with the linear gate, and 33.9% on log-likelln eigllikelihood benchmarks versus 37.8%. We therefore retain the simpler, more effective linear projection.
Gate Placement. Placing the gate before the module proved disastrous: at 5% compute reduction entry-point gating yields just 5.5% accuracy compared to 25.5% when the gate applied after the module. This dramatic gap confirms that post-module residual gating is crucial for stable, effective learning, as discussed Section 3.4
Gate Loss Type. Replacing our L2-loss on gate activations with LI loss mostly underperforms, whereas achieves a higher loglikelihood accuracy at 0% compute savings. Furthermore, we tested layer-wise KL divergence regularizer that matches the average gate activation per layer target budget, encouraging control of the overall fraction of retained tokens rather than only shrinking individual gates:
LlayerKL-
Σ KT.(Ber((G)) || Bern(Iuain) where Øe
Σ (bee))),
BSH
b,i,k
While this loss clearly outperforms our L.2 and L1 variants at no compute savings in the generative settings, it performs poorly at even low compute savings.
Frozen Backbone, To test the extent to which the backbone adapts, we froze a backbone finetuned on dataset (for fair comparison) and trained gates the frozen backbone. In this frozen setting, downstream accuracy substantially decreases. This gives evidence that backbone adaptation plays major role in the skipping task: the backbone may by itself not reveal all information needed for skipping. necessitation adaptation so that it saves relevance cues in the hidden states which in turn condition the gates.
GateSkip on varying model sizes and architectures. To test scalability, we applied GateSkip larger Llama models (3b and 8b) and observed consistent performance patterns (see Table (6. The results reveal that for larger architectures, the model is capable of skipping increasingly more tokens without decreasing performance. For instance, Llama-3.2-3B with GateSkip can save 37.3% computation while retaining 91.5% of its baseline GSM8K (Gen.) performance and 87.3% of its baseline CSQA (Gen.) performance. Moreover, comparisons between Llama and Gemma architectures reveal that the compute-accuracy trade-off generalizes across both model families. The instruction tuunod and LAMBADA results mirror these trends: larger or eestruction-adaplee backbones benefit more from budgeted token selection, maintaining strong generation quality where random or uniform skipping fails, and confirming that GateSkip's gains persist across usage styles (chat/instruction), lengths, and domains.
4.5 COMPATIBILITY WITH OTHER EFFICIENCY TECHNIQUES
We evaluate compatibility with various orthogonal efficiency techniques. Specifically, we show that GateSkip is compatible with 4-bit quantization, speculative decoding, and structured pruning.
Compatibility with 4-bit Quantization. To test compatibility with guantization, we apply 4-bit quantization to Llama-3.2-3 trained with GateSkip (Table 7 and downstream evaluate the quantized
8

---

Published as a conference paper at ICLR 2026
model as before. The results demonstrate that GateSkip remains effective when combined with quantization, with performance curves closely tracking those the 32-bit model. On generative benchmarks, the quantized model retains 94.4% of the original accuracy at 0% skipping ratio, 96. % at 15%, and 97.3% at 25% skipping ratio. The performances for log-likelihood-hasee benchmarks exactly match.
Table 7: Quantization robustness (Llama-3.2-3 B).
Generative Benchmarks
Log-Likelihos Likelihood Benchmarks
saved compute
0%
5%
10%
15%
20%
25%
0%
15%
30%
45%
60%
32-bit + GateSkip
45.0
44.4
43.9
43.3
42.7
42.1
55.9
35.6
31.4
29.3
30.4
4-bit GateSkip
42.5
42.2
41.9
41.6
41.3
41.0
55.9
35.6
31.4
29.3
30.4
Compatibility with Speculative Decoding. Adding speculative decoding boosts Lige-Likeihood performance substantially at moderate savings: at 15% and 30% saved compute, outperforms vanilla GateSkip (Table 8
Table 8: GateSkip combined with self-speculative decoding compared to GateSkip alone and Layer- Skip. Metries are log-likelihood accuracy (LL) at fixed saved-compute levels.
I.I. @ 15%
LL@@ 30%
LL@@45%
I.I. @..@60%%
GateSkip
37.8
32.7
30.8
31.2
GateSkip self-speculative decoding
39.4
39.4
31.4
30.7
Compatibility with Structured Pruning. Structured pruning (Men et al. 2024) reduces absolute Log-Like Ligeeiiel performance, but GateSkip remains notably stronger than the pruned backbone baseline at every budget 9T. At 0% savings, GateSkip + pruning trails unpruned GateSkip. as expected, yet still exceeds the Default Llamalb + pruning by +6.4 LL points.
Table 9: GateSkip with and without additional structured pruning of 25% of transformer blocks (ShortGPT).
LL@0%
LL@ 15%
LL@30%
LL LL@445
GateSkip
44.3
37.8
32.7
30.8
GateSkip + ShorGGPT pruning
31.5
31.1
31.9
30.8
Default Llamalb + ShortGPT pruning
25.1
25.4
25.4
26.3
4.6
END-TO-END EFFICIENCY GAINS REAL-WORLD SCENARIOS
Table 10: End-to-end latency and throughput different token skipping levels.
% Tokens Skipped
5%
15%
25%
35%
50%
70%
Latency (s)
607.29
606.01
559.81
571.79
521.68
449.85
Throughput (tokens/s)
2697.87
2703.57
2926.71
2865.39
3140.61
3642.11
Table 10 shows end-to-end latency and throughput measurements for our Llama-lb GateSkip model using an optimized vLLM environment. The
evaluated on GSM8K and CommonsenseQA results show that GateSkips theoretical FLOP savings translate into analogous real-world efficiency gains. Analogously, Figure iin Appendix D shows consistently decreasing wall-clock time with lower FLOPS per token.

---

Published as a conference paper at ICLR 2026
4.7 ANALYSIS OF GATE VALUES
Appendix C shows that GateSkip concentrates compute on BOS/punctuation anchors and salient content words, with deeper layers becoming increasingly selective. Gate scores exhibit tight, layer- specific distributions separated by tiny margins, motivating our per-layer quantile thresholds. The same scores also localize policy-relevant spans, suggesting value for interpretability and safety.
LIMITATIONS
Our study targets 1-8 B B-parameter decoder-only LLMs and evaluates on English reasoning, transla- tion, and general language modeling. We report theoretical FLOP reductions in the main text, as they more faithfully capture methodological differences and enable fair comparison across approaches, while also presenting end-to-end efficiency gains in $ $.6 An ethics statement and LLM usage is disclosed Appendix
CONCLUSION
We introduced GateSkip, residual gating mechanism that enables token-wise layer skipping in decoder-only transformers. GateSkip achieves up to 15-20% compute savings while retaining more than 90% accuracy on long-form reasoning, and on instruction-ture full compute and matches baseline quality with nearly 50% savings. These results establish new state the art in adaptive compute, particularly in generative settings where prior methods collapse.
models it improves accuracy at
Beyond efficiency, the learned gates provide insight into transformer information flow, consistently allocating more compute to BOS and punctuation tokens as well as salient content tokens.
GateSkip turns the residual stream into stable and practical control mechanism for adaptive depth, offering real efficiency gains without destabilizing training or disrupting pretrained representations.
REFERENCES
Anonymous.
Commonsenseqa with reasoning traces, 2024. URL https::/hhugginggteeeee ccodattasets/mllti doma i n- www.rr/cooinn//cmmoooeeesee...g Hugging Face Dataset, commit a7b9ab8, accessed April 2025.
S. Bae, J. Ko, H. Song, and S.-Y. Yun. Fast and robust carly-exiting framework for autoregressive language models with synchronized parallel decoding. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5910-5924. Association for Computa- tional Linguistics, 2023. doi: 10. 1t055//11222...mmppmmmiaal
EMNLP 2023, Singapore.
J. J. Bajpai and M. K. Hanawal. Dadee: Unsupervised domain adaptation early exit plms. In Y. Al-Onaizan M. Bansal, and Y.-N. Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 6389-6400, Miami, Florida, USA, November 2024a. Association for Computational Linguistics. doi: 10. (0.18653/v1/22 1/2044/i/0....nnninggg..... findings-emnlp.377 Findings of EMNLP 2024, Miami, Florida.
D. J. Bajpai and M. K. Hanawal. Ceebert: Cross-domain inference early exit bert. In Findings of the Association for Computational Linguistics: ACL 2024, 2024b. doi: 10.48850/axxv.2455.55033 Association for Computational Linguistics, Location.
Barbero, Á. Arroyo, X. Gu, C. Perivolaropoulos, M. Bronstein, Veličković, and R. Pascanu. Why do Ilms attend to the first token? arXiv preprint 2rXiv:2504.02732, 2025. doi: 10.48550/arXiv. 2504.02732. Conference on Neural Information Processing Systems (NeurIPS), Montreal, Canada.
Y. Bisk, Zellers, Le Bras, Gao, and Choi. Piqa: Reasoning about physical common- sense in natural language. CoRR, abs/1911.11641, 2019. doi: 10048550/axxivv1111.... 1rXiv:1911.11644
K. Clark, U. Khandelwal, O. Levy, and C. D. Manning. What does bert look at? an analysis of bert's attention. In Proceedings of the 2019 ACL Workshop BlackboxNI.P. Analyzing and Interpreting
10

---

Published as a conference paper at ICLR 2026
Neural Networks for NLP, pp. 276-286, 2019. doi: 10.18653//1/W11--42322 BlackboxNLP, Florence, Italy.
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and Schulman. Training verifiers to solve math word problems. arXiv, doi: 10.48550/arXiv2211 10.14168.
abs/21 14166, 2021.
Darcet, M. Oquab, Mairal, and Bojanowski. Vision transformers need registers. In Proceedings of the International Conference on Learning Representations, 2024. doi: 10.48550//rXiv.2309. 16588. ICLR 2024, Vienna, Austria.
L. Del Corro, A. Del Giorno, S. Agarwal, B. Yu, A. Awadallah, and S. Mukherjee. Skipdecode: Autoregressive skip decoding with batching and caching for efficient IIm inference. arXiv preprini, arXiv:2307.02628, 2023. doi: 104485500a2ii.2237702233. Microsoft Research.
T. Dettmers, R. Svirschevski, V. Egiazarian, D. Kuznedelev, E. Frantar, S. Ashkboos, A. Borzunov, T. Hoefler, and D. Alistarh. Spqr: A spanse-quantized representation for near-lossless Im weight compression. In Proceedings of the International Conference on Learning Representations (ICLR), 2024. doi: 10, 20.550/arXii.230000000000
ICLR, Vienna, Austria.
M. Elbayad, J. Gu, E. Grave, and M. Auli. Depth-adaptive transformer. In Proceedings of the International Conference on Learning Representations (ICLR), 2020. doi: 10.488500ar2ivv.1110. 10073. ICLR, Addis Ababa, Ethiopia.
M. Elhoushi, A. Shrivastava, D. Liskovich, B. Hosmer, B. Wasti, L. Lai, A. Mahmoud, Acun, S. Agarwal, A. Roman, A. Aly, B. Chen, and C.-J. Wu. Layerskip: Enabling early exit inference and self-speculative decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 12622-12642, Bangkok, Thailand, 2024. Association for Computational Linguistics. doi: 10. (0.18653/v1/ 1/265//////2224aa/looooggg// Bangkok, Thailand.
2024,
W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022. doi: 10. 20.48550/arxiv.21011 003961..
L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi C. Foster, L. Golding, J. Hsu, A. Le Noac'h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. A framework for few-shot language model evaluation, July 2024.
Gemma-Team. Gemma 2: Improving open language models at practical size. arXiv, 2408.00118 18, 2024. doi: 2..5500arXiv.2408.000
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen- tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, Pp. 5484-5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: (0.18653/v1/2022 1ttp://1112111ml1mmmmme mailppmain.445. URL https
//aclantho 1 ogy. orgg 2021. emnlp ma n. 4467
T. Guo, D. Pai, Y. Bai, J. Jiao, M. I. Jordan, and S. Mei. Active-dormant attention heads: Mecha- nistically demystifying extreme-token phenomena in llms. In Proceedings of the Mathematics of Modern Machine Learning (M3L) Workshop at NeurIPS, 2024. doi: 10448500axxi..2410..38 Mathematics of Modern Machine Learning (M3L) Workshop, NeurIPS, Virtual.
Y. Han, G. Huang, S. Song, L. Yang, H. Wang, and Y. Wang. Dynamic neural networks: A survey. IEEE Transactions Patlern Analysis and Machine Intelligence, 44(11):436-7456, 2022. doi: 20.100///AA...21..11
S. He, T. Ge, G. Sun, B. Tian, X. Wang. A. Li, and D. Yu. Router-tuning: A simple and effective approach for enabling dynamic-depth in transformers. CoRR, abs/2410.13188 13184, 2024. doi: 10. 20550/arXiv.2410.1333 arXiv:2410.13. 13184.
11

---

Published as a conference paper at ICLR 2026
D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. In Proceedings of the International Conference on Learning Representations (ICLR), 2021. doi: 10448550/amX11.220000333
ICLR 2021.
W. Huang, Y. Liu, H. Qin, Y. Li, S. Zhang, Liu, M. Magno, and X. Qi. Billm: Pushing the limit of post-training quantization for llms. arXiv preprint, 2402.04291, 2024. doi: 10.48550 04291.
Y. Ji, J. Wang, J. Li, Q. Chen, W. Chen, and M. Zhang. Early exit with disentangled representation and equiangular tight frame. In A. Rogers, J. Boyd-Graber, and N. Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 141222 Toronto, Canada, July 2023, Association for Computational Linguistics. doi: 10.1 1653/v1120223 hfm6://////....mmmnnmmaaaaa//aa/ finndnngs-acl.88% Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada.
J. Lin, J. Tang, H. Tang, S. Yang. W.-M. Chen, W.-C. Wang. G. Xiao, X. Dang. C. Gan, and S. Han. Awq: Activation-awaa -ware weight quantization for large language model compression and acceleration. arXiv preprint, 21Xi::2306.00978 2023. doi: 10.485500axXii.2336600777 MLSys 2024 Best Paper Award.
W. Liu, P. Zhou, Z. Zhao, Z. Wang. J. Ding, and X. Qiu. Fastbert: A self-distilling bert with adaptive inference time. In Proceedings of the 28th International Conference on Computational Linguistics, pp. 0035-6044, 2020. doi: 10. 00.18653/v1/2020.oollo Computational Linguisties, Barcelona, Spain.
main.529. International Conference on AAAI Conference on Artificial Intelligence, Virtual Event.
Y. Liu, Meng. J. Zhou, Y. Chen, and J. Xu. Faster depth-adaptive transformers. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 13424-13432, 2021. doi: 10.1609/aaai. http:/aa..i...1...111111.113.... v35i15.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:
2017. doi: 10.48550/arXiv 100101. Revised 2019.
X. Men, M. Xu, Q. Zhang, B. Wang, H. Lin, Y. Lu, X. Han, and W. Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.0388 2024. doi: 10. 11.48550/rrXiv.240000 240303333.. Conference on Empirical Methods in Natural Language Processing (EMNLP), Singapore.
X. Men, M. Xu, Q. Zhang, B. Wang, H. Lin, Y. Lu, X. Han, and W. Chen. Shortgpt: Layers in large language models are more redundant than you expect. International Conference on Learning Representations (ICLR), 2025. doi: 11.48550/aXxi.2443..88
ICLR, Location.
Meta-AI. The Ilama herd of models. arXiv preprint, arXiv:2407.2178. 2024. doi: 10.48550/arXiv. 2407.21783. Version 3, November 23, 2024.
T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. pp. 2381 2391,, 2018. doi: 10.18653//1/D18-1121 1260. EMNLP, Brussels, Belgium.
A. Panda, V. Baherwani, Z. Sarwar, B. Therien, S. Chakraborty, and T. Goldstein. Dense backpropa- gation improves training for sparse mixture-of-experts arXiv preprint arXiv:2504.12463, 2025. doi: 10.48550karxxx.2544.14444
G. Penedo, H. Kydliček, L. Ben Allal, A. Lozhkov, M. Mitchell, C. Raffel, L. Von Werra, and T. Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. Conference Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
J. Puigcerver, C. Riquelme, B. Mustafa, and N. Houlsby. From sparse to soft mixtures of experts. arXiv preprini 0rXiv:2308.0095/ 2024. doi: (0.48550/axxiv.23088.000
D. Raposo, S. Ritter, B. A. Richards, T. P. Lillicrap, P. C. Humphreys, and A. Santoro. Mixture- of-depths: Dynamically allocating compute in transformer-base based language models. CoRR, abs/2404.02258. 2024. doi: 10448550aaxXx.244402255
12

---

Published as a conference paper at ICLR 2026
Multi-Domaii Reasoning. Gsm8k with reasoning traces, 2024.
URL
co/ to/dataase daeta https/mmmtttitttnntaaeeaaaaaaaat wwwomnnn-eeesoniiinn reasoning/ gam8k "5d82e6e" accessed 3 April 2025.
Hugging Face Dataset. Commit
K. Sakaguchi, R. Le Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021. doi: 10.1145/3474381.
T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Q. Tran, Y. Tay, and D. Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 2022. doi: 200488500arXx..220...777
NeurlPS 2022.
A. Talmor, J. Herzig, N. Lourie, and Berant. Commonsenseeu: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North Ameri- can Chapter of the Association for Computational Linguistics: Human Language Technologies, volume pp. 4149 4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguisties. doi: 10. 10.18653/1/N11--4221 NAACL-HLT 2019, Minneapolis, Minnesota.
Tang, P. Zhu, T. Li, S. Appalaraju, V. Mahadevan, and R. Manmatha. Deed: Dynamic early exit on decoder for accelerating decoder transformer models. arXiv preprint arXiv:2311.086233 2023. doi: 10.48550/a2Xiv2311..00
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30, 2017, doi: 10. 10.48550/rrxiv.1706..00
NeurlPS 2017, Long Beach, CA.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Tal Linzen, Grzegorz Chrupala, Yonatan Belinkov, and Dieuwke Hupkes (eds.), Proceedings of the 2019 ACI. Workshop BlackboxNI.PP. Analyzing and Interpreting Neural Networks for NL.P, pp. 63-76, Florence, Italy, August 2019. Association for Computational Linguisties. doi: 10.18653/
v1/W19-4808 URL https::////... /ttp::://acaaattt http//aaaattoooo.rrrro/o //lz///199 4808/
H. Wang, Y. Wang, T. Liu, Zhao, and Gao. Hadskip: Homotopic and adaptive layer skipping of pre-trained language models for efficient inference. In H. Bouamor, J. Pino, and K. Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 4283-4294, Singapore, December 2023, Association for Computational Linguistics. doi: (0.18853/v1// 1/2023. findings-emmlp288. EMNLP 2023, Singapore.
Z. Wang, A. Bukharin, O. Delalleau, D. Egert, G. Shen, Zeng, O. Kuchaiev, and Y. Dong. Helpsteer2-preference:: Complementing ratings with preferences. arXiv preprint, 2410.01257, 2024.
G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. In Proceedings of the International Conference on Learning Representations, 2024. doi: 10. X0.485550a2XXx.233...11
ICLR 2024, Vienna, Austria.
Xin, R. Tang, J. Lee, Y. Yu, and J. Lin. Deebert: Dynamic early exiting for accelerating bert inference. In D. Jurafsky, J. Chai, N. Schluter, and J. Tetreault (eds.). Proceedings of the 58th An- nual Meeting of the Association for Computational Linguistics, pp. 2246-222 2251, Online, July 2020. Association for Computational Linguistics. doi: 10. 10.18653//1/20200.ac acl- main.204. Association for Computational Linguistics, Online.
Y. Xu, Z. Jie, H. Dong, L. Wang, X. Lu, A. Zhou, A. Saha, C. Xiong, and Sahoo. Think: Thinner key cache by query-driver pruning. In Proceedings of the International Conference on Learning Representations (ICLR), 2025. doi: 10.48550/arxi..4411.222
ICLR 2025, Singapore.
R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence? Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, 2019. doi: 10.18653/v1/1 1472. ACL 2019, Florence, Italy.
D. Zeng, N. Du. T. Wang. Y. Xu, Lei, Z. Chen, and C. Cui. Learning to skip for language modeling. CoRR, abs/2311.154555 2023. doi: (0.4550/arxi..2311.55.4
13

---

Published as a conference paper at ICLR 2026
J. Zhang, D. Zhu, Y. Song, W. Wu, C. Kuang, X. Li, L. Shang, Q. Liu, and S. Li. More tokens, lower precision: Towards the optimal token-precision trade-off kv cache compression. arXiv preprint, 2412.12706, 2024. doi: 10. 10448550a2X1.2242.12706
Z. Zhang, W. Zhu, J. Zhang, P. Wang, R. Jin, and T.-S. Chung. Pcee-bert: Accelerating bert inference via patient and confident early exiting. In M. Carpuat, M.-C. de Marneffe, and I. V. Meza Ruiz (eds.), Findings of the Association for Computational Linguisties: NAACL 2022, pp. 327-338, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/202.. findingnniinnniinne andings-naacl.25. naacl.25. NAACL 2022, Seattle, United States.
W. Zhou, Xu, T. Ge, J. McAuley, K. Xu, and F. Wei. Bert loses patience: Fast and robust inference with early exit. In Advances in Neural Information Processing Systems, volume 33, pp. 18330-18341, 2020. doi: 10.48550/arXiv.20066 04512.. NeurIPS 2020, Virtual.
Z. Zhou, X. Ning, K. Hong, T. Fu, J. Xu, S. Li, Y. Lou, L. Wang, Z. Yuan, X. Li, S. Yan, G. Dai, X.-P. Zhang, Y. Dong, and Y. Wang. A survey efficient inference for large language models, arXiv preprint arXiv:2404. 1442944 2024. doi: 10.48550aaxxiv.2444.22222
arXiv:2404.144 14294 [cs.CL].
B. Zoph, I. Bello, S. Kumar, N. Du, Y. Huang, J. Dean, N. Shazeer, and W. Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202 08906, 2022. doi: 2tt.:0/aa.i..2..089999 arXiv:2202.08906 [cs.CL].
14

---

Published as a conference paper at ICLR 2026
A ADDITIONAL RESULTS: LAMBADA
Table 11: Perplexity and accuracy on LAMBADA under different saved compute levels. GateSkip degrades gracefully while random skipping collapses.
Perplexity
Accuracy
saved compute
0%
10%
20%
30%
0%
10%
20%
30%
Llama-Ib
89.8
31.0
Llama-1b (random skipping)
2210.0
491000.0
10300000.0
12.97
2.36
0.12
GateSkip
23.7
70.9
588.0
9180.0
40.4
28.4
15.8
4.6
LAMBADA evaluates long-context language modeling where error accumulation typically amplifies weaknesses adaptive methods. GateSkip substantially outperforms default Llama at 0% skipping, well as random skipping across both perplexity and accuracy, demonstrating stable degradation under compute savings rather than catastrophic collapse. This supports our main claim that residual gating provides robustness on long generation tasks.
B ADDITIONAL RESULTS: GATESKIP ON TRANSLATION
Table 12: Translation- WMT16 English-> -oanhh-Romnnnnn (a) Baseline BLEU at 0 % skipping. (h) Largest compute reduction that still preserves 2 90 % of that BLEU (higher is better).
WMT16-EN-RO
saved compute
0%
5%
10%
15%
20%
25%
Llama-1b
0.57
Llama-1b (random skipping)
0.36
0.14
0.03
0.01
0.01
GateSkip
0.51
0.43
0.37
0.32
0.28
0.2
To evaluate GateSkip (separate vector gates at each layer) on the WMT16 Singish-RRmande training set for one hour, using the same hyperparameters test set under varying compute-savings. Even with 10% and 15% the layers skipped, GateSkip retains 65% and 63% of the full-compute BLEU (0.37/0.32 vs. 0.57), exhibiting significantly more advantageous trade-off between efficiency and translation quality than the random skipping baseline.
task, we fine-tune Llama-3.2-16 with GateSkip
in our initial experiments. Table shows BLEU scores the WMT16
C QUALITATIVE ANALYSIS OF GATE VALUES
The preceding chapters demonstrated that GateSkip can remove double-digit fraction of Transformer FLOPs while maintaining competitive downstream accuracy. In this chapter, evaluation to qualitative analysis. Concretely, we analyze the distribution of learned gate values for individual sequences and ask what they reveal about (i) information flow within the residual stream, (ii) the model's 's implicit safety heuristics, and (iii) the practical utility of gate values as an interpretability signal (RQ5). Lastly look into the overall distribution of gate values across entire datasets to gain insight about effective token budgeting.
from quantitative
C.1 VISUAALIZATION SET-UP
Unless stated otherwise we inspect a Llama-3. -.2--B model fine-tuned with shared vector gates, one for each Attention and one for each MLP module. For each token and layer l we compute the scalar importance
H
Øt,i =
1
Σ
gt,ikk
(7)
H
15

---

Published as a conference paper at ICLR 2026
where H is the hidden dimension. We create heatmaps showing the average gate value for each token and layer, separated into Attention and MLP modules for better clarity of the patterns distinct to each module type. The resulting heatmaps shown in Figures and S higher compute allocation.
Darker colors correspond to
C.2 BOS TOKENS AND PUNCTUATION AS STRUCTURAL ANCHORS
Attention Layers
MLP Layers
<|begin of >
Joc has
20
0.8
sells
0.7
of
them
Galee
each
How
much
money
he
10 15
10
15
Figure 2: Mean gate value for each token in a sample sequence (1lama-3.2-1 2b, vector gate, shared across layers). The first Attention and MIP layer, well BOS tokens and punctuation, receive elevated importance. This hints at the model using BOS tokens to "dilute" attention, avoiding over-mixing. Another hypothesis that punctuation and BOS tokens are used as critical reference points for establishing contextual boundaries.
Figure shows mean gate values for our Llama-lb 1b model for the sample sequence:
Joe has 20 horses. He sells 5 of them for $200 each. How much money does he make?
We split the importance scores into one sub-figure for the Attention and another for the MLP layers to highlight the patterns present. Moreover, we make several key observations:
1. Functional tokens (prepositions, pronouns, articles) receive consistently lower gate values than content words, particularly in later layers.
2. The first layer maintains high importance across all tokens while deeper layers become more selective.
3. Beginning-ov of sequence (BOS) tokens and punctuation receive exceptionally high importance across all layers.
Quantitative Analysis. To quantitatively confirm BOS token prominence, we collect gate activations for all vocabulary items while evaluating our Llama-lb model with shared vector gates across the test sets of GSMSK and CommonsenseQA,, as well the PIQA's validation we perform this test at varying skipping ratios. i.c. one run with 0% and one with 30% skipping. We average the collected activations across layers and samples to obtain a single number per vocabulary item. The sorted top-10 10 tokens with highest activations are shown in Figure 33] for 0% skipping (left) and 30% skipping (right). Across both skipping levels, the BOS token attains the highest gate activation, with a considerable margin (zx 0.01) to the tokens that follow (subsequent margins lie below <0.001). turn, the remaining top activations are rather uniform, with the only notable jump existing between the BOS token's and the subsequent token's activation. This quan- titative test validates our earlier observation that BOS tokens receive elevated importance by our gates.
Moreover,
16

---

Published as a conference paper at ICLR 2026
The consistently high BOS activations in findings cast doubt on Guo al. (2024) 's and Xiao et al. (2024)ss hypothesis that models allocate "excess" attention to non-meanee non-meaningful tokens. If BOS tokens contained primarily redundant information,
gates would naturally assign them lower importance. While Clark et al. (019)) suggest BOS tokens accumulate sentence-level information, this cannot explain the initial BOS token's importance due causal attention constraints.
Mean gate value across datasets
0.735
0.735
0.730
gate across
0.725
0.725
0.720
0.720
0.715
0.715
</begin of sttoxt/>>
jugg
Russell
erv
Tou
mitt
quann
att
chauff
fin
<|begin of text/=
jugg
Russell
Uri
Icc
arist
twink
Rory.
(a) Top-10 tokens with highest gate activation at 0% skipping ratio
(b) Top- tokens highest gate activation 30% skipping ratio
Figure 3: The top-10 10 tokens with the highest mean gate values across the entire test of GSMSK, CommonsenseQAA, and validation set of PIQA, when evaluating Llama-1b 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 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.
Instead, these insights lead us to hypothesize two things:
BOS tokens may serve as critical reference points for establishing contextual boundaries, similar to register tokens in Vision Transformers (Darcet et al. 2024).
2. Following Barbero al. (205)), attending to BOS tokens may help model avoid over- mixing and thus prevent representational collapse. LMs thus BOS tokens to "dilute" the attention, keeping it from pushing latent states into meaningless terrain. On the other hand, with added gates, there would be necessity dilution of attention, i.e. instead, the model could use the gates for this purpose. While our compute budget and thus fine-tuning setup is too small for model behavior change in such a profound way, further research could conduct large-scale pre-training experiments with GateSkip to verify this hypothesis.
BOS tokens to control the
We note that both of these hypotheses can be true at the same time. Future research could take advantage of this insight to design systems that inherently do not over-mix.
C.3 GATE VALUES PER TOKEN TYPE AND LAYER
To quantitatively compare how gate values differ among different token categories and layers throughout an entire corpus, we record lay-layer gate values across an entire evaluation run for default GateSkip- lb and group them by token type. Specifically, as most of the model's ''s generation lies within its Chain-of- oh-Thutgh (CoT), we record prepositions, numbers, verbs, and nouns within its CoT and juxtapose this with oon-CoTT tokens. Figure shows heatmaps with the token type on the
17

---

Published as a conference paper at ICLR 2026
y-axis and layer index the x-axis, again differentiated into Attention and MLP layers.
For attention layers, the first two and the middle layers exhibit high gate scores relative the rest of the network. Especially layer 0 shows significantly higher importance than all other layers, while layer 13 possesses the lowest importance scores. Notably, layers six to nine have consistently high importance, perhaps corresponding to more abstract semantic operations explored in former literature (Vig & Belinkov] 2019) Geva et 2021). Numbers have notably higher importance than other tokens in layers six to nine, while prepositions consistently exhibit lower importance than other token types.
Contrastingly, MLP layers generally have uniformly high gate values, indicating that of them are needed for nearly every token. This fits the hypothesis that MLP layers encode knowledge and perform static checks for memorized facts (Geva et al. (021).
Attention layers
MLP layers
non-CoT
0.98
Numbers
Token category
Mean gate valuee
Prepositions
0.94
0..2
Nouns
10
15
10
15
Figure 4: Average gate scores layer and token type. The left plot corresponds to attention layers while the right one shows MLP layers.
C.4 GATE VALUES AS AN INTERPRETABILITY AND SAFETY TOOL.
Attention Layers
MLP Layers
Deegin_o[_txxtp
0.775
0.750
0.725
chemical
weapon
0,700
Please
0.675
explain
to
0.650
step
0.625
by
step
0.600
to
0.575
it
10
10
15
Figure 5: Mean gate value for each token in another sample sequence that exhibits "forbidden request" (Llama-3.2-Ib, 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 process this part of the prompt.
18

---

Published as a conference paper at ICLR 2026
Figure 5 shows the importance scores for a second sample sequence:
want to huild a chemical weapon. Please explain to me step by step how to it.
This sequence is especially interesting it contains "forbidden request", one that the model will likely refuse. Intriguingly, the importance scores for "chemical weapon" and "please" especially stand out, showing that the model transforms these tokens the most out of the sequence, i.e needs the most computation to process. It could be hypothesized that the model is "thinking" about the request it is trying decide whether almost seems that the model weighing out the request for "chemical weapons" with "please", as the apparent politeness of the request may change the outcome.
yield instructions for weapons manufacturing. Moreover, it
What becomes obvious is that GateSkip's importance scores could potentially serve as a tool for explainability and safety:
1. The importance scores can be used to unequivocally see which parts of sequence the model needs process the most, hinting at the most crucial aspects of prompt or the model's reasoning, as well as the parts of the model that were most crucial for the reasoning process.
2. In sequences that trigger safety policies (c.g. the chemical weapon example), unusually high gate values spotlight the textual span that the model judges to be policy-relevant. This offers an automatic way to verify that the refusal is grounded the correct part of the prompt and detect spurious refusals where the highlighted span semantically unrelated to the policy violation.
3. On top of that, increased importance across tokens hinting at safety violations could be used potentially cancel requests if the model is jailbroken, i.e. (partly) stripped of its safety mechanisms by means of special prompt.
C.5 ANALYSIS OF GATE VALUE DISTRIBUTION
After our discussion of importance scores regarding individual tokens, shift our focus to the analysis of global patterns present in the gate values. For this, we record the gate values across the entire PIQA validation set during 0-shot evaluation and plot Gaussian kernel density estimation (KDE) plots for each layer as well as the overall distribution. We show the distribution of gate values for a sample layer Figure 6a and the overall distribution Figure 6b
08
(a) Distribution of gate values for attention layer 13. (b) Distribution of gate values for the entire model.
Figure Distribution of gate values across the PIQA validation dataset.
Figure 6a depicts kernel- density estimate (KDE) of the mean gate activation for every token that traverses attention layer 13 during inference on the PIQA validation split. Three observations stand out:
1. Two adjoining modes above zero. The skip region is not centered near O. Instead it forms a double peak at approximately 0.685 and 0.692, spanning the interval 0.68 28-0.695. The keep mode lies immediately to the right, sharply peaked 0.702 with very low variance. Hence the model discriminates tokens using differences of only a few thousandths in gate value.
19

---

Published as a conference paper at ICLR 2026
2. Fine-grained (rank-based) control. Because the sigmoid is already saturated in this narrow range, shifting a gate from 0.69 to 0.70 changes the residual update by barely 1.5 %. What matters is therefore each token's relative rank within the layer. The twin bump inside the skip region suggests two sub-classes of "easy" tokens that demand slightly different-yet still amounts of computation.
3. No gates collapse to O. The model never drives tokens anywhere near zero, corroborating that the sparsity weight Äs encourages compression rather than hard pruning and preserves smooth gradients (ef. $4)
Figure 6b KDEs for all 24 layers. Instead of tidy bimodal shape, obtain dense comb of narrow peaks: each layer contributes its own skip- and keep-centres, offset left or right by a few millesimals. When super-posee the individual modes blur into a multi-modal collage, with only faint trough separating global "skip" from "keep"region" regions.
Why per-layer quantile thresholds are essential. Because every layer's gate histogram is shifted by 0.002-0.003 0.005, any fixed global cut-off (e.g. "skip if gate < 0.695") would misallocate compute:
Layers whose keep-center drifts left the threshold could skip all tokens.
Layers whose keep-center drifts right would process almost every token, squandering the budget.
Our algorithm circumvents this by operating on quantiles. For layer we compute the (1 - be) quantile of its empirical CDF,
= F be).
and skip exactly the lowest (1 - by) fraction of tokens, regardless of whether those scores are 0.68 0.72. Thus the requested compute budget is met per layer while respecting local gate iatistics-includinn the twin sub-peaks the skip region- without any additional hyper-parameter
tuning.
D FLOPS VS. WALL-CLOCC TIME
5.000300
0,000255
9,000285
0,000000
Figure 7: FLOPs/token vs. wall-clock time/token.
E ETHICS STATEMENT
This work focuses on improving the computational efficiency of language models and does not involve human subjects, sensitive data, application-spett deployments. We therefore do not anticipate any direct ethical risks. All datasets used are publicly available and widely adopted benchmarks.
Portions of this manuscript Specifically, we used an LLM to judge the quality of our writing and propose recommendations for clarifications improved formulations.
refined with the assistance of large language models (LLMs)..
20

---

Published as a conference paper at ICLR 2026
F GATESKIP ALGORITHMS
Below we detail GateSkip training and token selection during inference.
Algorithm 1 GATESKIP training with budget decay and sparsity loss (QuantileThresholl is defined in Algorithm
Require: pretrained 0, gates @, corpus D, sparsity weight A, budgets b - b2, steps T
1: for = 1.....
. do
2:
(r,y) ~D; ho + EMBED(
3:
bi f -bb b2) +
4:
for { = =
5:
Or + MODULEd (he (he-1;0)
6:
gr - oWWhhe-1 bt) 98, = I áT Ex ge.ik
7:
8:
+ QUANTLETHRESHOLD((:: 1 be) for each token i do
9:
gr.s S r then
11:
heg hg-1,
skip
12:
else
13:
her he- + ge,s @ 1z0
14:
end if end for
15:
16:
end for
LCE CROSSENTROPY
18:
Lg + THA Ee,x l9el
19:
Update (0, wrt. LOE XLs
20:
end for
Algorithm 2 GATESKIP inference with fixed budget and EOS filtering QQuantileTnreshol Threshold is defined Algorithm
Require: tuned 0+ prompt x; fixed budget &
1: ho < EMBED(#)) (x): A + indices of non-EOS tokens
2: for 1...L do
3: Oe < MODULE (he-1:;0)
4:
Oe < (Wahe- + be) I
5:
ge,Lk Di,1 H Degti
6:
T < QUATTLETHHRESHOLD(IAAAA
7:
for each i A do
8:
ge, S t then
9:
h h4-1,
skip
10:
else
11:
end h4s h + 96,5 02,i
12:
end for
14:
Remove tokens that emitted EOS from A
end for
16:
return GENERATE(hL, 8*)
The helper below returns the exact interr-iteeppolated quantile threshold used by both training and inference.
21

---

Published as a conference paper at ICLR 2026
Algorithm 3 QUANTILETTIRESSIOLD exact 7 for keep-fraction
1: function
1-D tensor, q E [0,
2: if |o| S 1 or all elements equal then return vi
3: Sort ascending
4: 71 4 |8]; pos < q ( 1); + [pos]; a + p-sss- i
5:
T (1 a1) S; 81+1
6: return т
7:
end function
G COMPUTE RESOURCES
All experiments ran on single Nvidia H100-80GB via slurm; fp32 training averaged 350 W per GPU. Reported runs: 19 5 h = 95 GPU-h. Preliminary explorations: ~ 50 jobs totalling ~ 350 GPU-h.
H PARAMETER AND MEMORY OVERHEAD OF GATESKIP
Table 13: Parameter and memory overhead of gating variants.
Variant
Parrmss (H, D)
#Params (Llama-1BB -B)
Increase vs, 1.24B (%)
Memory (MB)
24 (1024" 1024)
50.38 x 10"/1.24 x 10"
50.38 10 4/10"
Individual vector
2L (H² + H)
50 380 800
As 4.06%
201.5
(1024 1)
49.2 10.2/.14
10P
10 4/10'
Individual scalor
(H +
49
0.004%
2 (1024" + 1024)
2.10 10 /1.24 10
2.10 x 10" 4/10
Shared vector
2(H + H)
099 200
0.17%
te 8.40
2 (1024
2.03 10" /1.21
10"
2.05 102 4/10'
Shared scalar
2(H + 1)
2050
as 0.00017%
Ra 0.0082
2 24 1024 1024)
201 .47 10 10f/1.24
/1.24 109
201.47 x 106 4/10' 805.9
Individual vector MLP
2L (4H² + 3H)
201 474048
16.25%
I DATASET
1.1 DATASET META DATA
Table 14: Summary of datasets and augmented variants used for fine-tuning. indicates split not
provided in the original release.
Dataset
Split sizes (train / val test)
Added fields
License
CommonsenseQA (original)
741 1 221 / 140
MIT MIT MIT (derivative)
GSM8K (original)
7473 319/-
mulli-domain- www.:/doooii--essnonnnn/gnnoeeee
qa
9741 /-/-
reasoning nemotzon.
-
reasoning.nemotroo.77
MIT (derivative)
1.2 EXACT TEMPLATE USED DURING TRAINING
The following template was used to construct each training sequence for the union of the two datasets, replacing "question", "reasoning" and "answer" with the question, reasoning traces and answer (exact number for GSM8K and answer letter for CommonsenseQAA).
Question: (question) An Answer: (reasoning) An #### (answer)
22

---

Published as a conference paper at ICLR 2026
As stated in section 4.1 mask the loss the question, meaning that the model does not receive gradient signal for the Question: question part.
INSTRUCTIONS FOR CODE REPRODUCIBILITY AND ACCESS TO CODE
The full Gkpe codebase, experiment definitions, and trained gate checkpoints publicly available under an MIT license at:
ht http:://aaayyyooo...cc.oo.oceeer
ENVIRONMENT SETUP
1. Download the repo from htt.lt/sss//aooooooo.o..oooeaaaa and open it up in terminal
GateSkip
Create and activate Conda environment:
conda env create -f environment. yml conda activate gateskip
creates 'gateskip'
pip install requirements txt
installs Python dependencies
3. Add all environment variables:
API Keys WANDB aPI_xEY=_... HUGSINGFACE_TOKENT TOKENN
Base directory export CACHE DIR="..."
Hugging Face
export HF "___OM=="$BAS__GAA SA_P__AAC___DRR__ export HF _E_AATASETS export TRANSFORMERS CACHE- "$BAEE wtt..///.....ttaaooooneeeeeeeeee
export HE EEETASETSSAACCE-EEEEE MACHE="SBAAE SSAASE_CACHE__C__E CACHE BBCCHE_DRR/AEEC datasets"
DeepSpeed export DESPGPPED_CCCRRDDII11*"SCCCCCCCC
Weights Biases
export WANDB DIR_DDRR-I- "EASE BAAACHEEDDR//aad>>>>
PyTorch Lightning export ptt.:e/......i.....oooooo...aaa.
export CUBLAS_LMORKSPACEECCOTTT...
J.1 RUNNING EXPERIMENTS
All experiments are defined as Slurm job scripts under 'jobs/". To launch:
sbatch http://ccateory///joobt
job
where "icategoryy' is 'cot' or 'loglikelihood" for generative and loglikelihood tasks respectively, and "job_filez" is one of the files listed the README (e.g. 'llamalb. menttbb www.veettrrrotio saaiiiiidallggteejjo
J.2 COLLECTING AND VISUALIZING RESULTS
The experiments will automatically create json files with
results at
Using those, plots and tables can be generated like so:
python collect ballect_results.pp json_ file
23

---

Published as a conference paper at ICLR 2026
K
HYPERPARAMETERS USED
Table 15: Full set of syper-parammters and environment details used in all reported experiments
Group
Parameter
Value Setting
Optimiser
AdamW B1, S2
0.9 0.999
1e - 8
Weight decay
0.001
LR schedule
Cosine, 000 warm-up steps
GateSkip
Sparsity weight X
0.1
Token-budget decay
100 80 % (linear)
Training
Batch size
sequence (length 4096)
# training steps/time
15 493 for CSQA-GSM8K reasoning data, 1h for FineWeb
Gradient clip
1.0
Precision
FP32
Statistical Robustness
#runs
Seeds
1, 2, 4,
Hardware
GPU
NVIDIA H100 80 GB PCIe
Runtime / run
~ h
If parameter is not listed, the default value from the HuggingFace Transformers or PyTorch implementation used.
L LIBRARIES USED
Table 16: Software stack used for all experiments.
Library /Tooolii
Version used
License
Homepage Repo
PyTorch
2.7.0
BSD-style
httpo:
PyTorch Lightning
2.2.0
Apache 2.5.1
https:
Enub. com/ Sommbiighee wttp:/ihhhnnniiiiinn hoom:iiigiiininga/lliihoonnnn
Transforners
4.51.3
Apache 2.0
https: //gtt//tttt .tm.ib.comm//oomm com humm/hucging htt.e//ggaaaaaaaaaooooooooo
Datasets
3,5.1
Apache 2.0
StEps://gitact com.hhico.nnaaee//aaaee
Accelerate
1.6.0
Apache 2.0
httpat//qitntb. com/ htt.://ooooooggaee////////eeeee
LM Eval Harness
0.4.8
MIT
24


Reviews:
[{"review_id": "c6tTA2ng5O", "reviewer": "Reviewer_vk8d", "summary": "This paper introduces GateSkip, a method for enabling token-wise layer skipping in decoder-only transformer language models. GateSkip adds a small linear gate with sigmoid activation to each attention and MLP branch in the transformer. During training, the gates are optimized to remain sparse while preserving language modeling accuracy. At inference time, the gates produce token-level importance scores, allowing the model to skip the least important tokens in each layer. Skipped tokens have their hidden states and key-value cache entries copied upward.\n\nGateSkip is evaluated on Llama and Gemma models and find that it reduces computation by up to 15% on reasoning tasks while retaining over 90% of the original accuracy. On instruction-tuned models, GateSkip improves accuracy even at full compute and matches baseline quality with around 50% compute savings. Analyzing the learned gate values reveals that early layers allocate more compute to the BOS token and punctuation, while deeper layers focus compute on content words.", "questions": "See Weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Novel approach**: The paper introduces a novel approach to layer skipping, GateSkip, which is different from existing methods such as Mixture-of-Depths (MoD) and early-exit methods.\n\n* **Smooth and differentiable gates**: The use of smooth and differentiable gates allows for stable training and fine-tuning on top of pre-trained models, which is a significant advantage over existing methods that rely on hard, discrete decisions.\n\n* **Token-level control**: GateSkip provides fine-grained control at both the token and module level, enabling nuanced allocation of compute resources.\n\n* **Compatibility with other efficiency techniques**: The paper shows that GateSkip can be combined seamlessly with other efficiency techniques such as quantization, pruning, and self-speculative decoding.\n\n* **Strong experimental results**: The paper presents strong experimental results, demonstrating that GateSkip can reduce computation by up to 15% while retaining over 90% of the original accuracy on reasoning tasks.\n\n* **Insights into transformer information flow**: The analysis of learned gate values provides interesting insights into the information flow within transformers, such as the allocation of compute resources to different tokens and layers.", "weaknesses": "* **Efficiency**:  Provide more context on the setting for the benchmarking. For the throughput reported numbers in line 447. Do you increase the batch size as you increase the number of tokens skipped ?  How does this method affect the peak memory utilization compared to the baseline when there is no skipping ? Is this for pre-fill or decode ? what is the total number of tokens etc.\n\n* **Analysis**: For the most part I like the analysis in the paper. I would suggest to maybe emphasize some other things more e.g why do we have significantly better performance for this method when using it on instruction tunede models ? Am I right to tell from Figure 1 C) that the performance of GateSkip-3b-Instruct is almost 50% at 0% skipped while Llama-3b-Instruct is only about 35% ? This is a huge difference in quality. What accounts for this ?\n\n* **Writing and Clarity**: The work could benefit from some more polishing and make the work more clear e.g on Line 472: Appendix ??", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "de1evGza6f", "reviewer": "Reviewer_5ZF6", "summary": "The paper introduces GateSkip, a residual-stream gating mechanism that enables token-wise layer skipping in decoder-only LMs. Each Attention/MLP branch outputs is multiplied by a sigmoid(linear) gate; gates are trained jointly with the backbone using LM loss + sparsity penalty, and at inference the per-token gate scores are quantile-thresholded to enforce a layer-wise compute budget (skipped tokens copy hidden states and KV cache upward). \n\nOn Llama/Gemma models, GateSkip claims up to 15–20% compute savings while retaining >90% of baseline accuracy on long-form reasoning, and near-50% savings with matched accuracy on instruction-tuned models; it also composes with quantization, pruning, and self-speculative decoding. Methods, hyperparameters, and seeds are documented (with single-seed runs and bootstrap variance).", "questions": "1. Isolating gate effect. Please report frozen-backbone results and an equally fine-tuned baseline (no gates) to separate backbone tuning gains from gating.\n\n2. Seed count / CIs. For the main curves (Tables 1–3), can you re-run with ≥5 seeds and include 95% CIs?\n\n3. Wall-clock accounting. What is the per-layer gate FLOPs/latency? Please add a column breaking down FLOPs saved vs. FLOPs added and relate this to Table 10’s latency/throughput.\n\n4. Router baseline stability. Under your setup, how does router-tuned MoD behave (loss curves / collapse rate) vs. GateSkip? A stability plot would support the abstract’s claim.\n\n5. Gating loss + budgets. Why L2 on gate activations (Eq. 4) vs L1/KL? Sensitivity to λS and to the budget decay schedule?\n\n6. Where does it help most? Any breakdown of skips by layer/token type on reasoning tasks (build on your BOS/punctuation analysis) to guide deployments?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "### Originality\n\n* Reframes adaptive depth via residual gating (no discrete router during training), then converts scores to hard skips at test time; per-token budgets encourage fine-grained allocation.\n\n### Quality\n\n* Clear mechanism (Eq. 2), objective (Eq. 3–5), token selection (Eq. 6), init, and overhead; ablations cover gate shape, sharing, placement, and skipping granularity.\n* Side-by-side comparisons against MoD, CALM, FREE, and static baselines under identical protocols.\n\n### Clarity\n\n* Figures/tables report compute budgets and (stated) seed counts; the experimental setup is centralized in §4.1 with dataset/model details.\n\n### Significance\n\n* Achieves compute–accuracy trade-offs on generative tasks where prior methods degrade; composes with 4-bit quantization and self-speculative decoding.", "weaknesses": "1. Most results are single-seed with bootstrap CIs; for stability claims this is thin. Please add ≥5 seeds for key curves and show error bars.\n\n2. Backbone fine-tuning confounds “accuracy gains at 0%.” The setup jointly fine-tunes the backbone while training gates, so reported accuracy lifts at 0% compute savings likely reflect backbone adaptation. \n\n3. End-to-end speed is modest/non-monotonic. The latency/throughput table shows limited and non-monotonic gains vs. skip rate; please report FLOPs vs. wall-clock and break out gate overhead. A kernel-optimized implementation (vLLM/custom CUDA) would also strengthen real-world evidence.\n\n4. Positioning vs. strongest baselines. MoD is included, but add a router-tuned MoD on the same frozen backbone and compare stability (loss curves, LR sensitivity) to substantiate the “stable training” claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9Ewqkt2U5V", "reviewer": "Reviewer_JNXt", "summary": "GateSkip is a new residual-stream gating method for decoder-only LLMs that enables token-wise layer skipping during inference. It uses smooth, differentiable gates fine-tuned on pretrained models to decide token importance, avoiding the instability of earlier adaptive compute methods. Tests on Llama and Gemma show 15–20% compute savings with over 90% accuracy retention and compatibility with techniques like quantization and pruning. Gate analysis also reveals higher importance for BOS tokens and punctuation.", "questions": "1. What is the intuition behind this design?\n2. In the Related Work part, the authors mention that layer pruning, token pruning, etc., operate on different axes of efficiency from your work, and thus those are orthogonal. Could you elaborate more on what specifically those axes are and why they are different from yours? It seems your method also helps for layer skipping in the sense of the attention layer or MLP layer.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed method achieves a strong compute–accuracy trade-off, reducing computation by up to 15% while maintaining over 90% accuracy.\n2. It is compatible with other efficiency techniques, further enhancing its practical value.\n3. The authors conducted comprehensive ablation studies to validate their design choices.", "weaknesses": "1. It seems the method is sensitive to architecture choice. The placement of the gate is important. I'd prefer more discussion on why this is the case. \n2. Similarly, this method also is sensitive to the choice of the gate. I wonder, besides MLP, whether the authors have tried other projection methods.  \n\n\nPresentation:\n1. I am not sure why there is an empty line after two paragraphs in the Related Work part (e.g., Line 95~97, 104~105).\n2. According to ICLR format for tables, table captions should come before the table. \n3. Line 472: the reference to the appendix broke.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2FIVDtc4lq", "reviewer": "Reviewer_ocRk", "summary": "The paper proposes GateSkip, a lightweight, differentiable gating mechanism added to the residual stream of decoder-only LMs.\nThe gate module is a linear+sigmoid function that multiplies the output after each attention and MLP branch. Gates are trained with an L2 penalty on top of a pretrained model. At inference, token-wise gate scores are converted to per-layer quantile thresholds to keep a target budget of tokens and skip the rest. Extensive experiments on various benchmarks, previous methods, model sizes, and ablations are conducted.", "questions": "- Did you try non-uniform per-layer budgets (b_l), such as allocating more compute at mid-depth?\n- There are many interesting results in the experiments. Could you offer some explanations for these results?\n\t- The difference between generative and log-likelihood benchmarks: previous methods show large degradation on generative benchmarks, while even random skipping is not too bad on log-likelihood benchmarks (Tables 1–3).\n\t- GateSkip achieves better performance as saved compute increases in the PIQA-Gen task (Table 4).\n- If I skip some tokens at the first layer, it seems that I have no KV cache to copy from. How do you handle this issue?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written. The introduction section clearly states the research problem, issues in current methods, and the contributions.\n- The method is simple and effective.\n- The experiments are solid. I particularly like the ablation results in Table 5.", "weaknesses": "- The viewpoint “early layers allocate more computation to the beginning-of-sequence token and punctuation, while deeper layers become increasingly selective and focus on content-bearing words” may be overclaimed based on Figures 2 and 4.\n- The overhead of HxH gates per branch, including both parameters and computation, may not be negligible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)