SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths

Speculative decoding通过利用一个更小更快的草稿模型来减少目标大型语言模型的推理延迟。其性能取决于一个超参数K，即每轮目标模型要验证的候选标记数。然而，以往的方法通常使用简单的启发式方法选择K，这可能导致次优性能。我们研究候选长度K的选择，并将其制定为马尔可夫决策过程。我们在理论上表明，该马尔可夫决策过程的最优策略采用阈值策略，即当获得拒绝的概率超过阈值时，当前的推测应该停止并进行验证。在这个理论的启发下，我们提出了SpecDec ++，这是一种增强版的推测解码，可以自适应地确定候选长度。我们使用一个训练好的接受预测头来增强草稿模型，以预测候选标记的条件接受概率。当预测到至少一个标记被拒绝的概率超过阈值时，SpecDec ++将停止当前的推测。我们实现了SpecDec ++并将其应用于llama-2-chat 7B＆70B模型对。我们的自适应方法在Alpaca数据集上实现了2.04倍的加速（比基线推测解码额外提高了7.2％）。在GSM8K和HumanEval数据集上，我们的方法分别实现了2.26倍的加速（9.4％的改进）和2.23倍的加速（11.1％的改进）。

本文旨在解决大型语言模型推理延迟问题，通过使用较小且更快的草稿模型进行推测解码来减少延迟。作者试图通过优化候选长度K来提高性能。

本文将候选长度K的选择问题建模为马尔可夫决策过程，并理论上证明了最优策略是阈值策略。作者提出了SpecDec ++，一种增强的推测解码方法，可以自适应地确定候选长度，并在草稿模型上添加了一个接受预测头来预测候选标记的条件接受概率。

实验使用了llama-2-chat 7B＆70B模型对SpecDec ++进行了测试，并在Alpaca数据集上实现了2.04倍的加速。在GSM8K和HumanEval数据集上，我们的方法分别实现了2.26倍和2.23倍的加速。作者提供了开源代码。

与本文相关的研究包括使用草稿模型进行推测解码的先前工作，以及将深度强化学习应用于语言模型的研究。相关论文包括Speculative BERT，Deep Reinforcement Learning for Non-Autoregressive Neural Machine Translation和Fast Decoding in Autoregressive Models with Augmented Restart.