用seq2seq加速Beam Search

Question

我有一个完全工作的seq2seq注意模型与光束搜索，它确实提供了改进的结果。但它需要＆gt; 1min用于推理（批量大小1024），k = 5（k是我的假设），因为它们都没有并行化。一切都是一次发生一次。

任务（简化）
目标是句子翻译，15个字朗A到15个字郎B.

• 编码器是一个RNN，它接收15个单词的句子并对其进行编码，给出[timestep，512]矩阵以及最终的隐藏状态。
• 解码器是另一个RNN，将编码器隐藏状态作为初始状态，使用[timestep，512]矩阵进行注意，并一次一次输出翻译后的单词[批次]。当然，到目前为止，存在某种形式的并行化。
• 在推理阶段，使用光束搜索。在解码器的每个时间步，而不是采用具有最高概率的预测单词，我采用k个最好的单词。并提供k个单词作为下一个时间步的输入，以便它可以预测句子中的下一个单词（算法的其余部分在下面给出）。在随后的时间步长中，算法使得解码不那么贪婪的预期结果具有更高的总概率。

for each element in the test-set
    calculate initial k (decoder-encoder step)
    for range(timesteps-1)
        for each prev k
            get hidden state
            obtain its best k
            save hidden state
        find new k from k*k possible ones
        ##update hypotheses based on new found k
        for element in k 
            copy hidden state
            change hypotheses if necessary
            append new k to hypotheses

有6个张量和2个列表可以跟踪和处理状态变化。这里有加速或并行化的空间吗？也许每个k可以同时通过enncode-decode？任何帮助深表感谢。