Чем регулируется размер вектора выхода сети transformer

ML: Градиентный метод

В этом документе продолжается обсуждение многомерной линейной модели. На её примере мы рассмотрим метод стохастического градиентного спуска и выбор оптимальных гиперпараметров. Приведеный код на Python с использованием библиотек Numpy, PyTorch и Keras, можно найти в файле: ML_Line_Model.ipynb.

Градиентный метод

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

Обозначим набор параметров вектором $\mathbf$. Ошибка является их функцией $L=L(\mathbf)$. В пространстве параметров существуют поверхности постоянного значения $L(\mathbf)=\mathrm$. Смещение $d\mathbf$ вдоль таких поверхностей не меняет $L$ и, следовательно, её диффернциал равен нулю:

Такими образом, вектор градиента $\mathbf$ перпендикулярен поверхностям $L=\mathrm$ и направлен в сторону увеличения $L$ (как и любая производная). При приближении к минимуму длина градиента, обычно уменьшается, а в точке минимума (экстремум) она равна нулю (выше второй рисунок).

Чтобы найти минимум $L$, необходимо двигаться в обратном к градиенту направлении (вдоль антиградиента), с шагом пропорциональным некоторому числу $\lambda$.
Этот гиперпараметр называется скоростью обучения:

Чем больше скорость обучения $\lambda$, тем быстрее параметры модели приближаются к оптимальным значениям. Однако при больших $\lambda$ существует риск проскочить минимум (несмотря на уменьшение длины градиента в его окрестности). Проиллюстрируем это примерами.

✍ В одномерном случае уравнение параболы с минимумом в точке $w=0$ и её «градиент» $g$ имеют вид: $$ L(w)=\frac\,\Gamma\,w^2,~~~~~~~~~~~~~g=\frac=\Gamma\,w, $$ где $\Gamma$ — константа. Спуск со скоростью $\lambda$ из точки $w^$, на $t$-той итерации приводит в точку: $$w^=(1-\lambda\,\Gamma)^t\,w^.$$ При $\lambda > 1/\Gamma$ метод становится неустойчивым: $w^ \to \pm \infty$. Однако, чем ближе $\lambda$ снизу к критическому значению $\lambda_c = 1/\Gamma$, тем быстрее достигается минимум.

✍ Аналогично велёт себя градиентный спуск для «многомерной параболы» с различными коэффициентами $\Gamma_i$ по каждой координате: $$ L = \frac\sum_i \Gamma_i\,w_i ^2,$$$g_i = \Gamma_i\, w_i,$$w^_i=(1-\lambda\Gamma_i)^t\,w^_i. $$ Если мы находимся в точке $\mathbf$, то кратчайший путь к минимуму $\mathbf_>=\mathbf$ будет в направлении вектора $-\mathbf$. Однако градиентный метод будет изменять параметры в отличающемся направлении $-\mathbf$ и поведёт к минимуму «окружным путём». Справа приведены траектории движения к минимуму $\$ в двумерном случае из различных начальных положений c $\Gamma=\$ и $\lambda=0.1$. Обратим внимание не только на кривизну траекторий, но и на заметное уменьшение шага по мере приближения к минимуму. Критическое значение $\lambda_c$ определяется наибольшим коэффициентом параболы: $\lambda_c = 1/\max \Gamma_i = 0.2$.

Stochastic Gradient Descent

SGD (Stochastic Gradient Descent) — несколько более продвинутый градиентный метод оптимизации. В нём ошибка $L$ вычисляется не по всем данным, а по выборке (пачка, batch) размера batch_size. Это ускоряет сходимость, т.к. для вычисления очередного шага требуются не все данные, а лишь небольшое их подмножество. После каждой подправки параметров, выборка меняется (поэтому метод называется «stochastic»). При прохождении через все N примеров (одна эпоха), происходит int(N/btach_size) итераций (шагов в пространстве признаков).

Так как данных для вычисления $\mathbf$ немного (обычно batch_size = 10-300), градиент от итерации к итерации может «вилять» в разные стороны. Чтобы уменьшить этот эффект, вектор $\mathbf$ усредняется $\langle\mathbf\rangle$ при помощи скользящего среднего. Это же усреднение помогает перемещаться по покрытой «мелкой рябью» поверхности. Степень сглаживания (усреднения) регулируется параметром $\beta=[0. 1)$. Итерации $t=0,1,2. $ метода SGD при $\langle\mathbf\rangle_=0$ имеют вид:

Величины batch_size и $\beta$ скоррелированы. Чем они больше, тем сильнее сглажен градиент. Увеличение batch_size замедляет обучение (за эпоху делается меньше итераций). Однако на графических карточках, благодаря быстрому перемножению больших матриц, увеличение batch_size может уменьшать время прохода по всем данным, даже в расчёте на одну итерацию.

☝ Скользящее среднее градиента имеет следующий смысл. С весами $\beta$ и $1-\beta$ (сумма которых равна 1) складываются среднее значение $\langle\mathbf\rangle^$ (полученное на предыдущей итерации) и текущий градиент $\mathbf^$. Чем ближе $\beta$ к 1, тем сильнее сглаживание $\langle\mathbf\rangle^\approx \langle\mathbf\rangle^$. При $\beta=0$ усреднения не происходит и $\langle\mathbf\rangle^ = \mathbf^$.

Ниже приведен пример сглаживания данных при помощи скользящего среднего. Синяя линия — это одномерное стохастическое случайное блуждание (подобное цене на акцию). Гладкие жёлтая и зеленая линии — это усреднение с различным параметром $\beta$.

Понятно, что чем больше усреднение (ближе $\beta$ к единице), тем сильнее отстаёт среднее от исходных данных.

Если усредняемая величина равна константе $\mathbf^=\mathbf=\text$, то её среднее, через достаточно большое число итераций, также выйдет на эту константу (сумма геометрческой прогрессии): $\langle\mathbf\rangle^ = \beta\,\mathbf + (1-\beta)\,\mathbf$, $$ \langle\mathbf\rangle^ =\beta\,\langle\mathbf\rangle^+ (1-\beta)\,\mathbf=(1+\beta)(1-\beta)\mathbf. ~~~~~~~~ \langle\mathbf\rangle^ = (1+\beta+. +\beta^)(1-\beta)\,\mathbf = (1-\beta^t)\,\mathbf. $$

☝ В библиотеках PyTorch и Keras параметр $\beta$ называется momentum, а скорость обучения lr — это $\lambda\,(1-\beta)$ и $\langle\mathbf\rangle^ \mapsto (1-\beta)\langle\mathbf\rangle^$. Кроме этого, метод SGD «includes support for momentum, learning rate decay, and Nesterov momentum«. Распад означает уменьшение lr с каждой эпохой iterations в соответствии с параметром распада decay (код из Keras):

lr_t = lr * (1. / (1. + decay * iterations ))

Далее происходят вычисления:

v = momentum * v - lr_t * grad if nesterov: params = params + momentum * v - lr_t * grad else: params = params + v

По умолчанию приняты следующие параметры: SGD(lr=0.01, momentum=0.0, decay=0.0, nesterov=False).

Вычисление градиента

Найдём градиент от среднеквадратичной ошибки (MSE) в линейной модели. Для этого воспользуемся правилом дифференцирования сложной функции. В данном случае $L$ зависит от $y_$, которые уже непосредственно зависят от параметров: $y_=\sum_\mu\,x_w_<\mu\gamma>+b_<\gamma>$. Поэтому для среднеквадратичной ошибки $L= \langle (\mathbf-\hat<\mathbf>)^2\rangle$, усреднённой по $N$ примерам (первый индекс) и $m$ выходам модели (второй индекс) имеем:

Аналогично для градиента по $\mathbf$:

В матричном виде это можно записать так:

Произведение $N\,m$ равно числу элементов массива $\mathbf$, которые в numpy находятся в атрибуте size.

Численный поиск минимума

Приведём алгоритм поиска оптимальных параметров методом SGD при помощи библиотеки numpy. Ошибку будем вычислять по batch_size примерам, проходя по всем обучающим примерам epochs раз:

def Loss(X,Y, W, B): # функция ошибки return np.mean((X @ W + B - Y) ** 2) def My_SGD(X, Y, lr, mo, batch_size, epochs): W = np.random.random ((nX, nY)) # случайные начальные значения B = np.random.random ((nY, ) ) # параметров модели W,B agW, agB, iters = 0, 0, int( len(X)/batch_size ) # средние градиенты, число батчей for epoch in range(epochs): # эпоха - проход по всем примерам idx = np.random.permutation( len(X) ) # перемешанный список индексов X, Y = X[idx], Y[idx] # перемешиваем данные for i in range(0, iters*batch_size, batch_size): # примеры разбиты на пачки xb = X[i: i+batch_size] # входы пачки yb = Y[i: i+batch_size] # выходы пачки (целевые значения) y = xb @ W + B # выходы модели gW = (2 / y.size) * xb.T @ (y - yb) # градиент по w gB = (2 / y.size) * np.sum(y - yb, axis=0) # градиент по b agW = mo*agW + (1-mo)*gW; W -= lr * agW # сглаживаем градиенты и agB = mo*agB + (1-mo)*gB; B -= lr * agB # подправляем параметры return W, B, Loss(X,Y, W, B) # параметры и ошибка

Обратим внимание, что обучающие примеры перед началом каждой эпохи случайно перемешиваются. Для этого, при помощи функции np.random.permutation создаётся перемешенный список целых чисел о 0 до len(X)-1. Следующая строчка собственно производит перемешивание. Перемешивать данные перед каждой эпохой в принципе не обязательно, но если данных немного это может улучшить сходимость к минимуму. Кроме этого, если число примеров N нацело не делятся на batch_size, не все данные попадут в вычисление градиента. Перемешивание ослабляет эту проблему.

Справа приведены графики ошибки при различных скоростях обучения и mo=0. Критическое значение lr находится в районе 1.75. При приближении к этому значению скорость обучения заметно возрастает.

Оптимальные гиперпараметры

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

Нарисуем карту высот ошибки Loss линейной модели, как функцию гиперпараметров $\lambda$ (=lr) и $\beta$ (=mo).
Синий цвет — минимум ошибки, коричневый — максимум. Количество эпох epochs = 10 и размер пачки batch_size = 10 фиксированы (точек N=100). Голубая линия с точками означает значение $\beta$ соответствующее минимальной ошибке при данном $\lambda$. Под картой высот приведен график ошибки как функции $\lambda$ (при оптимальном $\beta$ для данного $\lambda$). Коричневый цвет на карте высот соответствует области неустойчивости градиентного метода:

Как видно, для такой простой модели, при $\beta \sim 0.8$, значение $\lambda$ может меняться в широких пределах. Оптимальные гиперпараметры: mo=0.63, lr=5.31, loss=0.00013. Отметим, что подходящая параметризация гиперпараметров важна.
Например, если выбрать как в Keras lr=$\lambda\,(1-\beta)$, то получится менее устойчивая картина:

Оптимальные гиперпараметры в этом случае: mo=0.58, lr=2.25, loss=0.00013.

Градиент на графе в PyTorch

Библиотека PyTorch от Facebook является популярным инструментом при разработки сложных архитектур нейронных сетей. PyTorch создаёт динамический вычислительный граф (строится в процессе вычислений и не требует компиляции).

Воспроизведём на PyTorch простой градиентный метод. Параметры модели W,B являются терминальными узлами вычислительного графа, для которых будет вычисляться градиент (requires_grad=True):

Метод randn возвращает тензор нормально распределённых случайных чисел. По умолчанию они имеют тип float32. Чтобы он совпадал с типом обучающих данных X,Y, происходит явное задание типа аргументом dtype.

Затем идёт основной цикл вычислений:

import torch def grad_tourch(X, Y, lr=1, batch_size=10, epochs=10): W = torch.randn(nX, nY, dtype=torch.float64, requires_grad=True) B = torch.randn(nY, dtype=torch.float64, requires_grad=True) for epoch in range(epochs): # эпоха - проход по всем примерам idx = torch.randperm( len(X) ) # перемешанный список индексов X, Y = X[idx], Y[idx] # мешаем данные sumL, iters = 0, int( len(X)/batch_size) # суммарная ошибка и число батчей for i in range(0, iters*batch_size, batch_size): # примеры разбиты на пачки xb = torch.from_numpy(X[i: i+batch_size]) yb = torch.from_numpy(Y[i: i+batch_size]) y = xb.mm(W).add(B) # модель y = bx @ W + B loss = ((y-yb)**2).mean() # mse ошибка по батчу sumL += loss.data.item() loss.backward() # вычисление градиентов with torch.no_grad(): W.add_(- lr * W.grad) # без перестраивания графа B.add_(- lr * B.grad) W.grad.zero_() # обнуляем градиенты B.grad.zero_() print(f"() loss:") return W, B grad_tourch(X, Y)

Вначале формируются torch-тензоры пачек из numpy-тензоров (новая память при этом не выделяется). Затем строится вычислительный граф, корнем которого будет ошибка модели loss. Попутно происходят вычисления на графе. Затем для корня запускается метод backward, который возвращается по веткам и вычисляет градиенты.

Любое выражение с тензорами, имеющими свойство requires_grad=True приводит к динамической перестройке вычислительного графа, ведущего к его корню (ошибке модели loss). Чтобы этого не произошло, после backward устанавливается окружение no_grad(), которое блокирует изменение графа При этом, меняя параметры, мы не создаём новых тензоров (функция add_ это «инкрементация»). После изменения W,B необходимо сбросить в ноль градиенты (для следующей итерации).

Нейронная сеть на PyTorch

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

import torch.nn as nn model = nn.Sequential( nn.Linear(nX, nY) )

Создадим SGD оптимизатор (который будет подправлять параметры) и mse функцию ошибки:

optimizer = torch.optim.SGD(model.parameters(), lr=1, momentum=0.9) criterion = nn.MSELoss()

Цикл обучения выглядит следующим образом. На каждой итерации на сеть передают пачку примеров y=model(xb), в результате чего происходит прямое распространение (forward). Полученный выход сети y, вместе с «истинными» значениями yb передают функции ошибки. Для неё вызывают обратное распространение backward, в результате которого будут вычислены градиенты. Метод оптимизатора step, при помощи этих градиентов, подправляет параметры модели. Затем оптимизатор обнуляет градиенты (zero_grad):

iters = int( len(X)/batch_size ) for epoch in range(epochs): # эпоха - проход по всем примерам for i in range(0, iters*batch_size, batch_size): # примеры разбиты на пачки bx = torch.Tensor( X[i: i+batch_size] ) # по numpy массивам создаём by = torch.Tensor( Y[i: i+batch_size] ) # тензоры torch типа float32 y = model(bx) # прямое распространение loss = criterion(y, by) # вычисляем ошибку loss.backward() # вычисляем градиенты optimizer.step() # подправляем параметры optimizer.zero_grad() # обнуляем градиенты print('epoch: %d Loss: %.6f' % (epoch, loss) )

После обучения можно вывести параметры модели:

print(model[0].weight, model[0].bias)

Матрица весов weight, как и в sklearn, хранится в транспонированном виде.

Нейронная сеть на Keras

Библиотека Keras является простым в использовании инструментом для проектирования нейронных сетей. В настоящее время она является составной частью фреймворка tensorflow от Google.
Введение в Keras можно найти в документе NN_Base_Keras.

Линейная модель (один слой без активационной функции) в Keras реализован в полносвязном слое Dense. Создания нейронной сети с одним таким слоем делается следующим образом:

from tensorflow import keras # keras из tensorflow from keras.models import Sequential # способ формирования слоёв (стопка) from keras.layers import Dense # полносвязный слой from keras.optimizers import SGD # метод оптимизации model = Sequential() # линейная стопка слоёв model.add(Dense(units=nY, input_dim=nX))

После создания модели, она компилируется (compile) и запускается её обучением (fit). При обучении будем использовать SGD-оптимизатор и mse-ошибку:

model.compile(optimizer = SGD(lr=2, momentum=0.5), loss = 'mse') res = model.fit(X, Y, batch_size=batch_size, epochs=10, verbose=0 ) )

После обучения можно вывести график ошибок, как функцию числа эпох:

plt.plot(res.history['loss'], marker=".") plt.legend(["loss: %.5f" % ( res.history['loss'][-1] )]) plt.show()

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

for lr in model.layers: # по всем слоям w = lr.get_weights()[0] # веса синапсов нейронов b = lr.get_weights()[1] # смещения нейронов .

ViT — на кухне фаворит

Прошедший 2021-й год ознаменовался настоящей революцией в области компьютерного зрения.

Трансформеры, подобно новым штамма Ковида, вытеснившие конкурентов в области обработки естественного языка (NLP) и задачах, связанных с обработкой звука, добрались и до компьютерного зрения.

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

В данном обзоре я хотел бы рассказать о нескольких самых ярких прорывах и идеях в совершенствовании архитектур и обучении ViT-ов (Visual Transformers).

Введение

До сравнительно недавнего (если смотреть не по меркам DL) времени сверточные сети (CNN) безраздельно доминировали в области компьютерного зрения (Computer Vision). Сверки обладают рядом замечательных свойств — локальностью , позволяющей учитывать отношения близости между соседними пикселями, применением одних и тех же весов к каждому пикселю карты активации (feature map), построением иерархических представлений — от простых примитивов вроде границ и контуров до более сложных и составных понятий вроде кошек и собак (во всяком случае так утверждается многими).

Казалось бы, что можно вообще было бы придумать более подходящее и оптимальное с точки зрения использования параметров и вычислений среди возможных архитектур нейронной сети? Тем более, что за последние несколько лет было придумано множество наворотов и ухищрений для повышения качества сверточной нейронной сети, либо скорости работы.

В качестве самых значимых достижений можно вспомнить добавление разных видов skip-connections, depthwise сверток, inverted bottlenecks. Современные архитектуры вроде EfficientNet, NFNet прошли большой путь эволюции по сравнению с vanilla ResNetа-ми.

Но все же, сверточные сети несовершенны. Локальность операции свертки, преподнесенная выше как достоинство, является и недостатком. Пиксель в выходной карте активаций может зависеть лишь от области входной карты в пределах ядра свертки. Поэтому для сбора глобальной информации требуется большое количество слоев (при пулингах и свертках стандартного размера типа 2,3,5).

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

Трансформеры произвели настоящий фурор в области задач (NLP) обработки естественного языка, камня на камне не оставив от популярных ранее многослойных реккурентных сетей на LSTM и GRU, и вообще в задачах связанных с последовательностями.

Но как применить self-attention в задачах компьютерного зрения стало очевидно далеко не сразу. Первое, что могло бы прийти в голову — рассматривать каждый пиксель картинки, как слово, и считать attention между всеми пикселями внутри картинки. Проблема здесь в том, что вычислительная сложность и обьем используемой памяти в стандартном self-attention растет квадратично с длиной последовательности. Картинки на датасете больше игрушечных MNIST и CIFAR-10 имеют разрешение порядка сотен пикселей вдоль каждой размерности (скажем 224×224) и считать в лоб self-attention выходит слишкои накладно.

Были работы, которые считали его локально, но такой подход в каком-то смысле сродни сверткам. В DETR было предложено использовать feature map с нижнего слоя ResNet, где количество пикселей уже невелико, для self-attention и полученная конструкция сработала довольно неплохо в задаче детекции. Но в этих решениях основной рабочей лошадкой не был механизм внимания.

An image is worth 16×16 words

Настоящий триумф трансформеров в компьютерном зрении пришел с работой An image is worth 16×16 words.

Решение, позволившее добиться адекватной вычислительной стоимости и памяти для хранения, оказалось гениальным в своей простоте — использовать в качестве слов не отдельные пиксели, а кусочки картинки некоторого размера , тем самым уменьшив вычислительную сложность с до . Для стандартного разрешения на ImageNet — 224 и патча размера 16 выходит вполне себе подьемно (196 токенов).

Использованная архитектура является по существу цепочкой энкодеров а-ля BERT.

Для задачи классификации в дополнение к токенам, соответствующим отдельным патчам, добавляется дополнительный [CLS] токен для классификации.

На момент публикации самая большая версия полученной архитектуры — ViT-H/14 (H — Huge) установила новый SOTA (state-of-the-art) на ImageNet-1k . Здесь, правда, нужно отметить важный нюанс — для достижения такого высокого качества необходимо обучение на огромном количестве данных. В распоряжении исследователей Google был датасет JFT-300M . Без предобучения на большом количестве данных, даже с сильной регуляризацией ( weight_decay = 0.1) модель подвержена переобучению и работает заметно хуже ResNet-ов.

DeiT (Data-Efficient Image Transformer)

Тот же ViT, но лучше.

Необходимость предобучения на громадном количестве картинок могла бы ограничить применимость трансформеров в компьтерном зрении, но вскоре после вышеупомянутой работы вышла статья Training data-efficient image transformers & distillation through attention.

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

В статье авторы использовали мощный набор аугментаций и регуляризационных процедур:

• Label smoothing. Правильной метке дается вероятность , а остальная вероятность распределяется равномерно между остальными классами.
• Rand Augment. Выбирается некоторое множество преобразований, из которых случайным образом для каждого примера применяется какое-то количество из них с некоторой вероятностью и параметрами.
• Stochastic Depth. Так как в трансформерах есть skip-connections с некоторой вероятностью можно проигнорировать выход блока энкодера и подать просто выход прошлого слоя вперед.
• Mixup и CutMix. Mixup смешивает две картинки и соответствующие им целевые метки в классификации. CutMix вставляет уменьшенную версию одной картинки поверх другой и целевая метка классификации берется как смесь меток для каждого класса, причем доля класса пропорциональна занимаемой площади.
• Repeated Augmentation. Прогонять через аугментации можно не только лишь один, но и большее количество раз.
• Erasing. Из картинки вырезается некоторая область случайным образом.

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

Другим решением, дополнительно повысившим качество модели была дистилляция (knowledge distillation). Вкратце напомню, что идея дистилляции в том, чтобы кроме ground_truth меток подавать еще предсказания модели (учителя), хорошо обученной на рассматриваемом наборе данных.

Если в функцию потерь подаются вероятности (или логиты) то мы имеем дело с soft-distillation:

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

Если же подается предсказанный учителем класс (он может быть и ошибочным), то это hard-distillation.

Что занятно (и мне непонятно), второй способ сработал лучше.

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

С точки зрения архитектуры — DeiT ничем не отличается от ViT.

PVT (Pyramid Vision Transformer)

Интересное решение, позволившее использовать более мелкие патчи было предложено в статье Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions.

FPN (Feature Pyramid Network) и различные ее вариации довольно неплохо зарекомендовала себя в задачах сегментации и детекции. Признаки с верхних слоев фокусируются на извлечении мелких деталей и примитивов, в то время как более глубокие слои имеют представление о глобальной семантике. Использование признаков с разных слоев позволяет одновременно учитывать мелкие и крупные детали. В vanilla ViT все feature maps имеют один и тот же размер, поэтому нет разделения на мелкие и крупные признаки. Кроме того, крупные патчи не обеспечивают достаточного разрешения для разрешения мелких деталей.

В PVT было предложено использовать патчи размера 4×4 на первой стадии и затем последовательно уменьшать разрешение. На каждой стадии разрешение уменьшается вдвое с помощью strided свертки с увеличением размерности вектора embedding.

Тем не менее, на первых слоях при размере патча 4×4 все еще остается слишком много операций. Для того, чтобы уменьшить расход памяти на верхних слоях авторы предложили уменьшать длину последовательностей key и value.

Сложность вычисления произведения пропорциональна произведению длин последовательностей key — и query . Полученная матрица имеет размер . Если последовательность value имеет ту же длину, что и ключи, то возможно умножить матрицу внимания на и выход будет иметь ту же длину, что и query.

Уменьшение длины последовательностей key и query достигается следующим образом. Пусть и — количество патчей вдоль каждой из осей (высоты и ширины) а размерность эмбеддинга на — й стадии. Тогда:

• Входная последовательность длины и размерности эмбеддинга решейпится (звучит ужасно, знаю) в последовательность длины c размерностью эмбеддинга .
• Слой nn.Linear(R_i ** 2 * C_i, C_i) уменьшает размерность эмбеддинга до исходной (проектирует на подпространство).

После этого поступаем точно так же, как и в стандартном self-attention. В итоге получается экономия в в вычислительной сложности и памяти.

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

В первых слоях фактор довольно большой — 8, и уменьшается вдвое на каждой следующей стадии. На самой последней стадии . Кроме того, патч размера 2×2 c feature map с прошлой стадии используется в качестве пикселя (элементарной ячейки карты активации) на следующей стадии.

Наличие карт активации разного размера позволяет применить идею Feature Pyramid в PVT.

Полученная модель неплохо себя показывает на ImageNet.

Но по-настоящему польза от PVT становится заметной на детекции и сегментации.

Swin (Hierarchical Vision Transformer using Shifted Windows)

Основной проблемой при использовании ViT, особенно в Dense Prediction tasks — детекции и сегментации, является быстрый рост сложности с уменьшением размера патча. Патч размера 16×16 выходит слишком грубоватым для извлечения тонких деталей.

В статье Swin Transformer: Hierarchical Vision Transformer using Shifted Windows был предложен изящный способ уменьшить вычислительную сложность для feature map с большим количеством патчей. Как и PVT, подход в Swin мотивирован пирамидой признаков из CNN. Карта признаков на верхнем уровне составлена из мелких патчей (более конкретно, размера 4×4) и через некоторое количество слоев пространственная размерность уменьшается вдвое вдоль каждой оси (происходит слияние соседних патчей), а размерность эмбеддинга удваивается.

Но способ «удешевления» attention в верхних слоях другой. В верхних слоях attention считается только в пределах окна некоторого размера, причем количество токенов в окне постоянно во всех слоях сети. То есть, если на нижней стадии размер патча и attention захватывает для каждого токена все остальные токены, то на предыдущей стадии с размером патча attention локализован лишь на четверти входной картинки, а слое еще ниже (где патчи имеют размер на картинки. Благодаря этому становится возможным использование мелких патчей.

Сравним вычислительную сложность windowed self-attention c глобальным self-attention. Пусть ширина и высота feature map на данном слое — H и W, соответственно. Тогда при использовании окон, захватывающих области высотой H/R и шириной W/R потребуется вычислять self-attention для каждого из окон. Но так как вычислительная сложность операции внимания растет квадратично с длиной последовательности, то в силу имеем в конечном итоге выигрыш в раз по сравнению с исходной операцией.

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

• на четных слоях разбиваем на патчи одним способом (так чтобы верхний левый угол вернхего левого патча совпал с верхним левым углом всей картинки)
• на нечетных шагах сдвигаем разбиение на половину размера патча в данном слое

В остальном блоки трансформера в Swin повторяют ViT. Вычисление двух последовательных блоков в Swin имеет следующий вид:

В итоге получился очень сильный бэкбоун для задачи классификации и Dense predictions tasks (детекции, сегментации).

При сопоставимом количестве операций с плавающей точкой модели Swin значительно превосходят ViT и DeiT (но все же уступают наиболее совершенным CNN вроде EfficientNet).

Стандартные фреймворки детекции и сегментации состоят из backbone, который строит признаки и новое представление обьекта, и головы (head) для детекции и сегментации. Для того, чтобы сравнить качество извлекамых с помощью Swin признаков авторы статьи обучили модели с Cascade Mask R-CNN (голова для одновременной детекции и сегментации) на MS COCO.

Модели Swin заметно превзошли бейзлайны на основе ResNet-ов и DeiT с сопоставимыми характеристиками (числом параметов и операций) как в детекции, так и сегментации.

Использование shifted windows, как показывает ablation study, действительно важно для достижения хорошего результата, особенно для детекции и сегментации.

XCiT (Cross-Covariance Image Transformers)

Еще один подход побороть квадратичную зависимость от количества патчей был предложен в статье XCiT: Cross-Covariance Image Transformers от исследователей из Фейсбука (ныне Мета).

Идея состоит в том, чтобы транспонировать операцию attention.

В исходной операции self-attention c головами:

Сложность вычисления — , а расход по памяти .

Для транспонированного внимания (называемого в статье cross-covariance) операция имеет следующий вид:

где — некоторый параметр температуры. Квадратичная сложность переносится с длины последовательности на размерность эмбеддинга. Для cross-attention вычислительная сложность и расход памяти — . Поэтому вычислительная сложность для XCiT будет расти не так быстро, как для ViT, с уменьшением размера патчей или увеличением разрешения.

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

Для того, чтобы иметь явное взаимодействие между соседними патчами, авторы добавили так называемое локальное взаимодейсвтвие патчей (Local Patch Interaction). В качестве LPI используется последовательность двух depthwise сверток с батч-нормализацией и GeLU между ними. Последовательность токенов перед LPI разворачивается в 2d картинку, к этой картинке применяется описанная выше последовательность слоев, и картинка сворачивается обратно в последовательность токенов.

Приятным бонусом от XC-attention является меньшая чувствительность к изменению разрешения подаваемой картинки. Так как свертка при вычислении XC-attention проводится вдоль внутренней оси, размер матрицы внимания не меняется. Качество модели, обученной на разрешении проседает не так сильно при уменьшении разрешения, по сравнению с ResNet и DeiT, и даже заметно возрастает при увеличении разрешения до .

Бэкбоун получился очень даже замечательным. При сопоставимых размерах различные варианты XCiT оказываются эффективнее не только EfficientNet-ов и ранних ViT, но и сильных конкурентов вроде Swin-ов.

В задаче детекции и сегментации XCiT показал себя с хорошей стороны, превзойдя бэкбоуны на основе PVT и ViL (не затронутого в данном обзоре). XCiT-S12/8 превзошел даже Swin-T с похожими характеристиками, но более крупный свин таки подложил свинью в сравнении с XCiT-S24/8.

PS-ViT (Pooling and Attention Sharing)

В сверточных сетях обыкновенно карты признаков на верхних слоях обладают большим разрешением, и постепенно посредством pooling или strided-сверток разрешение уменьшается с увеличением числа каналов. Таким образом производится переход от локальных признаков к глобальным представлениям.

Разумно предположить, что аналогичный подход может хорошо сработать и для visual трансформеров.

И в работе Better Vision Transformer via Token Pooling and Attention Sharing была предложена архитектура такая архитектура, давшая существенный прирост качества на ImageNet по сравнению с DeiT при том же числе операций (6.6% для PSViT-2D-Tiny по сравнению с DeiT-Tiny ).

В качестве основных результатов данной статьи следует отметить:

• Механизм уменьшения количества токенов с увеличением глубины сети
• Переиспользование одного и того же attention в нескольких последовательных блоках трансформера

Pooling в PS-ViT

В статье авторы рассматривают разные стратегии модификации архитектуры (взяв за основу DeiT-Tiny) и сохраняя примерно то же количество FLOPs.

1. Увеличение глубины сети (количества блоков) при сохранении размерности эмбеддинга неизменной
2. Увеличение размерности эмбеддинга при том же количестве блоков

И то, и то сработало достаточно неплохо, но увеличение ширины несколько лучше. Кроме того, авторы рассматривают два варианта пулинга.

В первом случае, где классификация осуществляется через [CLS] токен, свертка 1×1 меняет размерность эмбеддинга, а затем проводится MaxPooling . Эта стратегия называется PSViT-1D .

В другом случае для классификации используется результат усреднения последней карты активации и для пулинга strided свертка с шагом 2. Этот подход, называемый PSViT-2D , работает даже немного лучше.

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

И последним по порядку, но по значению является оптимальный выбор расположения элементов и количества слоев в трансформере.

Полный перебор возможных вариантов расположения слоев с пулингом и размерностей эмбеддингов — слишком сложная комбинаторная задача, поэтому пространство поиска пришлось существенно ограничить. Размерности эмбеддинга и максимальное число блоков зафиксировано на каждой стадии (при фиксированном количестве токенов).

В каждом блоке есть 3 выбора:

1. Использовать обычный блок трансформера
2. Два последовательных блока с одним и тем же attention
3. Тождественную операцию (Identity)

На каждом проходе (forward pass) один из трех вариантов выбирается из равномерного распределения и при обратном проходе (backward pass) обновляются параметры для этого варианта (если это не Identity, конечно). Оптимальная архитектура определяется с помощью эволюционного алгоритма.

Работает это, по всей видимости, и правда неплохо:

VOLO (Vision Outlooker for Visual Recognition)

Довольно занятную вариацию внимания предложили в статье VOLO (если честно, я даже не понимаю, почему она работает так здорово).

Блок энкодера имеет стандартный вид:

Здесь — это LayerNorm, а вот что действительно интересно, так это операция Делается она следующим образом ( C — число каналов, K — размер ядра свертки):

1. Линейный слой nn.Linear(C, K ** 4) для каждого пикселя из feature map создает вектор размерности K ** 4 .
2. Полученный вектор решейпится (прошу прощения за англицизм) в матрицу K ** 2 x K ** 2 . Данная матрица играет роль матрицы внимания в пределах окна размера K x K . То есть матрица внимания предсказывается в один шаг, без создания ключей (keys) и запросов (queries) c последующим вычислением попарных скалярных произведений.
3. Линейный слой nn.Linear(C, С) выдает значения (values) для каждого токена (как в обычном трансформере).
4. Полученная на шаге 2 матрица attention перемножается на values и получается выходное представление.

Таким образом, получается некий trade-off между локальностью операции и вычислительной сложностью. В стандартном self-attention вычислительная сложность растет как поэтому использовать патчи размером меньше 16, особенно при большом разрешении довольно проблематично. В предложенном подходе же асимптотика линейна по количеству токенов . Размер ядра свертки K должен быть небольшим (в работе K = 3). Благодаря этому можно брать меньший патч (скажем 8) при большом разрешении (384×384, 512×512).

OutlookAttn — гибрид свертки и стандартного self-attention — локальный, но с большим receptive field. При таком подходе большой receptive field может быть достигнут при меньшем числе блоков, чем в типичной CNN и в то же время зашито понятие локальности и близости в саму архитектуру.

Получился классный бэкбоун, позволивший добиться впечатляющих результатов не только на ImageNet (87.1% без дополнительных данных), но и в задачах семантической сегментации на Cityscapes и ADE20K.

Заключение

Универсальность и гибкость архитектуры трансформера, способность улавливать глобальный контекст, оказалась полезной и в области компьютерного зрения.

За год с небольшим, прошедших с публикации An image is worth 16×16 words, трансформеры сильно изменили наши представления о том, как надо решать задачи компьютерного зрения, толкнули науку далеко вперед.

В данном обзоре я рассмотрел лишь отдельные работы из моря публикаций по этой теме за 2021 год. Многие другие интересные идеи, вроде Transformer in Transformer и CoAtNet не были затронуты в силу ограниченности обьема обзора. Кроме того, были рассмотрены только задачи классификации, детекции и сегментации картинок. ViT-ы показали впечатляющие задачи так же в мультимодальных задачах, при работе с видео и self-supervised, semi-supervised learning, генеративных моделях.

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

Но я уверен, что за развитием этой области будет очень интересно следить в 2022.

Список источников

Статьи

• An image is worth 16×16 words
• Training data-efficient image transformers & distillation through attention
• Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• Swin-Transformer
• PSViT: Better Vision Transformer via Token Pooling and Attention Sharing
• XCiT: Cross-Covariance Image Transformers
• VOLO
• Stand-Alone Self-Attention in Vision Models
• DETR