Masked Image Modeling

Введение

В нашей прошлой статье про Self-Supervised Learning был проведён общий разбор этого метода обучения нейронных сетей, а также высокоуровнево описаны некоторые подходы в SSL.

В новой статье мы подробно расскажем про подход Masked Image Modeling (MIM) и обсудим:

• историю возникновения;
• основную идею;
• способы применения;
• причины, из-за которых его не стоит путать с Image Inpainting.

Итак, начинаем! 🙂

История развития

BERT

Впервые маскирование токенов было осознанно использовано для предобучения модели BERT.

Идея маскирования в BERT заключается в обучении модели восстановлению исходного текста после замены слов на маскировочные токены — такой подход называется Masked Language Model (MLM). Он вносит в модель глубокие универсальные представления о структуре предложений и семантике слов, благодаря чему модель лучше показывает себя в downstream задачах.

ViT & BEiT

В 2020 году в Computer Vision приходят трансформеры со статьёй ViT, основная идея которой — представление оригинального изображения как набора патчей, словно токенов текста, и использование трансформерной архитектуры в качестве бэкбона.

Спустя год публикуется статья BEIT: BERT Pre-Training of Image Transformers, которая применяет наработки из BERT для предобучения ViT, адаптируя Masked Language Modeling под задачу компьютерного зрения и называя свой подход «Masked Image Modeling».

Masked Image Modeling

Подход Masked Image Modeling реализуется на неразмеченных данных и заключается в маскировании частей изображения и обучении модели восстанавливать недостающую информацию, используя остальную видимую часть изображения. Благодаря этому модель глубже понимает структуры объектов и контекст сцен на изображениях, что формирует информативные внутренние представления.

Masked Image Modeling можно разделить на два основных фреймворка — Reconstruction и Contrastive. Давайте рассмотрим их более подробно 🙂

Reconstruction

Метод Reconstruction представляет собой обучение модели восстанавливать исходное изображение по небольшому набору патчей (~25%).

Основной процесс выглядит так:

• Оригинальное изображение делится на патчи, случайным образом отбрасывается часть из них (~75%), а оставшиеся патчи подаются на вход трансформеру с соответствующим их положению position encoding.
• Видимые патчи переводятся в токены и попадают в encoder (ViT). На выходе получаются эмбеддинги патчей.
• Формируется матрица с исходным расположением патчей. Для видимых патчей используется выход encoder, а позиции маскированных патчей инициализируются как mask_token.
• Матрица из видимых и маскированных токенов подаётся в декодер, на выходе получается восстановленное изображение.
• Теперь мы, имея оригинальное изображение, разбитое на патчи, а также восстановленное изображение, считаем MSE loss между пикселями маскированных патчей.

Contrastive

Подход Contrastive сочетает в себе Masked Image Modeling, Self-Distilation из DINO и Contrastive Learning из CLIP (для большего понимая рекомендуется предварительно ознакомиться с этими статьями). Модель обучается за счёт contrastive loss, поэтому представления обладают хорошими семантическими высокоуровневыми признаками.

Фреймворк Contrastive MIM можно разделить на две ветви: teacher и student (левую и правую соответственно, см. Рисунок 6).

Основной процесс выглядит так:

• Student stage. К исходному изображению применяются сильные аугментации, оно делится на патчи равного размера, часть патчей отбрасывается, остальные подаются в student encoder, на выходе которого получаются эмбеддинги.
• Teacher stage. К исходному изображению применяются слабые аугментации, оно делится на патчи равного размера, патчи подаются в teacher encoder, на выходе которого получаем эмбеддинги. При этом веса teacher encoder — это Exponential Moving Average весов student encoder:

\( \theta_{\text{teacher}} \leftarrow \tau \cdot \theta_{\text{teacher}} + (1 — \tau) \cdot \theta_{\text{student}} \)

• Считаем Contrastive loss между эмбеддингами патчей маскированных патчей, полученных из модели ученика (\(\hat{H}\)), и эмбеддингами всех патчей, полученных из модели учителя (\(\tilde{H}\)):

\(\mathcal{L}(\varphi, \theta, M) \leftarrow -\sum_{i=1}^{n} \log \frac{    \exp\left( \langle \hat{H}[i], \tilde{H}[i] \rangle / \tau \right)}{    \sum_{j=1}^{n} \exp\left( \langle \hat{H}[j], \tilde{H}[i] \rangle / \tau \right)}\)

MIM vs Image Inpainting

На первый взгляд может показаться, что MIM — задача восстановления изображения Image Inpainting, и модель учится восстанавливать недостающие части изображения. Однако отметим отличия двух подходов:

• Конечная цель Image Inpainting заключается в правдоподобном восстановлении частей изображения, а в Masked Image Modeling восстановление изображения — это лишь способ получения качественных представлений о сценах и объектах.
• Из-за различия конечных целей также отличаются и лосс функции: для MIM могут применяться MSE, Cross-Entropy, а в Inpainting используются Perceptual loss, Adversarial loss.

Заключение

Большинство state-of-the-art подходов используют в обучении Self-Supervised Learning, в частности Masked Image Modeling.

Masked Image Modeling является мощным подходом Self-Supervised Learning, который продолжает развиваться до сих пор: подход из статьи Contrastive Masked Autoencoders are Stronger Vision Learners объединяет рассмотренные Reconstruction и Contrastive идеи.