Transformer 模型舍弃了 step by step 的时序循环,依靠 self-attention 的并行结构获得了远超同量级 LSTM 网络的训练速度。加入 target-mask 后的 decoder,可以做到像 encoder 一样的“完全并行”。
但是,在预测的时候,还是避免不了 auto-regressive —— 只有前一个词完全解码出来后,才能进行下一个词的预测。也就是说,“并行”的只是训练过程,由于网络结构复杂,预测过程反而比 LSTM 慢多了。
在进行上线部署前,解码端的加速就变成了一个必须要面对的问题。在不过分损失精度的前提下,有几种便于实施的方案,值得尝试。需要说明的是,这里不涉及编码端的优化。
1. Attention Cache
Attention Cache 是一种工程式的优化,不损失精度,不改变网络,因此可以看作是“必选项”。在官方代码 tensor2tensor 的 multihead_attention() 函数有一个“cache”输入参数,注释中这样说明:The caching works by saving all the previous key andvalue values so that you are able to send just the last query location to this attention function. I.e. if the cache dict is provided it assumes the query is of the shape [batch_size, 1, hidden_dim] rather than the full memory.
也就是说,把 self-attention 层的 K 和 V,缓存起来