Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks. However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is $\textit{speculative decoding}$: use a small model to sample a $\textit{draft}$ (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with $\textit{membership cost}$. This framework can be viewed as an extension of the well-known $\textit{maximal-coupling}$ problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of $k$ candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in $k$. We then propose a valid draft selection algorithm whose acceptance probability is $(1-1/e)$-optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token. Using this $new draft selection$ algorithm, we develop a new autoregressive sampling algorithm called $\textit{SpecTr}$, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

翻译：自回归采样从大型语言模型中生成文本，已在多个自然语言任务中取得最先进成果。然而，自回归采样每次仅生成一个token，导致速度缓慢，甚至在某些任务中无法实际应用。加速采样的一种方法是推测解码：使用小型模型采样一个草稿（token块或序列），然后通过大型语言模型并行对草稿中的所有token进行评分。基于统计方法，草稿中的部分token被接受（其余被拒绝），以确保最终输出遵循大型模型的分布。在本工作中，我们通过最优传输理论（带隶属成本）提供了对推测解码的原理性理解。该框架可视为经典的最大耦合问题的扩展。这一新公式使我们能够将推测解码方法泛化到token级别允许k个候选集的情形，从而优化隶属成本。我们证明，最优草稿选择算法（传输方案）可通过线性规划求解，但其已知最优运行时间随k呈指数增长。随后，我们提出一种有效的草稿选择算法，其接受概率在乘法意义上达到(1-1/e)-最优，且计算时间与单个token的域大小近似线性。利用这一新型草稿选择算法，我们开发了名为SpecTr的自回归采样算法，在确保解码输出质量不下降的同时实现加速。实验表明，对于最先进的大型语言模型，所提方法在标准基准上实现2.13倍的实时加速，比推测解码进一步加速1.37倍。