SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 4: Object detection and instance segmentation results on COCO dataset. † means using ViT-Adapter [35] networks. †† means that model pre-trained on ImageNet-22K dataset and ⋅ denotes using additional distillation teacher in pre-training.
Backbone	APbox	APmask	Params
with Mask R-CNN:
MAE-B [8]	49.8	44.3	110M
MAE-B† [8]	51.3	45.4	122M
RevColV2-B	52.4	46.2	119M
MAE-L [8]	53.1	47.1	323M
MAE-L† [8]	53.3	47.4	330M
RevColV2-L	54.0	47.8	363M
Backbone	APbox	APmask	Params
with Cascade Mask R-CNN:
Swin-B [49]	51.9	45.0	145M
ConvNeXt-B [41]	52.7	45.7	146M
RevCol [1]	53.0	45.9	196M
ViTDet-B [36]	54.0	46.7	141M
EVA-02-B †† [54]	55.5	47.1	141M
RevColV2-B	55.2	47.9	156M
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\mathrm{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\mathrm{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\mathrm{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
Table 5: Results of the data scaling. RevColV2-L with larger pre-training data and additional teacher achieves comparable performance than 2.1G parameter RevColV1-H and EVA-02-L.
Model	Params	Dataset	COCO (w/o Obj365)	ADE20K
Data	teacher	Detector	APbox	APmask	Segmenter	mIoU	+ms
MaskDistill-L	0.3 G	ImageNet1K	CLIP	-	-	-	UperNet	56.5	-
BEiTv2-L	0.3 G	ImageNet22K	CLIP	-	-	-	UperNet	57.5	-
EVA-02-L	0.3 G	merged 33M	EVA01-CLIP	Cascade	62.3	53.8	UperNet	60.1	-
RevCol-H	2.1 G	private 168M	semi-labeled	HTC++	61.1	53.0	Mask2Former	60.4	61.0
RevColV2-L	0.3 G	Laion400M+IN-1K	OpenCLIP	Cascade	62.1	53.2	Mask2Former	59.5	60.4
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.
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 \times 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.
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 low-level 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
Figure 3: Linear probing accuracy of each levels on ImageNet-1K after pre-training 300 epochs. The left is RevCol+MAE baseline, and the right is RevColV2 with reversible multi-column decoder.
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.
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 fine-tuning. 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
Table 6: Ablation results on ImageNet-1K. In different configurations of architectures, we keep the same overall FLOPs for fair comparison unless special descriptions.
(a) Decoder column depth.
Depth	ftBase	ftLarge
4	84.1	85.2
8	83.6	84.9
12	82.9	-
(b) Different schedules.
Epochs	ftBase	ftLarge
300	84.1	85.2
800	84.5	85.9
1600	84.7	86.3
(c) Type of decoder.
Decoder type	ft
MAE	83.5
UNet	83.4
Columns (w/o rev.)	83.7
Columns (w/ rev.)	84.1
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.
Disentangled feature learning. Disentangled representation refers to separating the factors of variation [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.
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.
[66] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748, 2018.
[67] 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.
[68] 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.
[69] 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.21\mathrm{x}$  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.
Table 7: Results of the inference latency, all models are tested with batch-size 32 on single A100 GPU.
Model	FLOPs	Latency	Model	FLOPs	Latency
RevCol-L	39G	61 ms	RevColV2-L	67G	51 ms
ViT-L	62G	42 ms	ViT-L (48 blocks)	64G	52 ms
Table 8: The throughput under different inference batch-size of RevColV2.
bs=16	bs=32	bs=64	bs=128	bs=256	bs=512
MAE-L [8]	730	754	786	811	820	823
RevColV2-L [8]	432	629	661	697	721	741
Speedup	0.591	0.834	0.841	0.859	0.879	0.900
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.
Table 9: The training time and memory cost comparison on a single A100 (80GB) x 8 machine.
Model	Time Cost	Memory (each GPU)
MAE-B [8]	220 s/epoch	43G
RevColV2-B [8]	249 s/epoch	49G
RevColV2-B + FlashAttn [8]	211 s/epoch	42G
RevColV2-B + FlashAttn + Reversible [8]	240 s/epoch	18G
- 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.
config	value
optimizer	AdamW
base learning rate	1.5e-4
weight decay	0.05
optimizer momentum	β1, β2=0.9, 0.95
batch size	4096
learning rate schedule	cosine decay
warmup epochs	40
training epochs	1600
augmentation	RandomResizedCrop
Table 10: MIM Pre-training settings.
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 intermediate 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
config	value
optimizer	AdamW
base learning rate	2.5e-4
weight decay	0.05
optimizer momentum	β1, β2=0.9, 0.999
layer-wise lr decay	0.8
batch size	4096
learning rate schedule	cosine decay
warmup epochs	5
training epochs	90
augmentation	RandAug (9, 0.5)
label smoothing	0.1
mixup	0.8
cutmix	1.0
drop path	0.1 (B), 0.2 (L)
head init	0.001
ema	None
Table 11: End-to-end IN-22K intermediate fine-tuning settings.
config	value
optimizer	AdamW
base learning rate	5e-4
weight decay	0.05
optimizer momentum	β1, β2=0.9, 0.999
layer-wise lr decay	0.75
batch size	1024
learning rate schedule	cosine decay
warmup epochs	5
training epochs	100 (B), 50 (L)
augmentation	RandAug (9, 0.5)
label smoothing	0.1
mixup	0.8
cutmix	1.0
drop path	0.1
head init	0.001
ema	0.9999
Table 12: End-to-end ImageNet-1K fine-tuning settings
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 R-CNN [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.
config	value
optimizer	AdamW
base learning rate	2.5e-5
weight decay	0.01
optimizer momentum	β1, β2=0.9, 0.999
layer-wise lr decay	0.9
batch size	512
learning rate schedule	cosine decay
warmup epochs	None
training epochs	30
augmentation	RandAug (9, 0.5)
label smoothing	0.1
mixup	None
cutmix	None
drop path	0.1(B), 0.2 (L)
head init	0.001
ema	0.9999
Table 13: End-to-end ImageNet-1K fine-tuning settings (after IN-22K intermediate fine-tuning).
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].