PRISM：面向投机采样草稿模型的参数化推理重构

摘要

大语言模型（LLM）受其自回归性质的约束，解码速度缓慢。投机解码方法已成为加速 LLM 解码的一种有前景的方案，吸引了系统社区与 AI 研究社区的关注。近来，对更好草稿质量的追求推动了草稿模型参数规模增大的趋势，这不可避免地引入了可观的计算开销。现有工作试图在预测精度与计算延迟之间进行权衡，而我们通过架构创新来化解这一根本困境。我们提出 PRISM，它将每个预测步骤的计算拆分到不同的参数集上，重构草稿模型的计算通路，成功地将模型容量与推理成本解耦。大量实验表明，PRISM 优于所有现有草稿架构，在保持极小草稿延迟的同时取得了出色的接受长度，带来卓越的端到端加速。我们还借助 PRISM 重新审视了缩放定律，发现与其他草稿架构相比，PRISM 随数据量扩大而扩展的效率更高。通过严格而公平的比较，我们表明 PRISM 将一个已经高度优化的推理引擎的解码吞吐量提升了 2.6 倍以上。

PRISM：面向投机采样草稿模型的参数化推理重构 — 图3：Prism 架构及其计算路径示意图。（a）目标模型预填充或验证多个 token，其最后一层 transformer 的输出隐藏状态被用作草稿模型的输入，并与输入 token 嵌入融合。融合后的特征被送入 Prism 草稿模型的第一个 transformer 进行预填充。（b）在融合上一个生成 token 的嵌入与之前生成的隐藏状态之后，将融合特征应用于第二个 transformer 层以执行第一个解码步骤。注意，transformer 1 的 KV 缓存会传递给 transformer 2。