Source-linked AI summary
LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, Carole-Jean Wu
TL;DR
LayerSkip addresses the high compute and memory cost of LLM inference and the difficulty of making early exits accurate. It trains models with layer dropout and shared early-exit supervision, then uses self-speculative decoding to verify and correct early predictions. The approach reports speedups across tasks, reaching 2.16× in the reported experiments.
Problem
LLM inference has high compute and memory requirements, while reducing layers can threaten accuracy and ordinary models often rely on later layers.
Method
LayerSkip combines depth-dependent layer dropout with shared early exit loss, then uses early layers to draft tokens and remaining layers to verify and correct them.
Results
1.34×–2.16× speedups are achieved depending on the task.
Takeaways & Limitations
The training recipe improves early-exit accuracy, while self-speculative decoding provides up to 1.86× speedup without a separate draft model.
Takeaways & Limitations
The method requires finetuning or recipe-based pretraining, and its hyperparameters require tuning to avoid reduced final-layer accuracy.
Abstract
from arXiv · showhide
We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.
1 Introduction
LayerSkip is an end-to-end approach to accelerate LLM inference by training models for earlier exits and correcting early predictions with later layers. It achieves task-dependent speedups while avoiding separate draft models.
- Motivation: Unlike conventional speculative decoding, LayerSkip avoids maintaining key-value caches for two different models, reducing memory and implementation complexity.Conventional speculative decoding pairs a main model with a faster draft model.
- Training and inference: LayerSkip combines layer dropout and early exit loss to support inference using subsets of transformer layers.The training recipe uses lower dropout rates earlier and higher rates later, with a shared exit.
- Training and inference: Its inference method exits at earlier layers, then verifies and corrects predictions using the remaining layers.This self-speculative decoding design uses the same model for drafting and verification.
- Results: 1.34×–2.16× speedups are reported across tasks.The reported range depends on the task.
2 Motivation
The motivation for LayerSkip is that tokens often require fewer layers than the full model, but ordinary LLMs distribute computation across layers and early exits can reduce accuracy. The paper addresses this by training earlier layers to predict reliably and by verifying early predictions with later layers.
- 2.1 Exiting Earlier in LLMs: A HumanEval example shows that tokens can require far fewer than all model layers, motivating earlier exits.In the example, a token required 23.45 of 32 layers on average, and even perfect prediction of exit timing would save at most 26% computation.
- 2.1 Exiting Earlier in LLMs: Even straightforward tokens may require all layers, indicating that standard LLMs are not trained to make final predictions early.The token starting a for-loop required all 32 layers to predict “for”.
- 2.1 Exiting Earlier in LLMs: Layer dropout uses lower dropout rates in earlier layers and higher rates in later layers to reduce reliance on later computation.The paper describes this as skipping layers during training so the model becomes less dependent on later layers.
- 2.1 Exiting Earlier in LLMs: Early exit loss trains a shared LM head to interpret embeddings from every transformer layer.Unlike approaches with dedicated heads or extra modules, LayerSkip directly supervises lower layers for language modeling.
- 2.2 Correcting if we Exit Too Early: Figure 3 presents layer dropout and early exit loss as producing a model equivalent to an ensemble of models with different depths.This summarizes the intended effect of making multiple layer prefixes usable for inference.
- 2.2 Correcting if we Exit Too Early: Because early exits can reduce accuracy, self-speculative decoding verifies groups of early-generated tokens in parallel with the remaining layers and corrects them.The approach uses early layers for autoregressive generation and later layers for verification.
3 Related Work
Related work spans dropout-based regularization, layer skipping, early exit, and speculative decoding for accelerating neural networks and language models. LayerSkip builds on these areas by applying layer dropout to decoder-only language models and using it for early-exit inference.
- Dropout: Dropout regularizes training by stochastically replacing portions of layer outputs with zeros, while transformer usage has varied with model scale and training regime.The passage distinguishes unstructured dropout from layer dropout and notes examples ranging from a 0.1 transformer rate to zero pretraining dropout in PaLM.
- Layer Dropout: Layer dropout, also called stochastic depth, skips layers stochastically during training and has improved robustness to pruning layers at inference.Prior work progressively increased layer-dropout rates across layers or iterations, but decoder-only language models and large-scale language-model training remained underexplored.
- Early Exit: LayerSkip differs from prior early-exit work by proposing layer dropout for decoder-only models and using it to improve early-exit inference.The authors describe this application as novel relative to the cited literature.
- Early Exit: Early-exit methods add branch modules or auxiliary losses so deep networks can terminate computation before their final layers.Prior language-model studies explored encoder-only and encoder-decoder architectures, static or dynamic exits, and token- or sequence-level exit granularities.
- Speculative Decoding: Speculative decoding uses a fast draft model to generate tokens and a slower main model to verify and correct them in parallel.Self-speculative variants can use the same model with intermediate layers skipped, while LayerSkip's stages reuse activations and KV cache because they execute shared early layers identically.
4 Proposed Solution
LayerSkip combines layer dropout and early-exit loss during training, then uses early exit and self-speculative decoding to reduce inference computation and memory while reusing one model.
- LayerSkip has three stages: training with layer dropout and early-exit loss, early-exit inference, and speculative verification and correction.
- Training using Layer Dropout & Early Exit Loss: Layer dropout applies a per-layer, per-iteration dropout rate, with rates increasing exponentially from the first layer toward the last.
- Training using Layer Dropout & Early Exit Loss: Early-exit loss directly supervises intermediate layers through the shared language-model head, whose normalized layer weights can penalize later layers more strongly.
- Inference using Early Exit: During inference, the model can run only its first E layers and send the resulting representation to the language-model head, reducing autoregressive cost from L layers per sample to E layers per token when E < L.
- Verification and Correction using Speculative Decoding: With layer dropout and early-exit loss, early exit speeds autoregressive generation but still incurs an accuracy cost relative to the full model, motivating verification and correction.
- Verification and Correction using Speculative Decoding: Self-speculative decoding drafts tokens with the first E layers, verifies them with the remaining L − E layers, and reuses KV-cache information and the exit query.
5 Experiments
The experiments evaluate LayerSkip across multiple training settings and model configurations, including continual pretraining, code-data finetuning, and task-specific finetuning.
- The evaluation covers different training types, including pretraining from scratch or finetuning.
- Continual Pretraining: Continual pretraining uses diverse natural-language and code data, including 52B tokens with pretrained Llama2 7B and 13B models.
- Finetuning on Code Data: The experiments include finetuning a pretrained Llama1 7B model on 5.2B tokens of CodeLlama data.
- The study compares layer dropout only, early-exit loss only, and their combination, denoted LD, EE, and LD+EE.
6 Results
LayerSkip improves early-exit accuracy across pretraining, continual pretraining, and finetuning settings, while self-speculative decoding provides substantial speedups by verifying early-layer predictions with later layers.
- Early Exit Inference: LayerSkip improves accuracy at earlier layers across continual pretraining and pretraining-from-scratch experiments.Continual pretraining shows clear gains at earlier layers with minimal last-layer loss, while scratch pretraining improves earlier-layer accuracy but can slightly reduce final-layer accuracy.
- Code Finetuning: LayerSkip was clearly better than the baseline at earlier layers on code-finetuned tasks, while the combined configuration had almost the same final-layer accuracy.This experiment used domain-specific finetuning with escale = 1.0.
- Task-Specific Finetuning: 77% exact match at layer 12 was achieved on TOPv2 semantic parsing with LayerSkip, whereas the baseline reached 0 EM when layers were removed.The LayerSkip configuration incurred a 3% regression at the final layer.
- Self-Speculative Decoding: 2.16× speedup was reached in scratch-pretrained models, exceeding traditional speculative decoding.The larger model achieved the bigger speedup, opposite to the continual-pretraining trend.
- Self-Speculative Decoding: 1.82× speedup was obtained on HumanEval with no accuracy drop for a code-finetuned Llama 7B model.The experiment used 12 speculations and exited at layer 6.
- Self-Speculative Decoding: 2.0× speedup on TOPv2 was accompanied by token acceptance rates of 76.0%, 97.2%, and 98.9% at exits E = 6, 12, and 18.The evaluation used eight speculations and greedy generation of the next 80 tokens.
7 Ablation Studies
Ablation studies examine how training scale and inference caching affect LayerSkip, showing that KV-cache reuse consistently reduces per-token latency.
- Scaling with Pretraining Tokens: Perplexity scaling experiments vary the number of GPUs and training tokens while tracking last- and middle-layer performance on The Stack.Each experiment used 50,000 steps, per-device batch size 4, and context window 4096.
- Generation Ablations: Generation ablations compare baseline, LayerSkip layer dropout, early-exit loss, and their combination across continual-pretraining, scratch-pretraining, code, and TOPv2 settings.The cited tables report generation results for these configurations.
- KV Cache: KV-cache reuse consistently saves 9–20 ms per token across tasks.The effect was measured under the same inference setup used for self-speculative decoding.
8 Limitations
LayerSkip’s benefits depend on modifying training and tuning recipe-specific hyperparameters, with additional learning-rate tuning required for scratch pretraining.
- Training and Tuning: LayerSkip requires finetuning or pretraining with its recipe, unlike the cited alternative self-speculative approach, which does not change model weights.The recipe’s pmax, escale, and R hyperparameters must be tuned to avoid reducing last-layer accuracy.
- Training and Tuning: Scratch pretraining with layer dropout requires increased learning rates, making optimal-accuracy tuning potentially difficult and time consuming.This is identified as a limitation of the training procedure.
9 Conclusion
LayerSkip combines layer dropout and early exit loss with curriculum to improve early-exit accuracy, then uses self-speculative decoding to verify and correct predictions with later layers. The approach achieves up to 1.86× speedup, while future work targets higher early-layer accuracy and dynamic exit selection.
- Combining layer dropout, early exit loss, and curriculum improves early-exit accuracy during inference.
- Self-speculative decoding uses early layers for generation and remaining layers for verification and correction, achieving up to 1.86× speedup.
- The authors suggest applying layer dropout and early exit loss to pretraining and finetuning, including alongside parameter-efficient methods such as LoRA.
- Future work aims to improve early-exit-layer accuracy and dynamically choose exit layers per token to increase self-speculative decoding speedups and token acceptance rates.
A.1 Experiment Details
The experiments specify training configurations, model architectures, evaluation tasks, and self-speculative decoding procedures. Results include near 2× CPU speedup on TOPv2 and higher early-layer accuracy from layer dropout under higher learning rates.
- Training configuration: Training experiments document hyperparameters, configurations, learning rates, and model architectures in Tables 8–10.
- Training configuration: When pretraining from scratch, layer dropout produces higher accuracy when trained with a higher learning rate.
- Evaluation tasks: Evaluation spans classification tasks, generation tasks, and perplexity on The Stack and Books datasets.
- Evaluation tasks: The classification evaluation includes common-sense reasoning benchmarks such as BoolQ, PIQA, SIQA, HellaSwag, Winogrande, ARC, OBQA, COPA, RACE, and MMLU.
- Evaluation tasks: The generation evaluation includes Natural Questions, Textbook Question Answering, MATH, GSM8K, HumanEval, and MBPP.
- Inference procedure: Near 2× CPU speedup is reported for TOPv2 using the first 100 test samples, seven speculations, and greedy generation of the next 50 tokens.
- Inference procedure: Self-speculative decoding drafts tokens with the first E layers, verifies them with the remaining L-E layers using cached states, and updates the cache after matching.