Прошедший 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 имеют разрешение порядка сотен пикселей вдоль каждой размерности (скажем 224x224) и считать в лоб self-attention выходит слишкои накладно.

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

An image is worth 16x16 words

Настоящий триумф трансформеров в компьютерном зрении пришел с работой An image is worth 16x16 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 было предложено использовать патчи размера 4x4 на первой стадии и затем последовательно уменьшать разрешение. На каждой стадии разрешение уменьшается вдвое с помощью strided свертки с увеличением размерности вектора embedding.

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

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

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

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

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

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

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

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

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

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

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

Swin (Hierarchical Vision Transformer using Shifted Windows)

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

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

Но способ "удешевления" 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] токен, свертка 1x1 меняет размерность эмбеддинга, а затем проводится 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) при большом разрешении (384x384, 512x512).

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

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

Заключение

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

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

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

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

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

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

Статьи

• An image is worth 16x16 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

Большой Перечень архитектур ViT

https://github.com/dk-liang/Awesome-Visual-Transformer

Ну и куда без Янника Килхера

https://www.youtube.com/c/YannicKilcher