Opening article
Self‑speculative decoding, introduced in the LayerSkip paper, lets a single large language model draft tokens with its early layers and verify them with deeper layers, reducing inference time and memory use. The method relies on a training recipe that adds early‑exit loss and progressive layer dropout, enabling the model’s LM head to unembed logits from intermediate layers. Hugging Face’s transformers library now supports this technique via the assistant_early_exit argument, and several Llama‑based checkpoints have been pretrained with the LayerSkip recipe.
Self-speculative decoding, proposed in LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding is a novel approach to text generation. It combines the strengths of speculative decoding with early exiting from a large language model (LLM). This method allows for efficient generation by using the same model's early layers for drafting tokens, and later layers for verification.
This technique not only speeds up text generation, but it also achieves significant memory savings and reduces computational latency. In order to obtain an end-to-end speedup, the output of the earlier layers need to be close enough to the last layer. This is achieved by a training recipe which, as described in the paper, can be applied during pretraining, and also while fine-tuning on a specific domain. Self-speculative decoding is especially efficient for real-world applications, enabling deployment on smaller GPUs and lowering the overall hardware footprint needed for large-scale inference .
In this blog post, we explore the concept of self-speculative decoding, its implementation, and practical applications using the 🤗 transformers library. You’ll learn about the technical underpinnings, including early exit layers , unembedding , and training modifications . To ground these concepts in practice, we offer code examples, benchmark comparisons with traditional speculative decoding, and insights into performance trade-offs.
Note: While the assistant_early_exit argument can potentially enable early-exit self-speculative decoding for any decoder-only transformer, the logits from the intermediate layers cannot be unembedded (process of decoding through LM Head, described later in the blog post) unless the model is specifically trained for that. You will also only obtain speedups for a checkpoint that was trained in such a way to increase the accuracy of earlier layers. The LayerSkip paper proposes a training recipe to achieve that (namely, applying early exit loss, and progressively increasing layer dropout rates). A collection of Llama2, Llama3, and Code Llama checkpoints that have been continually pretrained with the LayerSkip training recipe are provided here .
Unembedding can be performed at any transformer layer, turning early-exit into an efficient token-prediction mechanism. A natural question arises: how can the LM head be adapted to unembed logits from earlier layers when it was initially trained to work with the final layer only? This is where the training modifications come into play.
Once training is complete, we can apply self-speculative decoding during inference. The process begins with self-drafting , where tokens are generated by exiting early from some intermediate layer. The number of speculative tokens defines how many draft tokens are produced during this stage, and the layer we exit at defines how large and accurate is the draft stage. Both parameters can be specified at inference based on a trade-off between speed and accuracy of the draft stage .
So far, the 🤗 transformers library has implemented the first optimization (Shared Weights) in this pull request . As the number of models that use this method increases, we'll consider the additional optimizations. Feel free to open a PR if you're interested!
Dive straight into the following Hugging Face artifacts to know more about the method and try it out yourself:
Illustration of LayerSkip inference on facebook/layerskip-llama2-7B (Llama2 7B continually pretrained with the LayerSkip recipe).
Traditional speculative decoding uses two models: a smaller one (draft model) to generate a sequence of draft tokens, and a larger one (verification model) to verify the draft’s accuracy. The smaller model performs a significant portion of the generation, while the larger model refines the results. This increases text generation speed since the larger model verifies full sequences at once, instead of generating one draft at a time.
By enabling faster, lower‑memory text generation on smaller GPUs, self‑speculative decoding can broaden the deployment of large language models in real‑world applications and reduce operational costs.