Описание архитектуры Multi-Head Self-Attention

Рассмотренная ранее технология Self-Attention выявляет зависимости между объектами последовательности в некоем контексте и потом ранжирует их с использованием функции Softmax. Но в решении практических задач далеко не всегда можно однозначно дать такую оценку. Как правило, коэффициенты зависимости между объектами очень сильно меняются при изменении точки зрения или контекста анализируемого элемента. Окончательное решение о зависимости элементов — это всегда компромисс. Именно использование Multi-Head Self-Attention призвано помочь находить зависимости между элементами со всесторонним рассмотрением исходных данных. А вводимая дополнительная обучаемая матрица весовых коэффициентов поможет модели научиться находить этот компромисс.

Наверное, самым простым вариантом решения подобной задачи будет взять наш класс внимания CNeuronAttention и дополнить его массивом блоков Self-Attention. Такой подход возможен, но он иррациональный. Он ведет к увеличению числа объектов пропорционально увеличению количества голов внимания. Кроме того, последовательный вызов операций каждой головы внимания не дает нам возможности для организации одновременного параллельного вычисления внимания всех голов. Кроме того, последующая операция конкатенации результатов работы голов внимания тоже потребует затрат ресурсов и времени.

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

Давайте пройдем по алгоритму Self-Attention и подумаем о трансформации операций для реализации многоголового внимания Multi-Head Self-Attention.

1. Вначале вычисляем векторы Query (запрос), Key (ключ) и Value (значение). Указанные векторы получаются путем умножения каждого элемента исходной последовательности на соответствующую матрицу WQ, WK и WV.

Для организации Multi-Head Self-Attention нам необходимо повторить данную операцию по количеству голов внимания. Чтобы сильно не усложнять, возьмем, к примеру, три головы внимания.

Думаю, здесь все понятно и не вызывает вопросов.

Теперь посмотрим на размерности тензоров. Напомню, архитектурой модели предусмотрено одинаковое число элементов последовательности на всех этапах. Каждый элемент последовательности описывается неким вектором значений. А так как механизм Self-Attention одинаково применяется к каждому элементу последовательности, то для примера мы можем разобрать операции только с вектором описания одного элемента. При этом размер этого вектора одинаковый для тензоров исходных данных и значений. Но он может отличаться от размерности вектора описания одного элемента последовательности в тензорах запросов и ключей. Давайте обозначим через nI размер вектора исходных данных и через nK размер вектора ключей. Тогда тензоры будут иметь нижеследующие размеры.

Указанные размеры тензоров справедливы для всех голов внимания. А давайте попробуем объединить соответствующие матрицы весовых коэффициентов в одну большую.

Такие матрицы весовых коэффициентов будут иметь размер nI*3nK для матриц запросов WQC и WKC. Матрица WVC получит размер nI*3nI, где 3 — количество голов внимания.

Подставим конкатенированные матрицы в формулы определения векторов.

По правилам умножения матриц получим следующие размеры тензоров.

Сравните размеры тензоров в двух таблицах — они очень похожи. Единственное их отличие в умножении на количество голов внимания. Какую это имеет для нас практическую ценность? Тут все очень просто. Вместо создания нескольких экземпляров объектов для каждой головы внимания, мы можем создать всего по одному объекту для расчета каждой сущности. Как и при организации аналогичного процесса в механизме Self-Attention, мы можем воспользоваться нашими сверточными слоями, только размер окна результатов нужно будет увеличить пропорционально количеству голов внимания.

 

2. Далее определяем парные зависимости между элементами последовательности. Для этого перемножим вектор Query с векторами Key всех элементов последовательности. Данная итерация повторяется для вектора Query каждого элемента последовательности. В результате данной итерации получаем матрицу Score размером N*N, где N — размер последовательности.

В результате этой операции мы ожидаем получить по одному коэффициенту зависимости между парой элементов последовательности для каждой головы внимания. Но операция умножения двух конкатенированных векторов вернет нам только одно значение. Как и в случае одноголового Self-Attention.

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

Выход все же есть. Давайте вспомним правило умножения матриц.

Подставим сюда наши двумерные матрицы многоголового внимания. И не забудем, что вторая матрица транспонируется перед умножением.

Как можно заметить, вектор, который мы ожидали получить, образовывает диагональ матрицы результатов. А все остальные операции для нас только потеря ресурсов. Но мы можем разложить данную процедуру на операции. К примеру, не будем транспонировать матрицу ключей и воспользуемся адамарным произведением матриц (поэлементное умножение матриц).

После этого для получения ожидаемого результата нам достаточно лишь построчно сложить элементы матрицы.

Да, в итоге мы получили результат за две операции вместо одной. Но нужно отметить два момента:

• В формуле Self-Attention используется транспонированная матрица, а это тоже операция над матрицей, хотя и не выделена отдельно. И на ее выполнение тоже требуются ресурсы. При разложении на операции мы отказались от данной процедуры.
• Определение вектора коэффициентов осуществляется в две операции независимо от количества голов внимания.

3. Следующим этапом разделим полученные значение на квадратный корень из размерности вектора Key и нормализуем функцией Softmax в разрезе каждого Query. Таким образом, получаем коэффициенты попарной взаимозависимости между элементами последовательности.

В этом пункте мы не будем что-либо усложнять или упрощать. Деление матрицы на константу всегда выполняется поэлементно независимо от размера матрицы, а нормализовать данные нам придется в разрезе голов внимания.

4. Умножаем каждый вектор Value на соответствующий коэффициент взаимозависимости и получаем скорректированное значение элемента. Цель данной итерации — акцентировать внимание на релевантных элементах и снизить влияние не релевантных значений.

Для решения этой задачи воспользуемся приемами, применяемыми в пункте 2. Сначала изменим размерность вектора значений и приведем его к двумерной матрице. В ней строки будут соответствовать каждой отдельной голове внимания.

После этого мы можем воспользоваться поэлементным умножением вектора коэффициентов зависимости на матрицу значений.

5. Далее суммируем все скорректированные векторы Value для каждого элемента. Результатом данной операции и будет вектор выходных значений слоя Self-Attention.

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

После выполнения всех пунктов механизма Self-Attention в режиме нескольких голов внимания мы получаем по вектору результатов для каждой головы внимания. Соответственно, общий размер тензора результатов будет превышать размер тензора исходных данных пропорционально количеству голов. Для понижения размерности алгоритмом Multi-Head Self-Attention предусматривается умножение конкатенированного тензора результатов на дополнительную матрицу весов W0. Как вы понимаете, данная процедура очень напоминает операцию полносвязного нейронного слоя без функции активации. Аналогичные операции мы выполняли в пункте 1 для определения векторов Query, Key, Values. А значит, мы можем воспользоваться тем же решением и использовать ранее созданные сверточные слои.

Здесь можно отметить еще один момент. При описании работы блока Self-Attention мы обращали внимание на момент равенства размера векторов описания одного элемента последовательности тензоров значений Value и исходных данных. Такое требование исходило из необходимости последующего сложения тензоров результатов Self-Attention и исходных данных. В случае же многоголового внимания мы в любом случае получаем конкатенированный тензор результатов больше тензора исходных данных. Для приведения их в соответствие используется умножение тензора результатов на матрицу W0. Следовательно, с целью экономии ресурсов мы можем понизить размерность вектора описания одного элемента последовательности в тензоре значений Value без риска получения ошибки при последующей обработке данных.

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

Теперь, когда у нас есть полное представление о принципах реализации алгоритма, мы можем перейти к его реализации.