Tokens-to-Token ViT: обучение визуальных трансформеров с нуля на ImageNet

Трансформеры, популярные в языковом моделировании, недавно начали применяться для решения задач компьютерного зрения, например Vision Transformer (ViT) для классификации изображений. Модель ViT разбивает каждое изображение на последовательность токенов фиксированной длины и затем применяет несколько слоев трансформера для моделирования их глобальных взаимосвязей для задачи классификации. Однако ViT демонстрирует худшую производительность по сравнению с сверточными нейронными сетями (CNN), когда обучается с нуля на средних по размеру наборах данных, таких как ImageNet. Мы установили, что это происходит потому, что: (1) простая токенизация входных изображений не моделирует важную локальную структуру, такую как грани и линии между соседними пикселями, что приводит к низкой эффективности использования обучающих примеров; и (2) избыточная архитектура внимания в ViT ограничивает богатство признаков при фиксированных вычислительных ресурсах и ограниченном количестве обучающих примеров. Чтобы преодолеть эти ограничения, мы предлагаем новый Tokens-to-Token Vision Transformer (T2T-ViT), который включает: (1) покадровое преобразование Tokens-to-Token (T2T) для поэтапной структуризации изображения в токены путем рекурсивного аггрегирования соседних токенов в один токен, благодаря чему можно моделировать локальную структуру, представляемую окружением токенов, и уменьшать длину последовательности токенов; и (2) эффективный бэκбон с глубокой и узкой структурой для визуальных трансформеров, мотивированный дизайном CNN после эмпирического исследования. Примечательно, что T2T-ViT сокращает число параметров и MACs по сравнению с базовым ViT вдвое, при этом достигает более чем на 3.0% лучшей точности при обучении с нуля на ImageNet. Он также превосходит ResNet и демонстрирует сопоставимую производительность с MobileNet при прямом обучении на ImageNet. Например, T2T-ViT сопоставимого размера с ResNet-50 (21.5M параметров) может достигать 83.3% top-1 точности при разрешении изображения 384x384 на ImageNet.