SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qi Han*
MEGVII Technology
hqer@foxmail.com
Yuxuan Cai*
larryx.tsai@gmail.com
Xiangyu Zhang†
zhangxiangyu@megvii.com
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\mathrm{mIoU}$  on ADE20K semantic segmentation. With extra teacher and large scale dataset, RevColv2-L achieves  $62.1\mathrm{AP}_{box}$  on COCO detection and  $60.4\mathrm{mIoU}$  on ADE20K semantic segmentation.
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].
†Corresponding author. This work is supported by The National Key Research and Development Program of China (No. 2017YFA0700800).
(a) Supervised
(c) BEiT
(d) RevColV2
Figure 1: Illustration of the information distribution in different pre-training schemes. (a): low-level information is gradually changed to be semantic in supervised pre-training; (b) and (c): low-level information firstly changes to be semantic then recovers to be low-level or entangled in MAE and BEiT pre-training; (d): low-level and semantic information is gradually disentangled in RevColV2 architecture with MIM pre-training.
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<sup>3</sup> and avoids the mixture of low-level and semantic information, fully exploring the pre-training abilities. The re-designed new architecture is named RevColV2.
<sup>3</sup>Our decoder behaves differently from traditional 'decoder' which refers to the network performing merely image reconstruction. The top-down reversible column decoder also includes semantic information in our design.
We build 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\mathrm{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\mathrm{AP}_{box}$  on COCO detection. In ADE20K segmentation, RevColv2-L reaches  $60.4\mathrm{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.
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.
$$
F o r w a r d: x _ {i} ^ {l} = \boldsymbol {F} _ {i} ^ {l} \left(x _ {i - 1} ^ {l}, x _ {i + 1} ^ {l - 1}\right) + \gamma x _ {i} ^ {l - 1} \tag {1}
$$
$$
I n v e r s e: x _ {i} ^ {l - 1} = \gamma^ {- 1} \left[ x _ {i} ^ {l} - \boldsymbol {F} _ {i} ^ {l} \left(x _ {i - 1} ^ {l}, x _ {i + 1} ^ {l - 1}\right) \right],
$$
where  $x_{i}^{l}$  denotes feature maps of the  $i$ -th level in  $l$ -th column;  $F_{i}^{l}$  denotes an arbitrary non-linear operation in  $i$ -th level, analogous to those residual functions in standard ResNets;  $\gamma$  is a reversible operation (e.g. channel-wise scaling), whose inverse is denoted by  $\gamma^{-1}$ . The  $i$ -th level takes the features of the lower level at the same column  $x_{i-1}^{l}$  and the upper level at the previous column  $x_{i+1}^{l-1}$  as input. These two features are fused together and passed through a stack of building blocks. After that, feature  $x_{i}^{l-1}$  are added with reversible operation  $\gamma$  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
Figure 2: The pipeline of RevColV2 is a unified architecture between pre-training and fine-tuning. The bottom-up columns inherit the design of RevCol with masked image patches as inputs. In the pre-training, the top-down columns receive the outputs of bottom-up columns along with the mask tokens, and reconstruct unseen raw patches at the lowest level of the last column. In the fine-tuning, the output features of top-down columns are selectively used in downstream tasks, such as classification task using the top level features.
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_{i+1}^{l-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:
$$
F o r w a r d: x _ {i} ^ {l} = \boldsymbol {F} _ {i} ^ {l} \left(x _ {i + 1} ^ {l}, x _ {i - 1} ^ {l - 1}\right) + \gamma x _ {i} ^ {l - 1} \tag {2}
$$
$$
I n v e r s e: x _ {i} ^ {l - 1} = \gamma^ {- 1} [ x _ {i} ^ {l} - \boldsymbol {F} _ {i} ^ {l} (x _ {i + 1} ^ {l}, x _ {i - 1} ^ {l - 1}) ].
$$
Two inputs of  $\pmb{F}_i^l$  become  $x_{i + 1}^{l}$  (feature of the upper level in current column) and  $x_{i - 1}^{l - 1}$  (feature of the lower level in previous column). For the input (to the highest level) of each top-down column,  $x_{i + 1}^{l}$  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 added 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 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 existing 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, Base and Large, 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.
Table 1: RevColV2 architecture configurations for different model sizes. FLOPs are measured on classification task with  $224^2$  resolution input. BU and TD are short for bottom-up and top-down.
#BU columns.	BU depth	#TD columns.	TD depth	dim	head dim	Params.	FLOPs
Base	3	12	3	4	416	32	101M	19G
Large	4	12	4	4	672	56	342M	67G
Table 2: ImageNet-1K classification results for various sizes models. The left table shows the end-to-end fine-tune results after MIM pre-training, and the right table shows the results after intermediate fine-tuning on ImageNet-22K.
Model	Size	Target	Params	FLOPs	FT
ImageNet-1K pre-train:
BEiT-B [9]	2242	DALL-E	87M	18G	83.2
MAE-B [8]	2242	Pixel	87M	18G	83.6
CAE-B [14]	2242	DALL-E	87M	18G	83.9
SdAE-B [40]	2242	EMA	87M	18G	84.1
MaskFeat-B [32]	2242	HOG	87M	18G	84.0
ConvNeXt-B [41]	2242	Pixel	89M	15G	83.8
SimMIM-B [17]	2242	Pixel	88M	16G	84.0
HiViT-B [42]	2242	Pixel	66M	16G	84.2
DeiT III-B [43]	2242	Label	87M	18G	83.8
HorNetGF-B [44]	2242	Label	88M	16G	84.3
SwinV2-B [11]	2242	Label	88M	20G	84.2
ConvNeXt V2-B [12]	2242	Pixel	89M	15G	84.9
RevCol-B [1]	2242	Label	138M	17G	84.1
RevColV2-B	2242	Pixel	88M	19G	84.7
DeiT III-L [43]	2242	Label	304M	62G	84.9
BEiT-L [9]	2242	Pixel	307M	62G	85.2
MAE-L [8]	2242	Pixel	307M	62G	85.9
CAE-L [14]	2242	DALL-E	307M	62G	86.2
MaskFeat-L [32]	2242	HOG	307M	62G	85.7
SimMIM-L [17]	2242	Pixel	197M	35G	85.4
ConvNeXt V2-L [12]	2242	Pixel	198M	34G	85.8
RevColV2-L	2242	Pixel	327M	67G	86.3
Model	Size	Target	Params	FLOPs	FT
ImageNet-1K pre-train + 22K intermediate fine-tune:
ViT-B [41]	3842	Label	86M	55G	84.0
DeiT III-B [43]	2242	Label	87M	18G	85.7
ConvNeXt-B [41]	2242	Pixel	89M	15G	85.8
ConvNeXt-B [41]	3842	Pixel	89M	45G	86.8
RevCol-B [1]	2242	Label	138M	17G	85.6
RevCol-B [1]↑	3842	Label	138M	49G	86.7
RevColV2-B	2242	Pixel	88M	19G	86.2
RevColV2-B↑	3842	Pixel	88M	64G	87.3
RevColV2-B↑	5122	Pixel	88M	130G	87.5
ViT-L [29]	3842	Label	307M	191G	85.2
DeiT III-L [43]	2242	Label	304M	62G	87.0
MOAT-3 [45]	2242	Label	190M	45G	86.8
SwinV2-L [11]	2562	Label	197M	48G	86.9
SwinV2-L↑ [11]	3842	Label	197M	115G	87.6
ConvNeXt V2-L [12]	2242	Pixel	198M	34G	87.3
ConvNeXt V2-L↑ [12]	3842	Pixel	198M	103G	88.2
RevCol-L [1]	2242	Label	273M	39G	86.6
RevCol-L↑ [1]	3842	Label	273M	116G	87.6
RevColV2-L	2242	Pixel	327M	67G	87.4
RevColV2-L↑	3842	Pixel	327M	215G	88.3
RevColV2-L↑	5122	Pixel	327M	417G	88.4
3.1 MIM Pre-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 SwinV2-B[11], HorNet $_{\mathrm{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
Table 3: Semantic segmentation result on ADE20K dataset with UperNet and Mask2Former framework.  $M2F$  denotes using Mask2Former framework. We report mIoU with single/multi-scale test.
Backbone	mIoU (ss.)	mIoU (ms.)	Params
ImageNet-1K:
MAE-B [8]	48.1	N/A	156M
CAE-B [8]	50.2	N/A	156M
PeCo-B [48]	48.5	N/A	156M
Swin-B [49]	48.1	49.7	121M
ConvNeXt-B [41]	49.1	49.9	122M
ConvNeXt V2-B [12]	N/A	52.1	122M
InternImage-B [50]	50.8	51.3	128M
RevCol-B [1]	49.0	50.1	122M
RevColV2-B	51.3	52.3	121M
BEiT [9]	53.3	N/A	421M
MAE-L [8]	53.6	N/A	421M
ConvNeXt V2-L [12]	53.2	53.7	235M
RevColV2-L	54.0	54.4	399M
Backbone	mIoU (ss.)	mIoU (ms.)	Params
ImageNet-1K + ImageNet-22K:
Swin-B [49]	50.3	51.7	121M
ConvNeXt-B [41]	52.6	53.1	122M
RevCol-B [1]	52.7	53.3	122M
RevColV2-B	53.1	53.9	121M
RevColV2-B+M2F	54.9	55.8	325M
Swin-L [49]	52.1	53.5	234M
ConvNeXt-L [41]	53.2	53.7	235M
RepLKNet-L [51]	52.4	52.7	207M
Focal-L [52]	54.0	55.4	240M
CSwin-L [53]	54.0	55.7	208M
RevCol-L [1]	53.4	53.7	306M
RevColV2-L	54.8	55.7	399M
RevColV2-L+M2F	58.2	58.6	570M