L ayer S kip: Enabling Early Exit Inference and Self-Speculative Decoding (2024)

Abstract: 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 at https://github.com/facebookresearch/LayerSkip.

## introduction large language models (llms) have been deployed to many applications, yet their high compute and memory requirements lead to high financial and energy costs when deployed to gpu servers @xcite . acceleration solutions do exist to deploy to commodity gpus on laptops but they suffer from significant drop in accuracy @xcite . accelerating llms further to mobile or edge devices is still an active research * equal contribution. moreover, a popular research trend in llm acceleration is speculative decoding @xcite ; @xcite that has no drop in accuracy, where a large model, referred to as the main model, is accompanied with a faster model, referred to as the draft model. the advantage of speculative decoding is that it leads to faster inference compared to the main model, but requires a larger memory footprint and complexity in implementation to maintain key-value (kv) cache in two different models. in addition to exiting early, this paper also proposes combining exiting early with speculative decoding to propose a self-speculative decoding approach that does not require an additional model or auxiliary layers. the contribution of this paper is an end-to-end solution: ## proposed solution our approach has three different stages: ## experiments we would like to evaluate our training recipe on different types of training, whether pretraining from scratch or finetuning. to verify our approach, we run different types of training experiments: ## conclusion we show that combining layer dropout & early exit loss with curriculum, improves accuracy of early exit during inference, and developed a novel selfspeculative decoding solution that led upto 1.86× speedup. we hope this encourages researchers to adopt the proposed recipe in pretraining and finetuning. in the future, we can increase accuracy of earlier layers to obtain better speedups for selfspeculative decoding, e.g., by combining with dynamic conditions (like @xcite ). figure a1 shows sample generations exiting at different layers for llama2 7b and llama2 13b pretraining from scratch figure a3 shows sample generations exiting at different layers for llama2 7b pretrained with and without layerskip.