LLM2D
动态宽度推测束搜索解码:高效的 LLM 推理
Dynamic-Width Speculative Beam Decoding for Efficient LLM Inference
作者: Zongyue Qin, Zifan He, Neha Prakriya, Jason Cong, Yizhou Sun
发布日期: 9/26/2024
arXiv ID: oai:arXiv.org:2409.16560v1

摘要

大型语言模型(LLM)在众多现实世界任务中展现出卓越的性能。然而,这些模型的自回归特性使得推断过程缓慢且代价高昂。推测解码作为一种很有前景的解决方案出现,它利用一个较小的辅助模型来起草未来的标记,然后由较大的模型同时验证这些标记,从而实现 1-2 倍的加速。尽管推测解码与多项式采样匹配相同的分布,但多项式采样本身容易产生次优输出,而束搜索因在每一步维护多个候选序列而被广泛认可为产生更高质量的结果。本文探索了推测解码与束搜索的新型集成。然而,存在四个关键挑战:(1)如何在给定来自小型模型的草稿序列的情况下,从大型模型的分布中生成多个序列;(2)如何动态优化束的数量以平衡效率和准确性;(3)如何并行高效地验证多个草稿;(4)如何解决束搜索固有的额外内存成本。为了应对这些挑战,我们提出了动态宽度推测束解码(DSBD)。具体来说,我们首先引入了一种新颖的草稿和验证方案,该方案基于来自小型模型的束搜索轨迹,生成遵循大型模型分布的多个序列。然后,我们引入了一种自适应机制,根据上下文动态调整束的数量,优化效率和有效性。此外,我们扩展了基于树的并行验证以同时处理多个树,从而加速验证过程。最后,我们说明了对我们算法的简单修改,以减轻束搜索的内存开销。