Rethinking Patch Dependence for Masked Autoencoders

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 09:16:01

Abstract: In this work, we examine the impact of inter-patch dependencies in the decoder of masked autoencoders (MAE) on representation learning. We decompose the decoding mechanism for masked reconstruction into self-attention between mask tokens and cross-attention between masked and visible tokens. Our findings reveal that MAE reconstructs coherent images from visible patches not through interactions between patches in the decoder but by learning a global representation within the encoder. This discovery leads us to propose a simple visual pretraining framework: cross-attention masked autoencoders (CrossMAE). This framework employs only cross-attention in the decoder to independently read out reconstructions for a small subset of masked patches from encoder outputs, yet it achieves comparable or superior performance to traditional MAE across models ranging from ViT-S to ViT-H. By its design, CrossMAE challenges the necessity of interaction between mask tokens for effective masked pretraining. Code is available [here](https://anonymous.4open.science/r/mae-cross-anon-11EB/README.md).

Main paper: 
 introduction: 
 Masked image modeling [46, 30, 61, 4] has emerged as a pivotal unsupervised learning technique in computer vision. One such recent work following this paradigm is masked autoencoders (MAE):given only a small, random subset of visible image patches, the model is tasked to reconstruct the missing pixels. By operating mostly on this small subset of visible tokens, MAE can efficiently pre-train high-capacity models on large-scale vision datasets, demonstrating impressive results on a wide array of downstream tasks [33, 38, 49] .The MAE framework employs self-attention across the entire model for self-supervised reconstruction tasks. In this setup, both masked and visible tokens engage in self-attention, not just with each other but also with themselves, aiming to generate a holistic and context-aware representation. However, the masked tokens inherently lack information. Intuitively, facilitating information exchange among adjacent masked tokens should enable the model to synthesize a more coherent image, thereby accomplishing the task of masked reconstruction and improving representation learning. A question arises, though: Is this truly the case?We decompose the decoding process of each mask token into two parallel components: self-attention with other mask tokens, as well as cross-attention to the encoded visible tokens. If MAE relies on the self-attention with other mask tokens, its average should be on par with the cross-attention. Yet, the quantitative comparison in Figure 1 .(b) shows the magnitude of mask token-to-visible token cross-attention (1.42) in the MAE decoder evaluated over the entire ImageNet validation set far exceeds that of mask token-to-mask token self-attention (0.39). This initial observation prompts two questions: 1) Is the self-attention mechanism among mask tokens in the decoder necessary for effective representation learning? 2) If not, can each patch be Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute. For each set of 5 images, from left to right, are the original image, masked image with a mask ratio of 75%, MAE [30] , CrossMAE (trained to reconstruct 25% of image tokens, or 1/3 of the mask tokens), and CrossMAE (trained to reconstruct all masked tokens). Since CrossMAE does not reconstruct them, all model outputs have the visible patches overlaid. Intriguingly, CrossMAE, when trained for partial reconstruction, can decode all mask tokens in one forward pass (shown above), indicating that the encoder rather than the decoder effectively captures global image information in its output tokens. Its comparable reconstruction quality to full-image-trained models suggests that full-image reconstruction might not be essential for effective representation learning.independently read out from the encoder output, allowing the reconstruction of only a small subset of masked patches, which in turn, accelerates the pretraining without performance degradation?In addressing these questions, we introduce CrossMAE, which diverges from MAE in three ways:1. Cross-attention for decoding. Rather than passing a concatenation of mask and visible tokens to a self-attention decoder, CrossMAE uses mask tokens as queries to read out the masked reconstructions from the visible tokens in a cross-attention decoder. In this setting, mask tokens incorporate information from the visible tokens but do not interact with other mask tokens, thereby reducing the sequence length for the decoder and cutting down computational costs.2. Independent partial reconstruction. With self-attention removed, the decoding of each mask token, based on the encoded features from visible tokens, becomes conditionally independent. This enables the decoding of only a fraction of masked tokens rather than the entire image.3. Inter-block attention. Due to the separation of visible and mask tokens, we can use features from different encoder blocks for each decoder block. Empirically, we find solely relying on the last encoder feature map for reconstruction, the design present in MAE, hurts feature learning. We propose a lightweight inter-block attention mechanism that allows the CrossMAE decoder to leverage a mix of low-level and high-level feature maps from the encoder, improving the learned representation.The analysis performed on CrossMAE led to a novel way to understand MAE. Even though the patches to be reconstructed are independently decoded, our findings demonstrate that coherent reconstruction for each masked patch can be independently read out from the encoder output, without any interactions among masked tokens in the decoder for consistency (Figure 2 ). Furthermore, the downstream performance of the model remains robust even without these interactions (Figure 1 .(c), Tables 1 and 2 ). Both pieces of evidence confirm that the encoder's output features already encapsulate the necessary global context for image reconstruction, while the decoder simply performs a readout from the encoder output to reconstruct the pixels at the location of each patch.To sum up, our main contributions are the following:1. We present a novel understanding of MAE. Our findings show that MAE reconstructs coherent images from visible patches not through interactions between patches to be reconstructed in the decoder but by learning a global representation within the encoder. This is evidenced by the model's ability to generate coherent images and maintain robust downstream performance without such interactions, indicating the encoder effectively captures global image information.2. We advocate replacing self-attention layers with a simple cross-attention readout function.Given our discovery that the encoder in MAE already captures a comprehensive global representation, we propose replacing self-attention layers in the decoder with a more efficient information readout function. Specifically, we suggest utilizing cross-attention to aggregate the output tokens of the encoder into each input token within the decoder layers independently, thereby eliminating the need for token-to-token communication within the decoder.3. CrossMAE achieves comparable or superior performance with reduced computational costs in image classification and instance segmentation compared to MAE on vision transformer models ranging from ViT-S to ViT-H. Code is available here.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The claims in the abstract are justified in the method and the experiments section.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
self-supervised learning: 
 In self-supervised representation learning, a model trains on a pretext task where the supervision comes from the input data itself without labels. Contrastive learning methods learn representations by contrasting positive and negative samples, such as SimCLR [11] , CPC [44] , MoCo [29, 12, 13] ,CLD [59] and SwAV [7] . Additionally, in BYOL [26] , iBOT [65] , DINO [8] , DINOv2 [45] , and MaskAlign [62] make a student model to imitate a teacher model without negative pairs. Generative modeling, focusing on acquiring a generative model capable of capturing the underlying data distribution, is an alternative method for self-supervised learning. VAE/GAN [35] merges the strengths of variational autoencoders and generative adversarial networks to acquire disentangled representations of data. PixelCNN, PixelVAE, and PixelTransformer [55, 27, 54] generate images pixel by pixel, taking into account the context of previously generated pixels. Masked modeling, a large subclass of generative modeling, is discussed in the following subsection. After the pre-training stage, these generative models can be finetuned for many downstream applications.Masked modeling learns representations by reconstructing a masked portion of the input. Pioneering works in natural language processing (NLP) present various such pretraining objectives. BERT [19] and its extensions [41, 34] use a bidirectional transformer and present few-shot learning capabilities from masked language modeling. GPT [47, 48, 5] , uses autoregressive, causal masking and demonstrates multi-task, few-shot, and in-context learning capabilities.Early works in computer vision, such as Stacked Denoising Autoencoders [57] and Context Encoder [46] , investigated masked image modeling as a form of denoising or representation learning.Recently, with the widespread use of transformer [20] as a backbone vision architecture, where images are patchified and tokenized as sequences, researchers are interested in how to transfer the success in language sequence modeling to scale vision transformers. BEiT [3] , MAE [30] , and Sim-Figure 3 : MAE [30] concatenates all mask tokens with the visible patch features from a ViT encoder and passes them to a decoder with self-attention blocks to reconstruct the original image. Patches that correspond to visible tokens are then dropped, and an L2 loss is applied to the rest of the reconstruction as the pretraining objective. CrossMAE instead uses cross-attention blocks in the decoder to reconstruct only a subset of the masked tokens.MIM [61] are a few of the early works that explored BERT-style pretraining of vision transformers.Compared to works in NLP, both MAE and SimMIM [30, 61] find that a much higher mask ratio compared to works in NLP is necessary to learn good visual representation. Many recent works further extend masked pretraining to hierarchical architectures [61, 40] and study data the role of data augmentation [9, 21] . Many subsequent works present similar successes of masked pretraining for video [52, 58, 22, 28] , language-vision and multi-modal pretraining [1, 39, 23] and for learning both good representations and reconstruction capabilities [60, 37] .However, BERT-style pretraining requires heavy use of self-attention, which makes computational complexity scale as a polynomial of sequence length. PixelTransformer [54] and DiffMAE [60] both use cross-attention for masked image generation and representation learning. Siamese MAE [28] uses an asymmetric masking pattern and decodes frames of a video condition on an earlier frame. In these settings, all masked patches are reconstructed. In this work, we investigate if learning good features necessitates high reconstruction quality and if the entire image needs to be reconstructed to facilitate representation learning. PCAE [36] progressively discards redundant mask tokens through its network, leading to a few tokens for reconstruction. VideoMAEv2 [58] concatenates randomly sampled masked tokens with visible tokens and uses self-attention to reconstruct the masked patches.In comparison, we minimally modify MAE with a cross-attention-only decoder and masked tokens are decoded in a conditional independent way.In addition to the prevalent use of self-attention in computer vision, cross-attention has shown to be a cost-effective way to perform pooling from a large set of visible tokens. Intuitively, cross-attention can be seen as a parametric form of pooling, which learnably weighs different features. Touvron et al. [53] replace mean pooling with cross-attention pooling and find improvement in ImageNet classification performance. Jaegle et al. [32] uses cross-attention to efficiently process large volumes of multi-modal data. Cross-attention is also widely used for object detection. Carion et al. [6] utilizes query tokens as placeholders for potential objects in the scene. Cheng et al. [16, 15] further extend this concept by introducing additional query tokens to specifically tackle object segmentation in addition to the query tokens for object detection. Distinct from thes prior works, we are interested the role of cross-anttention for representation learning in a self-supervised manner.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: The limitations of the work have been discussed in the Discussion and Conclusion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
crossmae: 
 We start with an overview of vanilla masked autoencoders in Section 3.1. Next, in Section 3.2, we introduce the use of cross-attention in place of self-attention in the decoder for testing the necessity of interaction between mask tokens for representation learning. In Section 3.3, we discuss how eliminating self-attention in the decoding process enables us to reconstruct only a subset of masked tokens, leading to faster pretraining. Finally, Section 3.4 presents our inter-block attention mechanism, which allows decoder blocks to leverage varied encoder features.Masked Autoencoders (MAE) [30] pretrain Vision Transformers (ViTs) [20] . Each image input is first patchified, and then a random subset of the patches is selected as the visible patches. As depicted in Figure 3 , the visible patches, concatenated with a learnable class token [CLS], are subsequently fed into the ViT encoder, which outputs a set of feature latents. The latent vectors, concatenated with the sum of the positional embeddings of the masked patches and the learnable mask token, are passed into the MAE decoder. The decoder blocks share the same architecture as the encoder blocks (i.e., both are transformer blocks with self-attention layers). Note that the number of tokens fed into the decoder is the same length as the original input, and the decoding process assumes that the decoded tokens depend on both visible and masked tokens. Decoder outputs pass through a fully connected layer per patch for image reconstruction. After the reconstruction is generated, the loss is applied only to the masked positions, while the reconstructions for visible spatial locations are discarded.Recall in Sec. 1 we measure the mean attention value across all attention maps over the ImageNet validation set to study the properties of MAE. We grouped the attention values by cross-attention and self-attention between visible and masked tokens. We observed that in the decoding process of an MAE, mask tokens attend disproportionately to the class token and the visible tokens (see Figure 1.(b) ). This motivates us to make design decisions and conduct experiments specifically to answer the following question: Can we simplify the decoding process by eliminating self-attention among masked tokens without compromising the model's ability to generate coherent images and perform well on downstream tasks?To address this question, we substitute the self-attention mechanism in the decoder blocks with cross-attention, using it as a readout function to decode the latent embedding from the encoder to raw pixel values. Specifically, the decoder employs multi-head cross-attention where the queries are the output from previous decoder blocks (or the sum of position embedding of the masked patches and mask token for the first decoder block). The keys and values are from the encoded features.In the most basic CrossMAE, the output from the final encoder block is used as the key and value tokens for all layers of the decoder, as illustrated in Fig. 4(a) . Further exploration in Sec.3.4 reveals that utilizing a weighted mean of selected encoder feature maps can be beneficial. The residual connections in each decoder block enable iterative refinement of decoded tokens as they progress through decoder blocks.Diverging from the original transformer architecture [56] , our decoder omits the causal self-attention layer before the introduction of multi-head cross-attention. This elimination, coupled with the fact that layer normalization and residual connections are only applied along the feature axis but not the token axis, enables the independent decoding of tokens. This design choice is evaluated in the ablation study section to determine its impact on performance.Given the disparity in the dimensions of the encoder and decoder, MAE adapts the visible features to the decoder's latent space using an MLP. However, in CrossMAE, as encoder features are integrated at various decoder blocks, we embed the projection within the multi-head cross-attention module.Cross-attention layers serve as a readout function that decodes the global representation provided in the encoder's output tokens to the pixel values within each patch to be reconstructed. However, CrossMAE does not restrict the architecture to a single cross-attention block. Instead, we stack multiple cross-attention decoder blocks in a manner more akin to the traditional transformer [56] .The fact that CrossMAE uses cross-attention rather than self-attention in the decoder blocks brings an additional benefit over the original MAE architecture. Recall that mask tokens are decoded independently and thus there is no exchange of information between them, to obtain the reconstructions at a specific spatial location, CrossMAE only needs to pass the corresponding mask tokens to the cross-attention decoder. This allows partial reconstruction in contrast to the original full-image reconstruction in the MAE architecture which needs to pass all the masked tokens as the input of the decoder blocks due to the existence of self-attention in the decoder blocks.To address the second question in Sec. 3.1, rather than decoding the reconstruction for all masked locations, we only compute the reconstruction on a random subset of the locations and apply the loss to the decoded locations. Specifically, we name the ratio of predicted tokens to all image tokens as prediction ratio (γ), and the mask ratio (p). Then the prediction ratio is bounded between γ ∈ (0, p].Because we are sampling within the masked tokens uniformly at random and the reconstruction loss is a mean square error on the reconstructed patches, the expected loss is the same as in MAE, while the variance is (p/γ) times larger than the variance in MAE. Empirically, we find that scaling the learning rate of MAE (β) to match the variance (i.e. setting the learning rate as γβ/p)) helps with model performance. Since cross-attention has linear complexity with respect to the number of masked tokens, this partial reconstruction paradigm decreases computation complexity. Empirically, we find that the quality of the learned representations is not compromised by this approach.MAE combines the feature of the last encoder block with mask tokens as the input to the self-attention decoder, which creates an information bottleneck by making early encoder features inaccessible for the decoder. In contrast, CrossMAE's cross-attention decoder decouples queries from keys and values. This decoupling allows different cross-attention decoder blocks to take in feature maps from different encoder blocks. This added degree of flexibility comes with a design choice for selecting encoder features for each decoder block. One naive choice is to give the feature of the ith encoder block to the last ith decoder (e.g., feeding the feature of the first encoder to the last decoder), in a U-Net-like fashion. However, this assumes the decoder's depth matches the depth of the encoder, which is not the case for MAE or CrossMAE.Instead of manually matching each decoder block with an encoder feature map, we make the selection learnable and propose inter-block attention for feature fusion for each decoder block (Figure 4(c) ).Analogous to the inter-patch cross-attention that takes a weighted sum of the visible token embeddings across the patch dimensions to update the embeddings of masked tokens, inter-block attention takes a weighted sum of the visible token embeddings across different input blocks at the same spatial location to fuse the input features from multiple blocks into one feature map for each decoder block.Concretely, each decoder block takes a weighted linear combination of encoder feature maps {f i } as keys and values. Specifically, for each key/value token t k in decoder block k in a model with encoder depth n, we initialize a weight w k ∈ R n ∼ N (0, 1/n). Then t k is defined asEQUATIONIn addition to feature maps from different encoder blocks, we also include the inputs to the first encoder block to allow the decoder to leverage more low-level information to reconstruct the original image. We can select a subset of the feature maps from the encoder layers instead of all feature maps.This reduces the computation complexity of the system. We ablate this in Table 3d .We show that using the weighted features rather than simply using the features from the last block greatly improves the performance of CrossMAE. Intriguingly, in the process of learning to achieve better reconstructions, early decoder blocks tend to prioritize information from later encoder blocks, while later decoder blocks focus on earlier encoder block information, as demonstrated in Section 4.5.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments: 
 We perform self-supervised pretraining on ImageNet-1K, following MAE [30] 's hyperparameter settings, only modifying the learning rate and decoder depth. The hyperparameters were initially determined on ViT-Base and then directly applied to ViT-Small, ViT-Large, and ViT-Huge. Both CrossMAE and MAE are trained for 800 epochs. We provide implementation details and more experiments in the appendix.Setup. The model performance is evaluated with end-to-end fine-tuning, with top-1 accuracy used for comparison. Same as in Figure . 2, we compare two versions of CrossMAE: one with a prediction ratio of 25% (1/3 of the mask tokens) and another with 75% (all mask tokens). Both models are trained with a mask ratio of 75% and a decoder depth of 12.Results. As shown in Table 1 , CrossMAE outperforms vanilla MAE using the same ViT-B encoder in terms of fine-tuning accuracy. This shows that replacing the self-attention with cross-attention does not degrade the downstream classification performance of the pre-trained model. Moreover, CrossMAE outperforms other self-supervised and masked image modeling baselines, e.g., DINO [8] ,MoCo v3 [14] , BEiT [3] , and MultiMAE [2] .Setup. We additionally evaluate models pretrained with CrossMAE for object detection and instance segmentation, which require deeper spatial understanding than ImageNet classification. Specifically, we follow ViTDet [38] , a method that leverages a Vision Transformer backbone for object detection and instance segmentation. We report box AP for object detection and mask AP for instance segmentation, following MAE [30] . We compare against supervised pre-training, MoCo-v3 [14] ,BEiT [4] , and MAE [30] . Cross-Attention vs Self-Attention. As shown in Table 3a , CrossMAE, with its cross-attentiononly decoder, outperforms vanilla MAE in downstream tasks as noted in Section 4.1. Additionally, combining cross-attention with self-attention does not enhance fine-tuning performance, indicating that cross-attention alone is adequate for effective representation learning.Mask Ratio and Prediction Ratio. In our experiments with different mask and prediction ratios (i.e., the ratio of mask tokens to all tokens and the ratio of reconstructed tokens to all tokens, respectively) (see Table 3b and Table 3c ), we found that our method's performance is not significantly affected by variations in the number of masked tokens. Notably, CrossMAE effectively learns representations by reconstructing as few as 15% of tokens, compared to the 100% required by vanilla MAE, with minimal impact on downstream fine-tuning performance, which shows that partial reconstruction is sufficient for effective representation learning.Inter-block Attention. Our ablation study, detailed in Table 3d , explored the impact of varying the number of encoder feature maps in our inter-block attention mechanism. We found that using only the last feature map slightly lowers performance compared to using all 12. However, even a partial selection of feature maps improves CrossMAE's performance, with the best results obtained using 6 feature maps. This indicates that CrossMAE does not require all features for optimal performance.Decoder Depth. Table 3e shows that a 12-block decoder slightly improves performance compared to shallower ones. Remarkably, CrossMAE achieves similar results to MAE with just one decoder block, demonstrating its efficiency. Our experiments in Figure 7 that models with lower prediction ratios benefit more from deeper decoders.Input Resolution. We extend CrossMAE to longer token lengths by increasing the image resolution with constant patch size. Escalating the resolution from 224 to 448 increases the token length from 197 to 785, challenging the scalability of current approaches. Thus, we opt for a CrossMAE variant with a 25% prediction ratio. In Table 3f , we observe that the classification accuracy positively correlates with the input resolution, indicating that CrossMAE can scale to long input sequences.Due to partial reconstruction and confining attention to between mask tokens and visible tokens, CrossMAE improves pre-training efficiency over MAE. Results in Table 10 show that the FLOPs reduction does translate to an 1.54× training throughput and at least 50% reduction in GPU memory utilization compared to MAE.Visualizing Per-block Reconstruction. Rather than only visualizing the final reconstruction, we have two key observations that allow us to visualize the work performed by each decoder block:1) Transformer blocks have skip connections from their inputs to outputs.2) The final decoder block's output goes through a linear reconstruction head to produce the reconstruction. As detailed in Appendix D, we can factor out each block's contribution in the final reconstruction with linearity.This decomposition allows expressing the reconstruction as an image stack, where summing up all the levels gives us the final reconstruction. As shown in Figure 5 (a,b), we observe that different decoder blocks play different roles in reconstruction, with most details emerging at later decoder blocks. This justifies the need for low-level features from early encoder blocks, motivating inter-block attention.Visualizing Inter-block Attention Maps. As shown in the visualizations of the attention maps of inter-block attention in 5(c), CrossMAE naturally leverages the inter-block attention to allow the later decoder blocks to focus on earlier encoder features to achieve reconstruction and allow the earlier decoder blocks to focus on later encoder features. This underscores the necessity for different decoder blocks to attend to different encoder features, correlating with the performance improvements when inter-block attention is used. 
discussion and conclusion: 
 In our study, we present a novel understanding of MAE, demonstrating that coherent image reconstruction is achieved not through interactions between patches in the decoder but by learning a global representation within the encoder. Based on this insight, we propose replacing self-attention layers in the decoder with a simple readout function, specifically utilizing cross-attention to aggregate encoder outputs into each input token within the decoder layers independently. This approach, tested across models ranging from ViT-S to ViT-H, achieves comparable or better performance in image classification and instance segmentation with reduced computational requirements, showcasing the potential for more efficient and scalable visual pretraining methods.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We described the compute requirements in Appendix A.6. We do not use GPUs from a cloud provider.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: This paper aims to advance the field of self-supervised learning. Like other selfsupervised learning methods, our work may have various societal implications. However, we do not believe any specific consequences need to be highlighted in this context.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not pose such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The code and datasets used in this work follow the original MAE work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: To compare the attention values for mask tokens in vanilla MAE (Figure 1 ), we trained a ViT-B/16 MAE for 800 epochs using the default hyperparameters provided in [30] . For each image, we randomly generate a 75% binary mask (m) for all tokens, with m i = 1 representing a token being masked and m i = 0 otherwise. During the forward pass of the decoder, for each self-attention operation, the attention map is stored. This means that for the default MAE, a total of 8 attention maps, each with 16 attention heads are stored. Based on the mask pattern, we calculate the outer product (m • m ⊤ ) for the self-attention among mask tokens, and m • (1 -m ⊤ ) for the cross-attention from the mask token to the visible tokens. We then calculate the average across all feature maps and attention heads for self-attention and cross-attention to get the image average values. Lastly, we averaged across the entire ImageNet validation set to obtain the final values.We tried a few implementations for inter-block attention (IBA) and found the following implementation to be the fastest and most memory-efficient. In this implementation, we combine inter-block attention for all encoder layers as a single forward pass of a linear layer. For each decoder block, we index into the output tensor to extract the corresponding feature map, and a layer norm will be applied before the feature map is fed into the decoder block. Other alternatives we tried include 1)performing separate inter-block attentions before each decoder block, and 2) 1x1 convolution on the stacked encoder feature maps.In MAE, there exists a layer norm after the last encoder feature map before feeding into the decoder.In our implementation, we only add layer norm after inter-block attention. We find that adding an additional layer norm before inter-block attention to each encoder feature map does not lead to improvements in model performance but will significantly increase GPU memory usage.The pseudo-code of inter-block attention is the following:1 class InterBlockAttention () : Additionally, we further investigate the importance of using a cross-attention decoder, where each decoder block can use different feature maps from the encoder for decoding. In this experiment, we incorporated IBA into MAE, which uses only a self-attention decoder. Specifically, we concatenate the interblock attention features with the masked tokens. We then feed the combined features into MAE's self-attention decoder. We pre-trained the model and finetuned it for Imagenet classification.The results are presented in Table . 4, where all models are pre-trained for 400 epochs. We observe that inter-block attention has negligible performance improvements for MAE, potentially because MAE only takes in one feature map in its decoder. In contrast, inter-block attention allows cross-attention layers in CrossMAE to attend to features from different encoder blocks, thanks to its decoupling of queries with keys and values.In Section 4. can be used to improve the consistency for reconstruction. Specifically, we modify the formulation in the original transformer paper [56] , where the mask/query tokens are first passed through a multihead self-attention and a residual connection before being used in the multiheaded cross-attention with the features from the encoder. The primary difference with the vanilla transformer decoder implementation [56] is we do not perform casual masking in the multi-head self-attention. Please reference Figure 6 for a more visual presentation of the method. In Table 3d , the following cases are considered. 1 feature map (row 1) does not use inter-block attention. Each decoder block only takes the last feature map from the encoder as the keys and values.For scenarios where more than one feature map is used, the output of the patch embedding (input to the ViT) is also used.In addition to the simple design of inter-block attention proposed above, we also experimented with a variant of inter-block attention by further parameterizing the attention with linear projections. Specifically, rather than directly performing weighted sum aggregation to form the features for each cross-attention layer in the decoder, we added a linear projection for each encoder feature before the feature aggregation. We denote this variant as CrossMAE+LP. As shown in the Table. 5 (with ViT-B pre-trained for 400 epochs, consistent with the setting in Table . 3), adding a linear projection slightly improves the performance. This indicates that it is possible to design variants of readout functions, such as through improved inter-block attention, to improve the feature quality of CrossMAE.Acc. (%) CrossMAE 83.3 CrossMAE + LP 83.5 Pre-training: The default setting is in Table 6 , which is consistent with the official MAE [30] implementation. As mentioned in Sec. 3.4, we scale the learning rate by the ratio between mask ratio (p) and prediction ratio (γ) to ensure the variance of the loss is consistent with [30] . Additionally, we use the linear learning rate scaling rule [25] . This results in lr = γ * base_lr * batchsize/(256 * p).For Table 1 , we use 12 decoder blocks, with mask ratio and prediction ratio both 75%, and interblock attention takes in all encoder feature maps. For the 400 epochs experiments in Table 2 , we scale the warm-up epochs correspondingly. Other hyperparameters, such as decoder block width, are the same as MAE.Finetuning: We use the same hyperparameters as MAE finetuning. We use global average pooling for finetuning. In MAE, the layer norm for the last encoder feature map is removed for finetuning, which is consistent with our pretraining setup. Please refer to Table 7 Each of the pretraining and finetuning experiments is run on 2 or 4 NVIDIA A100 80GB GPUs. The batch size per GPU is scaled accordingly and we use gradient accumulation to avoid out-of-memory errors. ViTDet [38] experiments use a single machine equipped with 8 NVIDIA A100 (80GB) GPUs.We copy the datasets to the shared memory on the machines to accelerate dataloading. We use FlashAttention-2 [18] to accelerate attention calculation. [51] 0.1 mixup [64] 0.8 cutmix [63] 1.0 drop path [31] 0.1 Similar to MAE [30] , we here ablate the masking pattern. Instead of random masking, we perform grid-wise sampling that "keeps one of every four patches" (see MAE Figure 6 ). The finetuning performance is reported in Table. 9 for ViT-B (at 400 epochs), which shows that grid masking does not lead to additional improvements in downstream performance. All experiments in Table 10 are conducted on a server with 4 NVIDIA A100 (80GB) GPUs, with the standard hyperparameters provided above for pretraining. NVLink is equipped across the GPUs. We use the default setting for MAE and set the global batch size to 4096. For CrossMAE, we also use the default setting with a prediction ratio 0.25, and this takes around 41GB memory per GPU without gradient accumulation (i.e., local batch size is set to 1024 samples per GPU). However, the same local batch size results in out-of-memory (OOM), which indicates that the total memory requirement is larger than the available memory for each GPU (80GB). To run MAE on same hardware, we thus employ gradient accumulation with a local batch size of 512 to maintain the global batch size.The benchmark runs each method and measures the average per epoch runtime as well as the max memory allocation for 10 training epochs. Our experiments in Figure 7 show that models with lower prediction ratios benefit more from deeper decoders. Our model performs on par or better when compared to MAE, with up to 3.7× lower decoder FLOPS.We propose a more fine-grained visualization approach that allows us to precisely understand the effect and contribution of each decoder block.Two key observations enable per-block visualization: 1) Transformer blocks have residual connections from their inputs to outputs. Let f i be the output and g i (•) the residual function of decoder i, sof i = f i-1 + g i (f i-1 ).2) The final decoder block's output goes through a reconstruction head h, which is linear, consisting of a layer-norm and a linear layer, to produce the reconstruction. With D as the decoder depth, f 0 the initial input, and y the final output, y is recursively defined asy = h(f D-1 + g D (f D-1)), which simplifies due to the linearity of h:y = h(f 0 + g 1 (f 0 ) + • • • + g D (f D-1 )) = h(f 0 ) Pos Embed. + Mask Token + h(g 1 (f 0 )) Block 1 + • • • + h(g D (f D-1 )) Block DThis decomposition allows us to express the reconstruction as an image stack, where the sum of all the levels gives us the final reconstruction. We present the visualization in Figure 5 .NeurIPS Paper ChecklistJustification: The research conforms to the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: After the rebuttal and discussion period, the submission received mixed reviews. The AC carefully reviewed the paper and agrees with the concerns raised by Reviewer yYWM, particularly regarding the lack of experimental comparisons with state-of-the-art MAE approaches, especially those that also consider inter-patch relations during self-supervised learning. Simply comparing the model's performance with the original MAE is insufficient, as the presented results did not demonstrate a significant performance improvement.

Given this concern, the AC could not support acceptance. The authors are however strongly encouraged to revise the submission based on the feedback for consideration at a future venue.

Reviews: review: paper_summary: paper_summary: This paper presents CrossMAE, a methodology for improving pre-training efficiency over that of MAE for an encoder. The paper motivates its approach by presenting visual evidence that, in standard MAE pre-training, masked tokens attend to other masked tokens significantly less than to non-masked (aka, visible) tokens. Using this motivation, the paper then presents CrossMAE, which differs from MAE largely in that it replaces the MAE self-attention with cross-attention between the masked tokens and a learnable weighted combination of the encoder feature maps. This aspect decouples queries from keys and values (which is not the case in MAE), which the paper then exploits to allow only some (but not necessarily all) mask tokens to be used during reconstruction to pre-train the model. The paper presents an analysis of which encoder block features are optimal to cross attend with each decoder block, and it presents ablation studies on multiple design decisions. Finally, it presents visual and fine-tuning results showing comparable performance to MAE and similar methods.
strength_weakness: strength_weakness: 1) In Fig 1b, IIUC, for one particular mask token, the two $\mu$'s are the respective attention values averaged over all transformer blocks and all masked/non-masked tokens. If this is the case, my concern is that by averaging over all transformer blocks, variations in the attention is being hidden. Naively, I would think that for early blocks, the attention due to masked tokens would be small (as the paper concludes) but becomes larger for the later blocks (since now the masked tokens have actual useful signal in them). Did you consider this?

2) I do not follow why CrossMAE does not need an MLP at the end to convert final decoder tokens back to raw pixels. Line 218 says that the inputs to the first encoder block are included in the feature maps for cross attentions. Does this cause a final MLP to not be used?

3) Less critical:
  3a) Fig 1b should point the reader to Section A.1. I spent much of my reading confused about what $\mu$ is.
  3b) Fig 4a should have a different number of decoder layers than encoder layers. When I saw this figure, I immediately wondered why a decoder block wasn't being paired with feature maps from its "partner" encoder. I had to wait until lines 204-207 to get an explanation of why this doesn't work.
  3c) Line 187 references a "second question" in Sec 3.1, which doesn't exist as far as I can tell.
  3d) Fig 4a shows the "vanilla" version of Cross MAE, where the final encoder layer feature maps are attended with all decoder layers. But the paper presents results exclusively (?) on the version that uses a learned combination of the feature maps. Anyway, the figure confused me. Maybe I just didn't understand what the solid arrows vs dotted ones are supposed to represent.
questions: questions: See "Weaknesses".
limitations: limitations: No weaknesses are specifically addressed. But as this paper is essentially an optimization to MAE, I'm not sure this question is relevent.
paper_summary: paper_summary: The paper introduces a novel pre-training approach called CrossMAE. Instead of concatenating the masked and visible tokens for the decoder, the authors add cross-attention to decode the masked tokens by using them and the visible patch embeddings as separate inputs to the decoder. Further, the authors introduce a method to only partially reconstruct the masked patches, and leverage inter-bock attention to fuse feature across different layers.
strength_weakness: strength_weakness: - I feel like the paper is missing a more structure ablation of the individual contributions. I think the paper would benefit from having a simple table where all contributions are added sequentially to better identify the performance effect of the individual contributions as in:
	MAE X.X
	+ Cross-Attn X.X
	+ Partial Reconstruction X.X
	+ Inter-Block Attn X.X
- As can be observed from Table 3 c), the final setting (underlined) of the prediction ratio, 0.75, turns out to be exactly the same as the optimal masking ratio, 0.75. If I understood correctly, this means that in practice, CrossMAE works best when it predicts all tokens that were masked, not just a fraction of them. Only predicting part of the masked tokens was previously listed as a contribution. Therefore, I don’t understand how this additional hyper parameter provides any benefit for better downstream performance. Maybe I’m missing something and this be cleared up by answering the previous point.
- All models are only trained for 800 epochs. The original MAE reaches peak performance at 1600 epochs. For a thorough comparison, it would be necessary to also train CrossMAE for 1600 epochs and see if the performance gains sustain, or if performance has peaked at 800 epochs.
- Table 1 is missing the CrossMAE ViT-H with 75% masking ratio
- Contribution 2 and 3 don’t seem to be as well motivated in the introduction in comparison to Contribution 1
- Better performance is listed as a contribution. IMO this is not a contribution, rather a result of the technical contributions
questions: questions: I like the idea and motivation of the paper. It starts from an interesting observation of the vanilla MAE, and aims to improve this aspect. Unfortunately, it is not fully clear which of the proposed contributions actually have an impact on performance. Table 3a) shows that adding Cross-Attn improves downstream performance. But since the authors choose the masking ratio to be the same as the prediction ratio, there doesn’t seem to be an improvement resulting from the second contribution. The effect of improved computation complexity only exists if prediction ratio < masking ratio. Lastly, according to Table 3 d), with the right number of fused feature maps, inter-block attention CAN improve the model, but the authors choose 12 as their default number of fused feature maps, which doesn’t improve performance over just adding Cross-Attn.

Concretely, I think the following additions could improve the paper: 
- Introduce a comprehensive analysis of the individual contributions’ impact on performance, and also computational complexity if you want to highlight that, in a similar manner as proposed above
- Train both models for 1600 epochs and evaluate if the performance increase can be sustained 

I’m willing to increase my score if my concerns are adequately addressed, and/or if the other reviewers list further convincing arguments for accepting the paper.
limitations: limitations: The authors have sufficiently addressed the limitations of their approach.
paper_summary: paper_summary: This paper reveals the role of inter-patch dependencies in the decoder of MAE on representation learning. The paper shows that MAE achieves coherent image reconstruction through global representations learned in the encoder rather than interactions between patches in the decoder. Based on this, the authors propose CrossMAE, which only utilizes cross-attention in the decoder.
strength_weakness: strength_weakness: 1. Idea/Novelty
- The claim that MAE reconstruction is achieved through global representation learning within the encoder rather than interactions between patches needs more support. Recent studies linking MAE to contrastive learning have found that the receptive field of specific mask tokens in the decoder is relatively small. Could the role of mask tokens in the decoder be to capture local area information? This might explain the smaller attention magnitude of masked tokens compared to visible tokens in Figure 1(b). 
- There is a concern that without self-attention (i.e., with the proposed method), the observation that authors made on the vanilla MAE may no longer be valid. Additional explanation on this point is necessary as this observation is the main motivation for suggesting CrossMAE.

2. Additional justification
- Effectiveness of using a subset of mask tokens as queries: Unlike the traditional architecture, this method uses only a subset of mask tokens as queries. Detailed analysis and interpretation are needed on why this is effective. 
- Performance differences when using the entire set of mask tokens versus a subset (and what percentage of mask tokens is used) should be reported.

3. Experiment
- For a fair comparison, CrossMAE's performance should be evaluated using the same setting as the original MAE, especially regarding the fine-tuning recipe.
- The current experimental results do not convincingly demonstrate the effectiveness of the method. For classification tasks, only the linear-probing and fine-tuning results on IN1K are reported. Following the previous works, classification on various downstream datasets should be also considered.
- For generalizability, evaluation on another task like semantic segmentation (e.g. on ADE20K) would be useful to verify that the suggested method learns the generalizable feature representation.
questions: questions: Please refer to the weakness section.
limitations: limitations: The authors have not discussed limitations of this work except for the very last sentence of section 5, indicating that they have discussed limitations in this section in the questionnaire #2. It is strongly recommended to disclose more detailed limitations of the proposed work.

Comments: title: Gentle reminder - 1 day left for the author-reviewer discussion
Response to rebuttal
A gentle reminder - 2 days left for the author-reviewer discussion
A gentle reminder - 2 days left for the author-reviewer discussion
Additional experiment for CrossMAE ViT-H with 75% masking ratio
comment: Dear reviewer,

Please let us know if our response addressed your concerns and/or if you have any additional questions. We would be happy to answer them before the end of the discussion period.

Thanks!

Authors
Thank you for your detailed responses to my concerns. From W1, I conclude that my intuition is far from always correct. I appreciate the study you made.
Thank you for your tendency to increase your score! Please feel free to let us know if any other question arises before the rebuttal deadline, and we would be happy to answer them.
Dear authors,
Thanks for addressing my concerns regarding your paper.
For my final rating, I will take into consideration the arguments which arise during the upcoming AC-Reviewers discussion phase. Overall, I lean towards increasing my score.

 Best regards!
Dear reviewer,

We wanted to ask if you had a chance to check our response with additional details and clarifications, analyses on transfer learning to different datasets and task, and an updated limitations section.

Please also consider checking our general response, which includes a more detailed analysis of the improved efficiency of CrossMAE. 

Please let us know if we addressed your concerns and/or if any further information or experiments would be helpful. We would be happy to provide them.

Many thanks!
Authors
Dear reviewer,

We wanted to ask if you had a chance to check our response with additional details and clarifications, a more structured ablation study, a comparison of model performance at 1600 epochs, CrossMAE ViT-H with 75% masking ratio, and an updated list of contributions.

Please also consider checking our general response, which includes a more detailed analysis of the improved efficiency of CrossMAE. 

Please let us know if we addressed your concerns and/or if any further information or experiments would be helpful. We would be happy to provide them.

Many thanks!
Authors
> [W4] Table 1 is missing the CrossMAE ViT-H with 75% masking ratio

Our CrossMAE ViT-H run with 75% masking ratio just finished. The updated comparison for ViT-H is shown in the table below. **Our method, either with 25% or 75% masking ratio, surpasses MAE in terms of ImageNet fine-tuning performance on ViT-H, justifying the scalability of our method.** The result further justifies that reconstructing only 25% of the masked tokens brings large efficiency gains, as analyzed in the general response, but has a marginal impact on downstream performance.

| Method          | ViT-S | ViT-B | ViT-L | ViT-H |
|-----------------|-------|-------|-------|-------|
| Supervised | 79.0  | 82.3  | 82.6  | 83.1  |
| MAE        | 78.9  | 83.3  | **85.4**  | 85.8  |
| CrossMAE (25%)  | 79.2  | 83.5  | **85.4**  | 86.3  |
| CrossMAE (75%)  | **79.3**  | **83.7**  | **85.4**  | **86.4**     |