Rethinking Patch Dependence for Masked Autoencoders
Authors: Letian Fu
Long Lian
Renhao Wang
Baifeng Shi
XuDong Wang
Adam Yala
Trevor Darrell
Alexei A Efros
Ken Goldberg
Venue: Accepted by TMLR ()

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. This approach achieves comparable or superior performance to traditional MAE across models ranging from ViT-S to ViT-H and significantly reduces computational requirements. By its design, CrossMAE challenges the necessity of interaction between mask tokens for effective masked pretraining. Code and models are publicly available: https://crossmae.github.io/

Published in Transactions on Machine Learning Research (04/2025)
Rethinking Patch Dependence for Masked Autoencoders
Letian Fu" UC Berkeley
mww.huullckoo@@@eerkkee........
Long Lian" UC Berkeley
longliond@bereeley.ce edu
Renhao Wang UC Berkeley
rw435@@berkeley.com edu
Baifeng Shi UC Berkeley
edu
Xudong Wang UC Berkeley
udwanggeecs..eere werkseey..bbkk..............ee.. edu
Adam Yala² UC Berkeley, UCSP
yala@berkeley.c edu
Trevor Darrellt UC Berkeley
bresondarrell@barrkkeyy edu
Alexei A. Efros UC Berkeley
efros@eecs. berkkeey.sberrkeley........ edu
Ken Goldberg UC Berkeley
goldbery@bbereeeey.coo edu
Reviewed on OpenReview:
https: // openreview. net/ forum? id= JB2
Abstract
In this work, 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 protraining framework: cross-attention masked autoencoders (CrossMAE). This framework employs only cross-attention in the decoder independently read out reconstructions for small subset of masked patches from encoder outputs. This approach achieves comparable superior performance to traditional MAE across models ranging from VIT-S to VIT-H and significantly reduces computational requirements. By its design, CrossMAE challenges the necessity of interaction between mask tokens for effective masked pretraining. Code and models are publicly available: https: /etts://croosnee wit.aae.iittuuub.... io/
1 Introduction
Masked image modeling Pathhkk et al. 2016 He et al. 2022 Xie et al. 2022 Bao et al. 2022 has emerged
pivotal unsupervised learning technique in computer vision. One such recent work following this paradigm
Equual contribution. Equal advising.

---

Published in Transactions on Machine Learning Research (04/2025)
MAE Attention Visualization
Reconstruction Comparison
Input Image
Masked Input
Mask Tokens 0.39
Visible Tokens 1.42
MAE Acc: 83.3%
CrossMAE Acc: 83.5%
Figure 1: Method Overview. (A) Masked autoencoder (MAE) starts by masking random patches of the input image. (B) To reconstruct mask loken (marked by the blue star), MAE attends to both the masked lokens (B.Left) and the visible tokens (B.Right) A quantitative comparison over the ImageNet validation shows that the masked tokens in MAE disproportionally attend to the visible tokens (1.42 vs 0.39), questioning the necessity of attention within mask tokens. (C) We propose CrossMAF, the masked patches are reconstructed from only the cross attention between the masked tokems and the visible tokens. Surprisingly, CrossMAF attains the same or better performance than MAE on ImageNet classification and COCO instance segmentation.
is masked autoencoders (MAE): given only 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 (Kirillov et al. 2023 Li et al. 2022b Radosavovic et al. 2023
The MAE framework employs self-aftention across the entire model for self-supervised roconstruction 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 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 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 the encoded visible tokens. 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 (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 loken-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 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, 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- attenttonnio 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 dccoding 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 leverage a mix of low-level and high-level feature maps from the encoder, improving the learned representation.
2

---

Published in Transactions on Machine Learning Research (04/2025)
Original
CressMAE (25%)
Original
CrosEMAE (25%)
Figure 2: Example reconstructions of ImageNet validation images. For each set of 5 images, from left to right, are the original image, masked image with mask ratio 75%, MAE He et al. (2022), CrossMAE (trained 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, CrossMAB 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.
The analysis performed on CrossMAE led to novel way understand MAE. Even though the patches to be reconstructed are independently decoded, our findings demonstrate that coherent reconstruction for each masked patch 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 (c), Tables | 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
the following:
1. We present 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 represenlation 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 simple cross-attention readout function. Given discovery that the encoder in MAE already captures comprehensive global representation, we propose replacing self-attention layers in the decoder with more efficient information readout function. Specifically, we suggest utilizing cross-attention to aggregate the output tokens of the encoder into cach input token within the decoder layers independenlly, thereby eliminating the need for token-lo-token communication within the decoder.
3. CrossMAE achicves 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. H. Code available here
3

---

Published in Transactions on Machine Learning Research (04/2025)
2
Related Works
2.1
Self-Super 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 (Chen et al. 2020bb, CPC (Oord et al., 2018). MoCo (He
ct al. 2020 Chen et al. 2020c; 2021a), CLD (Wang et al. 2022)) and SwAV (Caron et al. 2020). Additionally,
in BYOL |rrill et al. 2020), iBOT (Zhou et al. 2021). DINO (Caron et al. 2021), DINOv2 (Oquab et al.,
2023), and MaskAlign Xue et al. 2023) make student model to imitate a teacher model without negative pairs.
Generative modeling, focusing on acquiring generative model capable of capturing the underlying data distribution, is an alternative method for self-supervised learning. VAE/GAN (Larsen et al. 2016) merges
the strengths of variational autoencoders and generative adversarial networks to acquire disentangled representations data. PixeICNN, PixeIVAE, and PixelTransformer (Van den Oord et al.) 2016 Gulrajani et al. 2016 Tulsiani & Gupta 2021) 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.
2.2 Masked Modeling
Masked modeling learns representations by reconstructing masked portion of the input. Pioncering works in natural language processing (NLP) present various such pretraining objectives. BERT (Devlin al.) 2019) and its extensions (Liu et al., 2019; Lan al. 2020) use bidirectional transformer and present. few-shot learning capabilities from masked language modeling. GPT (Radford et al. 2018 2019 Brown et al. 2020), uses autorcgressive, causal masking and demonstrates multi-task, few-shot, and in-context learning capabilities.
Early works in computer vision, such as Stacked Denoising Autoencoders (Vincent et al. 2010 and Context Encoder (Pathak et al., 2016), investigated masked image modeling as a form of denoising or representation learning. Recently, with the widespread use of transformer (Dosovitskiy et al., 2020) as backbone vision architocture, where images are patchified and tokenized as sequences, rescarchers are interested in how to transfer the success in language sequence modeling to scale vision transformers. BETT (Bao et al. 2021
MAE (He et al. 2022), and SimMIM (Xie et al.) 2022) are few of the early works that explored BERT style pretraining of vision transformers. Compared to works in NLP, both MAE and SimMIM (He et al.
2022 Xie
ct al. 2022 find that a much higher mask ratio compared to works in NLP is nocessary to learn good visual representation. Many recent works further extend masked pretraining to hierarchical architectures (Xie el al.
2022 Liu al. 2022) and study data the role of data augmentation (Chen al. 2003 Fang al.
2023). Many
subsoquent works present similar successes of masked pretraining for video Tong et al. 2022
Wang et al.
2023
Feichtenhofer et al. 2022
Gupta et al. 2023), language-vision and multi-modal pretraining (Bachmann
et al. 2022 Li et al.]
2023b
Geng et al.] 2022) and for learning both good representations and reconstruction
capabilities WWei et al. 2023 Li et al. 2022a).
However, BERT-style pretraining requires heavy use of self-attention, which makes computational complexity scale as a polynomial of sequence length. PixclTransformer (Tulsiani & Gupta 2021) and DiffMAE (eei et al. 2023 both use cross-altention for masked image generation and representation learning. Siamese MAE (Gupta et al. 2023) uses asymmetric masking pattern and decodes frames of video condition 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 (Li et al. 2023a) progressively discards redundant mask tokens through its network, leading to a few tokens for reconstruction. VideoMAEv2 (Wang et al. 2023 concatenates randomly sampled masked tokens with visible tokens and uses self-attention reconstruct the masked patches. In comparison, we minimally modify MAE with a cross-attentionnonly decoder and masked tokens are decoded in conditional independent way.
4

---

Published in Transactions on Machine Learning Research (04/2025)
VIT Encoder
MAE
CrossMAE
1 Cross-Atthn Decoder
Doo Doccesaattion
Decoder
Mask Toketts
Figure 3: MAE He et al. 2022 concatenates all mask tokens with the visible patch features from a VIT encoder and passes them to decoder with self-attention blocks reconstruct the original image. Patches that correspond to visible tokens are then dropped, and an L. loss is applied to the rest of the reconstruction as the pretraining objective. CrossMAE instead uses aoss-aattnntoon blocks in the decoder reconstruct only a subset of the masked tokens.
2.3 Applications of Cross-Attention
In addition to the prevalent use of self-attention in computer vision, cross-attention has shown to be a cost-effoctive 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. (2021) replace mean pooling with atos-atttntion pooling and find improvement in ImageNet classification performance. Jaegle et al. also widely used for object detection. Carion et al. (2020] utilizes query tokens placeholders for potential objects in the scene. Cheng al. (2022 2021 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.
(2021) cross-attention to efficiently process large volumes of multi-modal data. Cross-attention
3 CrossMAE
We start with an overview of vanilla masked autocncoders in Section 3.1 Next, in Section 3.2 the use of cross-altention in place of self-altention in the decoder for lesting 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.
we introduce
3.1 Preliminaries: Masked Autoencoders
Masked Autoencoders (MAE) (He al. (022) pretrain Vision Transformers (ViIss (Dosovitskiy al. 2020). Each image input is first patchified, and then random subset of the patches is selected as the visible patches. As depicted in Figure 3 the visible patches, concatenated with learnable class token [CLS], 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 dccoder blocks share the same architecture as the encoder blocks (i.e., both are transformer blocks with teltiontention layers). Note that the number of tokens fed into the decoder the same length as the original input, and the decoding process assumes that the docoded 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. we measure the mean atlention 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 the class token and the visible tokens (see Figure (b)). This motivates make design decisions and conduct experiments specifically to answer the following question: Can simplify
5

---

Published in Transactions on Machine Learning Research (04/2025)
Feed
Features (K,V)
Add
Intoo block Artention
VIT Blocki
Decoder Blockr
Mult-HHed
VT Block
WW
VIT Block;
Decoder Blocky
VIT Block
Randomly Selected Mask Tokens (Q)
VIT Blocks,
Decoder Blocki
VIT Blocky
Kv
Positional Embedding
Weighted Feature Map
Mesk Tokens
All Mask Tokens
Inter-block Figure 4: Overview of CrossMAE. (a) The vanilla version of CrossMAE uses the output of the last encoder block as the keys and queries for cross-attention. The first decoder block lakes sum of mask tokens and their corresponding positional embeddings as queries, and subsequent layers use the output of the previous decoder block queries to reconstruct the masked patches. (b) Unlike the decoder block Vaswani al. (2017). the cross-attention decoder block does not contain self-attention, decoupling the generation of different masked patches. (c) CrossMAF's decoder blocks can leverage low-level features for reconstruction via inter- attention. It weighs the intermediate feature maps, and the weighted of feature maps is used the key and value for each decoder block.
CrossMAE
Partial Masking
the decoding process by eliminating self-attention among masked tokens without compromising the model's ability to generate coherent images and perform well on dounstream tasks?
3.2 Reconstruction with Atoss-Attention
To address this question, substitute the self-attention mechanism in the decoder blocks with cross-attention, using it readout function to decode the latent embodding from the encoder raw pixel values. Specifically, the decoder employs multi-head cross-allention 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 CAossMAF, the output from the final encoder block is used as the key and value tokens for all layers of the decoder, illustrated in Fig. I()). Further exploration in 3.4 reveals that utilizing a weighted mean of selected encoder feature maps be beneficial. The residual connections in cach decoder block enable iterative refinement of decoded tokens they progress through decoder blocks.
Diverging from the original transformer architecture /Vaswani et al., 2017 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 evaluated in the ablation study section to determine its impact performance.
Given the disparity in dimensions between the encoder and decoder, MAE adapts the encoded 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 cach patch be reconstructed. However, CrossMAE does not restrict the architecture to single aros-attention block. Instead, we stack multiple cross-attention decoder blocks in manner more akin to the traditional transformer (Vaswani et al. 2017).

---

has linear complexity Published in Transactions on Machine Learning Research (04/2025)
3.3 Partial Reconstruction
The fact that CrossMAF, uses cross-attention rather than self-attention the decoder blocks brings an additional benefit over the original MAE architecture. Recall that mask tokens are decoded independently and thus there is exchange of information between them, to obtain the reconstructions at specific spatial location, CrossMAF, only needs to pass the corresponding mask tokens to the cross-altention 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 socond question in Soc. 3.1 rather than decoding the reconstruction for all masked locations, we only compule the reconstruction on 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 (7), and the mask ratio (p). Then the prediction ratio is bounded between E (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 MAF, while the variance is (p/7) times larger than the variance in MAE. Empirically, we find that scaling the learning rate of MAE (3) to match the variance (i.e. setting the learning rate as 7/3/p) helps with model performance. Since cross-attentio 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.
3.4 Inter-block Attention
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, Cross cross-attention decoder decouples queries from keys and values. This decoupling allows different cross-attention decoder blocks 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., foeding the feature of the first encoder to the last decoder), in 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 cach docoder block (Figure [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 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 {fe} as keys and values. Specifically, for each key/ sey/value token & in decoder block & in a model with encoder depth n, we initialize a weight wh E RD ~ NO, 1/n). Then tx is defined
4
W J3-
In 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 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 reconstruc- tions, early decoder blocks tend to prioritize information from later encoder blocks, while later decoder blocks focus on earlier encoder block information, demonstrated in Section 4.5
7

---

Published in Transactions on Machine Learning Research (04/2025)
Method
ViT-S
ViT-B
VIT-L
ViT-H
Supervised Steiner et al.
2022
79.0
82.3
82.6
83.1
DINO Caron et al. 2021
82.8
MoCo v3 Chen et al. 2021b)
81.4
83.2
84.1
BETT Bao et al. 2021
83.2
85.2
MultiMAE Bachmann et al.
2022
83.3
MixedAE Chen et al. (2023
83.5
CIM Fang et al. (2023
81.6
83.3
MAE He et al. 2022)
78.9
83.3
85.4
85.8
Cross MAE (25%)
79.2
83.5
85.4
86.3
Cross.AA (75%)
79.3
83.7
85.4
86.4
Table 1: ImageNel-11 classifccalion accuracy. CrosMMAE performs on par, if not better than MAE without self-attention in the decoder. All experiments are run with 800 epochs. The best results are in bold while the second best results underlined.
Apbox
APmask ViT-L
Method
VIT-B
VIT-L
ViT-B
Supervised Li et al. (2022b
47.6
49.6
42.4
43.8
MoCo v3 Chen et al.
(2021b)
47.9
49.3
42.7
44.0
BETT Bao et al. (2022
49.8
53.3
44.4
47.1
MixedAE Chen et al.
2023
50.3
43.5
MAE Li et al. 2022b
51.2
54.6
45.5
48.6
CrossMAE
52.1
54.9
46.3
48.8
Table 2: COCO instance segmentation. Compared to previous masked visual pretraining works, CrossMAB performs favorably on object detection and instance segmentation tasks.
4 Experiments
We perform self-supervised pretraining on TmageNett1IK, following MAR (He et al. 2022 '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. After pre-training we evaluate the pre-trained models by fine-tuning them for image classification and instance segmentation. We provide implementation details and more experiments on different datasets and downstream tasks (iNaturalist (Van Horn et al. 2018 Places365 (Zhou et al., 2017). and ADE20K (Zhou et. al. 2019)) in Appendix B.5
4.1 ImageNet Classification
Setup. The model performance evaluated with end-to-end fine-tuning, with top-1 accuracy used for comparison. Same in Figure. 2 we compare two versions of CrossMAE: one with prediction ratio of 25% (1/3 of the mask tokens) and another with 75% (all mask tokens). Both models are trained with mask ratio of 75% and decoder depth of 12.
Results. As shown in Table 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, CrossMAF, outperforms other self-supervised and masked image modeling baselines, e.g., DINO (Caron et al. 2021) MoCo v3 (Chen al,
2021bb
BETT
(Bao et al. 2021), and MultiMAE (Bachmann et al. 2022)
4.2 Object Detection and Instance Segmentation
Setup. We additionally evaluate models pretrained with CrossMAE for object detection and instance segmentation, which require deeper spatial understanding than ImageNet classification. Specifically, follow
8

---

Published in Transactions on Machine Learning Research (04/2025)
Method
Acc.
Mask
Acc. (%)
Pred.
Acc. (%)
(%)
Ratio
Ratio
MAE
83.0
65%
83.5
15%
83.1
CrossMAE
83.3
75%
83.3
25%
83.2
CrossMAE +
83.3
85%
83.3
75%
83.3
Self-Attn
(b) Mask ratio. CrossMAE has
(c) Prediction ratio. CrossMAE
(a)
Attention type in decoder blocks. Adding back self-attention between mask tokens does not im- prove performance.
consistent performance across high mask ratios.
performs well even when only a frac- tion of mask tokens are reconstructed.
# Feature
Acc. (%)
Decoder
Acc.
Image
Acc. (%)
Maps Fused
Depth
(%)
Resolution
82.9
83.0
224
83.2
83.3
83.1
448
84.6
6
83.5
8
83.1
12
83.3
83.3
(d) Inter-block attention. A com-
(e) Decoder depth. Cross-
(f) Input resolution. CrossMAE MAE performance scales with de- scales longer input sequences. coder depth.
bination of six select encoder feature maps is best.
Table 3: Ablations on CrossMAE. We report fine-tuning performance on ImageNet-1K classification with 400 epochs (i.e., half of the full experiments) with ViT-B/16. MAE performance is reproduced using the official MAE code. Underline indicates the default setting for CrossMAE. Bold indicates the best hyperparameter among the tested ones. feature map fused (row 1, Table 3(d)) indicates using only the feature from the last encoder block. We use 25% prediction ratio for both settings in Table 3(f) to accelerate training.
ViTDet (Li et al. 2022b), 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 on the COCO dataset (Lin et al.] 2014), following MAE (He et al.) 2022]. We compare against supervised pre-training, MoCo-v3 (Chen et al. 2021b), BEiT (Bao et al.) 2022). and MAE (He et al. 2022)
Results. As listed Table 2 CrossMAE, with the default 75% prediction ratio, performs better compared these baselines, including vanilla MAE. This suggests that similar MAE, CrossMAE performance ImageNet positively correlates with instance segmentation. Additionally, CrossMAE's downstream performance scales similarly to MAE the model capacity increases from ViT-B to ViT-L. This observation also supports our hypothesis that partial reconstruction is suprisingly sufficient for learning dense visual representation.
4.3 Ablations
Cross-Attention decoder, outperforms vanilla MAF 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.
As shown in Table 3a CrossMAE, with its cross-attention-onll
Mask Ratio and Prediction Ratio. In our experiments with different mask and prediction ratios (i.e., the ratio of mask tokens all tokens and the ratio of reconstructed tokens to all tokens, respectively) (sce Table 3b and Table 3cb, found that our method's performance 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 Table 3d explored the impact of varying the number of encoder feature maps our inter-block attention mechanism. We found that using only the last feature map slightly lowers performance compared to using all 12. However, partial selection of feature
9

---

Published in Transactions on Machine Learning Research (04/2025)
maps improves CrossMAE's performance, with the best results obtained using 6 feature maps. This indicates that CrossMAF 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 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 CrossMAE variant with 25% prediction ratio. In Table 31 we observe that the classification accuracy positively correlates with the input resolution, indicating that CrossMAE can scale to long input sequences.
4.4 Improving Training Throughput and Memory Utilization
In this section, we show tangible improvement that Cross provides: by implementing partial recon- struction and confining attention to interactions between masked and visible tokens, CrossMAE significantly improves pro-training efficiency compared to vanilla MAE.
Method (prediction ratio)
MAE (0.75) CrossMAE (0.75)
CrossMAE (0.25)
Accuracy (%)
83.0
83.3
83.2
Runtime (mins per epoch)
9.35
8.41
6.32
Memory (MB per GPU)
68386
37987
36805
Table 4: Comparison of MAE and CrossMAE with different prediction ratios. Each model is trained for 400 epochs.
Method
Accuracy (%)
Runtime (mins per epoch)
MAE
83.0
9.35
+Cross Attention (Section 3.2
82.9
7.38
+Partial Reconstruction (Section 3.3
83.2
6.32
+ Inter-block Attention (Section 3.4
83.3
8.41
Table 5: A structured ablation study of each contributing component of CrossMAE VTT-B architecture. Each model is trained for 400 epochs using 2 NVIDIA A100 80CB GPUs.
In Table we present a comparison ImageNett classification accuracy, runtime, and memory usage. Table we provide a more structured ablation of the different components of CrossMAF and their affect on runtime. The setup mirrors that of Table 3 with runtime measured on 2xA100 80GB GPUs, utilizing Flash-Attention 2 DDo 2023) across all models and gradient accumulation set to 2 for 400 epochs.
It is important to note that partial reconstruction relies on our proposed use of cross-attention decoder blocks, which is not applicable to vanilla MAE for efficiency improvements. In vanilla MAE, each masked token attends to other masked tokens within the decoder blocks. As a result, removing any masked tokens affects the decoding of the remaining ones. contrast, CrossMAF, reconstructs each masked token solely based on the visible tokens, without relying on interactions with other masked tokens. This allows the loss applied to the subset to serve as an unbiased estimate of the loss applied to all reconstructed tokens. This method reduces the computational load and runtime without sacrificing significant performance. Further discussion can be found in Appendix
Since vanilla MAE uses a self-attention decoder, the computational complexity of the decoder is quadratic with respect to the total number of tokens. Using cross-attention in CrossMAE reduces this complexity to be lincar with respect to the number of visible tokens and the number of masked tokens. By further varying the number of masked tokens through partial reconstruction, we can reduce the complexity by decoding only a subset. of the masked tokens.
10

---

Published in Transactions on Machine Learning Research (04/2025)
Patch Mean
Block
Block
Block
Block 10
Block 12
(a) Residual
stack of each
decoder block
(b) Cumulative output of each
decoder block
(c) Inter-block attention maps
Masked Input
Early enc. feats
Late enc. feats
Early enc. feats
Late enc. feats
Figure 5: We visualize the output each decoder block. (a-b) Different decoder blocks play different roles in the reconstruction, with most details emerging at later decoder blocks, which confirms the motivation for inter-block attention. (c) Visualizations of inter-blodd block-block attention shows that different decoder blocks indeed attend to feature from different encoder blocks, with later blocks focusing earlier encoder features to achieve reconstruction. The reconstructions are unnormalized w.r.t ground truth mean and std for each patch.
Formally, let the total number of visible tokens be N, and the latent dimension be d, with mask ratio p and prediction ratio The attention complexity in the MAE decoder is on the order of 22d. In CrossMAE, the visible tokens (of length (1 P)N serve as the queries, and the masked tokens (of length >N) serve as the keys and values. This results in the attention complexity in CrossMAF. decoder being (1 - V2d. For instance, Table 1, the CrossMAE variant that decodes all masked tokens, or CrossMAE (0.75), uses p= 0.75 and 7 0.75 (i.e., full prediction with p = 7), which gives 16 Ndd. On the other hand, CrossMAE (0:25) p 0.75 and 7 0.25, which gives 16 22d, further reducing the complexity. More discussion and comparisons be found in Appendix C
4.5 Visualizations
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 '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 (a.b), we observe that different decoder blocks play different roles in reconstruction, with most details emerging later decoder blocks. This justifies the need for lew-level features from early encoder blocks, motivating biter-block attention.
Visualizing Inter-block Attention Maps. As shown in the visualizations of the attention maps of inter-block allention in 5 c), CrossMAE naturally leverages the inter-block atlention to allow the later decoder blocks to focus on earlier encoder features to achieve reconstruction and allow the earlier decoder blocks focus on later encoder features. This underscores the nccessity for different decoder blocks attend different encoder features, correlating with the performance improvements when inter-bloce attention is used.
11

---

Published in Transactions on Machine Learning Research (04/2025)
5 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 ViT-H, achieves comparable better performance in image classification and instance segmentation with reduced computational requirements, showcasing the potential for more efficient and scalable visual pretraining methods. Our findings underscore the efficacy of the encoder's global representation learning, paving the way for streamlined decoder architectures in future MAE implementations. CrossMAE's efficiency and scalability demonstrate potential for large-scale visual pretraining, particularly on underutilized in-the-wild video datasets. However, our work has not yet explored scaling to models larger than VIT-H, the largest model examined in MAE, leaving this for future research.
12

---

Published in Transactions on Machine Learning Research (04/2025)
References
Roman Bachmann, David Mizrahi, Andrei Atanov, and Amir Zamir. Multimae: Multi-modal multi-task masked antocncoders. In European Conference on Computer Vision, pp. 348-367. Springer, 2022.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pro-training of image transformers. arXiv preprint arXie:2106.088 08254, 2021.
Hangbo Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. In ICLR, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. European conference on computer vision, pp. 213 229. Springer, 2020.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 33:9912-9999 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference compuler vision, pp. 9650 1000-6660-- 2021.
Kai Chen, Zhili Liu, Lanqing Hong, Hang Xu, Zhenguo Li, and Yan Yeung. Mixed autoencoder for sele-supervised visual representation learning. In Proceedings of the IEEE/CVP Conference on Computer Vision and Pottern Recognition (CVPR), pp. 22742 22751, June 2023.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and llya Sutskever. Generative pretraining from pixels. 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020b.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. ar preprint arXiv:2003.042: 04297, 2020c.
Xinlci Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers, 2021a.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. arXiv preprint arxiv:2104.00 02057, 2021b.
Bowen Cheng, Alexander G. Schwing, and Alexander Kirillov. Per-pixel classification not all you need for semantic segmentation. 2021.
Bowen Cheng, Ishan Misra, Alexander G Schwing, Alexander Kirillov, and Rohit Girdhar. Masked-attention mask transformer for universal image sogmentation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1290-1299, 2022.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V Le. Randaugment: Practical automated data augmentation with reduced search space. arxiv e-prints, page. arXiv preprind Xix:9:999.. 137199 4, 2019.
Tri Dao. FlashAllention-2:: Faster allention with better parallelism and work partitioning. 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. 2019.
13

---

Published in Transactions on Machine Learning Research (04/2025)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Yuxin Fang, Li Dong, Hangbo Bao, Xinggang Wang, and Furu Wei. Corrupted image modeling for self- supervised visual pre-training. In The Eleventh International Conference Leorning Representations, 2023. URL https://openrrvieew net/forumm1d-00NWcaaa httpp://ppparrrree...eeee/rrrtto
Christoph Feichtenhofer, Haoqi Fan, Yanghao Li, and Kaiming He. Masked autocncoders as spatiotemporal learners. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Aduances in Newral Information Processing Systems, 2022. URL https:://pppereeiiiitt
Xinyang Geng, Hao Liu, Lisa Lee, Dale Schnurams, Sergey Levine, and Pieter Abbeel. Multimodal masked autoencoders learn transferable representations. arXiv preprint orXiv2225514200 2022.
Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv: 02677, 2017a.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint ar Xiv: 02677, 2017b.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271 222224,,11,,4, 2020.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vazquez, and Aaron Courville. Pixelvac: A latent variable model for natural images. arXiv preprint arXio:1611.05011 2016.
Agrim Gupta, Jiajun Wu, Jia Deng, and Li Fei-Foi. Siamese masked autocncoders. arXiv preprint 24844 2023.
Kaiming He, Haoqi Pan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 16000-16009, June 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Aision-FCCC 2016: 14th European Conference, Amsterdam, The Netherlands, Oclober 11-14, 2016, Proceedings, Part IV 14, pp. 646 661. Springer, 2016.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin lonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Percciver io: A general architecture for structured inpuls & oulputs. arXiv preprint ar Xrxi::2177.1 14795, 2021.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tote Xiao, Spencer Whitehead, Alexander C. Berg, Wan- Yen Lo, Piotr Dollar, and Ross Girshick. Segment anything. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV), pp. 4015-4026, October 2023.
Zhenzhong Lan, Mingda Chen, Scbastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020. URL https://ppeaneeeeeeam////orttrrp
14

---

Published in Transactions on Machine Learning Research (04/2025)
Anders Boesen Lindbo Larsen, Søren Kaac Sonderby, Hugo Larochelle, and Ole Winther. Autoencoding beyond pixels using learned similarity metric. In International conference. on machine learning, pp- 1558 1566. PMLR, 2016.
Jin Li, Yaoming Wang, XIAOPENG ZHANG, Yabo Chen, Dongsheng Jiang, Wenrui Dai, Chenglin Li, Hongkai Xiong, and Qi Tian. Progressively compressed auto-encoder for self-supervised representation learning. In The Eleventh International Conference on Learning Representations, 2023a. URL https: : //openrevien www.ee...eeeeeeeee nwt/forummiiiieeee foo.:mmmidd8114mm2222W
Tianhong Li, Huiwen Chang, Shlok Kumar Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan. Mage: Masked generalive encoder lo unily representation learning and image synthesis. ar preprinl 0rXiv:2211.0911 2022a.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer backbones for object detection. In European Conference on Computer Vision, pp. 280 296. Springer, 2022b.
Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-image pre-training via masking. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 23390-23400 2023b.
Tsung- Yi Lin, Michael Maire, Serge J. Belongie, Lubomir D. Bourdev, Ross B. Girshick, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: common objects in context. Coll,, abs/1405.0311, 2014. URL http: //arxiv/ http:://arxiv.rrr/aaa//6666 0312
Jihao Liu, Xin Huang, Jinliang Zheng, Yu Liu, and Hongsheng Li. Mixmac: Mixed and masked autocncoder for elficient pretraining of hierarchical vision transformers. ar Ni::2205.13137, 2022.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zewlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907. 11692, 2019.
Byy Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. 2017a.
Πya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint xrX:::711100111,, 2017b.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint ar arXic:1807.0774 2018.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint 20X102334,07198,, 2023.
Decpak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexci A Efros. Context encoders: Feature learning by inpainting. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2536-236 2544,544, 2016.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Nly Sutskever. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. 2019.
Ilija Radosavovic, Tele Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, and Trevor Darrell. Real-world robot learning with masked visual pre-training. In Conference on Robot Learning, pp. 2023.
PMLR,
Andreas Peter Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https://openreriir: http:://pppmerrrroooe.////oooooo .nt//forum?id 4nPswriKeP
15

---

Published in Transactions on Machine Learning Research (04/2025)
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEFE conference on compuler vision and pattern recognition, pp. 2818 2826, 2016.
Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. VideoMAE: Masked autoencoders are data-efficient learners for self-supervised video pre-training. In Advances in Neural Information Processing Systems, 2022.
Hugo Touvron, Matthieu Cord, Alaaeldin El-Nouby, Piotr Bojanowski, Armand Joulin, Gabriel Synnaeve, and Hervé Jégou. Augmenting convolutional networks with attention-based aggregation, 2021.
Shubham Tulsiani and Abhinav Gupta. Pixcltransformer: Sample conditioned signal generation. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10455-10464. PMLR, 18-24 Jul 2021. URL html
https://proceeddngg mlr.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Conditional image generation with pixclenn decoders. Advances in neural information processing systems, 29, 2016.
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pictro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference compuler vision and pollern recognilion, pp. 1769-8775 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017.
Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou. Stacked denoising autoencoders: Learning useful representations in a deep network with local denoising criterion. Journal of machine learning research, 11(1)), 2010.
Limin Wang, Bingkun Huang, Zhiyu Zhao, Zhan Tong, Yinan He, Yi Wang, Yali Wang, and Yu Qiao. Vidcomac v2: Scaling video masked autocncoders with dual masking. In Proceedings of the IEEE/CV Conference on Computer Vision and Pattern Recognition, pp. 14549-14560. 2023.
Xudong Wang, Ziwei Liu, and Stella X Yu. Unsupervised feature learning by leves instance- discrimination. In Proceedings of the IEEE/CVE conference on computer vision and paltern recognilion, pp. 12586-12599 2021.
Chen Wei, Karttikeya Mangalam, Po- Yao Huang, Yanghao Li, Haoqi Fan, Hu Xu, Huiyu Wang, Cihang Xie, Alan Yuille, and Christoph Feichtenhofer. Diffusion models as masked autoencoder. In ICCV, 2023.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9653-9663, June 2022.
Hongwei Xue, Peng Gao, Hongyang Li, Yu Qiao, Hao Sun, Houqiang Li, and Jiebo Luo. Stare at what you see: Masked image modeling without reconstruction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 22732 22741, 2023.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 6023 6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018. URL https://openrevview
met/forumPiddrrDdptttBB
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE transactions on pattern analysis and machine intelligence, 40(6): 2017.
16

---

Published in Transactions on Machine Learning Research (04/2025)
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision, (27 302 321, 2019.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pro-training with online tokenizer. arXiv preprint or arXio:2111.0783 2021.
17

---

Published in Transactions on Machine Learning Research (04/2025)
A Implementation details
A.1 Attention Calculation
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 (He et al. 2022). For each image, we randomly generate 75% binary mask (m) for all tokens, with m, = representing token being masked and m, 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, total of 8 attention maps, each with 16 attention heads are stored. Based on the mask pattern, we calculate the outer product (m--n for the self-attention among mask tokens, and m (1-m) --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.
A.2 Attention to Masked Tokens Across Decoder Blocks
addition to comparing the overall attention values for masked tokens in vanilla MAE, we explore whether this pattern is shared by all attention blocks. To do so, we analyze the cross-attention and self-attention magnitudes at each block of the pre-trained MAE.
In the table below, report the attention to visible tokens and masked tokens across all decoder blocks:
Block
Attention to visible tokens
Attention to masked tokens
(closest to the encoder)
0.196
0.042
0.149
0.058
0.181
0.047
0.153
0.057
0.193
0.043
0.161
0.054
0.207
0.038
8 (closest to the reconstruction) Sum (reported in the paper)
0.179
0.048
1.418
0.388
Table 6: Atlention maguitudes for visible tokens and masked tokens to mask attention across decoder blocks.
Based on these results, we do not observe significant increase in the attention to masked tokens we progress through the later decoder blocks. The attention lo masked tokens remains relatively consistent across all blocks. Additionally, we observe that the magnitude of attention is consistently larger for masked tokens' cross-attention to visible tokens compared to their self-attention across all decoder layers. This pattern highlights the consistent interaction between masked tokens and visible tokens throughout the decoding process.
A.3 Inter-Block Attention
We tried a few implementations for inter-block attention (IBA) and found the following implementation be the fastest and most memory-efficient. In this implementation, we combine inter-block attention for all encoder layers a single forward pass of linear layer. For each decoder block, we index into the output tensor to extract the corresponding feature map, and 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.
18

---

Published in Transactions on Machine Learning Research (04/2025)
The pseudo-code of inter-block attention is the following:
class InterBlockAttentiion ():
def __init__ (self, san_feat_man decoder tecoder_depth); self linear Linear (num_feat__map decoder_depth. bias=False)
std_dev (self linear reight, mean=0., atd=std_dev)
sqrt
def forward (self, feature_maps list):
feature_mat maps: a list of length num feat_eaaaappeeeaeeeeeeeeeeeeapda _ape,, each with dimension Batch Size x Nun. Tokens x Embedding Dim.
10
11
12
faacked_fetture mape stack (featureeetaaaaeeeeeeeeeeeeeeeee _aaps, dim=-1) return self linear (stacked, _aaps)
Additionally, we further investigate the importance of using a cross-attention decoder, where each decoder block use different feature maps from the encoder for decoding. In this experiment, incorporated IBA into MAE, which uses only self-attention decoder. Specifically, 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. 7 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 attend features from different encoder blocks, thanks to its decoupling of queries with keys and values.
concatenate the interblock attention
Method MAE
Acc. (%)
83.0
MAE IBA
83.0
CrossMAE (25%)
83.2
CrossMAE (75%)
83.3
Table 7: For MAE, inter-block attention has very small differences in terms of finetuning performance, potentially due to the fact that MAE's decoder only takes in one set of features.
A.4 Ablation that Adds Self-Attention
In Section 4.3 (a); we propose adding self-attention back to CrossMAF. ablation. In that particular ablation study, analyze the effect of self-attention between the masked tokens, which can be used to improve the consistency for reconstruction. Specifically, we modify the formulation in the original transformer paper (Vaswani et al. (017), where the mask/ mask/query tokens are first. passed through multi-head self-attention and a residual connection before being used in the multihcaded cross-attention with the features from the encoder. The primary difference with the vanilla transformer decoder implemenntttio (Vaswani et. al.] 2017 is we do not perform casual masking in the multi-head self-attention. Please reference Figure for a more visual presentation of the method.
A.5 Ablation on Inter-block Attention
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 feature map is used, the output of the patch embedding (input to the ViT) is also used.
Additionally, we ablated an alternative to the linearly weighted inter-block attention. Instead, we enforce that the weights of the inter-block attention always sum to one. We do so by applying softmax on the weights before performing the linear combination. We run the experiment on ViT-B with a mask ratio p = 0.75 and a prediction ratio 7 = 0.25 with 400 pre-training epochs, 100 fine-tuning epochs. For this ablation, all encoder feature maps are used for block attention. The results are presented in Table R We found this alternative design to achieve comparable results the linearly weighted inter-block attention.
19

---

Published in Transactions on Machine Learning Research (04/2025)
Add Norm
Forward
Features (K.V)
Mult-Heed Cross Attention
aQ
Add Norm
Mult-Hoad Attention
okktvt
Randomly Selected
Mask Tokens (Q)
Figure 6: Modification for self-attention ablation
Method
Acc. (%)
CrossMAE (25%)
83.2
CrossMAE (25%) w/ sftmax weights
83.2
Table 8: Softmax weighted inter-bles block attention. We found that whether the weights in inter-block attention sums to one does not significantly impact model'ss performance.
In addition to the simple design of intcr-block attention, we experimented with 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 the decoder, we added linear projection for each encoder feature before the feature aggregation. We denote this variant. as LP... As shown in the Table. g (with VIT-B pre-trained for 400 epochs, consistent with the setting in Table. 3), adding linear projection slightly improves the performance. This indicates that it is possible to design variants of readout functions, such through improved inter-block attention, to improve the feature quality of CrossMAE.
Method
Acc. (%)
CrossMAE
83.3
CrossMAE + LP
83.5
Table 9: Improving inter-block attention by adding linear projections to the input features. The performance gain indicates that possible to design variants of readout functions to improve CrossMAE.
A.6 Hyperparameters
Pre-training: The default setting is in Table 10 which is consistent with the official MAE (H et al. 2022
implementation. As mentioned in Sec. 3.4 we scale the learning rate by the ratio between mask ratio (p) and prediction ratio (7) to ensure the variance of the loss is consistent with (He el al. 2022). Additionally, we use the linear learning rate scaling rule (Goyal et al. 2017b). This results b== 7 base br batchsize/ (256 * p). For Table 1 we use 12 decoder blocks, with mask ratio and prodiction 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, the same MAF.
20

---

Published in Transactions on Machine Learning Research (04/2025)
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 I1 for more detail.
Config optimizer
Value
AdamW
Loshchilov & Huttor 2017b
base learning rale
1.5e-4
learning rate schedule
cosine decay [Lohhchllov & Hutter 2017a
batch size
4096
weight decay
0.05
optimizer momentum
81.32 = 0.9, 0.95 (Chen et al. 2020a
warm epoch Gooya et al. 2017a total epochs
20, 40
400, 800
RandomBeeszeeCrop.
augmentation
RandommHorizontalFiii
Table 10: Pretraining Hyperparameters
Config
Value
optimizer
AdamW
base learning rate learning rate schedule batch size weight decay optimizer momentum warm up epoch total epochs augmentation
1e-3
cosine decay 1024
0.05
31, 82 0.9, 0.999
5
100 (B), 50 (L)
RandAug (9, 0.5) (Cubuk et al.
2019
label smoothing (Szegedy et al. 2016)
0.1
mixup (Zhang el. al. 2018
0.8
cutmix Yun et al.]
2019)
1.0
drop path
Huang
2016
0.1
Table Finetuning Hyperparameters
A.7 Compute Infrastructure
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 (Li
et al. 2022b) experiments use a single machine equipped with NVIDIA A100 (SOGB) GPUs. We copy the
datasets to the shared memory on the machines accelerate dataloading. We use Flasnnottt (Dao
2023 to accelerate attention calculation.
B Additional Experiments
B.1
Linear Probe
We provide linear probe comparisons (at 800 epochs) for ViT-Small and ViT-Base Table. 12 For both of
these experiments, we run CrossMAE with a prediction ratio of 75% (reconstruction of all masked patches). These results show that CrossMAF, achieves slightly better linear probe performance than vanilla MAF.
21

---

Published in Transactions on Machine Learning Research (04/2025)
Method
VIT-S
ViT-B
MAE 49.7
65.1
CrosSMAE
51.5
65.4
Table 12: Linear probe experiments of CrossMAE.
B.2 Quality of CrossMAE Reconstruction
In addition to the visual examples illustrated in Figure quantitatively evaluate the reconstruction quality of CrossMAB compared to the baseline MAE. Specifically, we measure the average Peak Signal-to-Noise Ratio (PSNR) for the reconstruction of masked patches (i.e., patches hidden from the model) on the ImageNet validation set. The results are summarized in Table
Our findings indicate that CrossMAE achieves reconstruction performance closely comparable to the baseline MAE. Particularly noteworthy is that when predicting 75% of the masked patches, CrossMAE's PSNR differs by only 0.3% from the baseline MAE. Similarly, predicting 25% of patches results in a marginal PSNR decrease of just 0.7%. These results demonstrate that CrossMAE maintains high reconstruction quality while significantly reducing computational demands.
Method
Average PSNR for Masked Patches
CrossMAE (75% pred)
20.05 (-0.3%)
CrossMAE (25% pred)
19.97 (-0.7%)
MAE
20.11
Table 13: Comparing the PSNR of masked patch reconstruction on the ImageNet validation set.
B.3 Variance Experiments
We pre-train and finetune the ViT-B with 25% prediction ratio (same setup as Table 3 i.c. pre-train for 400 epochs, finetune for 100 epochs), with four different seeds to compute a confidence interval of the result. We The results suggest that. CrossMAF has small variance between different.
present. the results in Table 14
experiments.
Seed 1
2
3
4
Avg + std crr
83.20
83.29
83.24
83.21
83.24 0.018
Table
Results across different random seeds.
B.4 Masking Strategy
Similar to MAE (He et al. 2022 we here ablate the masking pattern. Instead 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. 15 for ViT-B (at 400 epochs), which shows that grid masking does not lead additional improvements in downstream performance.
Method
Acc. (%)
Grid Masking 83.2
Random Masking
83.3
Table 15: Ablation of masking strategies.
B.5 Ablations on More Downstream Datasets
To further investigate the performance of CrossMAE, we provide results on additional downstream tasks, including classification iNaturalist2019, Places365, and semantic segmentation ADF20K. These
22

---

Published in Transactions on Machine Learning Research (04/2025)
experiments demonstrate that CrossMAE performs comparably to MAE for transfer learning while offering improved performance on specific tasks.
We first compare the performance of CrossMAE and MAE on the iNaturalist2019 and Places365 datasets pre-trained on ViT-B for 800 epochs.
Dataset
MAE
CrossMAE (0.25)
CrossMAE (0.75)
2N11
79.8
79.4
80.1
Places365
57.9
57.9
57.8
Table 16: Transfer learning results on iNaturalist2019 (Van Horn al. 2018 and Places365 Zhou et al. 2017
evaluated for image classification accuracy. CrossMAE performs comparably to MAE on these datasets.
To evaluate transfer learning to another task, in addition to the instance segmentation result in Table on the COCO dataset {Lin al.) 2014), we also report results on ADE20K (Zhou et al.) 2019) for semantic seymentation. These results further illustrate that CrossMAE learns comparable more genera iizale features compared to MAE.
Dataset ADE20K mloU
MAE
CrossMAE (0.25)
CrossMAE (0.75)
47.7
47.7
48.1
Table 17: Semantic segmentation results on ADE20K. CrosSMAE (0.75) shows an improvement over MAE and CrossMAE (0.25).
C Runtime and GPU Memory Comparisons with MAE
All experiments in Table 18 are conducted server with 4 NVIDIA A100 (80GB) GPUs, with the standard hyperparamnterss 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.c., local batch size is set to 1024 samples per GPU). However, the same local batch size results out-of-memory (OOM), which indicates that the total memory requirement is larger than the available memory for each GPU (80GB). To run MAE same hardware, 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.7x lower decoder FLOPS.
thus employ gradient accumulation with local batch size
Decoder FLOPS ImageNet Acc
83.7
MAE CrossMAE
Memory (MB/GPU)
Runtime (min/epoch)
Acc. (%)
(75%, 12)
Method
(75%, 8)
(25%,
MAE
0OM (>81920)
5.19*
83.3
CrossMAE
41177
3.38
83.5
(25%,
83.3
Table 18: CrossMAE greatly improves the training throughput and reduces the memory requirements, lowering the barrier for masked pretraining. Statistics measured on NVIDIA A100 80GB GPUs. Please refer to Appendix
GFLOPS
Figure 7: We compare VIT-B which pre-trained for 800 epochs with different variants of CrossMAB v.s. MAE, For CrossMAE, we vary the prediction ratio p and number of decoder blocks n,, and denote each as (p. n). While all experiments are run with inter-block altention, CrossAEE has lower decoder FLOPS than MAE He et al. 2022) and performs on par better.
C for comparison details. "MAE's default batch size exceeds the capacity of GPUs, requiring gradient accumulation for runtime measurement.
23

---

Published in Transactions on Machine Learning Research (04/2025)
D Visualizing the Contributions per Decoder Block
We propose a more fine-grained visualization approach that allows us to procisely 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; be the output and g:(+) the residual function of docoder so f; = 9f 2) The final decoder block's output goes through reconstruction head h, which is linear, consisting of layer-norm and linear layer, to produce the reconstruction. With D the decoder depth, fo the initial input, and y the final output, y is recursively defined h(fD-1 + (D(f-1)), which simplifies due to the linearity of h:
y = h(Jo + (fo) gD(JD-1))
=
h(fo)
+ +h(gi(fo)) + + h(g((fo11)
Pos Embed. Mask Token
Block
Block D
This decomposition allows to express the reconstruction an image stack, where the sum of all the levels gives us the final reconstruction. We present the visualization in Figure 5
24


Reviews:
[{"review_id": "TAKya73yaw", "reviewer": "Reviewer_mJNF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: Paper appears to offer a pareto improvement on an existing masked autoencoder architecture. No noted weaknesses.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Authors show that patches in the decoder of a masked autoencoder can be decoded independently without cross-attention between them. They show that this improves the computational efficiency of masked autoencoders without degrading their performance. Intuitively, this result makes sense to me.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "n/a", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JYxyPkNvg6", "reviewer": "Reviewer_vdW1", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- S1. The paper closely studies self-attention in MAE and the proposed cross-attention improvements are clear and intuitive. In hindsight, it does seem strange that the masked tokens attend to each other. \n- S2. Overall, the paper is well written, containing excellent diagrams and figures. \n- S3. Paper shows a wide variety of experiments that demonstrate that CrossMAE matches or outperforms MAE in across several tasks, including classification and segmentation. These experiments also include a variety of different ViT sizes and datasets, as well as a good ablation section. \n\n**Weaknesses**\n- W1. The author's claim that their CrossMAE encoder learns a global representation, however the existing evidence is not convincing and should be strengthened. The paper may verify this by (a) showing that MAE's encoder cannot learn global representations or (b) a direct analysis of encoder representations.\n- W2. Together, the proposed improvements yield only incremental gains, with a +0.3% increase in accuracy over MAE and a 10% reduction in training runtime according to Table 5. It seems that the inter-block attention counter-acts the improvement in runtime created by cross attention and partial reconstruction. \n- W3. For Inter-block Attention (Sec 3.4), the authors chose a weight sampled from $\\mathcal{N}(0,1/n)$. This choice could be better justified as there is a risk that the weights may all go to zero, particularly as $n$ grows. For example, why aren't the weights normalized to ensure they add to 1? Could this explain why accuracy decreases with higher number of features fused, as per Table 3d?\n\n**Questions**\n- Q1. The paper states that MAE cannot perform partial reconstruction, yet in Table 4, it shows MAE using a prediction ratio of 0.75. Can you clarify this?\n- Q2. Confusing sentence in Sec 3.2: \"MAE adapts 'visible features' to the decoder's latent space using an MLP.\" The term 'visible features' just isn't understandable from the context. Could you clarify?\n- Q3. See W3 - justification of weight sampling for inter-block attention.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces CrossMAE, a self-supervised learning architecture that improves upon Masked Autoencoder (MAE). Its contributions are as follow:\n1. Cross-Attention Decoder: CrossMAE replaces the self-attention mechanism in MAE with a cross-attention and decodes each masked tokens using information from visible tokens. This removes inter-masked token interactions which reduces computational overheads and maintains downstream performance. \n2. Partial Reconstruction: Cross-attention allows the user to reconstruct a subset of masked tokens instead of whole image. This reduces computational overhead again. \n3. Inter-block Attention: The authors proposes the decoder block take as input, the weighted linear combination of feature maps over $n$ encoder blocks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Stronger evidence that CrossMAE learns global representations, either by showing that MAE cannot learn global representations or direct analysis. \n2. Stronger justification for weight sampling for inter-block attention\n3. Remove duplicate references to \"multimae\"", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BRCBvGhSAy", "reviewer": "Reviewer_2PZ7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nI enjoyed reading the paper!\n- The paper is written and organized neatly.\n- Extensive ablations are designed to demonstrate the effect of the proposed modifications.\n- The authors provide hyperparameter values and a link to a repository with implementation of their method and experiments, facilitating reproducibility.\n\n**Room for improvement**\n\n1. Table 1: For clarification, were models other than CrossMAE and MAE trained from scratch or were the numbers taken from the corresponding papers? To confirm, do MAE and CrossMAE have the same number of decoder blocks?\n\n2. Figure 2: The visualizations are helpful to get a qualitative understanding of how well the models are reconstructing the missing information. I believe it would be beneficial to also include a quantitative evaluation, such as reporting PSNR (peak signal to noise ratio) between the original inputs and their reconstructions from the different models.\n\n3. Table 3: It would be interesting to report CrossMAE performance with 100% prediction ratio in 3(c).\n\n4. Tables 1,2,3: It would strengthen the paper’s argument to include confidence intervals for reported results. However, I recognize that this might be computationally prohibitive.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents CrossMAE - a modification of the original MAE model that removes self-attention in the decoder and replaces it with cross-attention between the masked tokens and the encoded visible tokens coming from the encoder. Additionally, this work proposes a modification to the MAE objective function so that only a fraction of the masked tokens are reconstructed. Finally, the authors propose the introduction of inner-block attention in which each decoder block obtains a weighted sum of the encoded features of visible tokens coming from different encoder blocks which share the same spatial location. When evaluated empirically through image classification and instance segmentation, CrossMAE pretrained on ImageNet achieves comparable or better performance than the original MAE model, while being more computationally efficient.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the points above that are intended to strengthen the work.", "broader_impact_concerns": "It would be beneficial for this work to include a brief discussion about possible ethical concerns that might be relevant such as potential biases coming from the training data, etc.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "arqOsXhENM", "reviewer": "Reviewer_XZMY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The proposed method outperforms MAE while reducing runtime.\n\n2. The authors provide detailed descriptions of each method and include ablation studies for the designs.\n\n3. The presentation is relatively clear, and the paper is well-structured.\n\n**Weaknesses**\n\n1. Section 3.2 introduces several minor designs for reconstruction with Cross-Attention, but lacks solid experimental evidence or convincing explanations for their importance to the Cross-Attention framework.\n\n2. Partial reconstruction and inter-block attention appear independent of Cross-Attention. These additions do not seem essential to CrossMAE, as they could also apply to standard MAE. Moreover, inter-block attention increases computational cost.\n\n3. The ablation studies are incomplete. Additional experiments should demonstrate the necessity of the weighted mean of selected encoder feature maps and its learnable version", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates improving the performance and training efficiency of masked auto-encoders (MAEs) by modifying the decoder design. Motivated by the observation that the attention value magnitude between mask tokens and visible tokens is higher than between mask tokens themselves, the authors propose CrossMAE, where mask tokens serve as queries and visible tokens as keys and values to reconstruct the masked regions. Additionally, the authors introduce several minor enhancements, including learnable weighted averaging of selected encoder feature maps, partial reconstruction, and inter-block attention, to further improve performance. Experiments on ImageNet using ViT demonstrate that CrossMAE outperforms MAE with reduced training cost. Ablation studies are also provided to analyze the contributions of each design.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Address the concerns outlined in the **weaknesses**.\n\n2. While the paper demonstrates that combining Cross-Attention, partial reconstruction, inter-block attention, and learnable feature selection benefits MAE, the improvements from inter-block attention and learnable feature selection appear to be mere additive effects with limited novelty. The authors should provide theoretical insights or empirical evidence to justify the inclusion of inter-block attention and learnable feature selection, which introduce additional hyperparameters and computational costs. These additions could also be combined with MAE alone (without Cross-Attention), which questions their necessity for CrossMAE and undermines the paper's stated goal of reducing runtime.\n\n3. The writing in Sections 3.2 and 3.4 is unclear. In Section 3.2, the authors list designs without prioritizing their significance or explaining which are critical and which are not. In Section 3.4, the structure is disorganized, making it difficult to discern the key aspects of the design. This lack of clarity undermines the overall presentation and weakens the argument for the proposed method.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces CrossMAE, a self-supervised learning architecture that improves upon Masked Autoencoder (MAE). It claims the following components: Cross-Attention Decoder, Partial Reconstruction, and Inter-block Attention. The components were validated by empirical results.

After the rebuttal, three reviewers are "leaning to accept" and 1 reviewer is "leaning to reject". As the authors' claims are validated by the empirical results, the AE considers it meets the expectation of TMLR on its claim and evidence provided. However, the authors are still recommended to give materials to further justify the effectiveness of the proposed components.

Best,
Action Editor
Decision: Accept (poster)