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

arXiv: https://arxiv.org/abs/2408.08470 · HTML: https://mlsanity.com/paper/2408.08470/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 27 of its 42 references that are in the map.
- First posted on arXiv 16 Aug 2024.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [Accelerating Large Language Model Decoding with Speculative Sampling](https://mlsanity.com/paper/2302.01318/) (2023)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [RouteLLM: Learning to Route LLMs with Preference Data](https://mlsanity.com/paper/2406.18665/) (2024)
- [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://mlsanity.com/paper/2401.15077/) (2024)
- [Blockwise Parallel Decoding for Deep Autoregressive Models](https://mlsanity.com/paper/1811.03115/) (2018)
- [Break the Sequential Dependency of LLM Inference Using Lookahead Decoding](https://mlsanity.com/paper/2402.02057/) (2024)
- [Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding](https://mlsanity.com/paper/2401.07851/) (2024)
- [Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation](https://mlsanity.com/paper/2203.16487/) (2022)
- [Speculative Decoding with Big Little Decoder](https://mlsanity.com/paper/2302.07863/) (2023)
- [Instantaneous Grammatical Error Correction with Shallow Aggressive Decoding](https://mlsanity.com/paper/2106.04970/) (2021)
- [On the Limitations of Reference-Free Evaluations of Generated Text](https://mlsanity.com/paper/2210.12563/) (2022)
- [SPEED: Speculative Pipelined Execution for Efficient Decoding](https://mlsanity.com/paper/2310.12072/) (2023)
- [Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decoding](https://mlsanity.com/paper/2307.05908/) (2023)
- [FiDO: Fusion-in-Decoder optimized for stronger performance and faster inference](https://mlsanity.com/paper/2212.08153/) (2022)
- [Sorted LLaMA: Unlocking the Potential of Intermediate Layers of Large Language Models for Dynamic Inference](https://mlsanity.com/paper/2309.08968/) (2023)

## Built on by (most cited first, 5 of 5)

- [Learning to Route LLMs with Confidence Tokens](https://mlsanity.com/paper/2410.13284/) (2024)
- [Mamba Modulation: On the Length Generalization of Mamba](https://mlsanity.com/paper/2509.19633/) (2025)
- [Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs](https://mlsanity.com/paper/2510.20064/) (2025)
- [Learning Adaptive LLM Decoding](https://mlsanity.com/paper/2603.09065/) (2026)
- [Exploring Jailbreak Vulnerability Across Depth-Based Dynamic Inference in Large Language Models](https://mlsanity.com/paper/2610.15207/) (2026)

## Abstract (arXiv)

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.