Masked Autoencoders (MAE)

Маскированные автоэнкодеры (MAE) — метод обучения без учителя для изображений: модель скрывает большую случайную долю патчей изображения, обычно 75%, и учится восстанавливать отсутствующие пиксели по остальным, видимым. Каймин Хэ и его коллеги из Meta AI представили этот метод в статье 2021 года Masked Autoencoders Are Scalable Vision Learners, адаптировав предсказание замаскированных слов из BERT для визуальных трансформеров. Поскольку энкодер обрабатывает только видимые патчи, предварительное обучение MAE проходило как минимум в 3 раза быстрее, чем обработка каждого патча, а модель ViT-Huge, обученная таким способом, достигла точности 87.8% на ImageNet — лучшего на тот момент результата среди методов, использовавших только данные ImageNet-1K. Затем предварительно обученный энкодер служит бэкбоном для классификации, обнаружения и сегментации.

Как работают маскированные автоэнкодеры#

MAE делит изображение на сетку неперекрывающихся патчей — 196 патчей размером 16x16 пикселей для входного изображения 224x224 — и случайным образом скрывает большинство из них. Энкодер и декодер намеренно устроены по-разному. Энкодер, стандартный ViT, обрабатывает только видимые патчи, сокращая вычисления и расход памяти примерно до четверти. Лёгкий декодер, по умолчанию состоящий из 8 блоков трансформера и требующий менее 10% вычислений энкодера на токен, получает закодированные видимые патчи и общий обучаемый токен маски для каждой скрытой позиции, а затем предсказывает значения пикселей каждого скрытого патча. Функция потерь — это среднеквадратичная ошибка только на скрытых патчах; использование нормализованных по патчам пикселей в качестве целевых значений дополнительно улучшило выученные признаки. После предварительного обучения декодер отбрасывается:

Высокая доля маскирования делает задачу полезной. Соседние патчи изображения сильно дублируют друг друга, поэтому, если скрыто всего несколько патчей, модель может заполнить каждую пропущенную область интерполяцией по соседним патчам. Если скрыто 75%, ей приходится выводить целые объекты и сцены по фрагментам, что заставляет её учиться распознавать формы, части и контекст. Для сравнения, BERT маскирует только 15% слов, потому что каждое слово уже несёт много смысла.

Ключевой этап MAE — маскирование. Ниже он показан на PyTorch для пакета из четырёх изображений:

import torch

torch.manual_seed(0)
images = torch.rand(4, 3, 224, 224)  # a batch of 4 RGB images

# Patchify: each 224x224 image becomes 196 patches of 16x16 pixels, flattened to 768 values
patches = images.unfold(2, 16, 16).unfold(3, 16, 16).permute(0, 2, 3, 1, 4, 5).reshape(4, 196, 768)

# Random masking as in the MAE reference code: rank patches by random noise and keep the first 25%
keep = int(196 * (1 - 0.75))
ids_keep = torch.argsort(torch.rand(4, 196), dim=1)[:, :keep]
visible = torch.gather(patches, 1, ids_keep.unsqueeze(-1).expand(-1, -1, 768))

# Mask of hidden patches (1 = hidden): the decoder predicts these, and the loss is averaged over them only
mask = torch.ones(4, 196).scatter(1, ids_keep, 0)
print(f"Encoder input: {visible.shape[1]} of {patches.shape[1]} patches per image")
print(f"Reconstruction targets: {int(mask[0].sum())} hidden patches per image")

Выводятся Encoder input: 49 of 196 patches per image и Reconstruction targets: 147 hidden patches per image. Поэтому слои трансформера энкодера обрабатывают последовательность вчетверо короче полного изображения, а поскольку механизм внимания сравнивает каждый токен со всеми остальными, матрица оценок внимания в шестнадцать раз меньше, чем для всех 196 патчей.

Результаты и варианты#

Признаки MAE лучше всего проявляют себя после дообучения. При линейном зондировании энкодер ViT-L, предварительно обученный с MAE, достиг точности 73.5% на ImageNet, уступив контрастивным методам. Однако после дообучения и переноса на задачи обнаружения объектов и сегментации на COCO предварительное обучение MAE превзошло обучение с учителем. Этот подход успешно адаптировали для задач, выходящих далеко за пределы статичных изображений: VideoMAE маскирует 90% или больше пространственно-временных патчей видеоклипа видео, SatMAE предварительно обучается на спутниковых снимках, а ConvNeXt V2 использует маскированное автокодирование в свёрточных сетях с разреженными свёртками.

Маскированное предварительное обучение лежит в основе и широко используемых фундаментальных моделей. Segment Anything Model (SAM) от Meta, доступная в пакете Ultralytics, использует для энкодера изображений ViT, предварительно обученный с MAE, а её версия ViT-B, sam_b.pt, — это модель масок по умолчанию в Ultralytics auto_annotate, которая преобразует рамки детектора YOLO26 в метки сегментации.

Маскированные автоэнкодеры и связанные понятия#

• Автоэнкодер: сжимает весь вход в компактный код и восстанавливает его. MAE получает неполный вход и должен предсказать части, которых он никогда не видел; выход энкодера при этом содержит по одному вектору на каждый видимый патч.
• Моделирование маскированных изображений: более широкое семейство методов, к которому относится MAE. BEiT предсказывает дискретные визуальные токены для замаскированных патчей, а SimMIM передаёт токены маски через энкодер; MAE же предсказывает пиксели и не подаёт токены маски в энкодер.
• Контрастивное обучение: обучает модель сопоставлять аугментированные представления одного изображения, не используя декодер. Полученные признаки лучше работают без дообучения, тогда как признаки MAE сильнее улучшаются после дообучения.
• Предиктивная архитектура с совместными эмбеддингами (JEPA): также скрывает области изображения, но предсказывает их эмбеддинги, а не пиксели, не позволяя модели тратить ресурсы на детали текстуры.