Rethinking Patch Dependence for Masked Autoencoders

Thank you for the review and we appreciate these suggestions. We have performed the requested experiments and will revise the paper accordingly.

[W1] “The claim that MAE reconstruction is achieved through global representation learning within the encoder rather than interactions between patches needs more support.”

The interpretation seems inconsistent with what was proposed in the work. We demonstrate that interactions between masked and visible patches are the primary contributors to MAE reconstruction, while interactions among masked patches contribute marginally. Figure 1(b) supports this, showing masked tokens attend much more strongly to visible tokens than to other masked tokens.

We further validated this with two test-time ablations:

1. Disabling self-attention among masked tokens still produces valid reconstructions (Figure I for Reviewer yYWM).
2. Disabling cross-attention fails to reconstruct the image (Figure II for Reviewer yYWM).

These results confirm that vanilla MAE primarily uses cross-attention for reconstructions, not self-attention between masked patches.

[W1] “Could the role of masked 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).”

We apologize for the potential misunderstanding. The hypothesis that each masked token attends to tokens nearby is also supported by our empirical analysis in Figure 1. However, this observation does not contradict our claims that 1) all the encoder features of visible patches collectively construct a representation of the whole image, including the parts that were masked out, and 2) the masked patches attend to visible patches more than they attend to other masked patches.

For the first claim, we note that each masked patch can be reconstructed by attending to the encoder features of the visible patches, indicating the encoder features for the visible parts collectively contain information for the whole image. In this way, we can independently decode each masked patch without considering the pixel values of other masked patches. We will revise the wording to reduce confusion.

For the second claim, although the masked tokens often attend to tokens nearby, they still attend to visible tokens much more strongly, justified by Figure 1(b).

[W2] “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.”

We observed that in vanilla MAE, self-attention among masked tokens is much weaker than cross-attention between masked and visible tokens. Converting the self-attention MAE to a cross-attention-only decoder changes the attention mask to allow only attention from masked tokens to visible tokens (i.e., equivalent to setting self-attention magnitude to 0), aligning with our previous observation.

[W3] “Detailed analysis and interpretation are needed on why (using only a subset of masked tokens as queries) is effective.”

The intention and justification for using a prediction ratio that is different from the mask ratio are presented in the general response.

[W4] Performance differences when using the entire set of masked tokens versus a subset (and what percentage of masked tokens is used) should be reported.

We will revise to emphasize that the downstream performance tradeoff is provided in Table 1 and ablated in Table 3(c). The FLOPS efficiency tradeoff is provided in the general response and will be added to our work.

[W5] For a fair comparison, CrossMAE's performance should be evaluated using the same setting as the original MAE, especially regarding the fine-tuning recipe.

We use the same fine-tuning and evaluation recipe as outlined in MAE. Hyperparameters such as decoder block width and the number of training epochs are the same as MAE for a fair comparison.

[W6] “Following the previous works, the classification on various downstream datasets should be also considered.”

We provide more experiments on iNaturalist2019 and Places365 (ViT-B, pretrained for 800 epochs). We find that CrossMAE performs comparably to MAE for transfer learning.

[W7] “For generalizability, evaluation on another task like semantic segmentation (e.g. on ADE20K).”

In addition to COCO instance segmentation, we provide semantic segmentation results on ADE20K, which further demonstrates that CrossMAE learns generalizable features.

[Limitations] “The authors have not discussed the limitations of this work except for the very last sentence of section 5”

We will add more limitations in the revision:

While CrossMAE improves the efficiency of pretraining, it still inherits the limitations of MAE. For example, although CrossMAE performs on par or better than MAE with the same hyperparameters, as ablated in Table 3, finding the optimal masking and prediction ratio can require more experimentation on new datasets. In addition, since CrossMAE still follows a masked reconstruction objective, the learned model predicts content based on the training dataset and will reflect biases in those data. Finally, both MAE and CrossMAE only work on vision transformers and their variants, while adapting them to CNN-based methods is non-trivial and may require custom CUDA kernels to maintain efficiency.

Given the clarifications we’ve provided and the promising results in many requested experiments, we kindly ask if you would consider raising your assessment score. Additionally, please let us know if there are any new concerns or further questions we can address for you!

We want to thank the reviewer for the detailed review. We provide responses via the discussion below.

The most critical concern that the reviewer had was outlined in Weakness 2:

[W2] I don’t understand how the prediction ratio provides any benefit for better downstream performance.

We apologize for the confusion regarding why we introduced a varying prediction ratio (partial reconstruction). Partial reconstruction does not intend to improve the performance in terms of downstream accuracy. Instead, it allows our user to control and significantly reduce the computation required during training in terms of runtime and GPU memory usage with minimal impact on downstream performance, as explained in L196-197. Additionally, in the general rebuttal above, we present a derivation of the computational complexity as well as runtime statistics.

In addition to the absolute downstream performance that the reviewer focused on throughout the review, one of the ways to improve model performance is making training more efficient (i.e. reaching a similar downstream classification accuracy and segmentation accuracy as baselines with less training time and compute).

Practically, the outlined improvement in efficiency not only allows longer and more exploratory experiments to be run for improved downstream accuracy within the same compute budget but also makes pre-training more feasible in settings where compute resources are more constrained (Table 10). We hope that the findings we had in the paper along with the improved pre-training efficiency can lead to more future research in improving visual pre-training.

Given this, we address each of the questions individually:

[W1/Q1] “missing a more structure ablation of the individual contributions”

We refactored Table 3, the individual contribution’s impact on performance. All time trials are averaged over 10 epochs with FlashAttention 2 enabled. Note that the runtime is benchmarked with 2x NVIDIA A100 GPUs while Table 10 uses 4x A100 (L 584), so the runtime roughly doubles compared to Table 10.

* Since in cross-attention, the keys and queries can take different values. This enables each decoder block to use different encoder features through inter-block attention. This would otherwise not be possible with self-attention decoder blocks in vanilla MAE.

We will incorporate this table to improve the clarity of our work’s contributions and advantages.

[W3/Q2] “Train both models for 1600 epochs”

We provide an experiment at 1600 Epoch in the table below. Both MAE and CrossMAE are trained with the same hyperparameters and on the ViT-B architecture. CrossMAE still performs comparably while being faster at pre-training.

[W4] Table 1 is missing the CrossMAE ViT-H with 75% masking ratio

Due to resource constraints and limited time for the rebuttal, we are not able to finish training ViT-H at 75% masking ratio. However, we did show that even at 25% masking ratio our ViT-H performance surpasses the performance of MAE (86.3% vs 85.8%), which is a setting requiring even less compute compared to 75% masking ratio.

[W5,6] “Contribution 2 and 3 don’t seem to be as well motivated in the introduction in comparison to Contribution 1; better performance (contribution 3) is a result of the technical contributions”

We have restructured the list of contributions as below:

1. (left unchanged) 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 and transmits global image information.

2. Given our discovery that the encoder in MAE already captures a comprehensive global representation, we propose replacing self-attention layers with 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. Finally, we leverage additional properties of cross-attention to achieve an even better performance-efficiency trade-off. CrossMAE's ability to independently reconstruct masked tokens allows us to process only a fraction of the masked patches during training, significantly improving efficiency. Furthermore, the use of cross-attention enables different decoder blocks to utilize distinct encoder features, enhancing the performance-compute trade-off through inter-block attention. This approach achieves comparable or superior results in image classification and instance segmentation tasks across various model sizes (from ViT-S to ViT-H) while reducing computational demands compared to MAE.

In light of the clarifications and analyses we provide, we would like to ask if you might be open to increasing your assessment score, and if there are any additional concerns or questions that we can address for you!

[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.

Thank you for your valuable questions and suggestions! We provide responses via the discussion below:

[W1] “By averaging over all transformer blocks, variations in the attention may be hidden. Naively, (the reviewer) would think that for early blocks, the attention due to masked tokens would be small but becomes larger for the later blocks.”

This is indeed an interesting hypothesis! We test this hypothesis on the pre-trained MAE and report the cross-attention and self-attention magnitude at each individual block in the table below.

Based on these results, we do not observe a significant increase in the magnitude of the attention to masked tokens for the later decoder blocks. In addition, our observed pattern that “the magnitude of attention is larger in masked tokens’ cross-attention to visible tokens than in masked tokens’ self-attention” is shared across all decoder layers.

[W2] “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?

Sorry for the confusion! We do need a final MLP to convert the final decoder token back to raw pixels, which is consistent with MAE.

Line 218 refers to the process of passing the aggregated feature of each encoder block to each decoder block rather than converting the final decoder features to raw pixel values, where we treat the input feature after patchfication and linear projection as one of the encoder features to be aggregated. We will update the paper to clarify this further.

[W3] 3.a-d

Thank you so much for pointing these out! We have made the following changes to the paper to provide clarifications.

For 3.a: “Fig 1b should point the reader to Section A.1. I spent much of my reading confused about what μ is.” We updated the third paragraph to clarify what μ is and updated the figure captions. Now the third paragraph and the caption reads as below (with main changes in bold):

Line 28-33: We decompose the decoding process of each masked token into two parallel components: self-attention with other masked tokens, as well as cross-attention to the encoded visible tokens. If MAE relies on self-attention with other masked tokens, its average should be on par with the cross-attention. Yet, the quantitative comparison in Figure 1.(b) shows the average magnitude of masked token-to-visible token cross-attention (μ=1.42) in the MAE decoder evaluated over the entire ImageNet validation set far exceeds that of masked token-to-masked token self-attention (μ=0.39). We describe the attention calculation in Section A.1.

Figure 1 Caption: (B) MAE reconstructs a masked token (marked by a blue star) by attending to both masked tokens (B.Left) and visible tokens (B.Right). A quantitative analysis of the ImageNet validation set reveals that masked tokens in MAE attend disproportionately to visible tokens compared to other masked tokens (average attention magnitude μ=1.42 vs μ=0.39, respectively). This observation raises questions about the necessity of attention among masked tokens themselves.

For 3.b/d: “Fig 4a: should have a different number of decoder layers than encoder layers … Fig 4a shows the "vanilla" version of CrossMAE, but the paper presents results exclusively on the version that uses a learned combination of the feature maps.”

Thanks for the suggestions. We have updated Figure 4a accordingly and attached the updated figure in the PDF in the general rebuttal (see the figure for Reviewer gKk3).

For 3.c: “Line 187 references a "second question" in Sec 3.1, which doesn't exist as far as I can tell.” Now Line 187 reads: “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.”

Thank you once again for your positive feedback! If you have any further questions, please don’t hesitate to ask!