Repenser la dépendance aux patchs pour les autoencodeurs masqués

Rethinking Patch Dependence for Masked Autoencoders

January 25, 2024

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

cs.AI

papers.abstract

Dans ce travail, nous réexaminons les dépendances inter-patch dans le mécanisme de décodage des autoencodeurs masqués (MAE). Nous décomposons ce mécanisme de décodage pour la reconstruction de patchs masqués dans les MAE en auto-attention et attention croisée. Nos investigations suggèrent que l'auto-attention entre les patchs masqués n'est pas essentielle pour l'apprentissage de bonnes représentations. À cette fin, nous proposons un nouveau cadre de pré-entraînement : les Autoencodeurs Masqués à Attention Croisée (CrossMAE). Le décodeur de CrossMAE exploite uniquement l'attention croisée entre les tokens masqués et visibles, sans dégradation des performances en aval. Cette conception permet également de décoder uniquement un petit sous-ensemble de tokens masqués, améliorant ainsi l'efficacité. De plus, chaque bloc de décodeur peut désormais exploiter différentes caractéristiques de l'encodeur, ce qui améliore l'apprentissage des représentations. CrossMAE atteint les performances des MAE avec 2,5 à 3,7 fois moins de calculs de décodage. Il surpasse également les MAE sur la classification ImageNet et la segmentation d'instances COCO avec la même quantité de calculs. Code et modèles : 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