В данной статье разберём работу Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, авторы которой предложили новую архитектуру Vision Transformer’a, позволяющую решать не только задачу классификации, но и более общие задачи компьютерного зрения, при этом делая вычислительную сложность операций линейной по размеру изображения.

Проблемы ViT

Вкратце работа Visual Transformer’а выглядит так:

1. Изображение нарезается на патчи 16x16.

2. Патчи вытягиваются в векторы и преобразуются к нужной размерности с помощью линейного слоя.

3. К полученным векторам прибавляются обучаемые позиционные 1D-эмбеддинги.

4. Также добавляется специальный обучаемый эмбеддинг класса — по аналогии с [CLS] токеном в BERT.

5. Полученная последовательность векторов подаётся в обычный энкодер трансформера.

6. Предсказывается класс изображения с помощью MLP над эмбеддингом класса.

Источник: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Среди проблем ViT можно выделить следующие:

1. Отсутствие некоторых полезных Inductive bias’ов (подробнее можно почитать здесь), которые есть в свёрточных нейронных сетях: translation invariance (из-за использования ядер свёртки, которые применяются для обработки всего изображения, и пулингов объекты в разных частях изображения могут успешно обрабатываться свёрточными сетями) и feature locality (ядра свёртки захватывают достаточно маленькую часть изображения, благодаря чему появляется предположение о том, что каждый предмет расположен на изображении компактно, + нескольких слоёв позволяют ориентироваться на разный receptive field).
2. ViT не может быть применён к таким задачам, как детекция или сегментация. В классической архитектуре используются фичи с одного и того же скейла, что мешает качественно обрабатывать объекты различных размеров.
3. Квадратичная по размеру изображения вычислительная сложность использующегося в архитектуре Self-Attention’а. Для изображений с высоким разрешением это становится существенным.

Swin Transformer

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

Источник: Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

Иерархичность достигается за счёт использования на первых слоях маленьких патчей и их увеличения на последующих слоях — благодаря этому появляется возможность перенимать свойства Feature Pyramid сетей и использовать модели для pixel-aware задач наподобие детекции и сегментации.