Context-Aware Assistant Selection for Improved Inference Acceleration with Large Language Models (2024)

Abstract: Despite their widespread adoption, large language models (LLMs) remain prohibitive to use under resource constraints, with their ever growing sizes only increasing the barrier for use. One particular issue stems from the high latency associated with auto-regressive generation in LLMs, rendering the largest LLMs difficult to use without advanced computing infrastructure. Assisted decoding, where a smaller draft model guides a larger expert model’s generation, has helped alleviate this concern, but remains dependent on alignment between the two models. Thus if the draft model is insufficiently capable on some domain of interest relative to the target model, performance can degrade. Alternatively, one can leverage multiple draft models to better cover the expertise of the target, but when multiple black-box draft models are available, selecting an assistant without details about its construction can be difficult. To better understand this decision making problem, we observe it as a contextual bandit, where a policy must choose a draft model based on a context. We show that even without prior knowledge of the draft models, creating an offline dataset from only outputs of independent draft/target models and training a policy over the alignment of these outputs can accelerate performance on multiple domains as long as an individual draft model is effective. We observe these results hold on various settings with multiple assisted decoding candidates, highlighting its flexibility and the advantageous role that such decision making can play.

## introduction with the introduction of the transformer @xcite has emerged the era of large language models @xcite @xcite and the development of llms capable of reasoning and acting in astonishingly human-like manner (kaplan et al., 2020; wei * correspondance to jerry.huang@mila.quebec @xcite . however, the use of resource intensive models and techniques remains a pre-requisite and accordingly, methods have been developed and applied to alleviate concerns relating to the practical usability of these models @xcite . one major area that has observed consistent improvement over time is the auto-regressive decoding aspect of text generation, where each generation of a new token requires a complete inference pass through the model, which under-utilizes the property of attention and the ability of modern accelerators (e.g. gpus, tpus) to parallelize computations @xcite . a growing approach towards addressing this is speculative decoding @xcite . in speculative decoding, latency is reduced by minimizing the amount of high-latency sequential computations and replacing them with cheaper ones. rather than sampling directly from the larger model, the sampling is approximated with samples from a smaller and cheaper model through accept-reject sampling. specifically, a small draft model auto-regressively generates text which is then verified by a larger target model in parallel @xcite . thus the large model does not need to generate text repeatedly but rather guides the small model by correcting outputs when it is truly incapable. this can reduce the number of calls to the large llm, saving both time and memory. however two models are required, along with some similarity in their generative abilities in order for this method to see signficant speedups. while approaches exist to circumvent some of these needs @xcite @xcite @xcite , these are often limited by the need for additional tuning @xcite @xcite , which is difficult in resource constrained settings, or quality degradation in generations @xcite . because of state-action-reward pairs candidate draft models o ine dataset expert draft model draft model training example expert output draft outputs draft output draft output greedy generation candidate draft models draft model draft model testing example expert draft model assisted output policy policy select a draft candidate from the policy . use to assist with generation. use query as policy input. figure 1: overview of our methodology. we first train a policy using offline data collected from greedily decoded output from each model, which are scored to produce reward samples. at test time, the policy takes in a query q ′ to select a draft candidate model, which is then used for assisted generation with the target model. the evident size-cost tradeoff, this is very efficient if the draft model is well aligned to the target. however, while one can ensure that the final output follows the target distribution @xcite , selecting an inadequate draft model can lead to a lack of acceleration due to the signficant number of rejections that will occur. while other methods allow for changes in the output distribution shift @xcite @xcite to further speed-up inference, such types of shifts can be problematic in many high-risk scenarios. from this perspective, the presence of multiple draft models, each suited for different settings, can be helpful for inference acceleration without degredations in generation quality. by dynamically choosing a draft model, speedups can be achieved on multiple domains with marginal additional costs. however this requires learning how to choose the best draft option given a context, introducing a decision making problem which needs to be solved. so how can this decision making process be learned? we start by observing this as a contextual bandits @xcite , where the goal is to have a policy select a draft model based on a given query. this also requires rewards from which the policy can learn to estimate and compare the ideality of different actions that can be taken. to this end, we use an offline process to estimate the contextual alignment between different draft models and the target model on a set of training examples (figure 1 ), enabling the construction of a dataset that defines the preference a target can have towards specific draft candidates. this enables us to train a policy that can take into account such preferences without knowing further details about the draft models. by deloying this policy at inference time it becomes possible to weigh these preferences, leading to speedups in generation with the target. we further show that this policy is useful when using self-speculative decoding, whereby the draft model is a subset of the target model parameters. to summarize our contributions: ## conclusion this work presents the first work at attempting to integrate assisted generation within a setting where multiple black-box draft candidates exist. when no a-priori knowledge of which draft candidate is best suited for assisting the decoding of a given example, the problem can be modeled as a contextual bandits problem, where the goal is to estimate the unknown reward from each drafting option. our work demonstrates that offline rl presents an efficient method for learning to distinguish the available options and provide accelerated decoding across various examples within this setting, with a logical way to collect offline data from models for learning. our results and ablations show that learning a policy with this approach can adapt to general preferences while accounting for more complex aspects of the decision making, highlighting its robustness. furthermore, such a method is scalable and robust to the introduction of more draft models or the removal of draft models, presenting a viable alternative to settings where a uniquely superior draft model may be unavailable. nevertheless, areas of further development exist. for example, learning in an online fashion may render this method more broadly applicable. alternatively, exploring how to dynamically choose drafters at every decoding step rather than per example, as well as combining this direction of work with that which attempts to adaptively choose the speculation length at every step, are feasible ways of combining our findings with concurrent work in the hopes of reaping the benefits of all methods.