CrossViT: мультишкальный визуальный трансформер с перекрёстным вниманием для классификации изображений

Недавно разработанный визуальный трансформер (ViT) показал многообещающие результаты в задаче классификации изображений по сравнению со сверточными нейронными сетями. Вдохновлённые этим, в данной работе мы изучаем, как обучать многомасштабные представления признаков в моделях трансформера для классификации изображений. Для этого мы предлагаем двухветвевый трансформер, который объединяет фрагменты изображения (т.е. токены в трансформере) разного размера для получения более сильных признаков изображения. Наш подход обрабатывает токены мелких и крупных патчей двумя отдельными ветвями с разной вычислительной сложностью, после чего эти токены несколько раз сливаются исключительно посредством механизма внимания, дополняя друг друга. Кроме того, для сокращения вычислений мы разрабатываем простой, но эффективный модуль слияния токенов на основе перекрёстного внимания, который использует один токен для каждой ветви в роли запроса для обмена информацией с другими ветвями. Предложенное перекрёстное внимание требует линейного времени как по вычислительной, так и по памяти сложности вместо квадратичной в противном случае. Широкий набор экспериментов показывает, что наш подход превосходит или сопоставим с несколькими современными работами по визуальным трансформерам, а также эффективными моделями CNN. Например, на наборе данных ImageNet1K при некоторых архитектурных изменениях наш подход опережает недавний DeiT на значительную величину в 2% при небольшом или умеренном увеличении FLOPs и числа параметров модели. Наши исходные коды и модели доступны по адресу https://github.com/IBM/CrossViT.