Self-Attention в CV

Если говорить про Self-Attention в картиночных моделях, то тут есть 2 варианта. Олдскульный  “давайте просто перевзвесим фичи” в разных вариантах: поканально, пространственно, в некоторой проекции. И новомодный "давайте обучим трансформер" с представлением патчей как визуальных слов. Первый подход рабочий, но не дает значительного улучшения в плане метрик. Второй подход слишком вычислительно сложный и часто заточен на размер картинок.

Подход коллег из ByteDance AI Lab и университета Пекина сильно отличается от этих крайностей и является переосмыслением Attention-механизма трансформеров в работе свёрток.

Чтобы понять, что такое революционное они придумали, начнём с простого: вспомним, что такое Self-Attention.

Self-Attention в NLP

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

Коротко о принципе работы Self-Attention.

Для каждого токена (слова или подслова) из словаря строится набор ключей K, V, Q, где:

• K (key) – значимость данного токена относительно других, если мы смотрим на него; 

• Q (query) – значимость данного токена относительно других, если мы смотрим из него;

• V (value) – это репрезентация токена. 

Проще говоря, если нам нужно сгенерить новый токен, то Q сгенерённого токена будет смотреть на K остальных токенов. 

Далее вычисляем значимость Q с каждым K и превращаем в вероятности с помощью Softmax, распределяя значения так, чтобы они находились в промежутке [0-1] и их сумма была равна 1. 

Далее умножаем каждый value на итоговую значимость каждого токена, суммируем их и получаем вектор с информацией о значимости каждого токена относительно исходного:

После всех преобразований умножаем нашу сумму на embedding токенов, тем самым получая их ненормированные вероятности (логиты). Далее выбираем топ-1 токен или даём модели выбрать токен.

Self-Attention сам по себе представляет только слой, но если несколько таких слоёв объединить в большую нейронную сеть, то получится архитектура Transformer, которая используется в лингвистических моделях GPT и BERT.

Существующие решения

Хорошо, с Self-Attention в текстах мы разобрались, но что если попробовать применить его к изображениями? Таких подходов было очень много, но мы сейчас посмотрим на самый простой

В статье “Self-Attention In Computer Vision” описывают идею использования механизма для нахождения аномалий на рентгеновских снимках. Self-Attention позволяет построить heatmap (тепловую карту) изображения.

Как это работает:

1. Проходимся окном по картинке и смотрим на центральный пиксель окна, вычисляем его query с помощью весов и берём это значение. Напомню, в nlp мы использовали токены, а тут – пиксели. 

2. Вычисляем key каждого пикселя в окне и умножаем их на query центрального пикселя. Тем самым, получаем матрицу комбинаций нашего query центрального пикселя с каждым key и превращаем в вероятности с помощью Softmax.

3. После вычисления матрицы values умножаем её на матрицу вероятностей и складываем все элементы в матрицы, получая итоговый пиксель.

А вот результат работы Self-Attention после обработки всего изображения. Исходная картинка превратилась в heatmap, с её помощью можно найти нужный сегмент на изображении.

Это решение практически идентично реализации Self-Attention в NLP. Одно важное отличие – отсутствие словаря, который обязательно нужен в задачах с текстом для получения embedding. В данном решении нам хватит итогового пикселя.

Involution

Теперь, когда мы вспомнили немного теории, приступим к разбору оригинальной статьи “Involution: Inverting the Inherence of Convolution for Visual Recognition”. Авторы предложили идею объединения Convolution и Self-Attention методов для более быстрого и простого использования.

Вспомним Convolution. У нас, как и раньше, есть плавающее окно и ядро свёртки (фильтр), умножая данное окно на фильтр, мы можем выделить нужные данные и распределить их по дополнительным каналам. Однако, если в стандартном алгоритме Convolution - ядро свертки – это обучаемые веса, то тут это сгенерированный из входа тензор.

Реализация Self-Attention. Как было упомянуто ранее, для работы алгоритма Self-Attention требуются значения (query, key, value), которые получаются путём умножения входного тензора на матрицы весов каждого значения. В данной работе авторы отказались от явных и тяжелых вычислений Q, K, V. Они предложили генерировать ядро свёртки, которое будет уникальным для каждого окна. В него включена взаимосвязь query и key каждого пикселя в окне. Таким образом модель сама находит связь между пикселями без прямого вычисления query и key каждого. С помощью матрицы взаимосвязи, умножив её на окно, мы получим результат Self-Attention-механизма, который осталось просуммировать и записать в итоговый результат.

Подробное описание реализации.

На вход подаётся тензор размерности (CxWxH), где C – количество каналов, W – ширина, H – высота. Допустим, на вход подаётся картинка (3xHxW), первым шагом генерируется ядро свёртки для окна: 

1. Изменение разрешения входного тензора

   1. AvgPool2d, если stride больше 1, то мы проходимся пулингом для уменьшения входного тензора.

   2. Иначе пропускаем тензор через слой Identity*.

2. Conv2d (слой сокращения), с его помощью мы уменьшаем количество каналов. Для этого проходимся фильтром размера 1x1 с входным количеством каналов C, выходным C//r, где r – коэффициент сокращения.

3. Conv2d (слой сборки ядра) – последний слой в генерации ядра свёртки, размер окна 1, входное количество каналов C//r, выходное K*K*G, где K – размер ядра свёртки, G – количество групп.

4. По итогу получаем тензор размерностью (K*K*G, H, W), то есть набор значений для каждого пикселя исходного изображения.

5. Далее делаем решейп (G,K*K,H,W) и получаем набор ядер свёрток для

*Зачем это делается я сам не понял, ведь можно просто ничего с ним не делать.

Реализация свертки с помощью сгенеренного ядра.

1. Так как входное и выходное количество слоёв будет одинаковое, мы просто делаем unfold входного тензора размерности (C,H,W) без свёрток, получаем (C, K*K, H*W), и делаем решейп в (G, C//G, K*K, H, W).

2. Далее с помощью простого умножения (mul) мы перемножаем входные тензоры на их ядра свёртки. Полученные тензоры размерностью (K,K,C) суммируем, получая тензор размерности (G,C//G,H,W) и делаем решейп, получая размерность (C,H,W).

Итоговый тензор – это не что иное, как Self-Attention на каждый пиксель исходного тензора. Как и в предыдущей статье, вход и результат идентичны, отличается алгоритм получения результата.

В статье “Involution: Inverting the Inherence of Convolution for Visual Recognition” описан ряд экспериментов, в которых модифицированные модели с Involution показывают результат не хуже, а то и лучше, чем оригинальные. Например, стандартная модель Faster R-CNN с RedNet-50 на 31.6 млн параметров и модель RedNet-50 с Involution с 29 млн параметров показывают результаты сопоставимые по качеству, но вторая требует почти в 2 раза меньше операций в секунду. Это новое слово в реализации Attention в Computer Vision. Возможно, в некоторых задачах эта идея придёт на замену свёрточным слоям, ведь облегчение нейронной сети и ускорение её обучения без потери качества – главный профит от замены Сonvolution на Involution.

Этот алгоритм можно использовать во всех задачах, где требуется heatmap: будь то object detection and classification, key point detector или просто сам heatmap.

Со всеми экспериментами авторов статьи можно ознакомиться в репозитории github.