Masked Autoencoders (MAE)

Gli autoencoder mascherati (MAE) sono un metodo di apprendimento auto-supervisionato per le immagini: il modello nasconde una grande porzione casuale delle patch di un'immagine, in genere il 75%, e impara a ricostruire i pixel mancanti a partire da quelli visibili. Kaiming He e i suoi colleghi di Meta AI hanno introdotto il metodo nel documento del 2021 Masked Autoencoders Are Scalable Vision Learners, adattando alla visione artificiale la predizione delle parole mascherate di BERT. Poiché l'encoder elabora solo le patch visibili, il preaddestramento MAE è stato almeno 3 volte più veloce dell'elaborazione di tutte le patch; un modello ViT-Huge preaddestrato in questo modo ha raggiunto un'accuratezza dell'87.8% su ImageNet, il miglior risultato dell'epoca tra i metodi che usavano solo dati ImageNet-1K. L'encoder preaddestrato funge poi da backbone per classificazione, rilevamento e segmentazione.

Come funzionano gli autoencoder mascherati#

MAE divide un'immagine in una griglia di patch non sovrapposte — 196 patch da 16x16 pixel per un input di 224x224 — e ne nasconde casualmente la maggior parte. L'encoder e il decoder sono deliberatamente diversi. L'encoder, un ViT standard, opera solo sulle patch visibili, riducendo il calcolo e la memoria a circa un quarto. Un decoder leggero, composto per impostazione predefinita da 8 blocchi Transformer e con meno del 10% del calcolo dell'encoder per token, riceve le patch visibili codificate e un token maschera apprendibile condiviso per ogni posizione nascosta, e predice i valori dei pixel di ciascuna patch nascosta. La funzione di perdita è l'errore quadratico medio calcolato solo sulle patch nascoste; usare come target i pixel normalizzati per patch ha migliorato ulteriormente le feature apprese. Dopo il preaddestramento, il decoder viene scartato:

È l'elevato rapporto di mascheramento a rendere utile il compito. Le patch adiacenti di un'immagine sono molto ridondanti, quindi se ne vengono nascoste poche il modello può colmare ogni lacuna interpolando dai vicini. Quando il 75% è nascosto, deve dedurre oggetti e scene interi da frammenti, il che lo costringe ad apprendere forme, parti e contesto. BERT, invece, maschera solo il 15% delle parole, perché ciascuna parola porta già con sé una grande quantità di significato.

Il fulcro di MAE è il passaggio di mascheramento, mostrato qui in PyTorch su un batch di quattro immagini:

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

Stampa Encoder input: 49 of 196 patches per image e Reconstruction targets: 147 hidden patches per image. Di conseguenza, i livelli Transformer dell'encoder elaborano una sequenza lunga un quarto rispetto all'immagine completa; poiché l'attenzione confronta ogni token con tutti gli altri, ogni matrice dei punteggi di attenzione è un sedicesimo delle dimensioni che avrebbe con tutte le 196 patch.

Risultati e varianti#

Le feature MAE sono più efficaci dopo il fine-tuning. Con il probing lineare, un encoder ViT-L preaddestrato con MAE ha raggiunto un'accuratezza ImageNet del 73.5%, inferiore a quella dei metodi contrastivi; tuttavia, dopo il fine-tuning e il trasferimento a rilevamento degli oggetti e segmentazione su COCO, il preaddestramento MAE ha superato quello supervisionato. Questo approccio è stato adattato ben oltre le immagini statiche: VideoMAE maschera almeno il 90% delle patch spazio-temporali di una clip video, SatMAE esegue il preaddestramento su immagini satellitari e ConvNeXt V2 porta l'autoencoding mascherato nelle reti convoluzionali con convoluzioni sparse.

Il preaddestramento mascherato è alla base anche dei diffusi modelli fondazionali. Il modello per segmentare qualsiasi cosa (SAM) di Meta, disponibile nel pacchetto Ultralytics, basa il proprio encoder di immagini su un ViT preaddestrato con MAE; la sua versione ViT-B, sam_b.pt, è il modello di maschere predefinito in Ultralytics auto_annotate, che converte i riquadri di delimitazione di un rilevatore YOLO26 in etichette di segmentazione.

Autoencoder mascherati a confronto con i termini correlati#

• Autoencoder: comprime l'intero input in un codice compatto e lo ricostruisce. Un MAE riceve un input incompleto e deve predire le parti che non ha mai visto; l'output del suo encoder mantiene un vettore per ogni patch visibile.
• Modellazione di immagini mascherate: la famiglia più ampia a cui appartiene MAE. BEiT predice token visivi discreti per le patch mascherate e SimMIM immette token maschera nell'encoder, mentre MAE predice i pixel ed esclude i token maschera dall'encoder.
• Apprendimento contrastivo: apprende abbinando versioni aumentate della stessa immagine e non ha bisogno di un decoder; produce feature che funzionano meglio quando restano congelate, mentre le feature MAE migliorano maggiormente con il fine-tuning.
• Architettura predittiva a embedding congiunti (JEPA): nasconde anch'essa regioni dell'immagine, ma ne predice gli embedding anziché i pixel, evitando così che il modello sprechi capacità sui dettagli delle texture.