Почему декодер генерирует все скрытые состояния во время вывода?

Вопрос или проблема

Кажется, что в ванильных трансформерах (например, AIAYN) во время инференса скрытые состояния генерируются для всех токенов во входной последовательности, но только последнее используется для предсказания следующего токена.

Мой вопрос: почему другие скрытые состояния вообще генерируются во время инференса? Разве не было бы эффективнее, если бы генерировалось только последнее скрытое состояние?

Например, здесь: https://nlp.seas.harvard.edu/annotated-transformer/#greedy-decoding

out = model.decode(
            memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
        )
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]

Я понимаю, что вы имеете в виду, что только скрытые состояния на позиции t – 1 декодера используются для предсказания токена на позиции t. Это не так. Чтобы вычислить скрытые состояния на позиции t слоя L, используются скрытое состояние/токен с позиций 1 до t – 1 слоя L – 1 (используются все предыдущие позиции из предыдущего слоя).

Пример, который вы привели, является частью реализации жадного декодирования. Это должно происходить в цикле, чтобы автогрессивно декодировать токен за токеном. Также пропущена строка из исходного кода, которая обновляет ys, соединяя его с предсказанным токеном, чтобы на следующем шаге декодирования модель использовала все ранее декодированные токены.

Обратите внимание, что в этой базовой реализации вы пересчитываете скрытые состояния первых позиций на каждой итерации цикла. Эти ненужные вычисления можно избежать, сохраняя эти скрытые состояния и используя их в декодировании для повторного использования.

Ответ или решение

Во время инференса в декодере трансформеров действительно генерируются скрытые состояния для всех токенов во входной последовательности, но только последнее скрытое состояние используется для прогнозирования следующего токена. Это может показаться неэффективным, и логичен вопрос: почему не генерировать только последнее скрытое состояние?

На самом деле, скрытые состояния для всех токенов необходимы по нескольким причинам:

Кросс-вагонный контекст: Хотя на первый взгляд может показаться, что для предсказания следующего токена достаточно только последнего скрытого состояния, следует учитывать, что предыдущие токены влияют на текущий. Архитектура трансформеров использует механизм внимания, который учитывает все предшествующие токены. Это значит, что для создания текущего скрытого состояния необходимо учитывать информацию из всех предыдущих состояний.
Параллельные вычисления: Генерация всех скрытых состояний позволяет использовать параллельные вычисления во время инференса. Хотя на каждой итерации используется только последнее скрытое состояние, все скрытые состояния могут быть вычислены одновременно, что ускоряет процесс инференса в рамках одной итерации и позволяет использовать архитектуру эффективно.
Общая структура модели: Архитектура трансформеров была разработана таким образом, чтобы быть модульной и простой в реализации. Если бы в процессе инференса использовались только последние скрытые состояния, это потребовало бы значительной переработки текущих механизмов декодирования, что могло бы привести к ошибкам и добавить сложности в код.
Кэширование скрытых состояний: Как вы правильно заметили, в базовой реализации (например, в приведённом вами примере с жадным декодированием) скрытые состояния для первых позиций пересчитываются на каждой итерации. Однако, во многих современных реализациях существует возможность кэширования этих скрытых состояний. Это позволяет сохранять вычисленные значения, что существенно повышает эффективность, так как повторно не вычисляются состояния, которые уже были определены.

Итак, несмотря на то, что во время инференса генерация всех скрытых состояний может показаться избыточной, она обусловлена архитектурными особенностями трансформеров, необходимостью обеспечения качественного контекстного анализа, а также возможностью оптимизации через кэширование.