中科院计算所副研究员冯洋:神经机器翻译的训练改进和解码提速

将门创投
关注

解码

1. CubePruning

下面介绍在解码方面的两个工作,第一个工作要解决的是beam search每一步要计算BeamSize*|V|的问题,这个计算量大大降低了inference时候解码的速度。

这是解码过程中每个步骤的时间消耗,对于GPU来说,大部分的时间消耗在的计算上,其它三个步骤比较节省时间,对于CPU来说,最耗费时间的是最后两个步骤,因为|V|比较大。

传统的方法使用的是Beam Search,传统的 Beam Search其实是一个二维的搜索方法。其中第一维就是已经生成的部分的译文,假设Beam Size = 4,那么就是四个译文。第二维度是这四个译文都要进行下一步的Token预测计算。总共就需要计算4*|V|的概率。因为|V|的个数通常是几千上万级别的,所以这个部分的计算量就非常大。

我们的做法是将二维的搜索扩展成三维的搜索,具体的做法分为以下几步:

1.  Beam分组:假设我们要解码第11步,我们就将第10步解码出来相同Token的候选序列归为一组。

2.  分组预测第11步的候选Token:只用每个组得分最高的哪个候选序列来计算当前的Token分布。

3.  近似组员的Token分布:由上一步已经知道本组最优的候选序列的下一个token的预测分布,对于组员来说,也将共享其老大计算出来的Token分布score,然后和自身的序列score相加,得到自身扩展一个Token后的score。这个score作为自身的近似分。

4.  查找Top-K:经过上面的计算之后,这样每个组就是得分其实是一个二维矩阵,我们将矩阵横轴作为每个组员,纵轴表示当前步预测的token,然后保证右上角score最大,往右,往下都是减小。这样便于我们查找Top-K。具体请看下一张slides。

对于近似的score这里有两个选择,1) 如果取到的candidate是预测的score,那么用真实的状态来重新计算一下这个score,这时候也顺便更新了一下自己的隐状态,2)直接用预测的score,不使用更新的方式,这时候和老大哥共享隐状态。

这个是GPU上的结果,横轴是速度,纵轴是BLEU值,可以看出在取得最优的BLEU值的情况下,我们的方法所用的时间是更短的。速度可以提升3.3倍。在CPU下,提速可以达到3.5倍。

在Beam Size=40的情况下,GPU上速度提升3.8倍,CPU上提升4.2倍。

声明: 本文由入驻OFweek维科号的作者撰写,观点仅代表作者本人,不代表OFweek立场。如有侵权或其他问题,请联系举报。
侵权投诉

下载OFweek,一手掌握高科技全行业资讯

还不是OFweek会员,马上注册
打开app,查看更多精彩资讯 >
  • 长按识别二维码
  • 进入OFweek阅读全文
长按图片进行保存