Ускорение поиска луча с помощью seq2seq

Question

У меня есть полностью рабочая модель внимания seq2seq с поиском луча, и она дает улучшенные результаты. Но это займет> 1 мин для логического вывода (размер партии 1024) с k=5 (мои гипотезы k), потому что ни один из них не распараллелен. Все происходит 1 образец за один раз.

Задача (упрощенная)
Цель - перевод предложения, 15 слов Lang A до 15 слов Lang B.

• Encoder - это 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 одновременно? Буду признателен за любую оказанную помощь.