Involution: Attention is not what you need, или Как скрестить Self-Attention из NLP и Convolution в задачах CV
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 (тепловую карту) изображения.
Как это работает:
Проходимся окном по картинке и смотрим на центральный пиксель окна, вычисляем его query с помощью весов и берём это значение. Напомню, в nlp мы использовали токены, а тут – пиксели.
Вычисляем key каждого пикселя в окне и умножаем их на query центрального пикселя. Тем самым, получаем матрицу комбинаций нашего query центрального пикселя с каждым key и превращаем в вероятности с помощью Softmax.
После вычисления матрицы 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), первым шагом генерируется ядро свёртки для окна:
Изменение разрешения входного тензора
AvgPool2d, если stride больше 1, то мы проходимся пулингом для уменьшения входного тензора.
Иначе пропускаем тензор через слой Identity*.
Conv2d (слой сокращения), с его помощью мы уменьшаем количество каналов. Для этого проходимся фильтром размера 1x1 с входным количеством каналов C, выходным C//r, где r – коэффициент сокращения.
Conv2d (слой сборки ядра) – последний слой в генерации ядра свёртки, размер окна 1, входное количество каналов C//r, выходное K*K*G, где K – размер ядра свёртки, G – количество групп.
По итогу получаем тензор размерностью (K*K*G, H, W), то есть набор значений для каждого пикселя исходного изображения.
Далее делаем решейп (G,K*K,H,W) и получаем набор ядер свёрток для
*Зачем это делается я сам не понял, ведь можно просто ничего с ним не делать.
Реализация свертки с помощью сгенеренного ядра.
Так как входное и выходное количество слоёв будет одинаковое, мы просто делаем unfold входного тензора размерности (C,H,W) без свёрток, получаем (C, K*K, H*W), и делаем решейп в (G, C//G, K*K, H, W).
Далее с помощью простого умножения (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.