arXiv:2410.18234cs.CLcs.DC2024-10ICLR被引 9

提出多草稿推测采样新方法,提升生成效率与准确性。

Multi-Draft Speculative Sampling: Canonical Decomposition and Theoretical Limits

  • 将最优采样分解为两步:先用重要性采样选中间词元,再单草稿生成输出。
  • 两份相同草稿模型下,可实现100%接受率,且给出最优接受概率公式。
  • 新方法在多种场景中显著提升区块效率和词元生成速率,适合大模型推理加速。

本文研究多草稿推测采样,其中多个草稿模型独立生成候选序列。每一步中,基于词元级别的选择机制接收一组有效词元输入,并输出一个分布与目标模型一致的词元。先前工作表明,最优选择方案可建模为线性规划求解。本文进一步证明,该最优方案可分解为两步:第一步采用重要性采样(IS)方式选出一个中间词元;第二步应用单草稿推测采样生成最终输出。对于两个相同草稿模型的情况,我们进一步推导出接受概率达到1的充要条件,并给出最优接受概率的显式表达式。理论分析还启发了一类基于加权重要性采样的新型词元选择策略。实验结果表明,在多种场景下,该方法持续提升块效率和词元生成速率,优于基线方案。

原文摘要 · Abstract (English)

We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models. At each step, a token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accepting one of the input tokens) can be cast as a solution to a linear program. In this work we show that the optimal scheme can be decomposed into a two-step solution: in the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token. For the case of two identical draft models we further 1) establish a necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability. Our theoretical analysis also motives a new class of token-level selection schemes based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number of scenarios.

推测采样生成效率重要性采样大模型推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。