LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding (2404.16710v4)

Abstract: We present LayerSkip, an end-to-end solution to speed-up inference of 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.

• The paper introduces a hybrid strategy combining self-speculative decoding and layer dropout to achieve early exit inference without a separate draft model.
• It applies early exit loss and layer dropout during training, enabling robust predictions from early layers and reducing dependence on deeper layers.
• This efficient framework speeds up LLM inference by up to 2.16x, making it ideal for deployment in resource-constrained environments.

Enhancing LLM Efficiency with Self-Speculative Decoding and Layer Dropout

Introduction to Self-Speculative Decoding and Layer Dropout

This research explores innovative methods for improving the efficiency of LLMs by focusing on two main techniques: self-speculative decoding and layer dropout. The paper proposes a hybrid approach that combines existing acceleration strategies with these new techniques to achieve substantial inference speedups without compromising the model's accuracy.

Combining Early Exit and Speculative Decoding

Speculative decoding has been identified as an effective strategy for enhancing LLM inference speeds. Traditionally, this involves using two models: a fast, less accurate draft model and a slower, more accurate main model for verification. This paper introduces a self-speculative decoding technique that eliminates the need for a separate draft model by utilizing early exits within a single model framework. The approach leverages early layers of the model for quick draft predictions and later layers for verification, thus optimizing memory usage and reducing complexity.

Training Techniques: Layer Dropout and Early Exit Loss

A significant contribution of this paper is its dual training strategy, which employs both layer dropout and early exit loss. This strategy is designed to make models less dependent on deeper layers, allowing for accurate early exits during inference:

1. Layer Dropout: Implemented by stochastically skipping layers during training to encourage model robustness and decrease reliance on later layers.
2. Early Exit Loss: Augments the training process by applying a loss function at various layers, training the model to make accurate predictions at earlier stages.

The combination of these methods not only improves inference speeds but also helps in maintaining high accuracy even when the model exits early. This leads to an end-to-end solution where models are trained to perform well under a truncated layer setup.

Practical Implications and Theoretical Advancements

The paper reports speedups ranging from 1.34x to 2.16x, depending on the task, without a notable drop in accuracy. These results are important for deploying LLMs in environments with limited computational resources, such as mobile or edge devices. Theoretically, the introduction of self-speculative decoding presents a new avenue in LLM research, focusing on the interplay between early layer accuracy and overall model efficiency.

Future Directions

While the proposed methods show promising results, further work is required to explore the full potential of these techniques. Future research could focus on dynamically choosing exit points based on token complexity or improving the early layers' predictive power directly through advanced training regimens. This could lead to even greater efficiency gains and open up new uses for LLMs in real-time applications.

Concluding Remarks

This paper presents a compelling method for enhancing the efficiency of LLMs through innovative training and inference techniques. By integrating layer dropout with self-speculative decoding, it sets the stage for more resource-efficient LLMs capable of maintaining high accuracy. These advancements are crucial for the wider adoption of LLM technologies in resource-constrained environments.