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

TL;DR

The paper addresses the latency of large language model inference by proposing a context-aware approach that selects among multiple draft models for assisted decoding. It reframes drafting as a contextual bandit problem and learns an offline policy from alignment-based rewards derived from draft–target outputs, enabling fast, domain-aware draft selection without exposing draft-model internals. Empirical results across translation and summarization show that the learned policy achieves speedups with minimal quality loss, can generalize to multi-task drafters, and remains effective even when self-drafting is used. The findings indicate that offline policy learning for draft selection is scalable and robust, offering a practical path to adaptive, resource-efficient LLM deployment, while highlighting areas for online learning and per-step adaptation as future work.

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 noted issue is the high latency associated with auto-regressive generation, rendering large LLMs use dependent on advanced computing infrastructure. Assisted decoding, where a smaller draft model guides a larger target model's generation, has helped alleviate this, but remains dependent on alignment between the two models. Thus if the draft model is insufficiently capable on some domain 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 provided the candidates are effective. Further results show this to hold on various settings with multiple assisted decoding candidates, highlighting its flexibility and the advantageous role that such decision making can play.

Figures (3)

• 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.
• Figure 2: Effect of varying the tradeoff between output alignment and draft model size (controlled through $\alpha$). Each compares the use of Flan-T5-Small as a draft model (red horizontal line). As $\alpha$ increases, the model increasingly uses the smallest draft model for decoding, demonstrating that the offline dataset is sufficient to learn how to balance the quality of the draft model's outputs and the cost of using it. All cases use speculative sampling/decoding.
• Figure 3: Decoding speed using a dynamic policy as a function of the number of examples used to train the policy, tested on IWSLT2017 En-De (left) and XSum (right). The number of examples is marked in $\log$-scale. Horizontal lines denote decoding speeds of individual drafting options.