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 токенов).
(Слева) Принцип работы ViT. (Справа) Блок трансформера в ViT.
Использованная архитектура является по существу цепочкой энкодеров а-ля BERT.
Для задачи классификации в дополнение к токенам, соответствующим отдельным патчам, добавляется дополнительный [CLS]
токен для классификации.
SOTA на ImageNet на момент публикации статьи
На момент публикации самая большая версия полученной архитектуры — ViT-H/14
(H — Huge) установила новый SOTA (state-of-the-art) на ImageNet-1k
. Здесь, правда, нужно отметить важный нюанс — для достижения такого высокого качества необходимо обучение на огромном количестве данных. В распоряжении исследователей Google был датасет JFT-300M
. Без предобучения на большом количестве данных, даже с сильной регуляризацией (weight_decay
= 0.1) модель подвержена переобучению и работает заметно хуже ResNet-ов.
Качество на ImageNet-1k в зависимости от датасета, на котором проходило предобучение. BiT — модификация ResNet.
DeiT (Data-Efficient Image Transformer)
Тот же ViT, но лучше.
Кривые Парето для разных моделей на ImageNet. По оси пропускная способность V100 по количеству картинок в секунду.
Необходимость предобучения на громадном количестве картинок могла бы ограничить применимость трансформеров в компьтерном зрении, но вскоре после вышеупомянутой работы вышла статья Training data-efficient image transformers & distillation through attention.
Так как основной проблемой трансформеров в исходной постановке является подверженность переобучению, то естественно было бы предложить более совершенную процедуру регуляризации, и аугментация является признанным и эффективным средством для эффективного увеличения размера данных и борьбы с переобучением. Вопрос в том — достаточно ли хороша она?
В статье авторы использовали мощный набор аугментаций и регуляризационных процедур:
Label smoothing. Правильной метке дается вероятность , а остальная вероятность распределяется равномерно между остальными классами.
Rand Augment. Выбирается некоторое множество преобразований, из которых случайным образом для каждого примера применяется какое-то количество из них с некоторой вероятностью и параметрами.
Stochastic Depth. Так как в трансформерах есть skip-connections с некоторой вероятностью можно проигнорировать выход блока энкодера и подать просто выход прошлого слоя вперед.
Mixup и CutMix. Mixup смешивает две картинки и соответствующие им целевые метки в классификации. CutMix вставляет уменьшенную версию одной картинки поверх другой и целевая метка классификации берется как смесь меток для каждого класса, причем доля класса пропорциональна занимаемой площади.
Repeated Augmentation. Прогонять через аугментации можно не только лишь один, но и большее количество раз.
Erasing. Из картинки вырезается некоторая область случайным образом.
Авторы провели основательный анализ важности тех или иных аугментаций для достижения хорошего качества классификации.
Ablation study для разных способов аугментации картинок из ImageNet.
Другим решением, дополнительно повысившим качество модели была дистилляция (knowledge distillation). Вкратце напомню, что идея дистилляции в том, чтобы кроме ground_truth меток подавать еще предсказания модели (учителя), хорошо обученной на рассматриваемом наборе данных.
Если в функцию потерь подаются вероятности (или логиты) то мы имеем дело с soft-distillation:
Здесь определяет вес лосса учителя ( — дивергенции Кульбака-Лейблера) по сравнению с кроссэнтропией между предсказанием и истинной меткой, а температура — регулирует уверенность моделей в предсказании.
Если же подается предсказанный учителем класс (он может быть и ошибочным), то это hard-distillation.
Что занятно (и мне непонятно), второй способ сработал лучше.
Сравнение различных моделей дистиляции для DeiT
В качестве учителя лучше всего себя показали RegNet-ы (сверточные сети), лучше, чем более крупная модель трансформера. По всей видимости, так как сверточные сети и трансформеры имеют различный способ построения признаков, то знание, переданное от CNN более ново и полезно, чем просто от более мощной модели той же структуры.
С точки зрения архитектуры — DeiT ничем не отличается от ViT.
PVT (Pyramid Vision Transformer)
Интересное решение, позволившее использовать более мелкие патчи было предложено в статье Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions.
(Слева) Извлечение пирамиды признаков в типичных CNN (Центр) ViT не обладает иерархией признаков (Справа) Извлечение пирамиды признаков в PS-ViT
FPN (Feature Pyramid Network) и различные ее вариации довольно неплохо зарекомендовала себя в задачах сегментации и детекции. Признаки с верхних слоев фокусируются на извлечении мелких деталей и примитивов, в то время как более глубокие слои имеют представление о глобальной семантике. Использование признаков с разных слоев позволяет одновременно учитывать мелкие и крупные детали. В vanilla ViT все feature maps имеют один и тот же размер, поэтому нет разделения на мелкие и крупные признаки. Кроме того, крупные патчи не обеспечивают достаточного разрешения для разрешения мелких деталей.
Архитектура PVT. На каждой стадии разрешение feature map уменьшается вдвое. Красный бледный прямугольник (Spatial Reduction) — одна из основных наработок статьи
В 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)
уменьшает размерность эмбеддинга до исходной (проектирует на подпространство).
Схема работы Attention с Spatial Reduction
После этого поступаем точно так же, как и в стандартном self-attention. В итоге получается экономия в в вычислительной сложности и памяти.
Данная модификация, несомненно ограничивает выразительности сети, но выбор архитектуры — почти всегда баланс между качеством и скоростью (размером).
В первых слоях фактор довольно большой — 8, и уменьшается вдвое на каждой следующей стадии. На самой последней стадии . Кроме того, патч размера 2×2 c feature map с прошлой стадии используется в качестве пикселя (элементарной ячейки карты активации) на следующей стадии.
Разные версии моделей PVT.
Наличие карт активации разного размера позволяет применить идею Feature Pyramid в PVT.
Полученная модель неплохо себя показывает на ImageNet.
PVT против других (порошков) моделей
Но по-настоящему польза от PVT становится заметной на детекции и сегментации.
Сравнение детекции + сегментации на MS COCO (val 2017) с использованием Mask R-CNN головы.Семантическая сегментация на ADE20K. Сравнение проводится по метрике mIoU (mean Intersection over Union)
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 для каждого из окон. Но так как вычислительная сложность операции внимания растет квадратично с длиной последовательности, то в силу имеем в конечном итоге выигрыш в раз по сравнению с исходной операцией.
Последоватеньное слияние патчей и увеличение размера окон для attention
Но при таком подходе токены из соседних окон не взаимодействуют друг с другом, что ограничивает выразительную способность сети. Взаимодействие с соседями реализовано в Swin следующим образом:
на четных слоях разбиваем на патчи одним способом (так чтобы верхний левый угол вернхего левого патча совпал с верхним левым углом всей картинки)
на нечетных шагах сдвигаем разбиение на половину размера патча в данном слое
Окна в Swin в двух последовательных блоках трансформера
В остальном блоки трансформера в Swin повторяют ViT. Вычисление двух последовательных блоков в Swin имеет следующий вид:
В итоге получился очень сильный бэкбоун для задачи классификации и Dense predictions tasks (детекции, сегментации).
Сравнение моделей на ImageNet-1k (без использования дополнительных данных)
При сопоставимом количестве операций с плавающей точкой модели Swin значительно превосходят ViT и DeiT и сравнимы с EfficientNet (эффективнее по числу операций, чуть уступают в плане эффективности по параметрам).
Стандартные фреймворки детекции и сегментации состоят из backbone, который строит признаки и новое представление обьекта, и головы (head) для детекции и сегментации. Для того, чтобы сравнить качество извлекамых с помощью Swin признаков авторы статьи обучили модели с Cascade Mask R-CNN (голова для одновременной детекции и сегментации) на MS COCO.
Модели Swin заметно превзошли бейзлайны на основе ResNet-ов и DeiT с сопоставимыми характеристиками (числом параметов и операций) как в детекции, так и сегментации.
Сравнение Swin-T и ResNet50 в качестве бэкбоуна для детекции с помощью Cascade MASK R-CNN. Сравнение метрик AP проводится на датасете MS COCO. Сравнение метрики mIoU (mean Intersection over Union) для сегментации на ADE20K.
Использование shifted windows, как показывает ablation study, действительно важно для достижения хорошего результата, особенно для детекции и сегментации.
Ablation study для испозования shifted windows. Сверху — без сдвига окон, снизу — со сдвигом окон.
XCiT (Cross-Covariance Image Transformers)
Еще один подход побороть квадратичную зависимость от количества патчей был предложен в статье XCiT: Cross-Covariance Image Transformers от исследователей из Фейсбука (ныне Мета).
Идея состоит в том, чтобы транспонировать операцию attention.
В исходной операции self-attention c головами:
Сложность вычисления — , а расход по памяти .
Для транспонированного внимания (называемого в статье cross-covariance) операция имеет следующий вид:
где — некоторый параметр температуры. Квадратичная сложность переносится с длины последовательности на размерность эмбеддинга. Для cross-attention вычислительная сложность и расход памяти — . Поэтому вычислительная сложность для XCiT будет расти не так быстро, как для ViT, с уменьшением размера патчей или увеличением разрешения.
XCA (Cross-covariance) блок в XCiT и различие между Self-Attention и Cross-Covariance Attention
XC-attention, как и Self-attention, позволяет агрегировать глобальный контекст. Но агрегация происходит несколько в менее явной форме, через свертку по внутренней размерности в вычислении .
Для того, чтобы иметь явное взаимодействие между соседними патчами, авторы добавили так называемое локальное взаимодейсвтвие патчей (Local Patch Interaction). В качестве LPI используется последовательность двух depthwise сверток с батч-нормализацией и GeLU
между ними. Последовательность токенов перед LPI разворачивается в 2d картинку, к этой картинке применяется описанная выше последовательность слоев, и картинка сворачивается обратно в последовательность токенов.
Приятным бонусом от XC-attention является меньшая чувствительность к изменению разрешения подаваемой картинки. Так как свертка при вычислении XC-attention проводится вдоль внутренней оси, размер матрицы внимания не меняется. Качество модели, обученной на разрешении проседает не так сильно при уменьшении разрешения, по сравнению с ResNet и DeiT, и даже заметно возрастает при увеличении разрешения до .
XCiT более устойчив к изменению разрешения входной картинки.
Бэкбоун получился очень даже замечательным. При сопоставимых размерах различные варианты XCiT оказываются эффективнее не только EfficientNet-ов и ранних ViT, но и сильных конкурентов вроде Swin-ов.
Сравнение различных моделей классификации на ImageNet cо сравнивыми характеристиками.
В задаче детекции и сегментации XCiT показал себя с хорошей стороны, превзойдя бэкбоуны на основе PVT и ViL (не затронутого в данном обзоре). XCiT-S12/8 превзошел даже Swin-T с похожими характеристиками, но более крупный свин таки подложил свинью в сравнении с XCiT-S24/8.
Сравнение разных моделей с MASK-RCNN головой на MS COCO в задаче детекции (сравнение проводится по AP^b) и instance сегментации (сравнение по AP^m).
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 в нескольких последовательных блоках трансформера
Структура PS-ViT. После каждого блока количество токенов уменьшается вдвое.
Pooling в PS-ViT
В статье авторы рассматривают разные стратегии модификации архитектуры (взяв за основу DeiT-Tiny) и сохраняя примерно то же количество FLOPs.
Увеличение глубины сети (количества блоков) при сохранении размерности эмбеддинга неизменной
Увеличение размерности эмбеддинга при том же количестве блоков
И то, и то сработало достаточно неплохо, но увеличение ширины несколько лучше. Кроме того, авторы рассматривают два варианта пулинга.
В первом случае, где классификация осуществляется через [CLS]
токен, свертка 1x1
меняет размерность эмбеддинга, а затем проводится MaxPooling
. Эта стратегия называется PSViT-1D
.
В другом случае для классификации используется результат усреднения последней карты активации и для пулинга strided свертка с шагом 2. Этот подход, называемый PSViT-2D
, работает даже немного лучше.
Сравнение глубокой и широкой сети по сравнению с исходным DeiT
Вторым важным наблюдением является то, что карты внимания (attention map) в соседних слоях сильно скоррелированы друг с другом. Так как вычисление attention является дорогостоящим по числу параметров и операций, переиспользование его в следующем слое дает серьезную экономию.
Корреляции attention maps в соседних слоях
И последним по порядку, но по значению является оптимальный выбор расположения элементов и количества слоев в трансформере.
Полный перебор возможных вариантов расположения слоев с пулингом и размерностей эмбеддингов — слишком сложная комбинаторная задача, поэтому пространство поиска пришлось существенно ограничить. Размерности эмбеддинга и максимальное число блоков зафиксировано на каждой стадии (при фиксированном количестве токенов).
В каждом блоке есть 3 выбора:
Использовать обычный блок трансформера
Два последовательных блока с одним и тем же attention
Тождественную операцию (Identity)
На каждом проходе (forward pass) один из трех вариантов выбирается из равномерного распределения и при обратном проходе (backward pass) обновляются параметры для этого варианта (если это не Identity, конечно). Оптимальная архитектура определяется с помощью эволюционного алгоритма.
Суперсеть из возможных конструктивных элементов в PS-ViT. Возможная глубина сети от 0 (только Identity) до 36 (везде два Sharing Layer) блоков трансформера.
Работает это, по всей видимости, и правда неплохо:
Сравнение моделек на ImageNet
VOLO (Vision Outlooker for Visual Recognition)
Довольно занятную вариацию внимания предложили в статье VOLO (если честно, я даже не понимаю, почему она работает так здорово).
Блок энкодера имеет стандартный вид:
Здесь — это LayerNorm, а вот что действительно интересно, так это операция Делается она следующим образом (C
— число каналов, K
— размер ядра свертки):
Линейный слой
nn.Linear(C, K ** 4)
для каждого пикселя из feature map создает вектор размерностиK ** 4
.Полученный вектор решейпится (прошу прощения за англицизм) в матрицу
K ** 2 x K ** 2
. Данная матрица играет роль матрицы внимания в пределах окна размераK x K
. То есть матрица внимания предсказывается в один шаг, без создания ключей (keys) и запросов (queries) c последующим вычислением попарных скалярных произведений.Линейный слой
nn.Linear(C, K ** 4)
выдает значения (values) для каждого токена.Полученная на шаге 2 матрица attention перемножается на values и получается выходное представление.
Как работает Outlook Attention
Таким образом, получается некий trade-off между локальностью операции и вычислительной сложностью. В стандартном self-attention вычислительная сложность растет как поэтому использовать патчи размером меньше 16, особенно при большом разрешении довольно проблематично. В предложенном подходе же асимптотика линейна по количеству токенов . Размер ядра свертки K
должен быть небольшим (в работе K
= 3). Благодаря этому можно брать меньший патч (скажем 8) при большом разрешении (384×384, 512×512).
Разные версии VOLO
OutlookAttn
— гибрид свертки и стандартного self-attention — локальный, но с большим receptive field. При таком подходе большой receptive field может быть достигнут при меньшем числе блоков, чем в типичной CNN и в то же время зашито понятие локальности и близости в саму архитектуру.
Кривые Парето для VOLO и еще нескольких современных моделей на ImageNet
Получился классный бэкбоун, позволивший добиться впечатляющих результатов не только на ImageNet (87.1% без дополнительных данных), но и в задачах семантической сегментации на Cityscapes и ADE20K.
(Слева) Качество сегментации (по mIoU) на CityScapes. (Справа) Качество сегментации (по mIoU) на ADE20K.
Заключение
Универсальность и гибкость архитектуры трансформера, способность улавливать глобальный контекст, оказалась полезной и в области компьютерного зрения.
За год с небольшим, прошедших с публикации An image is worth 16×16 words, трансформеры сильно изменили наши представления о том, как надо решать задачи компьютерного зрения, толкнули науку далеко вперед.
В данном обзоре я рассмотрел лишь отдельные работы из моря публикаций по этой теме за 2021 год. Многие другие интересные идеи, вроде Transformer in Transformer и CoAtNet не были затронуты в силу ограниченности обьема обзора. Кроме того, были рассмотрены только задачи классификации, детекции и сегментации картинок. ViT-ы показали впечатляющие задачи так же в мультимодальных задачах, при работе с видео и self-supervised, semi-supervised learning, генеративных моделях.
В настоящий момент сложно сказать, как будет развиваться эта область в будущем. Мне кажется, что в следующие несколько лет мы увидим последовательное развитие и улучшение архитектур Visual трансформеров, которое имело место для сверточных сетей. Будет ли архитектура на основе механизма внимания или ее гибрид со свертками конечным этапом развития нейронных сетей в компьютерном зрении или придет другая, еще более мощная и универсальная архитектура, не берусь судить.
Но я уверен, что за развитием этой области будет очень интересно следить в 2022.
Список источников
Статьи
Большой Перечень архитектур ViT
https://github.com/dk-liang/Awesome-Visual-Transformer
Ну и куда без Янника Килхера
https://www.youtube.com/c/YannicKilcher