Masked Autoencoders (MAE)

Los autocodificadores enmascarados (MAE) son un método de aprendizaje autosupervisado para imágenes: el modelo oculta una gran proporción aleatoria de los parches de una imagen, normalmente el 75 %, y aprende a reconstruir los píxeles que faltan a partir del resto visible. Kaiming He y sus colegas de Meta AI presentaron el método en el artículo de 2021 Los autocodificadores enmascarados son aprendices de visión escalables, adaptando la predicción de palabras enmascaradas de BERT a los transformadores de visión. Como el codificador procesa solo los parches visibles, el preentrenamiento de MAE fue tres veces más rápido —o incluso más— que procesar todos los parches, y un modelo ViT-Huge preentrenado de esta forma alcanzó una precisión del 87,8 % en ImageNet, el mejor resultado de la época entre los métodos que usaban únicamente datos de ImageNet-1K. Después, el codificador preentrenado sirve como backbone para clasificación, detección y segmentación.

Cómo funcionan los autocodificadores enmascarados#

MAE divide una imagen en una cuadrícula de parches no superpuestos —196 parches de 16x16 píxeles para una entrada de 224x224— y oculta aleatoriamente la mayoría. El codificador y el decodificador son deliberadamente desiguales. El codificador, un ViT estándar, funciona solo con los parches visibles, lo que reduce su cómputo y memoria a aproximadamente una cuarta parte. Un decodificador ligero —de forma predeterminada, 8 bloques de transformador con menos del 10 % del cómputo del codificador por token— recibe los parches visibles codificados y un token de máscara aprendible compartido para cada posición oculta, y predice los valores de píxel de cada parche oculto. La función de pérdida es el error cuadrático medio únicamente en los parches ocultos, y usar píxeles normalizados por parche como objetivos mejoró aún más las características aprendidas. Tras el preentrenamiento, se descarta el decodificador:

La elevada proporción de enmascaramiento es lo que hace útil la tarea. Los parches de imagen cercanos son muy redundantes, así que, si solo se ocultan unos pocos, el modelo puede rellenar cada hueco interpolando a partir de sus vecinos. Con un 75 % oculto, tiene que inferir objetos y escenas completos a partir de fragmentos, lo que le obliga a aprender formas, partes y contexto. En cambio, BERT enmascara solo el 15 % de las palabras, porque cada palabra ya contiene un significado denso.

El núcleo de MAE es su paso de enmascaramiento, que aquí se muestra en PyTorch con un lote de cuatro imágenes:

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

Imprime Encoder input: 49 of 196 patches per image y Reconstruction targets: 147 hidden patches per image. Por tanto, las capas de transformador del codificador procesan una secuencia que tiene una cuarta parte de la longitud de la imagen completa y, como la atención compara cada token con todos los demás, cada matriz de puntuaciones de atención tiene una dieciseisava parte del tamaño que tendría con los 196 parches.

Resultados y variantes#

Las características de MAE son más eficaces tras el ajuste fino. Con evaluación lineal, un codificador ViT-L preentrenado con MAE alcanzó una precisión del 73,5 % en ImageNet, por debajo de los métodos contrastivos; sin embargo, tras el ajuste fino y la transferencia a la detección de objetos y la segmentación en COCO, el preentrenamiento de MAE superó al preentrenamiento supervisado. La receta se ha adaptado mucho más allá de las imágenes estáticas: VideoMAE enmascara el 90 % o más de los parches espacio-temporales de un clip de vídeo, SatMAE preentrena con imágenes por satélite y ConvNeXt V2 aplica la codificación automática enmascarada a las redes convolucionales mediante convoluciones dispersas.

El preentrenamiento enmascarado también es la base de muchos modelos fundacionales ampliamente utilizados. El Segment Anything Model (SAM) de Meta, que funciona en el paquete Ultralytics, construye su codificador de imágenes sobre un ViT preentrenado con MAE, y su versión ViT-B, sam_b.pt, es el modelo de máscaras predeterminado en Ultralytics auto_annotate, que convierte las cajas de un detector YOLO26 en etiquetas de segmentación.

Autocodificadores enmascarados frente a términos relacionados#

• Autocodificador: Comprime toda la entrada en un código pequeño y la reconstruye. Un MAE recibe una entrada incompleta y debe predecir las partes que nunca ha visto; la salida de su codificador conserva un vector por cada parche visible.
• Modelado de imágenes enmascaradas: La familia más amplia a la que pertenece MAE. BEiT predice tokens visuales discretos para los parches enmascarados y SimMIM introduce tokens de máscara en el codificador, mientras que MAE predice píxeles y deja los tokens de máscara fuera del codificador.
• Aprendizaje contrastivo: Aprende asociando versiones aumentadas de una misma imagen y no necesita decodificador; produce características que funcionan mejor sin ajustar, mientras que las características de MAE mejoran más con el ajuste fino.
• Arquitectura predictiva de incrustaciones conjuntas (JEPA): También oculta regiones de la imagen, pero predice sus incrustaciones en lugar de sus píxeles, lo que evita que el modelo dedique capacidad a los detalles de textura.