Source-linked AI summary
Learning to Skip for Language Modeling
Dewen Zeng, Nan Du, Tao Wang, Yuanzhong Xu, Tao Lei, Zhifeng Chen, Claire Cui
TL;DR
Large language models spend substantial computation on every input, although inputs differ in complexity and importance. SkipLayer uses context-dependent routing to skip layers for selected tokens. Across 24 NLP tasks, it improves 1-shot performance over competitive baselines at mild extra inference cost.
Problem
Large language models activate all parameters for each input, while allocating uniform computation can waste effort on easy or redundant examples.
Source: §1 · p. 1Method
SkipLayer uses a context-dependent router to let tokens skip arbitrary wrapped layers, with differentiable training and discrete forward decisions.
Source: §1 · p. 1, §2.1 · p. 2Results
SkipLayer significantly improves 1-shot performance across 24 NLP tasks compared with competitive baselines, with only a small extra inference cost.
Source: §6 · p. 8Takeaways & Limitations
SkipLayer assigns heterogeneous computation to tokens, allowing more resources to be used to improve predictive quality for harder tokens.
Source: §6 · p. 8The Top-1 formulation did not precisely control model sparsity because its gate remained a soft approximation of a binary decision.
Source: §2.1 · p. 2
Abstract
from arXiv · showhide
Overparameterized large-scale language models have impressive generalization performance of in-context few-shot learning. However, most language models allocate the same amount of parameters or computation to each token, disregarding the complexity or importance of the input data. We argue that in language model pretraining, a variable amount of computation should be assigned to different tokens, and this can be efficiently achieved via a simple routing mechanism. Different from conventional early stopping techniques where tokens can early exit at only early layers, we propose a more general method that dynamically skips the execution of a layer (or module) for any input token with a binary router. In our extensive evaluation across 24 NLP tasks, we demonstrate that the proposed method can significantly improve the 1-shot performance compared to other competitive baselines only at mild extra cost for inference.
1. Introduction
SkipLayer targets the waste and cost of uniform computation in large language models by routing tokens through or around layers. The paper reports strong 1-shot performance with controllable computation tradeoffs against competitive baselines.
1. Introduction
- Activating all model parameters for every input makes large-scale language-model prediction computationally costly, with trillions of FLOPs per prediction.The passage also identifies training and serving these models as non-trivial.Source: §1 · p. 1
- Uniform computation may waste effort on easy or redundant examples, while harder inputs may warrant more computation.The introduction motivates allocating computation according to input hardness.Source: §1 · p. 1
- Conditional computation can reduce per-example computation, but discrete input-dependent activation decisions make efficient end-to-end differentiable training challenging.Source: §1 · p. 1
- SkipLayer lets an input skip any wrapped layer based on its contextual representation, while supporting differentiable end-to-end training and discrete forward decisions.The overview depicts a router activating or skipping embedded layer logic based on input context.Source: §1 · p. 1, Figure 1 · p. 2
- SkipLayer models show strong 1-shot performance with controllable computation tradeoffs compared with competitive baselines.The framework is evaluated across open-domain QA, reading comprehension, common-sense reasoning, and natural language inference tasks.Source: §1 · p. 2
2. Method
SkipLayer wraps a layer with a context-conditioned binary router so tokens either undergo its computation or bypass it. The method pairs budgeted routing with sparse FFN execution and Transformer-specific handling of attention context.
2.1. SkipLayer
- A context-conditioned binary router sends selected tokens through a wrapped layer and bypasses it for the rest.Straight-Through Gumbel-Softmax uses binary forward decisions and soft backward probabilities; capacity constraints prioritize hard examples while skipping easy ones.Source: §2.1 · pp. 2–3
2.2. Efficient Implementation
- Dynamic gather-and-scatter groups non-skipped tokens for sparse FFN computation, while skipped inputs pass directly to the output.Group size trades off wasted slots in oversized groups against reduced parallelism in undersized groups.Source: §2.2 · p. 3, Figure 2 · p. 3
2.3. SkipLayer for Transformer-based Models
- The method wraps an entire Transformer layer, allowing tokenwise FFN skipping while preserving key and value projections for attention context and omitting skipped tokens’ queries.During decoding, key and value projections are cached; query projection and attention are computed only when the router activates the layer.Source: §2.3 · p. 4
3. Experiment Setup
The experiments use a 1.6-trillion-token pretraining corpus and compare SkipLayer variants with several baselines. Models are evaluated under a 1-shot protocol across 24 NLG and NLU datasets.
3. Experiment Setup
- The pretraining corpus contains 1.6 trillion tokens from filtered webpages, books, Wikipedia, conversations, forums, and news.An in-house classifier estimates webpage content quality for filtering.Source: §3 · p. 5
- All models use a 1,536-dimensional model representation and 8×-width FFN, with shared learning hyperparameters and a 1,024-token training sequence length.Training uses batches of 256 sequences, Adafactor, and an auxiliary loss weight λ of 0.1.Source: §3 · p. 5, Table 1 · p. 5
- Each benchmark example is evaluated with one randomly selected training example as its sole demonstration, using the prompting format of Radford et al. (2018).The demonstration and evaluation example are concatenated with two newlines.Source: §3 · p. 5
- The 1-shot evaluation covers 24 datasets: four NLG tasks scored by exact match and 20 NLU tasks scored by length-normalized option likelihood.Overall NLG and NLU performance is averaged across datasets.Source: §3 · p. 5
- The baselines include dense standard Transformers, HighwayNet gating on the FFN, and random gating; WideFFN doubles FFN width and applies SkipLayer only to the FFN.At a 50% skip probability, WideFFN keeps prediction FLOPs unchanged because its FFN is twice as large.Source: §3 · pp. 5–6
4. Full Results
Across NLG and NLU, deeper, sparser SkipLayer models improve 1-shot performance at modest extra FLOPs. The results also compare decoding and training efficiency, learned versus random gating, and token-level skipping behavior.
4. Full Results
- 51.5% performance gain at 18.2% increased compute FLOPs for SkipLayer (24L, 25%) versus the 6L baseline; deeper, sparser models also improve NLU performance.Other reported gains are 28% and 22% versus the 12L and 24L baselines, respectively, at 19.2% and 20% increased FLOPs.Source: §4 · p. 6
- SkipLayer (96L, 25%) achieves 20% decoding efficiency with a 0.5% quality loss versus the 48L baseline; SkipLayer (12L, 50%) decodes in nearly the same time as the 6L baseline.SkipLayer (48L, 50%) gains 8% in 1-shot performance for a 6% increase in per-step decoding time versus the 24L baseline.Source: §4 · p. 6
- Compared with their respective full dense baselines, SkipLayer (24L, 50%) gains 18% training speed and SkipLayer (48L, 12.5%) gains 3x.Training speed generally decreases as FLOPs per prediction increase.Source: §4 · p. 7
- SkipLayer models achieve the best 1-shot performance among the compared methods across NLG and NLU, including as model scale increases.Figure 5 compares methods at 6 and 12 effective activated layers; for NLG, the 12L dense baseline scales better than WideFFN, HighwayNet, and random gating.Source: §4 · p. 7, Figure 5 · p. 7
- At the same model density, learned SkipLayer gating performs significantly better than random gating, despite both allowing per-token parameter switching.The random-gating baselines have similar compute FLOPs per token at the same density.Source: §4 · p. 7, Figure 6 · p. 7
- During TriviaQA decoding, functional tokens skip more layers, while independent words such as “No” and “Paris” skip fewer.The analysis covers 300 frequent tokens from 500 sampled questions using SkipLayer (12L, 50%); a token’s skipping pattern can vary by context.Source: §4 · p. 8, Figure 7 · p. 8
5. Related Work
SkipLayer differs from early-exit and learned-activation approaches by allowing input-dependent skipping across layers without soft approximations in the forward pass. It can also complement MoE layers and differs from random structural dropout by learning skipping patterns from data.
5. Related Work
- Unlike early-exit methods that activate computation in bottom layers, SkipLayer permits 2^L compute paths across L stacked layers.Early-exit methods use confidence-thresholded classifiers to leave the stack early.Source: §5 · p. 8
- SkipLayer avoids forward-pass soft approximations, yielding computation savings during both inference and costly language-model pretraining.Prior soft-approximation and reinforcement-learning approaches address discrete activation decisions.Source: §5 · p. 8
- Concurrent work, including CODA and CoLT5 (Ainslie et al., 2023), applies similar token selection to activate Transformer layers, but only conditionally activates encoder layers in encoder-decoder models such as T5.Source: §5 · p. 8
- SkipLayer is orthogonal to MoE and can wrap an MoE layer, while conditionally activating either self-attention or FFN rather than focusing mainly on FFN.MoE layers sparsely activate a subset of experts.Source: §5 · p. 8
- Unlike structural dropout’s random layer dropping and uniform inference computation, SkipLayer learns skipping patterns from data and can assign non-uniform computation per example.The paper reports better performance than the random-skipping baseline and describes non-uniform inference computation as potential.Source: §5 · p. 8
6. Conclusions
SkipLayer dynamically skips arbitrary layers based on input context, enabling heterogeneous computation for tokens. Across 24 NLP tasks, it improves 1-shot performance over competitive baselines with only a small extra inference cost.
6. Conclusions
- SkipLayer dynamically skips arbitrary layers according to input context using a routing algorithm.The method supports heterogeneous computation for tokens of different complexity or importance.Source: §6 · p. 8
- The approach assigns more computation resources to harder tokens to improve predictive quality.Source: §6 · p. 8
- Across 24 NLP tasks, SkipLayer improves 1-shot performance over competitive baselines with only a small extra inference cost.Source: §6 · p. 8