论文
WaveFront 解码:循环语言模型的并行自推测解码
WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models
摘要
循环语言模型重复应用权重共享块来增加有效深度而不增加参数数量,但每个生成的词元所产生的 T 个顺序循环块调用大大增加了解码延迟。为了解决这个问题,我们引入了 Wavefront Decoding (WFD),这是一种专为循环语言模型设计的免训练自推测解码框架。 WFD 利用这些架构的两个属性:中间循环输出提供有效的草稿预测,权重共享允许在一次批量循环块调用中处理不同位置和循环深度的词元状态。 WFD 将这些混合深度状态组织成对角波前,不断在浅深度绘制新位置,同时将早期位置推进到全深度验证。因此,与分阶段起草然后验证的时间表不同,WFD 在同一循环调用中同时进行起草和验证,同时使用全深度预测来纠正被拒绝的草稿。在六个 Spec-Bench 任务类别中,WFD 通过自回归解码在 Ouro-2.6B 上实现了 2.42 倍加速,在 Huginn-3.5B 上实现了 3.54 倍加速,始终优于“先草稿后验证”。交叉循环 KV 共享进一步减少了波前 KV 流量,并将 WFD 在 Huginn-3.5B 上的加速提高到 4.81 倍。