Все уже выучили идею Attention - берём матрицу Query (Nxd), Key^T (dxN), умножаем и применяем софтмакс - получаем распределение того, на какие токены смотрит каждый токен. Матрица получается NxN, и от этого квадратичного размера у всех подгорает. Потом результат умножается на Value (Nxd) и получается выход.
Некоторые внимательные товарищи заметили, что если убрать софтмакс, то вычисления почти не изменятся, но attention внезапно станет полностью линейной операцией - Q x K^T x V. В этом случае вычисление выхода можно пересобрать так, чтобы сложность стала линейной!
Грубо, в двух словах, мы провернули следующее:
В обычном трансформере мы считаем для каждого токена:
s_1(q_i x k_1) x v_1 + s_2(q_i x k_2) x v_2 + ...
Где s_1, s_2 - это несвязанные между собой применения софтмакса с учётом других query-токенов.Убирая софтмаксы, мы получаем:
q_i x k_1 x v_1 + q_i x k_2 x v_2 + ...
Что позволяет вынести q_i за скобки:q_i x (k_1 x v_1 + k_2 x v_2 + ...)
Но эта сумма в скобке-то для всех токенов одна и та же!Мы можем посчитать её один раз и не умножать матрицы размером Nxd и dxN друг на друга. В реальности всё чуть сложнее - мы считаем эту сумму налету, двигаясь слева направо, чтобы воспроизвести логику Causal Mask - когда токены смотрят только на предыдущие токены.
Согласно экспериментам, производительность близка к трансформерной, но работает несоизмеримо быстрее. С другой стороны, в теории эта нелинейность может оказаться необходимой - по той же причине, по которой мы не можем убрать нелинейности из MLP. Но авторы компенсируют это втыканием нелинейностей в другом месте. Может быть, это и есть оптимальное решение - перенести линейности туда, где они не создают боттлнеков в расчётах.
Будем следить за развитием событий!
@knowledge_accumulator