Draft on the Fly: Adaptive Self-Speculative Decoding using Cosine Similarity (2024)

Abstract: We present a simple on the fly method for faster inference of large language models. Unlike other (self-)speculative decoding techniques, our method does not require fine-tuning or black-box optimization to generate a fixed draft model, relying instead on simple rules to generate varying draft models adapted to the input context. We show empirically that our light-weight algorithm is competitive with the current SOTA for self-speculative decoding, while being a truly plug-and-play method.

## introduction the main bottlenecks of transformer-based large language model (llm) inference are the sequential generation of tokens, due to its autoregressive nature, and the need to reload the kv cache, causing inference to be memory bandwidth bound @xcite . both of these bottlenecks prevent the full use of computing resources. speculative decoding is a method to decrease the inference time of llms by taking advantage of this available compute capacity through the observation that the time of processing a single input token is roughly the same as processing a small number of tokens in parallel. after processing the input context and generating the first token t 1 using a model m , instead of sequentially generating tokens until completion, speculative decoding uses a smaller draft model d to generate a short sequence of tokens [ t2 , t3 , ...]. the tokens [t 1 , t2 , t3 , ...] are then processed in parallel by m for verification. the first rejected token tr is replaced by m 's predicted token t r . the process is then repeated, generating a new sequence of tokens [ tr+1 , tr+2 , ...] for verification. when using greedy decoding, the verification process is to simply ensure that @xmath0. when sampling is used, a clever method has been devised to accept tokens which ensures that the distribution they were sampled from matches that of m @xcite . simply put, speculative decoding has the ability to speed up inference without sacrificing accuracy. the ability to speed up inference also stems from the varying difficulty in predicting the next token over the generated sequence, where even quite small draft models are capable of predicting a percentage of tokens correctly. choosing a d to maximize speedup then requires balancing its speed of token generation with the accuracy of the tokens it generates. in this work we propose a self-speculative decoding algorithm, where d is chosen as a subnetwork of m , generated on the fly and adapted to the input context. compared to the current state of speculative decoding described in section 2, our method is simple to implement as a truly plugand-play method, avoiding the challenges of finding a fixed draft model using fine-tuning or blackbox optimization. in section 4 we observe that despite the simplicity of our proposed adaptive draft model generator (admg, algorithm 1), our method is competitive with the current sota of self-speculative decoding. ## related work speculative decoding was first developed in @xcite . these works proposed using a second smaller model to draft tokens from. acquiring such a draft model is in general challenging, and will typically require finetuning @xcite @xcite . sufficient memory must also be available to store a second model as well as save both models' kv caches. to alleviate the need for a separate draft model, self-speculative decoding was developed in draft & verify @xcite , where d is chosen as a subnetwork of m by removing attention and mlp layers. choosing which layers to remove to minimize the average inference time is a binary optimization problem (np-hard) of a black-box objective function. an approximate solution to this problem is found using bayesian optimization with a small number of training samples. this is a time consuming process (see section 4), with the solution depending on the model architecture, dataset, and computing environment. kangaroo @xcite can be seen as a hybrid approach between using a fine-tuned draft model and self-speculative decoding, where d consists of the first @xmath1, a fine-tuned adapter layer, and the lm head of m . medusa @xcite is an adjacent method to speculative decoding which uses multiple decoding heads to predict tokens in parallel, avoiding the challenges of obtaining an appropriate draft model. in its simplest form, these additional heads must be trained with the original weights of m kept frozen. eagle @xcite also does not use a separate draft model, but drafts tokens using an additional decoder layer which must be trained. the added decoder layer takes as input the embedding of the last generated token and the input to the lm head of the penultimate generated token, with its output fed into the lm head of m to predict the next token. compared to these related works, a main benefit of our method is that it does not require the use of any fine-tuning or black-box optimization. finetuning can be challenging when the training set of m is not available, as this can result in a shift in the output distribution of d compared to m , making the token verification stage of speculative decoding more challenging @xcite . similarly, for draft & verify, the chosen draft model can be sensitive to the calibration data used, and its performance may suffer if an appropriate dataset is not available, see section 4.1 for further discussion. ## adaptive self-speculative decoding in order to generate draft models on the fly, our main technique is to use a simple thresholding rule to remove layers from the original model based on the cosine similarities of the hidden states of the input context. ## experiments we compare our method to the self-speculative decoding algorithm of draft & verify, which is most similar to our work. we conducted the same (fine-tuned) llama-2-13b experiments found in the body of their paper without changing any of the chosen hyperparameters. the experiments cover three models and three datasets: llama-2-13b and llama-2-13b-chat @xcite evaluated on 1000 random test set samples of cnn/dm @xcite and xsum @xcite , and codellama-13b (rozière et al., 2023) evaluated on humaneval @xcite . cnn/dm and xsum are both summarization tasks whereas humaneval is a python code generation task. the test sets were sampled using the same random seed as in draft & verify's experiments. an important factor which determines the level of inference speedup is the number of tokens drafted before verifying their accuracy with m . we used the adaptive draft-exiting mechanism proposed in draft & verify @xcite , section 3.4), keeping the hyperparameters equal to their tuned values @xcite ) to match their experiments. we conducted all experiments using two v100-32gb gpus. draft & verify relies on bayesian optimization to find their draft model, which is sensitive to the computing environment. in order to get the best performance from their method, we reran their optimization code to generate draft models tuned to our environment. in total three draft models were generated, one for each model, which took on average 11 hours to generate per model. for our admg (algorithm 1), we used the hyperparameters in ## conclusion we have proposed a self-speculative decoding method to generate draft models on the fly which are adapted to the input context. our method uses cosine similarity thresholding with simple layer removal rules to generate draft models without the need for any fine-tuning, black-box optimization, or training data. despite our proposed method's simplicity, we found that it is competitive with the current self-speculative decoding sota, while being an easy to implement plug-and-play method.