Masked Autoencoders (MAE)

Autoencoders mascarados (MAE) são um método de aprendizado autossupervisionado para imagens: o modelo oculta uma grande parcela aleatória dos patches de uma imagem, geralmente 75%, e aprende a reconstruir os pixels ausentes a partir dos demais pixels visíveis. Kaiming He e colegas da Meta AI apresentaram o método no artigo de 2021 Masked Autoencoders Are Scalable Vision Learners, adaptando a previsão de palavras mascaradas do BERT aos transformers de visão. Como o codificador processa apenas os patches visíveis, o pré-treinamento MAE foi pelo menos 3 vezes mais rápido do que processar todos os patches, e um modelo ViT-Huge pré-treinado dessa forma alcançou 87,8% de acurácia no ImageNet, o melhor resultado da época entre os métodos que usavam apenas dados do ImageNet-1K. O codificador pré-treinado então serve como backbone para classificação, detecção e segmentação.

Como funcionam os codificadores automáticos mascarados#

O MAE divide uma imagem em uma grade de patches sem sobreposição — 196 patches de 16x16 pixels para uma entrada de 224x224 — e oculta aleatoriamente a maioria deles. O codificador e o decodificador são deliberadamente desiguais. O codificador, um ViT padrão, opera apenas nos patches visíveis, o que reduz seu uso de computação e memória a cerca de um quarto. Um decodificador leve, composto por padrão de 8 blocos de transformer, com menos de 10% da computação do codificador por token, recebe os patches visíveis codificados mais um token de máscara compartilhado e treinável para cada posição oculta, e prevê os valores dos pixels de cada patch oculto. A função de perda é o erro quadrático médio apenas nos patches ocultos, e usar pixels normalizados por patch como alvos melhorou ainda mais as representações aprendidas. Após o pré-treinamento, o decodificador é descartado:

A alta proporção de mascaramento é o que torna a tarefa útil. Patches de imagem vizinhos são altamente redundantes, então, quando apenas alguns patches estão ocultos, o modelo pode preencher cada lacuna interpolando a partir dos vizinhos. Com 75% ocultos, ele precisa inferir objetos e cenas inteiros a partir de fragmentos, o que o força a aprender formas, partes e contexto. Já o BERT mascara apenas 15% das palavras, porque cada palavra já carrega um significado denso.

O núcleo do MAE é sua etapa de mascaramento, mostrada aqui em PyTorch com um lote de quatro imagens:

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")

Ele imprime Encoder input: 49 of 196 patches per image e Reconstruction targets: 147 hidden patches per image. Portanto, as camadas de transformer do codificador processam uma sequência com um quarto do comprimento da imagem completa e, como a atenção compara cada token com todos os outros, cada matriz de pontuação da atenção tem um dezesseis avos do tamanho que teria para os 196 patches.

Resultados e variantes#

As representações do MAE são melhores após o ajuste fino. Na avaliação linear, um codificador ViT-L pré-treinado com MAE alcançou 73,5% de acurácia no ImageNet, abaixo dos métodos contrastivos; porém, após o ajuste fino e a transferência para detecção de objetos e segmentação no COCO, o pré-treinamento MAE superou o pré-treinamento supervisionado. A receita foi adaptada muito além de imagens estáticas: o VideoMAE mascara 90% ou mais dos patches espaço-temporais de um clipe de vídeo, o SatMAE faz pré-treinamento com imagens de satélite, e o ConvNeXt V2 leva a codificação automática mascarada às redes convolucionais com convoluções esparsas.

O pré-treinamento mascarado também está por trás de modelos fundacionais amplamente utilizados. O Modelo de Segmentação de Qualquer Coisa (SAM) da Meta funciona no pacote Ultralytics, usa um ViT pré-treinado com MAE como codificador de imagem e sua versão ViT-B, sam_b.pt, é o modelo de máscaras padrão no Ultralytics auto_annotate, que converte as caixas de um detector YOLO26 em rótulos de segmentação.

Autoencoders mascarados vs. termos relacionados#

• Autoencoder: comprime toda a entrada em um código pequeno e a reconstrói. Um MAE recebe uma entrada incompleta e precisa prever as partes que nunca viu, e a saída do codificador mantém um vetor por patch visível.
• Modelagem de imagens mascaradas: a família mais ampla à qual o MAE pertence. O BEiT prevê tokens visuais discretos para patches mascarados, e o SimMIM alimenta tokens de máscara ao codificador, enquanto o MAE prevê pixels e mantém os tokens de máscara fora do codificador.
• Aprendizado contrastivo: aprende combinando versões aumentadas da mesma imagem e não precisa de decodificador; produz representações que funcionam melhor congeladas, enquanto as representações do MAE melhoram mais com o ajuste fino.
• Arquitetura Preditiva de Embeddings Conjuntos (JEPA): também oculta regiões da imagem, mas prevê suas incorporações em vez de seus pixels, o que impede o modelo de gastar capacidade com detalhes de textura.