Vision Transformer (ViT)

Визуальный трансформер (ViT) — это нейронная сеть, которая применяет архитектуру Transformer, изначально разработанную для языкового моделирования, к изображениям. Он делит изображение на фрагменты фиксированного размера, преобразует каждый фрагмент в эмбеддинг токена и обрабатывает последовательность с помощью самовнимания, благодаря чему каждый фрагмент может учитывать все остальные уже с первого слоя. Исследователи Google представили эту архитектуру в статье «Изображение стоит 16×16 слов», показав, что чистый Transformer без слоёв свёртки может сравниться с современными свёрточными нейронными сетями (CNN) в задаче классификации изображений, если предварительно обучить его на достаточном объёме данных. Сегодня ViT выступают в роли энкодеров изображений в визуально-языковых моделях, базовых моделях сегментации и самообучаемых бэкбонах.

Как работают Vision Transformer#

ViT преобразует изображение в последовательность и обрабатывает её так же, как текст. Ниже описаны шаги для ViT-Base/16 на изображении 224×224 пикселя:

1. Встраивание фрагментов: Изображение делится на фрагменты размером 16×16 пикселей — сетку 14×14 из 196 фрагментов. Каждый фрагмент разворачивается в вектор и умножается на обучаемую матрицу, в результате чего получается 768-мерный эмбеддинг. Обычно эта операция реализована как свёртка 16×16 с шагом 16.
2. Токен класса и позиционные эмбеддинги: В начало последовательности добавляется обучаемый токен класса, а каждому токену назначается обучаемый позиционный эмбеддинг. Само внимание не учитывает порядок, поэтому именно эти эмбеддинги — единственный источник информации о том, откуда взялся каждый фрагмент.
3. Энкодер Transformer: Стек одинаковых блоков — в ViT-Base их 12 — чередует многоголовое самовнимание, которое смешивает информацию всех токенов, и двухслойную MLP с активациями GELU, которая преобразует каждый токен отдельно. Перед каждым подслоем применяется нормализация слоя, а сам подслой обёрнут в остаточное соединение.
4. Голова: Итоговый выход токена класса подаётся на линейный классификатор. Для плотных задач, таких как сегментация, вместо него используют токены фрагментов, преобразованные обратно в сетку 14×14.

Код на PyTorch ниже формирует последовательность токенов для одного изображения и передаёт её в один блок энкодера.

import torch
from torch import nn

# Параметры ViT-Base/16: фрагменты 16x16, токены размерности 768, 12 голов внимания
image = torch.randn(1, 3, 224, 224)
patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16)  # одна линейная проекция на каждый фрагмент

# Разбить изображение на сетку 14x14 из фрагментов и преобразовать её в 196 токенов
tokens = patch_embed(image).flatten(2).transpose(1, 2)

# Добавить в начало обучаемый токен класса и обучаемые позиционные эмбеддинги
cls_token = nn.Parameter(torch.zeros(1, 1, 768))
pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
tokens = torch.cat([cls_token, tokens], dim=1) + pos_embed

# Один блок энкодера с предварительной нормализацией, в котором каждый токен учитывает все остальные
block = nn.TransformerEncoderLayer(
    d_model=768, nhead=12, dim_feedforward=3072, activation="gelu", batch_first=True, norm_first=True
)
print(block(tokens).shape)

Выводится torch.Size([1, 197, 768]): 196 токенов фрагментов плюс токен класса, каждый размерности 768 — они готовы к передаче в следующий из 12 блоков ViT-Base.

Размеры и варианты моделей#

В названиях ViT указаны размер модели и размер фрагмента: ViT-B/16 — базовая модель с 86 миллионами параметров и фрагментами 16×16, а модели Large и Huge содержат 307 миллионов и 632 миллиона параметров соответственно. Меньшие фрагменты создают больше токенов и повышают детализацию, но требуют больше вычислительных ресурсов. В нескольких вариантах изменили способ обучения ViT или структуру модели:

• DeiT: Методика обучения с интенсивной аугментацией и токеном дистилляции, позволяющая обучать конкурентоспособные ViT только на ImageNet-1k, без дополнительных данных для предварительного обучения.
• Swin Transformer: Вычисляет внимание в локальных окнах со сдвигом и объединяет фрагменты между этапами. В результате получается иерархия, похожая на CNN и подходящая для обнаружения объектов и сегментации.
• Маскированные автоэнкодеры (MAE) и DINOv2: Методы самообучения, которые предварительно обучают ViT на неразмеченных изображениях: MAE восстанавливает замаскированные фрагменты, а DINOv2 сопоставляет выходы модели-учителя.
• CLIP и SigLIP: Энкодеры изображений ViT, обученные сопоставлять изображения с подписями к ним. Их широко используют в качестве визуального энкодера мультимодальных моделей.

Применение#

• Визуально-языковые модели: Энкодер ViT преобразует изображение в последовательность токенов, которую большая языковая модель обрабатывает вместе с текстом. Это позволяет создавать подписи и отвечать на визуальные вопросы.
• Плотное предсказание: Токены фрагментов из бэкбонов ViT используются для сегментации изображений и оценки глубины, где для каждого выходного пикселя нужен контекст из удалённых областей изображения.
• Понимание видео: Видеоверсии ViT делят клипы на пространственно-временные фрагменты, а модели самообучения на видео, такие как V-JEPA, извлекают признаки движения и внешнего вида из неразмеченных видеозаписей для таких задач, как распознавание действий.

Vision Transformer и CNN#

CNN закладывают в каждый слой два предположения: локальность, поскольку каждый фильтр видит небольшую окрестность, и эквивариантность относительно сдвига, поскольку один и тот же фильтр применяется во всех позициях. Помимо сетки фрагментов, ViT приходится самостоятельно изучать пространственную структуру по данным. В исходной статье ViT, обученные только на ImageNet (1,3 миллиона изображений), набрали на несколько процентных пунктов меньше, чем ResNet сопоставимого размера. Предварительное обучение на наборе JFT-300M из 300 миллионов изображений повысило точность ViT-H/14 на ImageNet до 88,55% по метрике top-1 accuracy. Второе отличие — вычислительная стоимость. Самовнимание сопоставляет каждый токен со всеми остальными, поэтому удвоение длины стороны изображения увеличивает число токенов вчетверо, а вычислительные затраты на внимание — в 16 раз. Flash Attention сокращает трафик памяти при этих вычислениях, а внимание в окнах ограничивает число самих сравнений. Ultralytics запускает модель сегментации всего от Meta: контрольная точка sam_b.pt кодирует каждое изображение 1024×1024 с помощью бэкбона ViT-B. Фрагменты 16×16 образуют сетку 64×64 из 4096 токенов, поэтому 8 из 12 блоков используют внимание только внутри окон 14×14, а 4 — глобальное внимание.