1 Introduction↩︎

Self-supervised learning (SSL) has made remarkable progress in language and computer vision. Masked Image Modeling (MIM) is an effective framework in image SSL, which boasts a simple training pipeline, a few handcrafted data augmentations, and high performance across downstream tasks. In the pioneering work of BEiT [1], 40% of the input image is masked, and the model is trained to capture the semantics of the masked patches by reconstructing the DALL-E [2] output features. To simplify pre-training and reduce computational overhead, MAE [3] only feeds the visible tokens into the encoder and encourages the decoder to reconstruct the raw pixels of masked patches. More recently, follow-up works have focused on adding auxiliary tasks or using large-scale pre-trained models to produce reconstruction targets. For instance, CMAE [4] explicitly adds a contrastive task and optimizes it in conjunction with the MIM task, while MILAN [5] and BEiT-v2 [6] employ multimodal pre-trained models such as CLIP [7] to generate the reconstruction features.

Among the various MIM methods available, pixel-based approaches such as MAE [3] are particularly interesting because of their simple pre-training pipeline and minimal computational overhead. However, these methods are typically biased towards capturing high-frequency details due to their emphasis on reconstructing raw pixels [1], [8]. As a result, they waste a significant amount of modeling capability that could be better utilized to capture low-frequency semantics. Our objective is to reduce this waste of modeling capacity, aiming for an improved quality of learned representation for downstream visual tasks. Toward this goal, we design two pilot experiments based on the representative work of MAE [3] to uncover its neglected design aspects.

1. Fuse Shallow Layers: Rather than solely using the output layer for pixel reconstruction, we implement a weight-average strategy to fuse the output layer with all previous layers. The weights assigned to each layer are normalized and dynamically updated during the pre-training process, with their absolute values indicating the significance of each layer for the reconstruction task. We track the changes in these weights and illustrate them in 1. As depicted in the figure, the model increasingly relies on the features of shallow layers as training progresses.

2. Frequency Analysis: To further understand the property of the representation learned with MAE, we analyze the frequency response of each layer’s feature. We adopt the tools proposed by [9] to transform the encoder features into the frequency domain and visualize the relative log amplitude of the transformed representation in 2. Typically, a higher amplitude indicates that the feature produced by one layer contains more high-frequency information. We empirically find that the shallow layers contain significantly more high-frequency components than deep layers, which are mostly related to low-level details (, textures).

Based on our analysis of the pilot experiments, we can conclude that the pixel reconstruction task exhibits a bias towards low-level details. This is evident from the low linear probing accuracy, which highlights the constraints of the pixel reconstruction task. Namely, it requires features that are semantically distinct enough to achieve linear separability. To address this limitation, we propose to explicitly incorporate low-level features obtained from shallow layers into the output layer for the pixel reconstruction task. By doing so, we alleviate the burden of the model having to focus excessively on these low-level details, allowing it to spend its modeling abilities to capture these high-level semantics.

We denote the proposed method as Multi-level Feature Fusion (MFF). Specifically, we extend the usage of the fusion strategy in the first pilot experiment and systematically investigate the design choices of multi-feature fusion, such as feature selection and fusion strategies. Despite the simplicity of the proposed method, it is a drop-in solution for unleashing the full modeling potential of pixel-based MIM approaches and has the following advantages:

1. Employing multi-level feature fusion can enhance the training efficiency of MAE by approximately \(\sim\)​5x, thus helping to reduce the carbon footprint. For example, by pre-training MAE with this strategy for only 300 epochs, we achieve semantic segmentation results that are on par with those obtained after 1600 epochs in the original paper.

2. We also consistently and significantly improve performance across all downstream tasks, including semi-supervised fine-tuning and linear probing. Notably, with a small model such as ViT-S, we outperform MAE by 2.8% on linear probing, 2.6% on semantic segmentation, and 1.2% on fine-tuning.

3. After evaluating our model on four out-of-distribution datasets, we observe that the approach with multi-level feature fusion exhibits greater robustness than the base method.

Furthermore, we conduct a thorough analysis to unveil how multi-feature fusion works for representation learning. Given the exploratory experiments from the perspective of latent features and optimization, we find that the fusion strategy attenuates high-frequency information in the latent features and flattens the loss landscapes. To summarize, our contributions are three-fold:

• Firstly, we develop a multi-level feature fusion strategy for isotropic backbones such as ViT, achieving superior results compared to various pixel-based MIM approaches.

• Secondly, we have conducted a thorough analysis of how this multi-level feature fusion strategy enhances the model from the perspectives of latent features and optimization. Our examination is meticulous and provides valuable insights.

• Lastly, we have performed extensive and rigorous ablation studies on the design details, which also strengthens the validity of our findings.

2 Related Works↩︎

Self-supervised Learning reviously, many works[10], [11] relied on abundant labeled datasets to achieve promising results. However, annotating this data requires a large number of human labors. Therefore, how to effectively capture useful semantics embedded in the abundance of data available on the Internet is currently a hot topic. In recent times, self-supervised learning has witnessed tremendous growth in computer vision, following remarkable achievements in natural language processing. These methods cater to diverse inputs, including images[1], [3], [12], [13], videos[14], [15], and multi-modality inputs[7], [16]. They capture rich semantic information by creating effective proxy tasks, such as contrastive learning and masked image modeling, in large amounts of unlabeled data. In comparison to supervised learning[11], [17], [18], these self-supervised learning approaches have gradually outperformed them in numerous downstream tasks and possess immense potential to become the principal pre-training paradigm.

Feature Pyramid Utilizing multi-level features has been extensively studied in previous years, and one of the most famous applications is the Feature Pyramid Network (FPN)[19]. This technique has been widely used in many dense tasks such as object detection and semantic segmentation to improve the model’s perception of objects of different scales. Incorporating FPN into existing designs in many works[20], [21] has led to significant improvements. However, the multi-level feature fusion module only accepts features of different scales as input, limiting its adaptation to isotropic architectures such as ViT[10], in which features from different layers are of the same scale. In masked image modeling, most approaches choose ViT as their encoder due to the masked patch prediction task. Therefore, there are few works exploring multi-level feature fusion in this domain. Even though some works [22], [23] aim to explore multi-level fusion in masked image modeling, their applications are still limited to the traditional hierarchical architecture and do not address the issue of being biased toward low-level details for these pixel-based methods.

3 Methods↩︎

We now introduce the proposed multi-level feature strategy for pixel-based MIM methods. In 3.1, we first give a short revisit to the universal framework of pixel-based MIM approaches. Then 3.2 describes how to insert multi-level feature fusion into pixel-based MIM approaches. Finally, two key components, the projection layer and fusion strategy, will be discussed in 3.3.

3.1 Introduction to Pixel-based MIM↩︎

We now introduce the unified formulation of recent pixel-based Masked Image Modeling(MIM) methods. This kind of MIM, using raw images as target, is a denoising autoencoder [24], and it follows a simple pipeline. Its primary objective is to predict raw pixel values of the original or post-processed images, such as the high-frequency filtered image in PixMIM [8]. When dealing with masked images, we can feed only visible tokens into the encoder or both visible and mask tokens. If only visible tokens are used for the encoder, both the mask token and the latent feature output by the encoder must be fed into the decoder.

3.2 MIM with Multi-level Feature Fusion↩︎

Multi-level feature fusion(MFF) can be incorporated into most existing pixel-based MIM approaches in a plug-and-play manner. 3 gives an overview of the whole framework. To keep the simplicity, we mainly focus on these steps relevant to MFF, leaving out other steps. Given an image \(\mathbf{I}\in \mathbb{R}^{H\times W\;\times3}\), we feed it into the encoder, \(\mathbb{E}\), to get the latent representations: \[\label{eq:extract} \mathrm{X} = \mathbb{E}(\mathbf{I})\tag{1}\] The latent representations, denoted by \(\mathrm{X} = \{x_0, x_1, ..., x_{N-2}, x_{N-1}\}\), correspond to the output feature from each transformer layer of the ViT, where \(N\) represents the depth of the encoder. For the pilot experiment in 1, we fuse all-level features from each layer of the encoder. However, indiscriminately fusing all of them may introduce redundancy or even makes the model much harder to be optimized. But finding the most effective layers to fuse induces a large search space. To simplify the layer selection procedure, we follow the guidelines below:

1. : We first conduct an ablation study to compare the results of fusing shallow layers or deep layers (as shown in 2, shallow layers contain low-level features, and deep layers contain high-level features), and more details are presented in 5.4. The results show that utilizing the features of the shallow layers performed significantly better than deep ones. Thus intuitive analysis and the quantitative numbers both indicate that the shallow layer should be selected for fusion.

2. : We then examine how many layers should be taken into consideration. In addition to the selected shallow layer and output layer, we also try to explore introducing different numbers of intermediate layers for fusion. We refer the reader to 5.4 for more details of this experiment.

We finally selected \(M=5\) additional layers besides the last layer (6 layers in total) and the output features from those layers are used for fusion. We define the indices for these selected layers as \(\mathcal{W}, |\mathcal{W}|=M\). After that, we apply a projection layer, \(\mathcal{P}_\text{i}\), to each of the additional \(M\) layers before fusion. \[\label{eq:proj} \tilde{\mathrm{X}} = \{\mathcal{P}_\text{i}(x_i)\}_{i \in \mathcal{W}} + \{x_{N-1}\}\tag{2}\] Adding a projection layer to align the feature space between different levels’ features is a common practice in self-supervised learning.

Finally, we introduce the fusion layer, \(\mathcal{F}\), to fuse multi-level features \(\tilde{\mathrm{X}}\): \[O = \mathcal{F}(\tilde{\mathrm{X}})\] \(O\) will be fed into the decoder for pixel reconstruction.

3.3 Instantiation of Projection and Fusion Layers↩︎

We further investigate the instantiation of the projection layer and the fusion layer.

Projection Layer: In terms of the projection layer, indicated as \(\mathcal{P}\), we focus on two popular options, namely linear projection and non-linear projection. Specifically, we instantiate the non-linear projection with the Linear-GELU-Linear structure. Our experiment has revealed that a simple linear layer is sufficient and effective within our framework.

Fusion Layer: The fusion layer aims to gather low-level information from the features of shallower layers feature. We evaluate two commonly employed fusion methods: weighted average pooling and self-attention-based fusion.

\[\label{eq:pooling} O = \sum_{i \in \mathcal{W}}w_i\mathcal{P}_i(x_i) + w_{N-1}x_{N-1}\tag{3}\]

The weighted average pooling fusion is illustrated in 3 . In this equation, \(w_i\) refers to the weight assigned to each of the \(M\) selected layers, while \(w_{N-1}\) is assigned to the output layer. During the training process, all these weights are dynamically updated and summed up to 1. As for the self-attention method, we use an off-the-shelf transformer layer.

\[\label{eq:attn} \hat{O} = \text{MultiHeadAttention}(([{\mathcal{P}_i(x_i)}_{i \in \mathcal{W}}, x_{N-1}])\tag{4}\]

After the multi-head attention layer, we extract the transformed tokens corresponding to \(x_{N-1}\) from \(\hat{O}\) to use in pixel reconstruction. Experimental results demonstrate that the weighted average pooling strategy is comparable to self-attention for this purpose while also being simpler and more computationally efficient.

Our method is intuitive and simple, and can be inserted into most of pixel-based MIM approaches without introducing non-negligible computational overhead. We evaluate it on MAE[3] and PixMIM[8], and more detailed results are shown in 5.2.

4 Analysis↩︎

In order to uncover the dark secret behind multi-level feature fusion, we first conduct a frequency analysis on the model (with and without MMF) in 4.1. In 4.2, we find multi-level feature fusion can be helpful to the optimization of the model, by flattening the loss landscape. Finally, in 4.3, we apply the first pilot experiment introduced in 1 to EVA [25] and supervised ViT [3], to investigate whether they too require low-level features and to confirm that the bias towards low-level details is the unique and inherent drawback of pixel-based MIM.

4.1 Frequency Analysis↩︎

We employ multi-level feature fusion to enhance MAE [3], resulting in MFF\(_\text{MAE}\). The aim of this fusion is to prevent the model from excessively focusing on low-level details. To investigate the change in frequency response before and after the fusion, we transform the feature from the last encoder layer into the frequency domain and calculate the amplitude of various frequency bands. As depicted in 4, multi-level feature fusion diminishes high-frequency responses and amplifies those that belong to the low-frequency range, which supports the efficacy of the fusion technique.

4.2 Optimization Analysis↩︎

Following [9], we analyze the Hessian max eigenvalue spectrum. Multi-level feature fusion has the additional benefit of reducing the magnitude of Hessian max eigenvalues. As shown in 5, the expected hessian max eigenvalue of MFF\(_\text{MAE}\) is smaller than that of MAE[3]. Hessian max eigenvalues represent the local curvature of the reconstruction loss function, and this result suggests that multi-level feature fusion flattens the loss landscape. Large eigenvalues can impede neural network training [26]. Therefore, multi-level feature fusion can help the model learn better representations by suppressing large Hessian eigenvalues.

4.3 Feature Bias of Different Pre-training Methods↩︎

In order to investigate whether being biased towards low-level features is the sole and inherent drawback of pixel-based MIM, we introduce multi-level feature fusion to EVA[25] and supervised ViT[3]. EVA is one of the representative works that focuses on regressing high-level features produced by CLIP[7], while supervised ViT is one of the works that require the model to map the input image to a semantic label. Both EVA and supervised ViT target high-level features that contain rich semantic information describing the input image. As shown in 6, unlike MAE, the weight of the last layer’s feature for both EVA and supervised ViT is significantly higher than that of the shallow layers. This observation suggests that the bias towards low-level features exhibited by these pixel-based MIM approaches is primarily caused by the raw-pixel reconstruction task.

5 Experiment↩︎

In 5.1, we present the experimental settings for pre-training and evaluation. Next, in 5.2, we apply MFF to two MIM baselines, namely MAE [3] and PixMIM [8], and show the improvements brought by such design. In addition, we also evaluate the effectiveness of MFF using a smaller model (, ViT-S) and fine-tune the pre-trained model under a low-shot setting. To evaluate the robustness of our proposed method, 5.3 includes additional analyses that assess the robustness of pre-trained models against out-of-distribution (OOD) ImageNet variants. Finally, 5.4 presents comprehensive ablation studies of our method.

5.1 Experiment Settings↩︎

To ensure the efficacy of our methods and design components, we conducted a series of extensive experiments on image classification using ImageNet-1K[27], object detection on COCO [28] and semantic segmentation on ADE20K[29]. Unless otherwise stated, our default settings are based on ViT-B.

ImageNet-1K [27] The ImageNet-1K dataset comprises 1.3 million images belonging to 1,000 categories and is divided into training and validation sets. To ensure the fairness of our experiments while applying our methods to MAE [3] and PixMIM[8], we strictly follow their original pre-training and evaluation settings on ImageNet-1K. This includes following the pre-training schedule, network architecture, learning rate setup, and fine-tuning protocols. Furthermore, in addition to the conventional fine-tuning protocol, we fine-tune the model using a low-shot setting, where only a fraction of the training set (% and 10%) is used. This approach is consistent with previous works[12] and we ensure that the low-shot fine-tuning setting also strictly follows that of the conventional fine-tuning.

ADE20K [29] To conduct the semantic segmentation experiments on ADE20K, we utilize the off-the-shelf settings from MAE[3]. With this approach, we fine-tune a UperNet[21] for 160k iterations with a batch size of 16 and initialize the relative position bias to zero.

COCO [28] For our object detection experiments on COCO, we adopt the Mask R-CNN approach[20] that enables simultaneous production of bounding boxes and instance masks, with the ViT serving as the backbone. As in MAE, we evaluate the box and mask AP as the metrics for this task. However, we note that there is no universal agreement for the setting of object detection fine-tuning epochs. We have chosen the commonly used 2\(\times\) setting, which fine-tunes the model for 25 epochs. Other settings strictly follow those in ViTDet [30].

Ablation studies We conduct all of our ablation studies based on the customary MAE settings [31], [32]. We pre-train all model variants on ImageNet-1K for 300 epochs and conduct a comprehensive performance comparison on linear probing, fine-tuning, and semantic segmentation. All other settings are consistent with those discussed previously.

5.2 Main Results↩︎

The application of multi-level feature fusion to MAE [3] and PixMIM [8] resulted in significant improvements in various downstream tasks, as shown in ¿tbl:tab:comparison?. After pre-training the model for 300 epochs, we achieve a 0.5%, 1.8%, and 3.6% improvement over MAE in fine-tuning, linear probing, and semantic segmentation, respectively. Additionally, our model exhibits scalability across pre-training epochs and consistently outperforms these base methods by a substantial margin. Compared to these methods, using an extra heavy tokenizer, CLIP, we also gradually close the performance gap with them. Although fine-tuning accuracy is often considered a reliable measure of the quality of non-linear features in a model, we find that it is not a sensitive metric, as compared to other metrics presented in ¿tbl:tab:comparison?. This may be attributed to pre-training and fine-tuning following the same data distribution, and the size of the training set and model capacity being sufficient to offset the performance gap between different methods. To address this limitation, we adopt the following two workarounds:

Low-shot fine-tuning. This protocol has also been adopted by many previous works, [12]. Rather than utilizing the entire training set, we fine-tune the pre-trained model end-to-end using only 1% and 10% of the training set. As indicated by ¿tbl:tab:comparison?, the performance gap between MFF and the base methods is much more prominent when using low-shot fine-tuning, which further verifies the effectiveness of MFF.

Pre-train with ViT-S. To mitigate the influence brought by model capacity, we pre-train MAE using ViT-S and compare their performance using fine-tuning, linear probing, and semantic segmentation. Since our objective is to evaluate the improvement brought by MFF to these base methods, we do not specifically tune hyper-parameters for the experiments with ViT-S to achieve state-of-the-art performance, but rather use the same settings as for ViT-B. Due to its smaller capacity compared to ViT-B, ViT-S requires a pre-training method that can effectively capture semantic features to perform well on downstream tasks. As demonstrated in 7, the method with MFF significantly outperforms their base method, further validating the effectiveness of MFF.

We also evaluate our pre-trained models with the object detection protocol and report the AP\(^\text{box}\) and AP\(^\text{mask}\). As shown in 1, MFF can still bring non-trivial improvements for object detection.

5.3 Robustness Evaluation↩︎

Robustness evaluation is a common practice in many previous works [3], [5], [37] to assess a model’s ability to handle different types of noise. In this study, we compare our pre-trained models with their corresponding baselines on four out-of-distribution (OOD) ImageNet variants: ImageNet-Corruption [38], ImageNet-Adversarial [39], ImageNet-Rendition [40], and ImageNet-Sketch [41]. These datasets introduce various domain shifts to the original ImageNet-1K and are widely used to evaluate a model’s robustness and generalization ability. As illustrated in 2, MFF significantly improves the robustness of MAE and PixMIM on all datasets by a clear margin. The enhanced robustness against domain shifts strengthens the value of our simple yet effective method.

5.4 Ablation Studies↩︎

Is shallow layer important? In order to determine the significance of low-level features from shallow layers, we explore the fusion of the output layer with either a deep or shallow layer. So we try to fuse the output layer with an extra shallow or deep layer, selected from the previous 11 layers of ViT-B[10]. And the specific index of the selected layer will be detailed in the appendix. As illustrated in [tbl:tab:shallow95layer], fusing the output layer with a deep layer only results in marginal improvements. However, incorporating low-level features directly from the shallow layer into the output layer leads to a significant performance boost, as it enables the model to focus on semantic information. Therefore, we have decided to use a shallow layer (the first layer) for multi-level feature fusion.

How many layers are used for fusion ? Aside from the output layer and the shallow layer picked in the previous selection, it is reasonable to consider using intermediate layers for fusion, as they may contain additional low-level features or high-level meanings that could assist in the reconstruction task. However, selecting these intermediate layers is a daunting task due to the large search space involved. To simplify the process, we pick an additional 1, 2, and 5 layers evenly spaced between the shallow layer and output layer selected in [tbl:tab:shallow95layer]. The specific indices for these selected layers are placed in the appendix. As shown in [tbl:tab:layers], introducing more layers brings consistent improvements because they may contain unique features, such as textures or colors, that help the model complete the reconstruction task. Nevertheless, when we fuse all these layers, we witness a performance drop in all downstream tasks. This drop may result from the difficulty of optimization, because of the redundancy in these layers.

Do the projection layer and fusion strategy matter? In 2 , we investigate the influence of the projection layer on the final results. Our findings indicate that a simple linear projection layer is sufficient to achieve satisfactory results, as compared to using no projection layer or a nonlinear projection layer. Incorporating a single linear projection layer offers benefits in mitigating the domain or distribution gap between different layers, as compared to using no projection layer. However, the addition of a nonlinear projection layer, which includes an extra linear projection and GELU activation before the linear projection, introduces computational overhead and is more challenging to optimize. As a result, the non-linear projection achieves sub-optimal performance. With regard to the fusion strategy, we found that the weight-average pooling strategy, which assigns a dynamic weight to each layer and then performs element-wise addition, achieves the best performance. Compared to attn, this strategy shares the merits of simplicity and smaller computational overhead.

6 Conclusion↩︎

In this study, we take the first step to systematically explore multi-level feature fusion for the isotropic architecture, such as ViT, in masked image modeling. Initially, we recognize that pixel-based MIM approaches tend to excessively rely on low-level features from shallow layers to complete the pixel value reconstruction task by a pilot experiment. We then apply a simple and intuitive multi-level feature fusion to two pixel-based MIM approaches, MAE and PixMIM, and observe significant improvements in both, gradually closing the performance gap with these approaches by using an extra heavy tokenizer. Finally, we conduct an extensive analysis of multi-level feature fusion and find that it can suppress high-frequency information and flatten the loss landscape. We believe that this work can provide the community with a fresh perspective on these pixel-based MIM approaches and continue to rejuvenate this kind of simple and efficient self-supervised learning paradigm.

7 Appendix↩︎

7.1 Pre-training↩︎

The settings for pre-training strictly follows those in MAE[3] and PixMIM[8], with details shown below:

7.2 Fine-tuning and linear probing↩︎

We also stick to the settings in MAE[3] for the ViT-B[10] model concerning fine-tuning and linear probing. Since our objective is to measure the enhancement brought by MFF and not attain the state-of-the-art (SOTA) performance, we employ the same settings as ViT-B without any specific adjustments for ViT-S.

7.3 Object detection and segmentation in COCO↩︎

All these settings also strictly follow those in MAE[3] but choose the commonly used 2\(\times\) settings, which fine-tunes the model on COCO[28] for 25 epochs.

7.4 Semantic segmentation in ADE20K↩︎

We stick to the settings used in MAE[3] and PixMIM[8], fine-tuning the pre-trained model end-to-end for 16k iterations with a batch size of 16.

7.5 Selected indices of the ablation study↩︎

Inspired by the results of the pilot experiment depicted in Figure 1 of the main paper, we choose layer\(_0\) as the shallow layer, and layer\(_{10}\) as the deep layer for the ablation experiment outlined in Table 3(a). Additionally, for ablation study in Table 3(b), we have selected additional two, four, and ten layers, evenly distributed between layer\(_0\) and the output layer (layer\(_{11}\)). The detailed indices for Table 3(b) is shown in the 6.

In addition, similar to the pilot experiment in Figure 1 of the main paper, we observe the weight for each layer of all experiments in Table 3(b) of the main paper. Just as shown in 8, no matter in which case, the model increasing relies on these shallow layers for the reconstruction tasks, indicating the significance of injecting low-level information into the output layer.

7.6 Transfer learning↩︎

We also study transfer learning where we pre-train on ImageNet-1K and fine-tune on several smaller datasets. We follow the training recipe and protocol in DINO[53]. MFF \(_\text{MAE}\) consistently outperforms MAE on CIFAR10, CIFAR100, and Stanford Cars. As shown in the following table, MFF \(_\text{MAE}\) consistently improves MAE on all datasets.

7.7 Feature-based MIM does not Suffer from being Biased toward Low-level Feature↩︎

To supplement the findings in 6, we apply multi-level feature fusion (MFF) to EVA[25] and MILAN[5], and evaluate their performance with linear probing, fine-tuning and semantic segmentation. Detailed results are shown below:

As shown in the table above, MFF brings marginal improvements to feature-based MIMs, consistent with the findings in 6.

7.8 The Effect of Deep Supervision↩︎

To exclude the influence of deep supervision[54], we detach all shallow layers before fusing with the last layer (MFF\(_\text{MAE}^\text{detach}\)), ensuring that gradients do not propagate through these shortcuts to the shallow layers. As shown in the table below, deep supervision alone does not improve MAE, and MFF’s improvements come from alleviating the problem of being biased toward high-freq components.

References↩︎

---

1. https://github.com/open-mmlab/mmpretrain↩︎