论文收藏夹
🗒️分块并行解码(Blockwise Parallel Decoding)
00 分钟
2023-10-9
2024-1-20
type
status
date
slug
summary
tags
category
icon
password

摘要

近年来,深度自回归序列到序列模型在各种任务中表现出了令人印象深刻的性能。虽然递归、卷积和自注意网络等常见架构类别在每层所需的计算量和训练时关键路径的长度之间做出了不同的权衡,但生成仍然是一个固有的顺序过程。为了克服这一局限性,我们提出了一种新颖的分块并行解码方案,在该方案中,我们对多个时间步骤进行并行预测,然后回退到由评分模型验证的最长前缀。当应用于可并行处理输出序列的架构时,这使得生成速度在理论上有了大幅提高。我们通过一系列实验验证了我们的方法,这些实验使用了最先进的机器翻译和图像超分辨率自注意模型,与基线贪婪解码器相比,迭代次数最多减少了 2 倍,而质量却没有下降,或者最多减少了 7 倍,但性能却略有下降。就挂钟时间而言,我们的最快模型比标准贪婪解码器的实时速度最多提高了 4 倍。

分块并行解码

该方法主要解决逐个生成 token 的限制,通过训练一些辅助模型,使用贪婪解码策略并行生成候选结果。

原理

具体来说:原始模型为 ,令 ,辅助模型为 ,给定前 个 token,第 个 token 的概率分布为 。通过以下三步在保证结果的同时使用步骤减少至 个( 是原始方法需要的步骤):
  1. 获得 个块预测结果: for
  1. 用原始模型验证分块推理结果,获得尽可能长的推理结果:找到能够使 for 成立的 ,也就是论文中所说的最长前缀
  1. 接受结果

示例

notion image
在预测步骤中,原始模型和辅助模型独立且并行预测出”in“,”the“和”bus“。在验证步骤中,原始模型对每一个预测进行评分。当使用 Transformer 或卷积序列到序列模型时,这三个计算可以并行完成。可以看到,”in“和”the“是正确结果,最终被接受。
💡
该方法能提高解码性能的关键在于能够并行预测&并行验证
但是,问题到这里还没有结束。当我们使用 transformer 时,每一次产生结果需要两次模型调用:第一次是使用 进行预测,第二次是使用 进行验证。这意味着即使我们有最完美的辅助模型,模型调用次数也只能从 减少到

混合使用原始模型和辅助模型

如果将第 步的验证和第 步的预测合并能将调用次数减少到
notion image