D yna M o: Accelerating Language Model Inference with Dynamic Multi-Token Sampling (2024)

Abstract: Traditional language models operate autoregressively, i.e., they predict one token at a time. Rapid explosion in model sizes has resulted in high inference times. In this work, we propose DynaMo, a suite of multi-token prediction language models that reduce net inference times. Our models *dynamically* predict multiple tokens based on their confidence in the predicted joint probability distribution. We propose a lightweighttechnique to train these models, leveraging the weights of traditional autoregressive counterparts. Moreover, we propose novel ways to enhance the estimated joint probability to improve text generation quality, namely co-occurrence weighted masking and adaptive thresholding. We also propose systematic qualitative and quantitative methods to rigorously test the quality of generated text for non-autoregressive generation. One of the models in our suite, DynaMo-7.3B-T3, achieves same-quality generated text as the baseline (Pythia-6.9B) while achieving 2.57 \times speed-up with only 5.87% and 2.67% parameter and training time overheads, respectively.

## introduction recent research has demonstrated the tremendous promise of large language models (llms) as competent artificial intelligence (ai) assistants @xcite . this has led to their rapid and widespread adoption as chatbots in diverse applications, e.g., healthcare, e-commerce, education, etc. however, the high computational requirements of llm training and inference and the use of massive closed-source corpora have restricted their development to a few laboratories. the increasing number of open-source llms, including pythia @xcite and @xcite , democratizes research in natural language processing (nlp). for instance, vicuna-13b @xcite , an instruction-finetuned * work done as an intern at samsung research america. llama model @xcite , has gained significant interest among researchers due to its exceptional instruction-following capabilities for its relatively compact size. nevertheless, access and study of llms remain limited due to challenges involved in their efficient evaluation on resourceconstrained devices. ## background and related works previous research explores various approaches to reduce token prediction latency in llms. it includes distillation @xcite , complexity reduction @xcite , sparsification @xcite , quantization @xcite , etc., to reduce model size or complexity, leveraging specialized hardware @xcite . other engineering solutions include flash attention @xcite that reduces memory reads/writes. recently, skeleton-of-thought decoding @xcite was proposed, wherein the llm first generates the skeleton of the answer and then conducts batched decoding to complete the contents of each skeleton point in parallel. speculative decoding @xcite ) is yet another approach that has gained recent prominence. it leverages a small draft model (which can be combined with the main model, @xcite to anticipate the main model predictions and queries it for batch verification. the batch size depends on the targeted number of token positions in the future, for draft prediction, and the number of top-k samples at each position. despite attempts at improving inference efficiency @xcite , such methods incur high computational overhead due to high-batch operations and result in poor compute utilization (e.g., sparse tree attention used by @xcite . for the greedy decoding scheme, such methods enable up to n× speed-up, however, at the cost of at least n× the compute. instead, in this work, we propose a low-compute approach that directly maps the joint probability distribution and implements cooccurrence weighted masking and adaptive thresholding, obviating the need for batched verification. further, medusa @xcite exploits simple feed-forward layers for draft prediction. this work explores various architectural modifications for draft prediction. nevertheless, the abovementioned approaches are orthogonal to the proposed method and can be used in conjunction to further boost performance. ## method in this section, we discuss the implementation details of multi-token prediction in the dynamo suite. ## experiments in this section, we present experimental results and comparisons of the proposed approach with the pythia baseline, which we used to instantiate the dynamo models. we provide test results for architectural and training variations in multi-token prediction in appendix c.2. ## discussion in this section, we discuss the implications of the proposed dynamo suite of multi-token prediction models and future work directions. dynamo-77m-t3 17.3 ±1.1 41.0 ±1.0 55.7 ±0.9 56.0 ±5.0 26.9 ±0.4 14.7 ±1.6 59.8 ±1.1 49.8 ±1.4 ## conclusion in this work, we presented dynamo, a suite of multi-token prediction language models. we trained the proposed model suite efficiently by reusing weights of existing pre-trained llms. we proposed novel ways to dynamically predict multiple tokens for a given context. the dynamo models dynamically back off to lower-order n-gram prediction based on a threshold. we also proposed adaptive thresholding and co-occurrence weighted masking on the modeled joint probability distribution to improve text generation quality. one of our proposed models, dynamo-7.3b-t3, achieved the same-quality generated text as the baseline (pythia-6.9b) while achieving 2.57× speed-up with only 5.87% and 2.67% parameter and training time overheads (see appendix a.2).