Abstract
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.7$\times$ 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
Overview
-
The paper introduces CrossMAE, a new framework for masked autoencoders in vision tasks that focuses on decoding efficiency and representation quality.
-
CrossMAE questions the role of self-attention among masked patches, suggesting self-attention to visible tokens might be sufficient.
-
By leveraging only cross-attention, CrossMAE saves computational resources and shows no decrease in performance on downstream tasks.
-
Empirical tests show CrossMAE matches or exceeds existing MAE models in performance while reducing decoding computation by 2.5 to 3.7 times.
-
The findings suggest potential for scaling and efficiency improvements in pre-trained vision models, encouraging a re-evaluation of self-attention's role.
Introduction
Masked Autoencoders (MAE) have gained prominence in unsupervised learning for computer vision, offering efficient pre-training of large-scale models. These models traditionally employ multi-headed self-attention throughout, wherein both visible and masked tokens exchange information. However, recent empirical evidence raises questions about the necessity of this self-attention paradigm, specifically the role of self-attention among the masked patches. CrossMAE, introduced in the reviewed paper, critically analyses and restructures the decoding mechanism of MAE, with a focus on efficiency and representation quality.
Analysis of Mask Tokens
The core inquiry starts with the attention mechanism applied in the decoder of MAE, distinguishing between self-attention among masked tokens and cross-attention where masked tokens attend to visible ones. CrossMAE's preliminary results indicate that masked tokens disproportionately attend to visible tokens rather than other masked tokens, suggesting that the self-attention among masked patches may not contribute significantly to the quality of the learned representation. This observation raises important questions: is self-attention among mask tokens necessary for effective representation learning, and can decoders be designed to reconstruct only a partial set of masked patches without diminishing performance?
CrossMAE Design
CrossMAE proposes a novel framework leveraging only cross-attention for masked patch reconstruction, removing the necessity for masked tokens to attend to each other. This adjustment significantly reduces computational resources in the decoding process, while empirically no decrease in downstream task performance is observed. In addition to restricting attention, CrossMAE introduces partial reconstruction, which allows independent decoding of each masked token, creating the option to reconstruct a subset of masked patches. Moreover, CrossMAE's decoder operates dynamically, utilizing different features from various encoder blocks for each decoding block, which contrasts MAE's static usage of the last encoder feature map. These combined strategies lead to substantial improvements in decoding efficiency and suggest an enhanced capacity for representation learning.
Empirical Validation and Efficiency Gains
The CrossMAE framework, with its evolved architecture, undergoes rigorous empirical validation. It demonstrates its capability to match or surpass MAE in performance metrics while simultaneously achieving a reduction in decoding computation by a factor of 2.5 to 3.7. On the ImageNet classification and COCO instance segmentation tasks, CrossMAE shows superior performance under identical computational constraints. Additionally, models trained under the CrossMAE framework scale favorably, suggesting that the disentangled design has implications for further scalability improvements in representation learning tasks.
Conclusion
The CrossMAE study presents compelling evidence that the decoding process of masked autoencoders can be made more computationally efficient without compromising the representational abilities of the model. The findings encourage revisiting the role of self-attention in vision pre-training. The improved design lends itself to scaling across longer input sequences, therefore expanding the realm of tasks and datasets that can be effectively processed. With CrossMAE, we step into an era where efficient pre-training of visual representations is not only desirable but eminently achievable.
Create an account to read this summary for free: