An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Разбор статьи

Проблема и мотивация

Свёрточные архитектуры (например, ResNet) долгое время оставались фундаментом современного компьютерного зрения благодаря заложенным в архитектуру индуктивным смещениям, таким как эквивариантность к сдвигу и двумерная локальность. В то время как Transformers стали стандартом в обработке естественного языка (NLP) благодаря вычислительной эффективности и масштабной масштабируемости, ранние попытки перенести механизм self-attention в зрение обычно сочетали его со свёрточным бэкбоном либо использовали специализированные паттерны локального или аксиального внимания, которые сложно эффективно масштабировать на аппаратных ускорителях. В этой статье исследуется, может ли чистая архитектура Transformer, применяемая напрямую к последовательностям патчей изображений с минимальными модификациями, конкурировать с современными CNN или превосходить их.

Ключевые идеи

• Изображения как последовательности слов: Изображение $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$ разбивается на неперекрывающиеся патчи размера $P\times P$. Преобразованные в одномерный вектор патчи обрабатываются аналогично NLP-токенам.
• Минимальное индуктивное смещение: В отличие от CNN, где пространственная локальность и трансляционная эквивариантность заложены в каждый слой, ViT использует глобальный self-attention между патчами. Двумерная структура учитывается только при нарезке патчей и при интерполяции позиционных эмбеддингов для файнтюнинга на более высоком разрешении.
• Масштаб важнее индуктивного смещения: Без сильного индуктивного смещения ViT уступает ResNet на датасетах малого и среднего размера (таких как стандартный ImageNet без сильной регуляризации). Однако при предобучении на больших наборах данных (ImageNet-21k с 14 млн изображений или JFT-300M с 303 млн изображений) ViT не уступает или превосходит CNN, затрачивая существенно меньше FLOPs при обучении.

Метод

1. Эмбеддинг патчей и формирование входа: При заданном размере патча $(P,P)$ изображение делится на $N=HW/P^2$ патчей. Каждый вытянутый в вектор патч ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^2\cdot C}$ линейно проецируется матрицей $\mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D}$ в размерность $D$. В начало последовательности добавляется обучаемый токен классификации ${\mathbf{x}}_{\text{class}}$, а затем суммируются одномерные обучаемые позиционные эмбеддинги ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}$: $${\mathbf{z}}_0=[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos}$$
2. Слои энкодера Transformer: Последовательность проходит через $L$ идентичных блоков, чередующих multi-head self-attention (MSA) и блоки MLP с использованием pre-LayerNorm (LN) и остаточных связей (residual connections): $${\mathbf{z}}_{\ell }^{\prime }=\mathrm{MSA}(\mathrm{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -1},\ell =1\dots L$$ $${\mathbf{z}}_{\ell }=\mathrm{MLP}(\mathrm{LN}({\mathbf{z}}_{\ell }^{\prime }))+{\mathbf{z}}_{\ell }^{\prime },\ell =1\dots L$$
3. Представление и классификационная голова: Итоговое представление извлекается из первого выходного токена: $$\mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)$$ Небольшой MLP (с одним скрытым слоем) служит классификационной головой на этапе предобучения, а при файнтюнинге он заменяется единичной линейной проекцией ($D\times K$).
4. Тонкая настройка на более высоком разрешении: При файнтюнинге на изображениях большего разрешения (например, $384\times 384$ или $512\times 512$) размер патча $P$ остаётся неизменным, что увеличивает длину последовательности $N$. Предобученные 1D-позиционные эмбеддинги корректируются с помощью двумерной бикубической интерполяции в соответствии с их исходным расположением на пространственной сетке.

Результаты

• Сравнение с SOTA:
  • Модель ViT-H/14, предобученная на JFT-300M, достигла 88,55% на ImageNet, 90,72% на ImageNet ReaL, 94,55% на CIFAR-100 и 77,63% на наборе из 19 задач VTAB, потребовав 2,5 тыс. TPUv3-core-days.
  • Модель ViT-L/16, предобученная на JFT-300M, показала результат 87,76% на ImageNet всего за 0,68 тыс. TPUv3-core-days, превзойдя BiT-L (ResNet152x4, 87,54%, 9,9 тыс. TPUv3-core-days) и приблизившись к Noisy Student EfficientNet-L2 (88,4/88,5%, 12,3 тыс. TPUv3-core-days).
  • Предобучение на открытых данных: ViT-L/16, обученная на ImageNet-21k, достигла 85,30% на ImageNet.
• Компромисс между вычислениями и качеством: В экспериментах по масштабированию на JFT-300M чистый ViT доминировал над ResNet на границе Парето, требуя в $2–4$ раза меньше вычислительных ресурсов для достижения аналогичной точности.
• Внутренние представления: В нижних слоях присутствуют головы внимания, охватывающие как всё изображение целиком, так и головы с малой локализованной дистанцией внимания. В целом дистанция внимания монотонно возрастает по мере увеличения глубины.

Ограничения и открытые вопросы

• Потребность в данных (Data Hunger): ViT плохо обобщается на небольших датасетах без масштабного предобучения (ImageNet-21k или JFT-300M) либо без агрессивной регуляризации.
• Отставание самообучения (Self-Supervised Gap): Первые эксперименты с предсказанием маскированных патчей дали 79,9% на ImageNet (ViT-B/16), сохраняя разрыв в 4% по сравнению с контролируемым обучением.
• Применение к задачам плотного прогнозирования (Dense Prediction): Исходная работа сосредоточена исключительно на классификации; адаптация чистого патчевого трансформера к задачам плотного пространственного предсказания, таким как детекция объектов и семантическая сегментация, остаётся актуальной исследовательской задачей.