Layer Skip: 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


Characterizing Solar Center-to-Limb Radial-Velocity Variability with SDO
    [
    April 29, 2024
========================================================================





§ 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 <cit.>. Acceleration solutions do exist to deploy to commodity GPUs on laptops but they suffer from significant drop in accuracy <cit.>. Accelerating LLMs further to mobile or edge devices is still an active research area <cit.>. While a large portion of LLM acceleration approaches reduce number of non-zero weights <cit.> (a.k.a. sparsity), number of bits per weight <cit.> (a.k.a. quantization), number of heads per layer <cit.> (a.k.a. head pruning), a smaller portion of approaches focus on reducing number of layers <cit.>. In this paper, we explore reducing the number of layers required for each token by exiting early during inference. Unlike quantization or sparsity, acceleration by reducing number of layers does not require specialized hardware or software kernels.



Moreover, a popular research trend in LLM acceleration is speculative decoding <cit.> 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:

    
  * a training recipe that combines layer dropout and early exit loss, that leads to,
    
  * inference that is more robust to exiting at earlier layers of the model, essentially creating different sized sub-models within the same model, and
    
  * a self-speculative decoding solution that decodes with earlier layers and verifies and corrects with later layers.


The solution achieves speedups between 1.34× and 2.16× depending on the task. We provide an overview of the solution in Figure <ref>.



§ MOTIVATION





 §.§ Exiting Earlier in LLMs


To motivate our approach, we investigate, with an example prompt, what happens in each layer in a LLM. In Figure <ref>, we provide the first prompt from the HumanEval coding dataset <cit.> to a pretrained Llama1 7B model <cit.>. The prompt consists of a Python function header and a docstring, and the model autocompletes it by defining the function body. When generating each token, we probe each transformer layer in the LLM by projecting its output embeddings on the language model (LM) head (that consists of the model's final layer normalization and linear layer), applying softmax, and then obtaining the index of the output element with highest value. The resulting index corresponds to the predicted token at this layer. This operation is referred to in some literature as the unembedding operation <cit.>, as it converts an embedding to an index. Unembedding at each layer is equivalent to early-exit at that layer, i.e., it is equivalent to skipping the remaining transformer layers to the model's LM head.




The token predictions across layers in Figure <ref> illustrate the evolution of embeddings from an input token fed to the model to the predicted next token by the model. When analyzing the token prediction in each layer in Figure <ref>, we make a few observations. First, token predictions in earlier layers appear to be irrelevant as they correspond to the previous token projected on the model's embedding layer's weights, which are different from the weights of the LM head. In later layers, token predictions converge to the final prediction. Second, we do not always need all the layers to predict the correct token. In fact, most of the time, the final token prediction is predicted fewer layers  before the end. We also notice that intermediate layers are sometimes hesitant and “change their minds”, e.g., for Token 05, the model was predicting “range” as early as Layer 07, but changed its mind between Layer 22 and Layer 26, before settling again on “range”.

Similar analysis was done in <cit.> on a GPT2 model <cit.> as it developed predictors to estimate when prediction saturates to exit early. For the particular example we present in Figure <ref>, we find, on average, a token requires 23.45 layers out of the model's 32 layers. Hence, even if we have a perfect predictor that has zero compute overhead, we can only save up to 26% of computation. Therefore, there is a need to make LLM models require fewer layers to predict each token, and spend less compute being hesitant or “changing its mind”. By default, deep learning models are not motivated to predict their final output early and instead spread their compute across all layers <cit.>. We see in Figure <ref>, that tokens we would consider easy or straightforward to predict, e.g., Token 02 that starts a for-loop, required all 32 layers to predict “for”. We would like our model to be less reliant on later layers and only use later layers for harder tokens. We would like our models to be more reliant on earlier layers than later layers. To do that, we propose skipping layers during training, which we refer to as layer dropout. However, we use higher dropout rates for later layers and lower dropout rates for earlier layers, to make the model less reliant on later layers. 

Moreover, LM heads in LLMs are trained to unembed embeddings from the last transformer layer. They were not trained to unembed from earlier layers. Therefore, our solution also adds a loss function during training to make LM heads better “understand” embeddings of earlier layers. While most papers that explored early exit <cit.> trained a dedicated LM head for each transformer layer, and some have introduced additional modules for each early exit <cit.>, we chose to have a shared LM head for all transformer layers in the model. This makes training faster, require less memory consumption for both training and inference, and eases deployment and maintenance. Hence, as shown in Figure <ref>, we train a deep learning model that is equivalent to an ensemble of models of various depths, capable of skipping from different transfomer layers to the LM head.





 §.§ Correcting if we Exit Too Early



Regardless if we use heuristics or predictors (as  <cit.>) to exit early, or if we modify the training procedure to make models predict early (as <cit.> and this paper as well), it is likely that exiting early during inference will lead to a reduction in accuracy. It will be ideal if there is a way to verify if an early prediction is accurate, and correct it by executing remaining layers. Some approaches like <cit.> proposed a confidence heuristic to decide after executing an early exit if the remaining layers are needed. Here, we leverage speculative decoding techniques to verify the early exit prediction and correct it. Speculative decoding benefits from the fact that verifying the prediction of a group of tokens is faster than generating each token auto-regressively. Hence, we present a self-speculative decoding approach where we use early exit to generate each token auto-regressively, and use the remaining layers to verify a group of tokens in parallel, and correct them.



§ RELATED WORK



Dropout: Dropout was first introduced by <cit.> and involved stochastically replacing a portion of output elements of fully-connected layers with zeros during training. We refer to this variant of dropout as unstructured dropout. It presented a regularization effect for training, with the purpose of reducing over-fitting. Unstructured dropout was commonly used in convolutional neural networks (CNNs) before batch normalization <cit.> replaced it as a means to improve generalization. However, the introduction of transformers brought it back to light as <cit.> used a dropout rate of 0.1. However, dropout faded again when pretraining dataset sizes increased, e.g., large scale models like Llama <cit.> and GPT3 <cit.> do not mention dropout in their papers, and PaLM <cit.> used a dropout of 0 for pre-training on its large corpus but 0.1 for fine-tuning on small datasets.

The intuition behind dropout’s regularization effect in <cit.> was that it enabled training to learn across an ensemble of many models, and avoiding co-adaptation between the model’s nodes. Another intuition is that dropout induces multiplicative noise into node’s activations <cit.>, and training a model with noise makes it more robust to the distribution shift between training and validation data.  <cit.> showed that dropout induces noise to mini-batch gradients and makes them more aligned to the full dataset gradient in the early epochs of training, and hence showed that applying dropout either early in training or late in training is better than applying dropout throughout training.

Layer Dropout: Skipping layers stochastically during training is referred to in literature with different terms such as stochastic depth or layer dropout. It was first explored in ResNets by <cit.>. ConvNext <cit.> used higher layer dropout rates for larger models: 0.1/0.4/0.5/0.5 for ConvNeXt-T/S/B/L respectively when trained on ImageNet <cit.>. However, when training on the larger ImageNet-22K dataset, ConvNeXt used smaller layer dropout rates: 0.0/0.0/0.1/0.1/0.2. In language models, LayerDrop <cit.> applied dropout to every other transformer layer, which increased its robustness to pruning layers at inference time.  <cit.> increased the pretraining speed of BERT by applying a dropout rate that progressively increased every iteration as well as every layer. To the best of our knowledge, layer dropout for training decoder-only models, or scaling language models to large model sizes or large datasets has not been explored. Moreover, our paper is the first to propose using layer dropout to improve early exit inference.



Early Exit: Exiting early in deep learning has first been explored in CNNs <cit.>. They added branch modules at different exit points in a deep learning network and introduced additional loss functions during training to improve the accuracies of those early exits. In language models, early exit was explored in encoder-only models by BERxiT <cit.> and BE3R <cit.>, in encoder-decoder models by  <cit.> and CALM <cit.>, and in decoder-only models by SkipDecode <cit.>.

 <cit.> added a dedicated LM head for each decoder layer in an encoder-decoder translation model, and explored different early exit granularities (per token as well as per sequence) in both training and inference, and explored exiting early both statically and dynamically, using auxiliary layers that learned during training when to exit. CALM <cit.> built upon  <cit.> and started with a model pretrained with early exit losses, and focused on finding optimal criteria to decide which layer to exit at during inference. SkipDecode <cit.> fine-tuned a decoder-only model to skip more layers for later tokens within a sequence.  <cit.> started with pretrained models and finetuned auxiliary fully-connected layers to map the embeddings outputted by earlier layers to later layers. In our proposed solution, we do not introduce any additional modules or linear layers for early exit, and instead used a shared exit for all layers.

Speculative Decoding: Speculative decoding <cit.> is a popular acceleration technique for language models. It is based on the fact that auto-regressive decoding of decoder models are slow as they generate one token a time, while measuring the likelihood of a group of generated tokens in parallel is faster. It uses a fast, less accurate model, referred to as the draft model, to generate multiple tokens auto-regressively, and a large, slower, more accurate main model, to verify the tokens in parallel, and correct them when needed. The draft model could have the same or different architecture as the main model, or could be a compressed version of the model.  <cit.> recently proposed a self-speculative decoding approach where the draft model is the same as the main model, but with a group of intermediate attention and feed forward network (FFN) layers skipped. The advantage of our proposed solution compared to <cit.> is that verification and correction stages can reuse the activation and KV cache from the draft stage as both stages execute the same early layers in the same order, while  <cit.> can not reuse them as it skips intermediate layers.  <cit.> used shared transformer layer groups and a shared LM head to exit each token at a different layer and execute different layer groups in a pipeline fashion.



§ PROPOSED SOLUTION



Our approach has three different stages: 


    
  * Training using Layer Dropout & Early Exit Loss
    
  * Inference using Early Exit 
    
  * Verification and Correction using Speculative Decoding


We explain each stage in the following sub-sections.



 §.§ Training using Layer Dropout & Early Exit Loss













We denote the input tokens to a transformer model as X and its output as Y, with an embedding layer that maps the token indices to token embeddings, x_0, and a transformer model with L transformer layers, where transformer layer l evolves embeddings outputted from its previous layer, x_l+1 = x_l + f_l(x_l), and a final LM head that maps the embedding outputs of the last layer, x_L to logits, e_L=g(x_L). We denote the binary cross entropy loss function that is usually used to train language models as J_BCE(e_L,Y).



  §.§.§ Layer Dropout



The first modification we apply to common training recipes, is to apply layer dropout. Hence the transformer layer operation at layer l and training iteration t changes to:

    x_l+1,t = x_l,t + M(p_l,t) f_l(x_l,t)

where p_l,t is the dropout rate of layer l at iteration t, M(p) is a Bernoulli function that returns 0 with probability p and returns 1 with probability 1-p. We apply the dropout operation on each sample separately within a batch. We remove the dropped samples from a batch, apply the transformer operation f_l on the remaining samples, and then concatenate the output with the dropped samples. To ensure higher speedup during training, we seed the random number generator for each GPU with the same seed, so that each transformer layer at each iteration will drop the same number of samples. 

The dropout rate can be different at each layer l and training iteration t, p_l,t: 

    p_l,t = S(t)D(l)p_max

where p_max is a hyperparameter that sets the maximum dropout rate in the model during training, D(l) is a per-layer scaling function, and S(t) is a per-time step scaling function. We found that the best per-layer scaling is to increase dropout rate exponentially across layers from 0.0 in layer 0, to 1.0 in last layer, L-1:

    D(l) = e^lln2/L-1 - 1

For scaling across time, S(t), we found that if we start with a pre-trained model and perform continual pre-training or finetuning, it is best to not scale across time and hence set S(t)=1. However, for pretraining from scratch, we found that an exponential curriculum, S_exp(t), lead to best accuracies for T training steps:

    S_exp(t) = e^tln2/T-1 - 1




  §.§.§ Early Exit Loss



To boost prediction accuracy of lower layers, we need to ensure that the model's LM head, g, is capable of unembedding outputs of different layers. Hence, during training, we augment layer dropout with early exit loss at each layer. During training we supervise the model directly to connect the early exit layers to the LM head, this enables us to directly supervise the lower layers for the language modeling task. The total loss of the model at iteration t is:

    J(X,Y,t) = ∑_l=0^l=L-1ẽ(t,l)J_BCE(g(x_l+1),Y)

Where ẽ(t,l) is a normalized per-layer loss scale, whose sum across all layers is equal to 1:

    ẽ(t,l) = C(t,l)e(l)/∑_i=0^i=L-1C(t,i)e(i)


C(t,l) is a binary curriculum function that determines if we enable early exit of layer l at iteration t. We build upon  <cit.> and set a scale that increases across layers, such as the scale at one layer is proportional to the sum of the scales of all previous layers:

    e(l)= 
    
        e_scale∑_i=0^i=li,                if  0 ≤ l < L-1
    
        L-1+e_scale∑_i=0^i=L-2i,            if  l = L-1


This way, we penalize later layers with quadratically higher weight, as predicting in later layers is easier. 0 ≤ e_scale≤ 1 is a hyperparameter that controls the scale of early exit loss.

Note that we do not add additional LM heads as proposed in other early exit papers <cit.>, as we essentially use the same LM head for all layers.

Early Exit Loss Curriculum: We find that adding early exit loss of all layers at all iterations during training slows down training and reduces the accuracy of the last layer. To overcome this, we introduce a curriculum, C(t,l). We have explored 2 different curricula. First, we explored a rotational early exit curriculum, C_rot,R, where we enable early exit at every R layers, and perform circular rotation at each iteration. This way, the early exit at each layer is enabled once every R iterations. Hence, at each training iteration, only ⌈ L/R ⌉ unembedding operations are applied. Second, we explored a gradual early exit curriculum, C_grad, where we gradually enable early exit loss from layers L-1 to 0, one layer at a time every T/2L iterations.


Overall, to summarize, the hyperparameters of our training recipe:

    
  * Layer Dropout:
    
        
  * p_max: maximum dropout rate of last layer of the model,
        
  * S(t): layer dropout curriculum. We use either no curriculum S(t)=1 for finetuning or continual pretraining, or an exponential curriculum, S(t)=S_exp(t) for pretraining from scratch,
    
    
  * Early Exit Loss:
    
        
  * e_scale: scalar scale of loss of earlier layers,
        
  * C(t,l): early exit loss curriculum, either rotational, C_rot,R(t,l), or gradual, C_grad(t,l) 
        
            
  * R: is a dilation across layers for rotational early exit loss curriculum
        
    




 §.§ Inference using Early Exit





When generating each token during autoregressive decoding, we run the first E transformer layers in a model, and skip to the model's LM head, i.e., the model's final output becomes g(x_E). We explore with different values of E and provide the accuracies in the Results section. 



 §.§ Inference using Self-Speculative Decoding





With layer dropout and early exit loss in training, we show it is possible to speedup autoregressive generation by exiting early, but this comes at an accuracy cost compared to using the full model. Speculative decoding <cit.> is able to leverage a faster yet less accurate model to speedup generation without accuracy cost. However, this requires storing and training 2 models, increasing memory and training requirements.

We introduce a novel self-speculative decoding algorithm built on top of early exit, enabling us to reduce memory through the use of a single model and latency of traditional speculative decoding through re-using hidden states in draft and verify steps. As shown in Figure <ref>, our self-speculation algorithm consists of 2 key steps (1) Self-Drafting, using the early exit to draft tokens from the same model (2) Self-Verification, using the remaining layers to validate the prediction. To enable re-use in (1) and (2), we develop a novel Cache Reuse technique that unifies the KV cache and storing the exit query. We provide a high level description of the algorithm in sections <ref> and <ref> and provide pseudo code in <ref>.



  §.§.§ Self-Drafting


The first step in speculative decoding is to define a set of draft tokens D_0...d-1. In our algorithm, we compute the first d draft tokens through early exit. We refer to d as the number of speculations. We leverage a subset of the LLM and conduct auto-regressive inference exiting at layer E.

Our training recipe enabled us to train the model once to get an ensemble of different candidate draft models at each layer depth. We can evaluate exiting at different layers and observe a trade off between latency and accuracy.



  §.§.§ Self-Verification



The next step in speculative decoding is verification. Verification leverages the full LLM to predict the next token for each draft token in a single forward pass. We then assess to see where the draft tokens and verified tokens agree. All the draft tokens up till the disagreement point are added to the output along with the next verified token and the process continues from the draft stage.

In our self-speculative decoding algorithm, the self-verification stage critically only requires computing the remaining layers of the model that were not used in the draft stage. For a model with L layers, the number of verification layers is L - E. In order to re-use the first E layers from the draft stage we employ some modifications to the KV cache as we show in the subsequent subsection.




  §.§.§ Reusing the Cache



In autoregressive transformers, the KV cache is a critical component of efficient generation, allowing us to avoid recomputing prior KV pairs in each transformer layer. 

As our draft stage uses the first E layers of the model and the verification stage uses the remaining L - E layers, we are able to re-use a significant amount of compute between the 2 stages:

    
  * Single KV Cache  As the draft model and verification model operate on the same model using the same order of layers, the first E layers are shared in both steps. Hence, in the draft stage, the KV cache in the first E layers are already computed, so we are able to effectively maintain a single KV cache for the draft and verify steps, reducing memory and latency.
    
  * Exit Query Cache: To further reduce computation of the first E layers, we introduce an exit query cache that saves the query vector of exit layer E-1 for verification to directly continue from layer E to last layer L. Critically note that we need to save only the query for the exit layer. We term the union of the KV cache and the exit query as KVQ cache.
































§ 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:

    
  * Continual Pretraining: start with a pretrained model and continue pretraining on 52B tokens from a corpus of diverse data containing natural language text and code. We experiment using pretrained Llama2 7B (32 layers), with p_max=0.1, e_scale=0.2, C_rot,R=8, and Llama2 13B (40 layers), with p_max=0.1, e_scale=0.1, C_rot,R=39.
    
  * Pretraining from Scratch: start with randomly initialized model and pretrain on 26B tokens from a corpus of diverse data containing natural language text and code. We experiment with Llama2 1.5B (a custom small Llama-like model with 24 layers) (see  <ref> for architecture details) with p_max=0.1, e_scale=0.2, C_rot,R=23 and Llama2 7B (32 layers) with p_max=0.2, e_scale=0.2, C_rot,R=31. Following  <cit.> we use higher learning rates when layer dropout is greater than 0.0.
    
  * Finetuning on Code Data: start with pretrained Llama1 7B model <cit.> and finetune on 5.2B tokens of CodeLlama <cit.> data mix. We use p_max=0.1, e_scale=1.0, C_rot,R=16.
    
    
  * Finetuning on Task-Specific Dataset: start with a pretrained Llama 1.5B (24 layers) and finetune on TOPv2 <cit.>, a multi-domain task-oriented compositional semantic parsing dataset. We post processed the dataset into a JSON format to be more aligned with code pre-training. We report our results on the TOPv2 evaluation set. We use p_max=0.2, e_scale=1.0, C_grad.

We try different variants of : layer dropout only (LD), early exit loss only (EE), and both layer dropout and early exit loss (LD+EE). We provide more details about training hyperparameters in Appendix <ref>.



§ RESULTS





 §.§ Early Exit Inference Results


After training each model configuration, we evaluate accuracy of exiting early at different layers. 

Continual Pretraining In Figure <ref>, we present our results for Llama2 7B and 13B on a diverse set of evaluation tasks (see  <ref> for task details) and compare with the baseline model from  <cit.>. In Table <ref> we zoom in and show the specific values of accuracies for the last layer and middle layer of each model. In Figure <ref> we show sample text generations for exiting at earlier layers for both models with and without continual pretraining with . Overall, for earlier layers,  is clearly better than the baseline. For last layer accuracy,  has minimal drop in accuracy compared to baseline.

It is noteworthy that some “classification” tasks, i.e., multiple choice question or true/false question tasks, maintain relatively decent accuracy on earlier layers on the baseline model, while open-ended “generation” tasks drop drastically. Surprisingly, MMLU <cit.> which is considered a challenging task, only drops from 55.2% to 49.2% on Llama2 13B baseline from the last to the middle layer. This could be because classification tasks are evaluated on generating one token only while generation tasks are evaluated on the accuracy of many tokens, and an error in one token may have a compounding effect when generating later tokens. Moreover, classification tasks evaluate a token out of 4 or 2 possible outcomes, while generation tasks evaluate each token out of thousands of possible entries in the LLM's dictionary. We observe 's significant importance on generation tasks, e.g., NaturalQuestions <cit.> drops from 25.1% to 0% when exiting in middle layers of Llama2 7B, but jump to 4% when using . 







Pretraining from Scratch In Figure <ref>, we present our results for Llama2 1.5B and 7B pretrained from scratch on 26B tokens using  on a diverse set of evaluation tasks (see  <ref> for task details) and compare with the same models pretrained on the same number of tokens from scratch without . In Figure <ref> we show sample text generations for exiting at earlier layers. The results show that introducing our proposed layer dropout configuration and/or early exit loss leads to higher accuracy than the baseline on earlier layers. On the last layer, in some downstream tasks, we do see a slight drop in accuracy compared to baseline, while in other tasks we see our layer dropout and/or early exit configuration leading to higher accuracy. Note that since the models were pretrained on relatively small number of tokens, the evaluation results on some classification tasks were close to random guesses, whether we used or not, and hence we have removed them from the results.








Finetuning on Code Data In Figure <ref>, we present our results on 2 coding tasks and compare accuracy to Llama1 7B finetuned on the same number of code tokens without . For earlier layers,  is clearly better than the baseline, with layer dropout combined with early exit loss showing a big improvement on one of the 2 tasks. For last layer accuracy,  with both layer dropout and early exit loss has almost the same accuracy as baseline. Note that since this experiment finetuned on specific domain data, we were able to increase e_scale to 1.0 (as opposed to e_scale=0.1 or 0.2 in the previous two configurations).

Finetuning on Task-Specific Dataset In Figure <ref>, we compare results of fine-tuning our Llama 1.5B model on TOPv2 training set with and without . In semantic parsing, correctness requires an exact match (EM) between generated sequence and annotated parse. We find that when removing layers from the baseline model, the model is not able to generate complete and accurate parses resulting in 0 EM. However, with , early exit inference improves to 77% at layer 12. We notice a regression in the final layer reducing accuracy by 3%. Again, as this configuration finetuned data on a specific task, we were able to set e_scale=1.0.









 §.§ Self-Speculative Decoding Results



We evaluate the self-speculative decoding algorithm introduced in <ref> on the different trained models. We report quality metrics, EM (exact match) and ROUGE-2 <cit.> implemented in torchmetrics <cit.>, token acceptance rate for the self speculation algorithm (how often verification accepts each of the draft tokens), throughput measured as tokens per token averaged over the sampled dataset, and finally a speed up compared to autoregressive decoding. For our early exit and our self-speculative decoding experiments, we specify the layer we exit at in the E column in each table below. We compare with Draft & Verify <cit.> on common models and tasks evaluated in both papers. All experiments were performed with greedy decoding and generated a maximum of 512 tokens for each sample. Following  <cit.>, speedup is calculated as the acceleration of average inference time per token compared to the “Autoregressive” baseline on the same setting. “Autoregressive” experiments use baseline models that were pretrained or finetuned without , while “Early Exit” and “Self Speculative” experiments use our models trained or finetuned with . Our implementation leverages HuggingFace <cit.>.

Continual Pretraining  In Table <ref>, we evaluate the continual pre-training of Llama2 7B and 13B with and without  on various tasks: CNN/DM <cit.>, XSUM <cit.> abstractive summarization tasks, and HumanEval <cit.> coding task. The experiments were performed on NVIDIA H100 GPUs. The number of speculations, i.e., the  number of tokens generated in the draft stage in self-speculation, is specified in the d column. We follow <cit.> and perform CNN/DM with 1-shot abstractive summarization, and XSUM with 0-shot. We obtain speedups between 1.34× and and 2.16× depending on the model or the task. In general, we observe higher speedups for the smaller Llama2 7B compared to the larger Llama2 13B model. When comparing with Draft & Verify <cit.> on the common model and tasks of both approaches, we are significantly faster on CNN/DM (1.81× vs. 1.5×) and slightly slower on XSUM (1.34× vs. 1.48×).
 



Pretraining from Scratch Our results are presented in Table <ref>. The experiments were performed on NVIDIA H100 GPUs. We found that an opposite tried to continual pretraining: the bigger model has a bigger speedup than the smaller model, reaching 2.16× speedup, which exceeds traditional speculative decoding. 

Finetuning on Code Data In Table <ref>, we evaluate our code-finetuned Llama1 7B on HumanEval using 12 speculations, and exit at layer 6 for self speculation & early exit. The experiments were performed on NVIDIA A100 GPUs. We show speedup of upto 1.82× with no accuracy drop.

Finetuning on Task-Specific Dataset In Table <ref> we show results for Llama 1.5B finetuned on TOPv2's training dataset and evaluated on TOPv2 test set. The experiments were performed on NVIDIA H100 GPUs. We present the EM (exact match) on the fully TOPv2 test set, further we sample 1000 samples for latency experiments where we leverage 8 speculations, and generate the next 80 tokens with greedy decoding. With self-speculation, the model was able to achieve high token acceptance rate, (E=6: 76.0%, E=12: 97.2%, E=18: 98.9%) reaching 2.0× speedup.












§ ABLATION STUDIES








Layer Dropout Configurations

In Figure <ref> we show that our layer dropout configuration leads to lower loss compared to a constant layer dropout across all layers with the same average value.





 §.§ Shared FC vs non-Shared FC

    


Scaling with Pretraining Tokens

In order to understand how the accuracy of last and middle layers change across time when pretraining from scratch, we ran 3 training experiments with different number of tokens on Llama 1.5B and show the results in Figure <ref>. Each experiment trained for 50,000 steps, per device batch size of 4, context window of 4096, but changed the number of GPUs to 32, 64, 128. We plotted the perplexity of a held out split of The Stack on the last layer (layer 24) and the middle layer (layer 12). As expected, perplexity on last layer decreases as we train on more tokens. However, surprisingly, we discover that perplexity on last layer increases drastically by default in training, unless we apply early exit loss. Layer dropout reduces the increase as well. This could open the door to more research on the dynamics of transformers and the evolution of embeddings in earlier layers to understand why embeddings across layers are close to each other early on in training but diverge drastically as training progresses. This could also present a motivation for our training recipe that has minimal drop in last layer accuracy while significantly improves accuracy of earlier layers.



KV Cache in Self-Speculation

In <ref> we introduced the re-use of KV cache as a method for improving model generation speed. We measure its effect in Table <ref>. We follow the same inference setup as described in <ref>. We find that the use of KV cache is able to consistently save us 9-20 ms per token depending on the task.












§ LIMITATIONS



    
  * Our self-speculative decoding solution requires finetuning a model or pretraining it with our recipe, while the self-speculative decoding approach propoposed in  <cit.> does not require changing a model's weights.
    
  * The introduced hyperparameters, p_max for layer dropout, e_scale and R for early exit, requires tuning in order to avoid a drop in last layer accuracy.
    
  * When pretraining with layer dropout from scratch, increasing the learning rate is required to maintain accuracy, and tuning learning rate to get optimal accuracy could be tricky and time consuming.
    




§ CONCLUSION










We have shown that combining layer dropout, early exit loss with curriculum, improves accuracy of early exit during inference. We then developed a novel self-speculative decoding solution that led upto 1.86× speedup.

We hope this paper will encourage researchers and engineers to adopt the proposed layer dropout and/or early exit loss in pretraining as well as finetuning recipes. For pretraining from scratch, layer dropout can lead to increase in training speed. While early exit loss may introduce an overhead for pretraining, using a large R can make that minimal. For finetuning, we hope that our proposed techniques may be combined with parameter efficient techniques suchs as LoRA <cit.>.

In the future, we would like to increase the accuracy of early-exit layers in order to obtain better speedups during self-speculative decoding. We can also explore dynamic conditions to determine a different exit layer for each token (like  <cit.>) and hence improve token acceptance rate of self-speculative decoding.



§ ACKNOWLEDGEMENTS

We would like to thank:

    
  * Volker Seeker, Artem Korenev, and Ilia Kulikov for logistic support.
    
  * FAIR's clusters' support team members, especially, Henry Estela, Hongsheng Song, Shubho Sengupta, and Nabib Ahmed, for their help in maintaing our clusters. Without them, we cannot run our research. 
    
  * Fabian Gloeckle, Andrey Gromov, and Aaditya Singh for discussions.



assets/plainnat






§ APPENDIX





 §.§ Experiment Details


We provide details of training configuration and hyperparameters for each of our experiments in Table <ref>.


When pretraining from scratch, layer dropout leads to higher accuracy when trained on higher learning rate  <cit.>. Therefore, we show learning rates of each experiment with and without layer dropout separately in Table <ref>.





  §.§.§ Model Architectures


We provide details of architectures of different models in Table <ref>.





  §.§.§ Evaluation Tasks


We have evaluated our language models on a wide range of tasks. For the sake of discussions in  <ref>, we categorize the tasks into:

    
  * “Classification” Tasks: where the model responds with one out of pre-defined answers, e.g., multiple-choice questions, or questions whose answers are either “True” or “False”:
    
        
  * Common Sense Reasoning Tasks
        
            
  * BoolQ <cit.>
            
  * PIQA (Physical Interaction Question Answering) <cit.>
            
  * SIQA (Social Interaction Question Answering) <cit.>
            
  * HellaSwag <cit.>
            
  * Winogrande 1.1 <cit.>
            
  * ARC (Abstraction and Reasoning Corpus) <cit.>
            
                
  * ARC Challenge
                
  * ARC Easy
            
            
  * OBQA (Open Book Question Answers) <cit.>
            
  * COPA (Choice Of Plausible Alternatives) <cit.>
        
        
  * RACE (ReAding Comprehension dataset from Examinations) <cit.>
        
            
  * RACE Middle
            
  * RACE High
        
        
  * MMLU (Massive Multitask Language Understanding) <cit.>
    
    
  * “Generation” Tasks: where the model responds with an open-ended sequence of tokens and we evaluate either the exact match of the tokens with the reference answer, or, in the case of code, run and test it on a compiler or interpreter.
    
        
  * Question Answering
        
            
  * NQ (Natural Questions) <cit.>
            
  * TQA (Textbook Question Answering) <cit.>
        
        
  * Mathematics
        
            
  * MATH <cit.>
            
  * GSM8K <cit.>
        
        
  * Code Generation
        
            
  * HumanEval <cit.>
            
  * MBPP (Mostly Basic Python Problems Dataset) <cit.>
        
    


We also evaluate perplexity on held out test sets on the following datasets:

    
  * The Stack, a coding dataset <cit.>
    
  * Books  <cit.>
    
  * Wikipedia







 §.§ Additional Results





  §.§.§ Self-Speculative Decoding Results



CPU Inference Experiments
We conduct our task specific fine-tuning on Llama 1.5B to measure decoding performance on CPU as well, showing a near 2× speed up on CPU as well, presented in Table <ref>. We conduct our experiments using the first 100 samples from the TOPv2 test set, leveraging 7 speculations, generating the next 50 tokens with greedy decoding. 





 §.§ Ablation Studies


Layer Dropout Configurations: In Figure <ref> we show that our layer dropout configuration leads to lower loss compared to a constant layer dropout across all layers with the same average value.
















 §.§ Self Speculation Pseudo Code


Below we share pseudo code for implementing self speculation





basicstyle=,breaklines=true
framextopmargin=50pt,frame=bottomline
[language=python]selfspeculation.py






