BanditSpec: Decodificação Especulativa Adaptativa via Algoritmos Bandit

BanditSpec: Adaptive Speculative Decoding via Bandit Algorithms

May 21, 2025

Autores: Yunlong Hou, Fengzhuo Zhang, Cunxiao Du, Xuan Zhang, Jiachun Pan, Tianyu Pang, Chao Du, Vincent Y. F. Tan, Zhuoran Yang

cs.AI

Resumo

A decodificação especulativa emergiu como um método popular para acelerar a inferência de Modelos de Linguagem de Grande Escala (LLMs, na sigla em inglês) enquanto mantém seu desempenho superior na geração de texto. Métodos anteriores adotam uma configuração fixa de decodificação especulativa, independentemente dos tokens de prefixo, ou treinam modelos de rascunho de forma offline ou online para alinhá-los ao contexto. Este artigo propõe uma estrutura de aprendizado online sem treinamento para escolher adaptativamente a configuração dos hiperparâmetros da decodificação especulativa à medida que o texto é gerado. Primeiro, formulamos esse problema de seleção de hiperparâmetros como um problema de Multi-Armed Bandit e fornecemos uma estrutura geral de decodificação especulativa chamada BanditSpec. Além disso, dois algoritmos de seleção de hiperparâmetros baseados em bandit, UCBSpec e EXP3Spec, são projetados e analisados em termos de uma nova quantidade, o arrependimento do tempo de parada. Limitamos superiormente esse arrependimento tanto em cenários de recompensa estocástica quanto adversariais. Ao derivar um resultado de impossibilidade teórico da informação, demonstra-se que o desempenho de arrependimento do UCBSpec é ótimo até constantes universais. Por fim, extensos experimentos empíricos com LLaMA3 e Qwen2 mostram que nossos algoritmos são eficazes em comparação com métodos existentes, e a taxa de transferência se aproxima do melhor hiperparâmetro oráculo em cenários simulados de serviço de LLMs na vida real com prompts de entrada diversos.

English

Speculative decoding has emerged as a popular method to accelerate the inference of Large Language Models (LLMs) while retaining their superior text generation performance. Previous methods either adopt a fixed speculative decoding configuration regardless of the prefix tokens, or train draft models in an offline or online manner to align them with the context. This paper proposes a training-free online learning framework to adaptively choose the configuration of the hyperparameters for speculative decoding as text is being generated. We first formulate this hyperparameter selection problem as a Multi-Armed Bandit problem and provide a general speculative decoding framework BanditSpec. Furthermore, two bandit-based hyperparameter selection algorithms, UCBSpec and EXP3Spec, are designed and analyzed in terms of a novel quantity, the stopping time regret. We upper bound this regret under both stochastic and adversarial reward settings. By deriving an information-theoretic impossibility result, it is shown that the regret performance of UCBSpec is optimal up to universal constants. Finally, extensive empirical experiments with LLaMA3 and Qwen2 demonstrate that our algorithms are effective compared to existing methods, and the throughput is close to the oracle best hyperparameter in simulated real-life LLM serving scenarios with diverse input prompts.