Title: RevColV2: Exploring Disentangled Representations in Masked Image Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: architecture design

representation learning

masked image modeling

self-supervised learning
Abstract: Masked image modeling (MIM) has become a prevalent pre-training setup for vision foundation models and attains promising performance. Despite its success, existing MIM methods discard the decoder network during downstream applica- tions, resulting in inconsistent representations between pre-training and fine-tuning and can hamper downstream task performance. In this paper, we propose a new architecture, RevColV2, which tackles this issue by keeping the entire autoen- coder architecture during both pre-training and fine-tuning. The main body of RevColV2 contains bottom-up columns and top-down columns, between which information is reversibly propagated and gradually disentangled. Such design enables our architecture with the nice property: maintaining disentangled low-level and semantic information at the end of the network in MIM pre-training. Our experimental results suggest that a foundation model with decoupled features can achieve competitive performance across multiple downstream vision tasks such as image classification, semantic segmentation and object detection. For exam- ple, after intermediate fine-tuning on ImageNet-22K dataset, RevColV2-L attains 88.4\% top-1 accuracy on ImageNet-1K classification and 58.6 mIoU on ADE20K semantic segmentation. With extra teacher and large scale dataset, RevColv2-L achieves 62.1 APbox on COCO detection and 60.4 mIoU on ADE20K semantic segmentation.

RevColV2: Exploring Disentangled Representations in Masked Image Modeling

Abstract

Masked image modeling (MIM) has become a prevalent pre-training setup for vision foundation models and attains promising performance. Despite its success, existing MIM methods discard the decoder network during downstream applications, resulting in inconsistent representations between pre-training and fine-tuning and can hamper downstream task performance. In this paper, we propose a new architecture, RevColV2, which tackles this issue by keeping the entire autoencoder architecture during both pre-training and fine-tuning. The main body of RevColV2 contains bottom-up columns and top-down columns, between which information is reversibly propagated and gradually disentangled. Such design enables our architecture with the nice property: maintaining disentangled low-level and semantic information at the end of the network in MIM pre-training. Our experimental results suggest that a foundation model with decoupled features can achieve competitive performance across multiple downstream vision tasks such as image classification, semantic segmentation and object detection. For example, after intermediate fine-tuning on ImageNet-22K dataset, RevColV2-L attains 88.4% top-1 accuracy on ImageNet-1K classification and 58.6 mIoU on ADE20K semantic segmentation. With extra teacher and large scale dataset, RevColv2-L achieves 62.1 AP box on COCO detection and 60.4 mIoU on ADE20K semantic segmentation.

1 Introduction

Pre-trained vision foundation models attract more attention in vision community [1, 2, 3, 4]. A key component in pre-training is how to learn generalizable features which meet the demands of various visual applications [5]. Typical series of methods focus on self-supervised learning, such as contrastive learning [6, 7] and masked image modeling (MIM) [8, 9]. The latter obtains promising results recently in most scenarios by learning occlusion invariant features [10], and becomes a commonly used approach in vision pre-training especially for large models [11, 4, 12, 13].

Dog

(a) Supervised

(b) MAE

(c) BEiT

43

78

91

55

67

14

89

35

24

Discrete Token

(d) RevColV2

Low-Level

Semantic

A typical MIM method, masked autoencoders (MAE) [8], employs an encoder to embed the masked images into semantic features and a decoder to reconstruct unseen patches, as shown in Figure 1 (b). Under such pre-training paradigm, features are rich in low-level information in both input and output. The semantic features, which are desired for downstream tasks, are reserved inside network. A common method to utilize such semantic feature is to manually partition the encoder and decoder based on the amount of semantic information in features and discard the decoder during downstream fine-tuning [8, 14, 15, 16]. Even so, discarding parts of the pre-trained network could incur information loss when transferring to downstream visual tasks. Existing works alleviate this defect by utilizing encoder only architectures [9, 17] and jointly modeling the un-masked patches and mask tokens. Even with extra computation cost in pre-training, these methods often show inferior generalization abilities, because the low-level information used for reconstructing images and the semantic information still appears in an entangled form during pre-training, as shown in Figure 1 (c).

In this paper, we tackle this problem in terms of architecture design. Rather than discarding the decoder, we keep the entire autoencoder architecture in both pre-training and fine-tuning. Similar architecture has proven to be successful in masked language modeling [18]. Nevertheless, vision tasks are naturally different from language tasks, due to their inconsistent output space between pre-training and fine-tuning [19, 20]. In order to obtain better transfer ability as well as keep the unified encoder-decoder architecture, it is important to separate the low-level and semantic information during the image reconstruction process in pre-training. RevCol [1] is a pioneer that combines the idea of reversible network [21, 22] and multi-column architecture [23, 24] to learn disentangled representation in label-guided supervised pre-training. A straightforward attempt is directly combining RevCol with the decoder used in MAE to perform MIM pre-training. However, in order to reconstruct raw images, both low-level and semantic information is needed for reasoning the unseen content and detail appearance, resulting in entangled information. This not only harms the downstream tasks but also destroys the disentangled learning objective of RevCol.

As shown in Figure 2, we re-design the architecture of RevCol to fit the MIM pre-training target. The new architecture contains a bottom-up reversible column encoder and a top-down reversible column decoder. The bottom-up columns and top-down columns are totally symmetric with masked images and encoder embedding as input. During MIM pre-training, the raw image reconstruction loss is connected to the end of the last column in decoder. Hence low-level information primarily sinks to the bottom level and semantic information moves upwards to other stages based on lossless propagation, as shown in Figure 1 (d). Between bottom-up columns and top-down columns, reversible connections are added to ensure the disentangle feature learning object of the whole network. Benefited from the disentangled representations, during downstream fine-tuning, features in the last top-down columns can be quickly adapted to various tasks. Thus the unified architecture does not have to discard the decoder network 3 and avoids the mixture of low-level and semantic information, fully exploring the pre-training abilities. The re-designed new architecture is named RevColV2 .

Webuild various RevColV2 models with different number of parameters and computation budgets and evaluate them on downstream ImageNet [25] image recognition, ADE20K [26] semantic segmentation and COCO [27] object detection after MIM pre-training. The experimental results show consistent improvements over RevCol(V1). Compared with other counterparts, RevColV2 achieves comparable performance to state-of-the-art models pre-trained on pure ImageNet-1K dataset, verifying the effectiveness of the new RevColV2 architecture. After intermediately fine-tuning on ImageNet-22K dataset, RevColV2-L achieves 88.4% top-1 accuracy on ImageNet-1K and 58.7 mIoU on ADE20K segmentation. Moreover, we design a new pre-training scheme which joint modeling the masked semantic features in the top-level, as well as the low-level image pixels in the bottom. Under such pre-training scheme and large scale dataset Laion400M[28], RevColv2-L achieves 62.1 AP box on COCO detection. In ADE20K segmentation, RevColv2-L reaches 60.4 mIoU with multi-scale test. We also conduct analytical experiments showing that RevColV2 with bottom-up reversible columns and top-down reversible columns can learn disentangled representation in MIM pre-training. This property benefits RevColV2 to share unified architecture between pre-training and fine-tuning, fully exploiting the potential of vision pre-training.

2 RevCol V2

In this section, we introduce the newly designed RevColV2, which is a pure isotropic transformer architecture [29]. The core idea in RevColV2, learning disentangled representations during MIM pre-training, is accomplished by the proposed symmetrical reversible encoder-decoder columns and the unified architecture in pre-training and fine-tuning.

2.1 Preliminary

RevCol [1] is a reversible column architecture which learns disentangled representation in supervised learning. The main body of RevCol is composed of multiple subnetworks named columns respectively. Each column contains several basic residual blocks and could be partitioned into four levels with corresponding feature maps. Between each column, reversible connections are introduced to keep lossless information propagation. In Equation 1, we summarize the forward and inverse propagation function in RevCol.

where x l i denotes feature maps of the i -th level in l -th column; F l i denotes an arbitrary non-linear operation in i -th level, analogous to those residual functions in standard ResNets ; γ is a reversible operation ( e.g. channel-wise scaling), whose inverse is denoted by γ -1 . The i -th level takes the features of the lower level at the same column x l i -1 and the upper level at the previous column x l -1 i +1 as input. These two features are fused together and passed through a stack of building blocks. After that, feature x l -1 i are added with reversible operation γ to get the final output. Given the features within one column, we can compute the features in other columns recursively during forward and backward propagation according to Equation 1. At the back-propagation, we can reconstruct the required activations on the fly from the last column to the first, which means we only need to maintain activations from one column in memory during training.

With this reversible property, RevCol can learn disentangled representations in supervised pre-training: the most semantic information is maintained in the last level according to the class label target, and the irrelevant low-level information is kept in other levels based on lossless propagation. In RevColV2, we inherit the RevCol design paradigm but make some improvements for MIM pre-training.

2.2 Symmetrical Encoder and Decoder with Reversible Columns

Figure 2 gives a sketch of the RevColV2 top-level architecture design. The bottom-up column encoder follows RevCol [1] in macro design. The input image is first split into non-overlapping patches by a patch-embedding module. In MIM pre-training, the image patches are randomly masked and the un-masked patches are input into each bottom-up column. The forward and inverse function still follows Equation 1, but we change the operation inside F . We select hardware-friendly vanilla transformer block [29] which includes pre-LayerNorm, Self-Attention and Feed-Forward network (FFN) as the basic building blocks in each column. The blocks are evenly split into four levels for

concat masked patches

Level 3

Level 2

Level 1

Level 0

Level 3

Level 2

Level 1

Level 0

Level 3

Level 2

Level 1

Level 0

Level 3

Level 2

Level 1

Level 0

Reversible

Connection

Intra-Column

Connection

Inter-Column

Connection

Dog

γ

γ

-1

Layer

Norm

Layer

Norm

Linear

Bottom-up

Columns

Top-down

Columns

十

easy implementation of downstream tasks. Two input feature maps of each level are first normalized by a LayerNorm module and then summed together. Consider one of the input x l -1 i +1 is from the previous column, we add another linear transformation on this inter-column branch to project the input space into current column's, as shown in Figure 2.

The top-down column decoder is symmetric to the bottom-up column encoder. Similar as MAE [8], output features of the last bottom-up column are first normalized and concatenated with learnable mask tokens, then input into the top-down columns. As the information propagation is in the opposite direction, the input features of each level are also opposite to the encoder. Equation 1 becomes:

Two inputs of F l i becomes x l i +1 (feature of the upper level in current column) and x l -1 i -1 (feature of the lower level in previous column). For the input (to the highest level) of each top-down column, x l i +1 is actually the last level's output feature of the encoder, which follows the design of image patches input to each bottom-up column. At the lowest level of the last top-down column, a pixel-level loss is add to reconstruct the unseen image patches. Obviously, the top-down column is also reversible. Apart from the input, all other calculations remain the same as the encoder, and that's why we describe the encoder and decoder as fully-symmetrical.

Compared with existing architectures [1, 29], the symmetrical reversible column architecture in MIM pre-training has the following advantages:

The disentangled feature learning object in RevCol is no longer restricted to supervised learning. Although the image reconstruction loss at the end of decoder requires mostly low-level information, the semantic information can still be maintained in other levels and gradually refined through reversible propagation. Therefore, we can learn disentangled representations in MIM pre-training without labels.

Low-level and semantic information is retained in both bottom-up columns and top-down columns. Thus there is no need to discard parts of the network during fine-tuning. The bottom-up encoder and top-down decoder can serve as a unified architecture that yields consistent representation during pre-training and fine-tuning.

2.3 The Unified Architecture between Pre-training and Fine-tuning

MIM Pre-training. As shown in Figure 2, in MIM pre-training, masked image patches are fed into the bottom-up columns and reconstruct unseen patches through top-down columns. At the lowest level of the last top-down column, we use mean-square error (MSE) loss to reconstruct raw images, similar to [8]. As described in Section 2.2, the low-level information is mainly gathered in the bottom levels (yellow regions in Figure 1(d)) with the constraint of MSE loss. Under the constraint of lossless propagation, the semantic information is accordingly decoupled to the top levels (blue regions in Figure 1(d)). Under such circumstances, the top-down column decoder learns not only reconstruction details, but also semantic features, indicating that downstream tasks can directly take advantage of these decoder outputs.

Joint Pre-training. Considering the semantic information will be gathered in the top levels during MIM training, we make a further step. We explicitly joint model the masked semantic features in the top-level, as well as the low-level image pixels in the bottom. Specifically, we introduce CLIP [30] as as an extra teacher to model the semantic features and use cosine similarity as loss function. We also keep the pixel reconstruction MSE loss. Thanks to the disentangled feature in the last column, we can apply two different loss separately at the top and bottom. Our method is different from raw pixel reconstruction methods like MAE [8] or mask distillation works like MaskDistill [31], which modeling homogeneous features at the output of the network.

Downstream Fine-tuning. Rather than discard the decoder during fine-tuning, we leverage the entire encoder-decoder architecture. We come up with two fine-tuning methods in downstream tasks, according to different characteristics. For classification tasks, which require highly semantic features only, we apply classification heads on top of the last top-down column. Since the whole network is optimized by raw pixel reconstruction pre-training target, low-level information sinks to the bottom level features and semantic information is preserved at the top. For dense prediction tasks, which require of both semantic and low-level information, we take all level's features of the last top-down column, then directly connect them to task oriented dense prediction heads.

Compared with exists autoencoder architectures [8, 14, 32, 33, 34] that have to drop the decoder in downstream applications, the unified RevColV2 architecture has consistent representations during pre-training and fine-tuning, fully taking advantage of vision pre-training. Meanwhile, encoder only architectures [9, 17, 4, 31] often entangle the low-level and semantic information in the final output. However, RevColV2 can learn disentangled representations during pre-training, naturally avoids these problems. Compared with CNN-based backbones, ViT backbones are cumbersome to transfer in dense prediction tasks, often with an adapter like ViT-Adapter [35] or SimpleFPN in ViTDet [36]. Thanks to the multi-level embedding output, RevColV2 is free from any extra feature pyramid structures (eg, FPN [37], BiFPN[38], etc) or adapters. We simply interpolate the output features in the last column to multiple resolutions, then leverage the hierarchical embedding in dense prediction tasks. The simple and effective solution for dense prediction tasks makes the ViT-based RevColV2 easy to transfer, like previous CNN-based backbones.

2.4 Model Variants

We provide two variants of RevColV2 models, B ase and L arge, as shown in Table 1. The depth of bottom-up columns and top-down columns are fixed to 12 and 4. The number of bottom-up columns and top-down columns is changed from 3 to 4 according to different model sizes. The model parameters include the whole bottom-up columns and top-down columns. However, in the downstream task such as image classification, not all levels are involved in computation (as shown in Figure 2), such the number of parameters is slightly lower (as shown in Table 2). We do not use intermediate supervision for each column as it needs careful tuning.

3 Experiments

3.1 MIMPre-training

3.1.1 Pre-training Details

We pre-train RevColV2 on ImageNet-1K [25] dataset. Hyper-parameters generally follow [8]. The mask ratio is set as 75% with random sampling strategy and the reconstruction target is the normalized raw pixel from the original image. We pre-train 1600 epochs for RevColV2 models. The pre-training image size is 224 2 and the pre-training optimization parameters are: batch size 4096, base learning rate 1.5e-4 for 256 batch-size and linear scaled up, AdamW with weight decay 0.05. We do not use stochastic depth strategy in pre-training. More details can be found in supplementary material.

3.1.2 ImageNet Classification

Setup. For image classification, we evaluate top-1 accuracy on ImageNet-1K [25]. We initialize weights using MIM pre-trained models, and fine-tune on ImageNet-1K with class label, similar to [8, 12, 9]. To fully exploit the potential of RevColV2, we intermediately fine-tune the models on ImageNet-22K [39] following [34, 12] after MIM pre-training. The intermediate fine-tuning hyper-parameters are almost the same as [12] and shown in supplementary material.

Results. Table 2 shows the results on ImageNet-1K classification. The MIM pre-train epochs for RevColV2 is 1600, and other methods are reported with their longest schedules. RevColV2-B achieves 84.7 top-1 accuracy with pure ImageNet-1K data, outperforms ViT [29] architecture pretrained with MIM [8, 9, 14, 40, 32] by a large margin using only raw pixels as reconstruction target. As a pure transformer isotropic architecture, RevColV2-B also achieves comparable performance with state-of-the-art hierarchical architectures, i.e. RevColV2-B reaches higher performance than SwinV2B[11], HorNet GF -B [44] and other counterparts in Table 2. For large size models, RevColV2-L with 86.3% top-1 accuracy outperforms ConNeXt V2 [12], MAE [8], and CAE [14] counterparts.

We initialize RevColV2 models with the MIM pre-trained model weights and then use larger dataset ImageNet-22K and supervised training methods to test the scaling up ability. Under this training

schedule, RevColV2-B reaches 87.5% top-1 accuracy on ImageNet-1K with 512 2 input resolution. For large size models, the performance of RevColV2 achieves 88.4 % top-1 accuracy, surpassing hierarchical counterparts of SwinV2 [11] and ConvNeXt V2 [12]. Compared with over RevCol(V1) [1], all sizes RevColV2s show consistent and significant improvements on classification tasks, further illustrating the effectiveness of our design.

3.1.3 Semantic Segmentation

Setup. For semantic segmentation tasks, we evaluate RevColV2 backbones on ADE20K benchmarks [26] with UperNet [46] and Mask2Former [47] framework. We interpolate the position embedding to a fixed input size for both bottom-up and top-down columns. Following previous works [11, 12], we initialize weights using ImageNet-1k classification fine-tuned models. Training hyper-parameters are available in supplementary material.

Results. Table 3 shows the head-to-head comparison results on ADE20K semantic segmentation with various backbones and UperNet segmentation head. RevColV2 models gain competitive performance over single-scale and multi-scale mIoU across different pure transformer and CNN architectures. RevCol-B/L pre-trained on ImageNet-1K achieve 52.3 and 54.4 mIoU on ADE20K with 512 2 input resolution respectively, outperforming other counterparts such as ConvNeXt V2 [12]. After intermediate fine-tuning on ImageNet-22K, RevColV2-B/L reach 53.9 and 55.7 mIoU on ADE20K benchmark with 640 2 resolution, exceeding other competitors such as Focal transformer [52] and CSwin [53]. With stronger segmentation framework Mask2Former [47], our RevCol-B/L backbones achieve 55.8 and 58.6 mIoU, further demonstrating the effectiveness of RevColV2 backbones.

3.1.4 Object Detection

Setup. For object detection and instance segmentation task, we evaluate RevColV2 backbones on COCO [27] dataset with Mask R-CNN [55] and Cascade Mask R-CNN [56] framework. We follow

the setting in [36] that using window attention in RevColV2 backbones. We train models with 1024 2 resolution crops using large scale jittering augmentations [57].

Results. Table 4 left shows the experiment results for RevColV2 backbones compared with MAE [8] baseline using Mask R-CNN detector. For MAE baseline, we reproduce the results using the same training configuration optionally with ViT-Adapter [35]. RevColV2 series achieve 52.4 and 54.0 box AP for base and large models, outperforming MAE series. In Table 4 right, we compare RevColV2 backbones with other architectures using Cascade Mask R-CNN detector. RevColV2-B achieve 55.2 box AP and 47.9 mask AP, outperforming RevCol(V1)-B [1], ViTDet-B [36] and other competitors.

Note that in all dense prediction tasks, we do not use any extra feature pyramid structures (eg, FPN, BiFPN, etc.) or adapters like ViT-Adapters in RevColV2 backbones implementation.

3.2 Joint Pre-training and Data Scaling

3.2.1 Pre-training Details

In joint pre-training, we use OpenCLIP-L as the teacher to represent the semantic features similar to MaskDistill and EVA. Except for the additional teacher, we use a larger dataset Laion400M[28], which contains about 400M unlabeled images in pre-training. Note that we do not use datasets such as COCO, ADE20K, Object365, etc. in pre-training to avoid artificial fitting to specific distribution (this is different from EVA-02 which uses a merged dataset that has overlapped data in the downstream task). We use 800 ImageNet-1k epochs on Laion400M dataset and then 300 epochs on ImageNet-1k dataset during pre-training.

3.2.2 Results

Then we evaluate our model on downstream tasks such as ImageNet1K classification, COCO detection with cascade Mask-RCNN, and ADE20K semantic segmentation with Mask2Former. The newly trained RevColv2-L achieves 87.7% Top-1 accuracy in ImageNet-1k classification with 224 × 224 input resolution. The larger dataset and the extra teacher lead to better performance compared with purely IN-1k MIM pre-training (86.3%) and IN-1k MIM + IN-22k intermediate fine-tuning (87.4%). As shown in Table 5, the performance gain is more prominent on dense prediction tasks. RevColv2-L achieves 62.1 box AP and 52.3 mask AP in COCO instance detection and segmentation. In ADE20K segmentation, RevColv2-L reaches 60.4 mIoU with multi-scale test.

3.3 Analysis

Experiments in this section are based on basic MIM pre-training.

RevColV2 can learn disentangle representation. Linear probing is a useful tool to evaluate the sparsity of features. We assume the semantic information tends to be more sparse, while the lowlevel information is more abundant. So, we evaluate the linear probing accuracy of each level to visualize the distribution of semantic and low-level information. We construct two experiments: 1) RevCol(V1) [1]+MAE [8] baseline in which the RevCol(V1) encoder has 3 columns and MAE decoder has 8 ViT blocks; 2) RevColV2 with 3 bottom-up columns encoder and 3 top-down columns decoder in which each column contains 12 blocks. The number of model parameters is maintained at the same level. Figure 3 shows the evaluation results on ImageNet-1K. The output level of the baseline encoder only reaches 47.2% accuracy, which is far less than the previous level 60.1%. The

Level 1

Level 2

Level 3

Level 4

Train with MAE decoder

Train with reversible multi-column decoder

Cor

Cor

Col

58.0 %51.9 %47.2%

51.0 %60.1%49.5%

17.9 %25.5%33.6%

Corr

Cor

Cors

Cor"

Cor-

Col

%79%919%99%69%9%9

%769%69%199%679%L1%067

18.3 %25.8 %31.2 %53.8 %53.9 %52.6 %

%8%04%0%10%%21%601

baseline method will entangle the low-level and semantic information during column propagation, because the MAE decoder needs both low-level and semantic features to reconstruct unseen patches. Conversely, RevColV2 with reversible top-down columns can learn the disentangled representation as described in Section 2. As shown in Figure 3, the last level of encoder and decoder reaches high linear probing accuracy (64.9% and 62.5%), and the accuracy decreases at the lower levels. Thus low-level and semantic information lies in a disentangled manner between the bottom and top levels in the last column.

We also show the linear probing accuracy of various sizes RevColV2 models in supplementary material. The separated semantic information significantly boosts the linear probing accuracy compared to other encoder only methods such as SimMIM [17].

RevColV2 matters in unified representation fine-tune. We experimentally verify the effectiveness of RevColV2 with unified architecture between pre-training and fine-tuning. We first compare the bottom-up encoder only variant and the entire autoencoder variant of RevColV2 on ImageNet-1K. The latter has about 0.9% (83.8% v.s. 84.7%) top-1 accuracy gain, showing the top-down columns are critical in downstream fine-tuning. To eliminate the influence of model capacity (#params. and FLOPs), we also build one experiment with the same number of parameters and FLOPs for comparison. The experimental results show the accuracy of this encoder only variant (83.9%) is still lower than the entire auto-encoder variant(84.7%).

Similarly, we conduct experiments for ViT [29] + MAE [8] that without dropping decoder in the fine-tuning stage. The accuracy on ImageNet-1K drops about 0.4% (from 83.6% to 83.2%). Give more model capacity does not increase the performance, indicating that the vanilla ViT decoder is useless in downstream fine-tuning.

We also make an ablation that only utilize the encoder's pre-trained weights while initializing the decoder weights randomly with RevColV2-B. This variant achieves 84.4% (-0.3%) top-1 accuracy on ImageNet-1K dataset and 50.7 (-0.6) mIoU on ADE20K dataset, with only ImageNet-1K MIM pre-trained encoder weights. These experimental results draw the same conclusion that the pre-trained decoder is necessary for RevColV2.

3.4 Ablation Study

Experiments in this section are based on basic MIM pre-training.

Depth of top-down decoder columns is analyzed similar to the decoder depth analysis in MAE [8], as shown in Table 6 (a). In MAE [8], the fine-tuning accuracy does not vary much over the decoder depth. However, the performance of RevColv2 drops with deeper decoder. This is because the top-down column decoder also takes up model parameters as we do not discard decoder during finetuning. We keep the total parameter the same, and deeper decoder will lead to shallower (narrower) encoder, which could give rise to degraded performance. This experiment also explains why we use shallower decoder (4 blocks decoder compared to 12 blocks encoder) in our RevColV2 variants.

Different type of decoder . The decoder in RevColV2 is a symmetric architecture with encoder, that includes multiple reversible top-down columns. To verify the effectiveness of these top-down columns, we compare this design with other variants: MAE decoder with a stack of ViT blocks; single

column decoder with shortcuts from encoder similar to UNet [58]; top-down multi-columns without shortcuts and reversible connections. Experiment results in Table 6 (c) draw the same conclusion with Section 2 that reversible column decoder achieves the best performance. This is because top-down reversible multi-columns can help feature disentangling while others not, as described in Section 2 and the analysis in Section 3.3.

Different pre-training schedules . We also conduct experiments with shorter schedules compared with the default 1600 epochs training, i.e. 300/800 epochs. Table 6 (b) shows the fine-tuning top-1 accuracy on ImageNet-1K. The results indicate larger models may need longer iterations (+0.6% from 300 to 1600 epochs for base size model, +1.1% from 300 to 1600 epochs for large size model) in MIM pre-training to fully activate the potential of models.

4 Related Works

Disentangled feature learning. Disentangled representation refers to separating the factors of variantion [5]. Traditional methods [59, 60, 61] mainly focus on generative models to learn disentangled representation. Hinton proposes GLOM [23] to build a general network that learning the part-whole hierarchies. Cai et al. [1] integrate the idea of GLOM with reversible networks [22, 21] to learn disentangle representation in supervised learning paradigm, namely RevCol. Although RevCol can learn disentangled representation, it needs large-scale labeled data to maintain this property that can not directly benefit from MIM. In this paper, We seek the advances of re-design RevCol with MIM to further enhance the representation abilities while maintaining the decoupled feature learning ability.

Masked image modeling. Self-supervised learning has a long history in computer vision research community [62, 63, 64, 65, 66, 7]. The recent method masked image modeling (MIM) erases the masked image patches and then predicts the unseen contents, a representative of which is masked autoencoders [8]. Following MIM pipeline, researchers make efforts on designing different reconstruction targets, such as DALL-E [9, 14], HOG [32], VQGAN [48], frequency [33, 67], to learn occlusion invariant features [10].

Some methods try to use consistent architecture during pre-training and fine-tuning, such as SimMIM [17] which utilizes an encoder only network. As a result, low-level and semantic information is entangled at the end of the network, resulting in degraded performance, especially for the linear probing accuracy. The above methods usually directly adopt vanilla ViT as backbones, ignoring the mutual promotion between architectures and MIM pipelines. Recent works [12, 68] aim to co-design the architecture with MIM, and focus on the CNN architectures [41]. In this paper, we explore the new research direction: learning disentangled representation during MIM. We design a unified architecture between pre-training and fine-tuning based on RevCol [1] to learn disentangled features and naturally avoid the above information mixture problem.

5 Conclusion

In this paper, we design a new architecture named RevColV2 which learns disentangled representations during MIM pre-training and keep a unified autoencoder architecture when transferring to downstream tasks. RevColV2 extends the reversible columns network which is previously limited in supervised learning to MIM training, and bridges the gap between pre-training and fine-tuning. These unified architectures improve the performance across various downstream tasks, including image classification, object detection and semantic segmentation, without using additional task specific adapters. By these impressive results, we hope to stimulate more research in learning generalizable features and help foundation model pre-training.

References

Yuxuan Cai, Yizhuang Zhou, Qi Han, Jianjian Sun, Xiangwen Kong, Jun Li, and Xiangyu Zhang. Reversible column networks. arXiv preprint arXiv:2212.11696 , 2022.

Hongyu Wang, Shuming Ma, Shaohan Huang, Li Dong, Wenhui Wang, Zhiliang Peng, Yu Wu, Payal Bajaj, Saksham Singhal, Alon Benhaim, et al. Foundation transformers. arXiv preprint arXiv:2210.06423 , 2022.

Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorck, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, et al. Image as a foreign language: Beit pretraining for all vision and vision-language tasks. arXiv preprint arXiv:2208.10442 , 2022.

Yuxin Fang, Wen Wang, Binhui Xie, Quan Sun, Ledell Wu, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao. Eva: Exploring the limits of masked visual representation learning at scale. arXiv preprint arXiv:2211.07636 , 2022.

Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence , 35(8):17981828, 2013.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 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/CVF Conference on Computer Vision and Pattern Recognition , pages 16000-16009, 2022.

Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint arXiv:2106.08254 , 2021.

Xiangwen Kong and Xiangyu Zhang. Understanding masked image modeling via learning occlusion invariant feature. arXiv preprint arXiv:2208.04164 , 2022.

Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12009-12019, 2022.

Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie. Convnext v2: Co-designing and scaling convnets with masked autoencoders. arXiv preprint arXiv:2301.00808 , 2023.

Mannat Singh, Quentin Duval, Kalyan Vasudev Alwala, Haoqi Fan, Vaibhav Aggarwal, Aaron Adcock, Armand Joulin, Piotr Dollár, Christoph Feichtenhofer, Ross Girshick, et al. The effectiveness of mae pre-pretraining for billion-scale pretraining. arXiv preprint arXiv:2303.13496 , 2023.

Xiaokang Chen, Mingyu Ding, Xiaodi Wang, Ying Xin, Shentong Mo, Yunhao Wang, Shumin Han, Ping Luo, Gang Zeng, and Jingdong Wang. Context autoencoder for self-supervised representation learning. arXiv preprint arXiv:2202.03026 , 2022.

Christoph Feichtenhofer, Yanghao Li, Kaiming He, et al. Masked autoencoders as spatiotemporal learners. Advances in neural information processing systems , 35:35946-35958, 2022.

Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training. arXiv preprint arXiv:2203.12602 , 2022.

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 , pages 9653-9663, 2022.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Xinlong Wang, Wen Wang, Yue Cao, Chunhua Shen, and Tiejun Huang. Images speak in images: A generalist painter for in-context visual learning. arXiv preprint arXiv:2212.02499 , 2022.

Jia Ning, Chen Li, Zheng Zhang, Zigang Geng, Qi Dai, Kun He, and Han Hu. All in tokens: Unifying output space of visual tasks via soft token. arXiv preprint arXiv:2301.02229 , 2023.

Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. The reversible residual network: Backpropagation without storing activations. Advances in neural information processing systems , 30, 2017.

Jörn-Henrik Jacobsen, Arnold Smeulders, and Edouard Oyallon. i-revnet: Deep invertible networks. arXiv preprint arXiv:1802.07088 , 2018.

Geoffrey Hinton. How to represent part-whole hierarchies in a neural network. Neural Computation , pages 1-40, 2022.

Jingdong Wang, Ke Sun, Tianheng Cheng, Borui Jiang, Chaorui Deng, Yang Zhao, Dong Liu, Yadong Mu, Mingkui Tan, Xinggang Wang, et al. Deep high-resolution representation learning for visual recognition. IEEE transactions on pattern analysis and machine intelligence , 43(10):3349-3364, 2020.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 633-641, 2017.

Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13 , pages 740-755. Springer, 2014.

Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. arXiv preprint arXiv:2111.02114 , 2021.

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. arXiv preprint arXiv:2010.11929 , 2020.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. A unified view of masked image modeling. arXiv preprint arXiv:2210.10615 , 2022.

Chen Wei, Haoqi Fan, Saining Xie, Chao-Yuan Wu, Alan Yuille, and Christoph Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14668-14678, 2022.

Jiahao Xie, Wei Li, Xiaohang Zhan, Ziwei Liu, Yew Soon Ong, and Chen Change Loy. Masked frequency modeling for self-supervised visual pre-training. arXiv preprint arXiv:2206.07706 , 2022.

Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. Beit v2: Masked image modeling with vector-quantized visual tokenizers. arXiv preprint arXiv:2208.06366 , 2022.

Zhe Chen, Yuchen Duan, Wenhai Wang, Junjun He, Tong Lu, Jifeng Dai, and Yu Qiao. Vision transformer adapter for dense predictions. arXiv preprint arXiv:2205.08534 , 2022.

Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer backbones for object detection. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part IX , pages 280-296. Springer, 2022.

Tsung-Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan, and Serge Belongie. Feature pyramid networks for object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2117-2125, 2017.

Mingxing Tan, Ruoming Pang, and Quoc V Le. Efficientdet: Scalable and efficient object detection. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10781-10790, 2020.

Tal Ridnik, Emanuel Ben-Baruch, Asaf Noy, and Lihi Zelnik-Manor. Imagenet-21k pretraining for the masses. arXiv preprint arXiv:2104.10972 , 2021.

Yabo Chen, Yuchen Liu, Dongsheng Jiang, Xiaopeng Zhang, Wenrui Dai, Hongkai Xiong, and Qi Tian. Sdae: Self-distillated masked autoencoder. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXX , pages 108-124. Springer, 2022.

Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11976-11986, 2022.

Xiaosong Zhang, Yunjie Tian, Lingxi Xie, Wei Huang, Qi Dai, Qixiang Ye, and Qi Tian. Hivit: A simpler and more efficient design of hierarchical vision transformer. In The Eleventh International Conference on Learning Representations , 2023.

Hugo Touvron, Matthieu Cord, and Hervé Jégou. Deit iii: Revenge of the vit. In European Conference on Computer Vision , pages 516-533. Springer, 2022.

Yongming Rao, Wenliang Zhao, Yansong Tang, Jie Zhou, Ser-Nam Lim, and Jiwen Lu. Hornet: Efficient high-order spatial interactions with recursive gated convolutions. arXiv preprint arXiv:2207.14284 , 2022.

Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan, Yukun Zhu, Alan Yuille, Hartwig Adam, and Liang-Chieh Chen. Moat: Alternating mobile convolution and attention brings strong vision models. arXiv preprint arXiv:2210.01820 , 2022.

Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing for scene understanding. In Proceedings of the European conference on computer vision (ECCV) , pages 418-434, 2018.

Bowen Cheng, Anwesa Choudhuri, Ishan Misra, Alexander Kirillov, Rohit Girdhar, and Alexander G Schwing. Mask2former for video instance segmentation. arXiv preprint arXiv:2112.10764 , 2021.

Xiaoyi Dong, Jianmin Bao, Ting Zhang, Dongdong Chen, Weiming Zhang, Lu Yuan, Dong Chen, Fang Wen, and Nenghai Yu. Peco: Perceptual codebook for bert pre-training of vision transformers. arXiv preprint arXiv:2111.12710 , 2021.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Wenhai Wang, Jifeng Dai, Zhe Chen, Zhenhang Huang, Zhiqi Li, Xizhou Zhu, Xiaowei Hu, Tong Lu, Lewei Lu, Hongsheng Li, et al. Internimage: Exploring large-scale vision foundation models with deformable convolutions. arXiv preprint arXiv:2211.05778 , 2022.

Xiaohan Ding, Xiangyu Zhang, Jungong Han, and Guiguang Ding. Scaling up your kernels to 31x31: Revisiting large kernel design in cnns. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11963-11975, 2022.

Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal self-attention for local-global interactions in vision transformers. arXiv preprint arXiv:2107.00641 , 2021.

Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12124-12134, 2022.

Yuxin Fang, Quan Sun, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao. Eva-02: A visual representation for neon genesis. arXiv preprint arXiv:2303.11331 , 2023.

Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision , pages 2961-2969, 2017.

Zhaowei Cai and Nuno Vasconcelos. Cascade r-cnn: Delving into high quality object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 6154-6162, 2018.

Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung-Yi Lin, Ekin D Cubuk, Quoc V Le, and Barret Zoph. Simple copy-paste is a strong data augmentation method for instance segmentation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 2918-2928, 2021.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Guillaume Desjardins, Aaron Courville, and Yoshua Bengio. Disentangling factors of variation via generative entangling. arXiv preprint arXiv:1210.5474 , 2012.

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations , 2017.

Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems , 29, 2016.

Carl Doersch, Abhinav Gupta, and Alexei A Efros. Unsupervised visual representation learning by context prediction. In Proceedings of the IEEE international conference on computer vision , pages 1422-1430, 2015.

Mehdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part VI , pages 69-84. Springer, 2016.

Deepak Pathak, Ross Girshick, Piotr Dollár, Trevor Darrell, and Bharath Hariharan. Learning features by watching objects move. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2701-2710, 2017.

Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. arXiv preprint arXiv:1803.07728 , 2018.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Hao Liu, Xinghua Jiang, Xin Li, Antai Guo, Deqiang Jiang, and Bo Ren. The devil is in the frequency: Geminated gestalt autoencoder for self-supervised visual pre-training. arXiv preprint arXiv:2204.08227 , 2022.

Keyu Tian, Yi Jiang, Qishuai Diao, Chen Lin, Liwei Wang, and Zehuan Yuan. Designing bert for convolutional networks: Sparse and hierarchical masked modeling. arXiv preprint arXiv:2301.03580 , 2023.

Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14 , pages 646-661. Springer, 2016.

A Analysis of Speed and Memory cost

Inference. The current model variants of RevColV2 introduce more latency compared with other works of the similar number of parameters and FLOPs, such as ViT. We test the inference latency of variant models in Table 7. As described RevColV1 [1], fragmented memory access takes a large part of latency. In RevColV2, we make some improvements: 1) remove the up-sample and down-sample operation in RevColV1; 2) reduce the number of total blocks; 3) hard-ware friendly architecture without hierarchy. As shown in Table 7, RevColV2 has lower latency than the V1 version during inference, but is still 1.21x higher than ViT. This is because of the large number of building blocks in RevColV2-L (about twice of ViT-L). Though we reduce the total number of blocks, the multi-column RevColv2 still requires at least 12 blocks in each column in the encoder. Shallower column leads to coarse representation which could harm the performance. On the other hand, if we make the ViT model deeper and maintain the same FLOPs, ViT-L-deeper (48 blocks) and RevColV2-L (48 blocks) have similar latency.

We also analyze the impact of batch size. We show throughput (#image/s) under the different batch size of RevColV2-L and ViT-L on a single A100 GPU.The results in Table 8 show that with the increase in batch size, the inference speed gap between RevCol-L and ViT-L is closing because the fragment memory access time can be distributed to each sample. Although the speed of RevColV2 is lower than vanilla ViT, we think it can be solved by advanced techniques.

Training We make some further analysis on the per-training speed and its memory cost, and compare them with the popular used ViT-MAE baseline. We take RevColV2-B and ViT-B for comparisons. We test training speed and memory cost on a single A100 (80GB) x 8 machine, with the same data-loader (implemented for our cluster). We use our own implementation for RevColV2 and the official implementation for MAE. Table 9 shows the training cost with batch size 4096 for one epoch. To speed up training and save memory, we equip RevColV2 with Flash Attention. We only use data parallel in this testing.

Table 9 shows that the vanilla implementation of RevColV2 pre-training has a little slower (249s vs. 220s) than ViT. Equipped with FlashAttn, RevColV2 achieves comparable pre-training cost (211s vs. 220s and 42G vs. 43G). We further analyze the impact of reversible propagation. We test the pre-training cost of the Reversible version of RevColV2 (re-compute the intermediate features during backward according to the last column outputs, rather than the vanilla autograd function in PyTorch. It is the key component of reversible column networks [1]). Results show that RevColV2-B can use extremely few GPU memory (only 18G) during pre-training with a total batch size 4096. This allows RevColV2 can be pre-trained with limited resources, such as RTX3090 GPU.

In addition to the above comparison, the fragmented access of memory can be optimized by some techniques which can be further investigated in further work. Here, we give two ways that may be further studied:

Kernel fusion. This can reduce the frequent access of the memory caused by a large number of blocks.

Model parallel. Before the calculation of previous columns is finished, parts of the current column can be calculated in parallel. This is the nature of the multi-column network and can be further studied to speed up the inference and training.

B More Training Details

This section gives more training details on MIM pre-training and fine-tuning on downstream tasks, such as ImageNet classification, COCO detection, and ADE20K segmentation. For ImageNet experiments, the base learning rate is based on batch size 256.

B.1 Training Details on MIM pre-training.

We use the same setting for different sizes RevCol models on MIM pre-training. The detail hyperparameters are shown in Table 10. Following exists works [8, 12], we do not use stochastic depth [69] and other regularization strategies in MIM pre-training.

B.2 Details on Image22K intermediate fine-tuning.

We further intermediately fine-tune RevColV2 models on ImageNet-22K dataset. The fine-tuning details is shown in Table 11. The hyper-parameters generally follow [1, 12].

B.3 End-to-end fine-tuning details on ImageNet-1K.

We end-to-end fine-tune RevCol variants on ImageNet-1K after MIM pre-training and intermediately fine-tuning on ImageNet-22K. Table 12 shows the detail training settings after MIM pre-training.

We also show training settings on ImageNet-1K after ImageNet-22K fine-tuning. Table 13 gives the detailed hyper-parameters.

B.4 Details on ADE20K semantic segmentation

For semantic segmentation, we evaluate different backbones on ADE20K dataset. We fine-tune the pre-trained networks on ADE20K with 160,000 iterations. For UperNet framework [46], the learning rate is 4e-5 with batch size 16, using AdamW optimizer. The layer-wise learning rate decay rate is set as 0.65 for both base and large size models. The drop path rate is 0.1. For Mask2Former

framework [47], the learning rate is 2e-5 with batch size 16. The drop path rate is set as 0.3 and the layer-wise learning decay rate is 0.9.

B.5 Details on COCO object detection and instance segmentation

For object detection and instance segmentation, we evaluate RevColV2 backbones with Mask RCNN [55] and Cascade Mask R-CNN [56] detectors. We use ImageNet-1K MIM pre-trained weights as initialization and fine-tune the models with 50 epochs and a batch size of 32, learning rate 1e-4 for Mask R-CNN framework. The large scale jittering data augmentation strategy is used with scale range [0.1, 2.0]. The drop path rates for RevCOlV2 are set as 0.2 (base) and 0.3 (large) and the layer-wise learning rate decay rates are set as 0.9. For Cascade Mask R-CNN framework, we train models with 100 epochs following [36] with large scale jittering augmentation strategy. The learning rate is 1e-4 with batch size 64. We do not use soft-NMS in our experiments.

C More Results

C.1 Compared with supervised baseline

To verify the effectiveness of RevColV2 architecture with MIM pre-train, we compare the performance on ImageNet-1K fine-tune using MIM pre-trained model weights and random initialization. We use the same setting with [8] in this supervised baseline, except additional 0.999 EMA strategy. The base/large models achieve 83.1% and 82.6% top-1 accuracy on ImageNet-1K. The MIM pre-trained RevColV2 models outperform supervised baseline by a large margin ( +1.6% and +3.7% ).

C.2 Linear probing results

We report the linear probing results on ImageNet-1K after pre-training for RevColV2 models and other counterparts on Table 14. Following [8], we fix the pre-trained backbone models and train a classification head for 90 epochs with LARS optimizer. We append this classification head on the last level of bottom-up columns in RevColV2. The linear probing performance of RevColV2 models surpasses other encoder only models such as SimMIM [17] and autoencoder models such as MAE [8].

Peer reviews:
comments: SUMMARY:
In this paper, the authors propose a new backbone RevColv2, which is suitable to the MIM pretraining and could learn disentangled representations during the pretraining. The strong experiment results show its effectiveness.

STRENGTHS:
Please refer to Questions

WEAKNESSES:
Please refer to Questions

QUESTIONS:
### strength
1. The paper is well-written and easy to follow
2. The proposed method is a novel combination of Mask Image Modeling and RevCol
3. The experiment part shows strong results.

### weakness
1. Some details are unclear, line 103 says 'the un-masked patches are input into each bottom-up column'. While line 134 and Fig.2 say 'masked image patches are fed into the bottom-up columns and reconstruct unseen patches through top-down column'.
Which one is the correct way?

2. Compared with vanilla ViT, the proposed model is quite complex, so I'm wondering about the speed of both the pretrain and inference.

3. Fig.3 compare the MAE decoder with RevColv2 decoder and claims that the proposed decoder can disentangle the low-level and high-level feature better. But for the linear accuracy difference between Level1 and Level4, the difference of MAE is 14.6, while the proposed method is 18.7, the difference seems not significant.

4. PeCo is a strong ImageNet-1k classification baseline while only segmentation result on ADE20K is reported in the paper, why?

LIMITATIONS:
Please refer to Questions
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed novel architecture to explore disentangled representations with masked image modeling. Different from previous MAE-like methods, this paper design a unified network and do not drop the decoder in downstream task. This paper showed that the disentangled representation is learned in different network levels. Experiments are done on ImageNet, MS-COCO and ADE20k for base and large model sizes.\n\nSTRENGTHS:\n1. The idea of combining MIM with disentangled representation learning is novel.\n2. The idea of keeping the entire autoencoder tackles the problem of inconsistent representation between pre-training and fine-tuning.\n3. The performance in downstream vision tasks is competitive.\n\nWEAKNESSES:\n1. Figure 2 is misleading, the pre-training target on ImageNet-1k is single MIM or combined MIM with image labels?\n2. The network parameters in Table 1 and Table 2 is not consistent. E.g. for base-size, 101M in Table 1 and 88M in Table 2.\n3. As described in line 204, the initialized weight for semantic segmentation is ImageNet-1k classification fine-tuned model and not MIM pre-trained model. It is not fair to compare it with a bunch of MIM pre-trained models.\n4. The DropPath op seems to be conflict to the idea of disentangle representation in each level.\n5. In the supplementary materials, the linear probing experiments is based on the bottom-up columns which is conflict to the idea of keeping the entire autoencoder.\n6. This paper had several typos and grammar mistakes. E.g.\nline 1: per-training -> pre-training\nline 13: performances -> performance\nline 15: intermediately -> intermediate\nline 66: per-trained -> pre-trained\nline 211: fine-turning -> fine-tuning \n7. There are some factual errors in this paper. E.g. in Table 2, the target of ConvNext-B is label not pixel and the target of BEiT-L is DALL-E not pixel.\n\nQUESTIONS:\nQuestion:\n1. Speed. It seems like the training and inference speed of revcol architecture is slower than deeper plain ViT. What is the training and inference speed of RevColV2 in pre-training and fine-tuning?\n2. What is the performance of dropping the decoder and only using the encoder for downstream tasks?\n3. How to show the scalability of RevColV2?\n4. Is there more evidence to show the representation difference of RevColV2 and RevCol? E.g. attention distance, KL divergence between different attention heads, etc.\n5. What is the performance of not using DropPath?\n6. Why using bottom-up columns representation for linear probing?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper proposes to keep the entire auto-encoder architecture during both pre-training and fine-tuning based on RevCoI. It contains bottom-up columns and top-down columns, and the information is reversibly propagated and gradually disentangled between them. Better results are achieved on ImageNet-1k and downstream tasks.

STRENGTHS:
1. Good motivation to keep the same structure for both pretraining and finetuning. 
2. Better results are achieved compared with RevCol.

WEAKNESSES:
1. This paper is a little bit hard to follow, and I do not think the figures help a lot for understanding this paper. Maybe better visualizations/figures are needed. 
2. Tiny and small size models are also suggested to exploit. Or the authors needs strong arguments why they did not do so. 
3. Some papers achieves better results than this method, but they did not present them and compare with them.

QUESTIONS:
1. What the performances for tiny and small size model?

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a revised version of RevCol, referred to as RevColV2, which is applicable for MAE training. RevColV2 consists of an encoder-decoder framework. The encoder is the same as RevCol, while the decoder uses reversed column connections. The paper also uses a unified fine-tuning framework utilizing a decoder for downstream tasks. RevColV2 demonstrates impressive performance on diverse vision tasks.\n\nSTRENGTHS:\n- RevColV2 decoder presents an interesting design with plenty of novelty. In particular, reversed column connection between enc-dec is an innovative architectural approach for MAE training.\n- RevColV2 achieves meaningful performance improvements on diverse tasks.\n- Multi-column architecture is different from the mainstream of transformer architecture, which enhances the novelty of RevColV2\n\nWEAKNESSES:\n- Basic component of the architecture is the same as RevCol. RevColV2 is an improved version, not a new architecture. Although it is interesting, the contribution of V2 paper is limited.\n- The paper's contribution is similar to ConvNeXt V2, enhancing existing architecture with MAE training and minor architecture revision. I think shedding light on existing architecture can be a contribution. But, similarity with ConvNeXt V2 paper might decrease the impact of this paper.\n- There are no reports for the throughput or latency. FLOPs numbers for detection and segmentation are omitted. I think throughput and FLOPs are necessary for architecture research in 2023. I strongly recommend authors to report those numbers.\n\nQUESTIONS:\n- Although the inverse pass is well described in the main section, I don't know where it is used in the experiments. What is the role of the inverse pass in RevCol V2? \n- IN22k intermediate fine-tuning is an interesting case. It would be better if RevCol V2 is compared with IN21k training in the following papers.\n  - [1] BEIT V2: Masked Image Modeling with Vector-Quantized Visual Tokenizers\n  - [2] DeiT III: Revenge of the ViT\n\nLIMITATIONS:\n.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces \"RevColv2,\" an advancement over the RevCol model, enabling compatibility with MIM training. The authors propose the new architecture comprising a bottom-up reversible column encoder and a top-down decoder, facilitating MIM compatibility while preserving disentangled low-level and semantic information throughout the network. The authors conduct experiments on ImageNet, detection, and segmentation tasks and the results show it achieve strong results on ImageNet and ADE20K.\n\nSTRENGTHS:\n•\tThe idea of RevColv2 to make RevCol compatible with MIM and the design of top-down column decoder is intuitive. It’s nice to see some downstream tasks could benefit from both pre-trained encoder and decoder.\n\n•\tThe Illustration of key motivation of maintaining disentangled low-level and semantic information is clear and further verified by analysis in Figure 3.\n\nWEAKNESSES:\n* The results on ImageNet-1K are strong; however, there is a lack of speed comparison with other methods. It would be valuable to assess the runtime speed for both pre-training and fine-tuning stages, particularly considering the impact of the reversible column network design on speed.\n* One benefit for reversible networks is memory-saving (at the cost of some speed). It would be beneficial to discuss whether this holds true for RevColv2. Exploring the trade-off between memory usage and speed for RevColv2 will add valuable insights to the paper.\n* In figure 2, the sequence length is different for the same level in the encoder and decoder. It seems unclear about the strategy used to handle this dimension change when connecting the encoder to the decoder.\n* For dense prediction tasks, RevColv2 utilizes both encoder and decoder pre-trained weights. To verify the effectiveness of this approach, one missing ablation is to compare to a variant that employs the same encoder and decoder during downstream fine-tuning but only utilizes the encoder's pre-trained weights while initializing the decoder weights randomly.\n* While the COCO detection results for the base model are strong, the performance of RevColv2-L appears to lag behind ViTDet-L using Mask R-CNN (54.0 vs. 55.6, citing the results from the ViTDet paper). Additionally, no results for RevColv2-L with Cascade Mask R-CNN are reported. It would be insightful to discussions the scaling results for RevColv2 on detection and provide some intuition on the potential reasons behind the observed performance differences.\n\nQUESTIONS:\nOverall, the authors propose RevColV2 to make RevCol network compatible with MIM training and it achieves consistent better results than RevCol. My main concerns are about some missing analysis/discussions for some ablations and results as listed in the weakness.\n\nLIMITATIONS:\nNo limitations have been discussed. Suggestions on potential limitation that worth to discuss could be about the speed issue and further scaling the model.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}