Самоконтролируемое зрение с маскированным автоэнкодером

Обучение с учителем — это традиционный способ обучения моделей машинного обучения, который требует меток для каждого наблюдения за данными. Данные относительно легко собирать, учитывая, насколько все доступно через Интернет. Например, если мы хотим собрать изображения собак, мы можем удалить фотографии из Интернета, используя только скрипт. Но маркировка часто является узким местом, замедляющим процесс сбора данных.

Что, если у нас есть способ обучать модели машинного обучения, не заморачиваясь сбором меток? Что, если вместо этого мы извлечем метки из тех же данных, которые мы собрали? Этот тип алгоритма обучения известен как обучение с самоконтролем. Этот подход хорошо зарекомендовал себя при обработке естественного языка. Одним из примеров является BERT¹, который Google использует для своей поисковой системы с 2019 года. К сожалению, этого нельзя сказать о компьютерном зрении.

Тем не менее, может произойти прорыв в самоконтролируемом зрении. В своей недавней статье He et al. от Facebook AI Research предложила маскированный автоматический кодировщик (MAE)², который следует архитектуре Vision Transformer (ViT)³. Авторы утверждали, что их подход работал лучше на ImageNet, чем Vision Transformer, обученный с нуля. В этой статье мы углубимся в их подход и посмотрим, как его можно реализовать в коде.

Маскированный автоэнкодер

Как и в других методах с самоконтролем, в MAE есть предлоговая задача: восстановить отсутствующие участки изображения. Изображение разделено на фрагменты, часть которых замаскирована. Остальные патчи преобразуются во вложения токенов посредством линейной проекции, аналогично ViT. Затем часть внедренных токенов удаляется (маскируется) перед отправкой кодировщику. После этого все эти замаскированные вложения токенов преобразуются в одно и то же встраивание маски, а затем добавляются к выходным данным кодировщика в соответствии с их исходным положением. После линейной проекции выходные данные кодировщика вместе с вложениями маски отправляются в декодер, который предсказывает значения пикселей замаскированных патчей.

Этот этап называется предварительной подготовкой, так как модель MAE позже будет использоваться для последующих задач, таких как классификация изображений. Выполнение модели на предлоговых заданиях не имеет решающего значения в самоконтроле. Смысл этих задач в том, чтобы позволить модели изучить промежуточное представление, которое, как ожидается, будет содержать хорошие семантические значения. После этапа предварительной подготовки декодер будет заменен головкой многослойного персептрона (MLP) или линейным слоем, который действует как классификатор, выдающий прогноз для последующих задач.

Архитектура модели

Кодер

Кодер представляет собой ViT. Он принимает изображение с тензорной формой (batch_size, RGB_channels, height, width). Вложение патча получается для каждого патча путем выполнения линейной проекции. Это можно сделать с помощью 2D-сверточного слоя. Затем тензор сглаживается в последнем измерении, чтобы стать (batch_size, encoder_embed_dim, num_visible_patches). Транспонировать в тензор формы (batch_size, num_visible_patches, encoder_embed_dim).

Как упоминалось в оригинальной статье Transformer, позиционное кодирование добавляет информацию о позиции каждого патча. Авторы используют версию «синус-косинус» вместо обучаемых позиционных вложений. Эта реализация ниже является 1-D версией.

Как и в Transformer, каждый блок состоит из нормальных слоев, модуля внимания с несколькими головками и слоя прямой связи. Промежуточная выходная форма — (batch_size, num_visible_patches, encoder_embed_dim). Ниже приведен код модуля внимания с несколькими головками:

Код для блока Transformer:

Эта часть предназначена только для тонкой настройки. Авторы следуют архитектуре ViT, которая имеет токен класса (патч) для классификации. Следовательно, они добавили фиктивный токен. Но они утверждали, что их метод хорошо работал и без него, поскольку они выполняли средний пул на других токенах. Здесь я включил среднюю версию реализации пула. После этого в качестве классификатора добавляется линейный слой. Окончательная форма тензора — (batch_size, num_classes).

Хотите получать самые свежие новости обо всем, что связано с глубоким обучением? Подпишитесь на еженедельную рассылку Deep Learning Weekly и получайте последние новости прямо на свой почтовый ящик.

Собрав все вместе, кодировщик реализован следующим образом:

Декодер

Подобно энкодеру, декодер состоит из ряда блоков преобразователя. В конце декодера находится классификатор, состоящий из уровня нормы и уровня прямой связи. Входной тензор имеет форму batch_size, num_patches, decoder_embed_dim), тогда как окончательный выходной тензор имеет форму (batch_size, num_patches, 3 * patch_size ** 2).

Соедините все вместе — архитектура MAE

MAE используется для предварительной подготовки на замаскированных изображениях. Сначала замаскированные входы подаются на кодировщик. Затем они передаются на уровень прямой связи, чтобы изменить размер внедрения, чтобы он соответствовал декодеру. Перед передачей в декодер маскированные патчи вставляются обратно. Кодирование положения снова применяется ко всему набору фрагментов изображения, как видимых, так и замаскированных. В статье авторы перетасовывают список, содержащий все патчи, чтобы маскированные патчи вставлялись правильно. Здесь этого не делается, так как это непросто реализовать на PyTorch. Вместо этого позиционные кодировки перетасовываются соответствующим образом перед добавлением в патчи.

Процедура обучения

Для самоконтролируемого предварительного обучения авторы обнаружили, что простая средняя абсолютная потеря по пикселям хорошо работает в качестве целевой функции. Набор данных, который они использовали, — это обучающий набор ImageNet-1K.

На этапе тонкой настройки декодер удаляется, а кодировщик обучается на том же наборе данных. Данные немного отличаются от данных предварительной подготовки, поскольку кодер теперь получает полный набор патчей изображения (без маскирования). В результате количество патчей теперь отличается от этапа предобучения.

Если вам интересно, можно ли использовать модель, используемую для предварительного обучения, для тонкой настройки, ответ — да. Позвольте мне объяснить. Кодер в основном состоит из модуля внимания, уровня нормы и уровня прямой связи. Чтобы проверить, влияет ли изменение количества патчей (второе измерение, индекс 1) на прямой проход, нам нужно увидеть форму тензоров параметров каждого слоя.

Параметр в слое нормы имеет форму (batch, 1, encoder_embed_dim). Он может транслировать по измерению патча (индекс 1) во время прямого прохода, поэтому он не зависит от размера измерения патча.
Слой прямой связи имеет матрицу весов формы (in_channels, out_channels) и матрицу смещения формы (out_channels,), обе из которых не зависят от количества патчей.
Модуль внимания, по сути, выполняет серию линейных проекций. Итак, по той же причине количество патчей не влияет на форму тензора параметров.

Однако количество исправлений в пакете по-прежнему должно быть постоянным из-за того, что параллельная обработка позволяет помещать данные в пакеты.

Полученные результаты

Давайте посмотрим на изображения реконструкции на этапе предварительной подготовки, как сообщается в исходной статье. Похоже, MAE неплохо справился с восстановлением изображений, даже когда 80% пикселей замаскированы.

MAE также хорошо справляется с тонко настроенными последующими задачами, такими как классификация изображений в наборе данных ImageNet-1K. Базовый уровень, который использовали авторы, ViT-Large, на самом деле работал лучше при обучении с использованием предварительного обучения MAE по сравнению с контролируемым способом.

В документ также включены контрольные результаты экспериментов по переносу обучения для последующих задач и различных исследований абляции. Я рекомендую вам ознакомиться с документом для более подробной информации.

Обсуждение

Если вы знакомы с BERT, вы можете заметить сходство между подходами BERT и MAE. При предварительном обучении BERT мы маскируем часть текста, и перед моделью ставится задача предсказать их. Более того, поскольку сейчас мы используем архитектуру на основе Transformer, было бы уместно предположить, что этот подход является эквивалентом BERT в области зрения.

Однако авторы статьи утверждали, что такой подход предшествует BERT. Например, в прошлых попытках самоконтроля изображения в качестве предтекстовой задачи использовались автоэнкодеры с шумоподавлением и закрашивание изображения. Сам MAE также использует автоэнкодер в качестве модели и предтекстовую задачу, аналогичную рисованию изображения.

Если это так, что заставляет MAE работать, а предыдущие работы не работали? Я считаю, что ключ лежит в архитектуре ViT. В своей статье авторы упомянули, что у сверточных нейронных сетей есть проблема с интеграцией в них «индикаторов», таких как токены маски и позиционные вложения, тогда как ViT устраняет этот архитектурный пробел. Если это так, то мы видим еще одну идею, разработанную в области обработки естественного языка, успешно реализованную в компьютерном зрении. До этого это был механизм внимания, затем концепция Transformer была заимствована в компьютерное зрение в виде Vision Transformers, а теперь и весь процесс предобучения BERT.

Заключение

Я в восторге от того, что самоконтролируемое зрение может предложить в будущем. Учитывая, насколько успешен BERT для обработки естественного языка, я склонен думать, что подход к моделированию с использованием масок, такой как MAE, принесет пользу компьютерному зрению. Как я упоминал в начале этого поста, данные изображения легко получить, но их маркировка может занять много времени. При таком подходе можно потенциально масштабировать процесс предварительной подготовки, курируя гораздо больший набор данных, чем ImageNet, не беспокоясь о маркировке. Потенциал безграничен. Станем ли мы свидетелями очередного возрождения компьютерного зрения, покажет только время.