AI论文精选
⏶52

Set Block Decoding 是一种语言模型推理加速器

发表
由 Itai GatItai Gat 提交
作者: Itai Gat, Heli Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, Yaron Lipman

摘要

AI 生成总结
Set Block Decoding 通过整合下一个 token 预测和掩码 token 预测来加速语言模型生成，从而实现未来 token 的并行采样，并在不牺牲准确性的情况下降低计算成本。
自回归下一词元预测语言模型具有强大的功能，但在实际部署中面临严峻的挑战，特别是由于推理（尤其是解码阶段）的高计算和内存成本。我们提出了一种简单而灵活的范式——集合块解码（SBD），该范式通过在单一架构中集成标准的下一词元预测（NTP）和掩码词元预测（MATP）来加速生成。SBD允许模型并行采样多个、不一定连续的未来词元，这是与先前加速方法的一个关键区别。这种灵活性允许使用离散扩散文献中的先进求解器，在不牺牲准确性的情况下显著提高速度。SBD不需要架构更改或额外的训练超参数，与精确的KV缓存兼容，并且可以通过微调现有的下一词元预测模型来实现。通过微调Llama-3.1 8B和Qwen-3 8B，我们证明SBD在生成过程中所需的正向传递次数减少了3-5倍，同时实现了与同等NTP训练相同的性能。
查看 arXiv 页面查看 PDF
Set Block Decoding 是一种语言模型推理加速器

评论

Itai GatItai Gat
论文提交者

我们提出了一种称为集合块解码的新范例，结合了下一个 token 预测和掩码（或离散扩散）模型，无需任何架构更改即可实现并行解码（速度提升 3-5 倍），并具有精确的 KV 缓存。性能与 NTP 相当！

Tom SpisTom Spis

取得了绝对惊人的结果，各位——祝贺，也感谢你们的了不起的工作并与世界分享！
一点小建议：图1是一个数据可视化图表上的罪过。横轴从1开始意味着你实际上无法直观地比较SBD和NTP——应该从0开始。此外，“Speedup”（加速）作为横轴标签是一个相当糟糕的选择，因为这个词可以被解释为增量改进——更好的选择是诸如“Speed Factor”（速度因子）或“Speed wrpt/NTP”（相对于NTP的速度）之类的词。

Deepak KhetwalDeepak Khetwal

做得好

Jon WangJon Wang

太棒了！训练和推理代码会开源吗？

Aman GuptaAman Gupta

有可以下载的SBD微调版llama-3b吗？如论文中所述？

David RauDavid Rau

@itaigat 你能不能帮忙实现一下？https://github.com/huggingface/transformers/issues/41084