arXiv:1906.05909

Статья сильно не новая, но полезная. В ней attention подход переносится из сетей, применяемых для работы с текстами (перевода), на сети, которые призваны решать задачи компьютерного зрения.

Обычно для работы с изображениями используются свёрточные нейронные сети, они преобразуют входной тензор $x \in \mathbb{R}^{h \times w \times d_{in}}$, сворачивая окрестность $\mathcal{N_k}$ размера $k \times k$ вокруг каждой точки с претренированными весами $W \in \mathbb{R}^{h \times w \times d_{out} \times d_{in}}$ и получают новый тензор $y \in \mathbb{R}^{h \times w \times d_{out}}$ по формуле:

\[y_{ij} = \sum_{a,b \in \mathcal{N}_k(i,j)} W_{i-a, j-b}x_{ab}\]

здесь суммирование выполняется по окрестности $\mathcal{N}_k(i,j) = \left\{a,b \vert \vert a-i \vert \le k/2, \vert b-j \vert \le k/2 \right\}$.

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

В статье предлагается заменить свёртки на self-attention блоки. Подход этот был изначально предложен для задач перевода текстов, а здесь обобщен для обработки изображений. Для пикселя $x_{ij} \in \mathbb{R}^{d_{in}}$ вновь рассматривается его окрестность, но результат получается не при помощи свёртки пикселей окрестности, а по формуле:

\[y_{ij} = \sum_{a,b \in \mathcal{N}_k(i,j)} {\rm softmax}_{ab}\left(q^T_{ij}k_{ab}\right)v_{ab}\]

здесь $q_{ij}=W_Qx_{ij}$ - запрос, $k_{ab}=W_Kx_{ab}$ - ключи и $v_{ab}=W_vx_{ab}$ - значения. Весовые матрицы $W_Q, W_K, W_V \in \mathbb{R}^{d_{out} \times d_{in}}$ получаются в процессе обучения. Т.е. значения в пикселе $x_{ij}$ преобразуются в некий запрос, а пиксели в окрестности данного - в набор пар ключ-значение. Результат получается как взвешенная сумма значений, с учетом близости запроса и соответствующих значениям ключей.

На практике обычно входной тензор разрезается на $N$ групп равномерно по оси глубины, преобразуется в набор выходных векторов и затем эти выходные вектора склеиваются. Т.е преобразования осуществляется для каждой части вектор $x^n_{ij} \in \mathbb{R}^{d_{in} / N}$, весовые тензоры соответственно $W^n_Q, W^n_K, W^n_V \in \mathbb{R}^{d_{out} / N \times d_{in} / N}$, и выход есть объединение полученных групп $y_{ij} \in \mathbb{R}^{d_{out}}$.

Пока формула для преобразования не подразумевает какой-либо информации о позиции пикселей, даже относительной, и значит “перемешав” пиксели в окрестности получим тот же результат, это не очень хорошо для задач компьютерного зрения. Поэтому предлагается дополнить преобразование информацией об относительной позиции пикселей. Для каждого пикселя $ab \in \mathcal{N_k} \left(i,j\right)$ определим пару относительных смещений от пикселя $ij$: по вертикали $a-i$ и по горизонтали $b-j$ и для каждого смещения будем тренировать представление этой позиции $r_{a-i}, r_{b-j} \in \mathbb{R}^{d_{out} / 2}$. И, наконец, объединим представление смещений в один вектор размерности $d_{out}$ и добавим в формулу:

\[y_{ij} = \sum_{a,b \in \mathcal{N}_k(i,j)} {\rm softmax}_{ab}\left(q^T_{ij}k_{ab} + q^T_{ij}r_{a-i,b-j} \right)v_{ab}\]

Таким образом на выбор веса с которым добавляем значение, соответствующее пикселю $ab$ влияет близость “ключа” $k_{ab}$ в этой точке к “запросу” $q_{ij}$ и позиция пикселя $ab$ относительно центра ячейки.

Плюс нового self-attention подхода в том, что он сокращает число параметров которые надо тренировать, а соответственно, понижает требования к вычислительным ресурсам. Авторы предлагают заменить в ResNet сетке bottleneck блоки на новые attention блоки. И действительно, сократив количество параметров процентов на 25-30 они получают такое же или чуть лучшее качество классификации на ImageNet датасете в сравнении с такой же глубины сеткой ResNet.

Так же авторы предлагают заменить первые слои (с большим ядром свёртки), которые работают непосредственно с изображением на вариант self-attention (не совсем такой как описан выше), но качество на ImageNet всё таки лучше, если сначала применить свёртку.