От трансформеров к исполнителям: приблизительное внимание

Математический трюк для ускорения трансформаторов

Несколько недель назад исследователи из Google, Кембриджского университета, DeepMind и Института Алана Тьюринга выпустили статью Переосмысление внимания с исполнителями, в которой делается попытка найти решение проблемы узкого места softmax в трансформаторах [1]. В их подходе используется хитрый математический трюк, который я объясню в этой статье.

Предварительные требования:

  • Некоторые знания о трансформаторах
  • Функции ядра

Охваченные темы:

  • Почему трансформаторы?
  • Проблема с трансформаторами
  • Обход узких мест softmax

Почему трансформаторы?

По сути, Transformer - это модель, предназначенная для эффективной работы с последовательными данными, и на самом деле она активно используется в задачах обработки естественного языка (NLP), которые требуют обработки последовательностей слов / букв. В отличие от других последовательных моделей, преобразователь использует механизмы внимания для параллельной обработки последовательных данных (то есть: нет необходимости обрабатывать одно слово / ввод за раз) [2].

Если вы не знакомы с трансформаторами, я рекомендую прочитать Внимание - все, что вам нужно - статья, в которой они были представлены в 2017 году, очень доступная - или Иллюстрированный трансформатор для более удобного для начинающих введения. [3] [4]

Проблема с трансформаторами

Трансформаторы основаны на внимании, которое рассчитывается следующим образом:

где Q, K, V (размеры L x d) - это матрицы запросов, ключей и значений, L - длина последовательности, а d - (произвольное) измерение векторов запросов, ключей и значений.

Примечание. Если вы не знаете, откуда это взялось, я настоятельно рекомендую прочитать документ, упомянутый ранее (Внимание - все, что вам нужно), хотя для целей этой статьи не обязательно знать, какие запросы, ключи и ценности и почему внимание рассчитывается таким образом (или что это такое).

Проблема с трансформатором возникает из-за функции softmax, давайте посмотрим, почему.

Сложность времени внимания

Напоминаем, что временная сложность умножения двух матриц размеров n x m и m x p равна O (nmp).

Если мы посмотрим на уравнение для внимания, мы увидим, что мы умножаем три матрицы: Q (с размерами L x d), K ^ T (d x L) и V (L x d). Мы получим разную сложность в зависимости от того, в каком порядке мы их умножаем.
Игнорируя softmax и знаменатель sqrt (d) на мгновение (который является просто скаляром), мы можем увидеть это, умножив QK ^ T сначала получаем сложность O (L²d), а если умножить K ^ T V сначала получаем сложность
O (d²L).

Мы предпочли бы O (d²L), поскольку d - это параметр, который мы можем выбрать, и обычно мы можем иметь d ‹L. Однако на самом деле мы не можем выполнить умножение в таком порядке, потому что Q K ^ T «застрял» внутри softmax, и нет простого способа избавиться от него. Это означает, что мы вынуждены иметь дело с временной сложностью O (L²d), которая квадратична по длине последовательности (поэтому обработка более длинных последовательностей становится все более затратной в вычислительном отношении).
Итак, softmax - это узкое место в трансформаторах, и мы хотели бы найти способ исправить это.

Обход узких мест softmax

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

По сути, можем ли мы найти такие матрицы Q ’ и K’, что

Цель достаточно проста для понимания, но детали того, как ее достичь, немного более запутаны.

Во-первых, давайте вспомним, что softmax - это функция, которая для вектора z длины n нормализует все элементы как:

Учитывая это, обратите внимание, что мы можем переписать softmax в уравнении внимания как:

где экспонента в A применяется поэлементно, 1 _L - это единичный вектор длины L, а D - диагональная матрица. с элементами A1 _L. D дает знаменатель softmax (на самом деле A1 _L - это просто вектор длины L, полученный суммированием столбцов А).

A с его поэлементной экспонентой представляет собой настоящую проблему, поэтому наша цель - как-то разложить его на части. Мы можем игнорировать скалярный знаменатель sqrt (d), поскольку он используется только для нормализации, но мы можем эквивалентно нормализовать запросы и ключи. Это означает, что наша цель - найти такие Q ’ и K’, что:

Нахождение ядра Softmax через гауссово ядро

Здесь в игру вступают ядра. Мы знаем, что ядра - это функции, которые эквивалентны скалярному произведению некоторой карты признаков φ:

Обычно, учитывая некоторую многомерную карту признаков φ, мы заинтересованы в поиске эквивалентной функции K, которая позволит нам избежать вычислений в многомерном пространстве φ. Однако в нашем случае мы фактически пойдем по противоположному пути: если мы предположим, что A - это матрица ядра с элементами A (i, j ) = K (q _i, k _j) = exp (q_i k _j ^ T) (где q_i и k _j - векторы-строки в Q и K), можем ли мы найти функцию map φ, чтобы помочь нам разложить A?

Теперь большинство ядер можно аппроксимировать картой характеристик вида

где h и f₁,…, f_l - некоторые детерминированные функции, а w₁,…, w_m - случайные значения, взятые из распределения D (поэтому φ (x) - вектор с элементами l x m). [5]

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

То есть, если мы извлечем w из нормального распределения со средним значением 0 и единичной дисперсией, мы можем получить гауссово ядро, используя карту характеристик:

Обратите внимание, что гауссово ядро ​​с единичной дисперсией определяется следующим образом:

Теперь помните, что мы хотим найти ядро ​​softmax:

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

Это означает, что мы можем фактически переписать ядро ​​softmax как

и мы можем повторно использовать карту функций, которая ведет к гауссовскому ядру, изменив функцию h с h (x) = 1 на:

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

Поиск более стабильного ядра Softmax

Исследователи обнаружили, что ядро ​​softmax также можно переписать как:

(Доказательство того, что это на самом деле ядро ​​softmax, можно найти в приложении к статье.)

Таким образом, мы можем просто взять форму карты объектов ранее и установить

Сделав это, мы увидим, что все значения будут положительными, поскольку мы используем exp, что решает проблему, которая у нас была раньше.

Авторы также предлагают альтернативную карту функций, которая ведет к тому же ядру, я рекомендую прочитать статью, если вам интересно.

Использование ядра softmax для поиска Q ’и V’

Подведем итоги.
Мы начали с уравнения внимания

и обнаружили, что можем переписать его как

Затем мы нашли карту функций для ядра softmax, которую можно использовать для аппроксимации матрицы A:

Итак, теперь мы можем заменить элементы в A, используя карту функций:

Обратите внимание, что мы переходим от векторов q_k длины L к векторам φ (q _i) и φ (k _j) длиной м.

Теперь мы можем разложить A на Q ' и K', где элементы Q ' и K ' - это φ (q _i) и φ (k _j). Наконец, мы можем изменить порядок умножения матриц и уменьшить временную сложность с O (L²d) до O (Lmd), тем самым получая линейную сложность. - вместо квадратичной - по длине последовательности.

И готово!

Заключение и несколько заключительных замечаний

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

Обратите внимание на несколько интересных моментов:

  • Хотя этот метод был разработан с учетом трансформаторов, его можно применить к любой модели, требующей softmax. Будет интересно посмотреть, где это окажется полезным.
  • Авторы отмечают, что этот метод не только быстрее, но и эффективнее с точки зрения памяти. В этом можно убедиться, посмотрев на размеры матриц, которые необходимо сохранить.

использованная литература

[1] Чороманский и др. «Переосмысление внимания с исполнителями, 30 сентября 2020 г. »
[2] Джей Аламмар. «Визуализация модели нейронного машинного перевода (механика моделей Seq2seq с вниманием), 9 мая 2018 г. »
[3] Vaswani et al. «Внимание - все, что вам нужно, 6 декабря 2017 г. »
[4] Джей Аламмар. «Иллюстрированный трансформер, 27 июня 2018 г. »
[5] Лю и др. «Случайные функции для аппроксимации ядра: обзор алгоритмов, теории и не только, 4 июля 2020 г. »

Первоначально опубликовано на https://chiaracampagnola.io 29 октября 2020 г.