arXiv:2501.19309cs.LGcs.CL2025-01ICLR被引 61

让大模型当裁判,提升快速生成速度且不丢质量。

Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

  • 用小模型判断草稿是否合理,突破对齐限制
  • 80亿模型实现9倍加速,700亿模型达141词/秒
  • 适合追求高效推理的开发者和部署团队

大型语言模型性能与其规模密切相关,导致模型越来越大、推理越来越慢。推测解码技术通过快速草稿模型提出候选词,并在目标模型下并行验证,虽能保证输出一致,但大量高质量草稿仍被拒绝,严重限制加速潜力。我们发现,即使强大如GPT-4o或人类文本,在标准验证下也难以获得高接受率。因此提出新思路:让目标模型自身作为‘裁判’,通过在嵌入层上训练紧凑模块,生成对当前延续的‘判断’。该方法在Llama-3.1系列上实现80亿/4050亿模型对比下的9倍加速,保持多基准测试质量。优化框架中,8B/70B-Judge达141词/秒,8B/405B-Judge达129词/秒,分别在2块和8块H100上运行。

原文摘要 · Abstract (English)

The performance of large language models (LLMs) is closely linked to their underlying size, leading to ever-growing networks and hence slower inference. Speculative decoding has been proposed as a technique to accelerate autoregressive generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based on their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as GPT-4o, as well as human text cannot achieve high acceptance rates under the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelmingly likely when solely relying on alignment of draft and target. We thus ask the following question: Can we adapt verification to recognize correct, but non-aligned replies? To this end, we draw inspiration from the LLM-as-a-judge framework, which demonstrated that LLMs are able to rate answers in a versatile way. We carefully design a dataset to elicit the same capability in the target model by training a compact module on top of the embeddings to produce ``judgements" of the current continuation. We showcase our strategy on the Llama-3.1 family, where our 8b/405B-Judge achieves a speedup of 9x over Llama-405B, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up to 141 tokens/s for 8B/70B-Judge and 129 tokens/s for 8B/405B on 2 and 8 H100s respectively.

推理加速模型判断推测解码

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