论文

LLM解码深度探索

Depth Exploration for LLM Decoding

模型推理推理加速

摘要

自回归 LLM 解码通过整个层堆栈评估每个生成的词元,即使许多词元在中间深度变得可预测。现有的无损深度自适应方法通过选择单个非最终退出深度并使用最终深度模型验证其预测来利用这种冗余。然而,我们的测量表明,这种基于选择的策略留下了巨大的空间:选择太晚的退出会浪费计算,而选择太早的退出会触发回退并丢弃相关草稿。我们提出深度探索解码(DEX),这是一种无损解码算法,用多个候选深度上的并行探索取代单深度选择。在每个提交位置,DEX 根据最终深度参考验证候选者,准确提交最终深度标记,并折叠探索网格以仅保留可重用的分支状态。这种扩展-提交-折叠过程保留了与标准自回归解码的等效性,同时降低了提交每个词元的成本。在早期退出训练和标准 LLM 中,DEX 的性能优于代表性深度选择基线,并实现了与推测性和分布式解码方法相比具有竞争力的端到端吞吐量。此外,DEX 随着探索的深度变得更精细而改进,这表明并行深度探索提供了一种可扩展的方式来利用 LLM 解码的未充分利用的深度轴。