Ripensare la dipendenza dalle patch per gli autoencoder mascherati

Rethinking Patch Dependence for Masked Autoencoders

January 25, 2024

Autori: Letian Fu, Long Lian, Renhao Wang, Baifeng Shi, Xudong Wang, Adam Yala, Trevor Darrell, Alexei A. Efros, Ken Goldberg

cs.AI

Abstract

In questo lavoro, riesaminiamo le dipendenze inter-patch nel meccanismo di decodifica degli autoencoder mascherati (MAE). Scomponiamo questo meccanismo di decodifica per la ricostruzione di patch mascherate in MAE in self-attention e cross-attention. Le nostre indagini suggeriscono che la self-attention tra le patch mascherate non è essenziale per apprendere rappresentazioni di qualità. A tal fine, proponiamo un nuovo framework di pre-addestramento: Cross-Attention Masked Autoencoders (CrossMAE). Il decoder di CrossMAE sfrutta esclusivamente la cross-attention tra token mascherati e visibili, senza alcuna degradazione delle prestazioni downstream. Questo design consente inoltre di decodificare solo un piccolo sottoinsieme di token mascherati, migliorando l'efficienza. Inoltre, ogni blocco del decoder può ora sfruttare diverse caratteristiche dell'encoder, portando a un miglioramento nell'apprendimento delle rappresentazioni. CrossMAE eguaglia le prestazioni di MAE con un calcolo di decodifica da 2,5 a 3,7 volte inferiore. Supera inoltre MAE nella classificazione su ImageNet e nella segmentazione di istanze su COCO a parità di calcolo. Codice e modelli: https://crossmae.github.io

English

In this work, we re-examine inter-patch dependencies in the decoding mechanism of masked autoencoders (MAE). We decompose this decoding mechanism for masked patch reconstruction in MAE into self-attention and cross-attention. Our investigations suggest that self-attention between mask patches is not essential for learning good representations. To this end, we propose a novel pretraining framework: Cross-Attention Masked Autoencoders (CrossMAE). CrossMAE's decoder leverages only cross-attention between masked and visible tokens, with no degradation in downstream performance. This design also enables decoding only a small subset of mask tokens, boosting efficiency. Furthermore, each decoder block can now leverage different encoder features, resulting in improved representation learning. CrossMAE matches MAE in performance with 2.5 to 3.7times less decoding compute. It also surpasses MAE on ImageNet classification and COCO instance segmentation under the same compute. Code and models: https://crossmae.github.io