R EVERSIBLE C OLUMN N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a new neural network design paradigm Reversible Column Network (RevCol) .
The main body of RevCol is composed of multiple copies of subnetworks, named columns respectively, between which multi-level reversible connections are employed.
Such architectural scheme attributes RevCol very different behavior from conventional networks: during forward propagation, features in RevCol are learned to be gradually disentangled when passing through each column, whose total information is maintained rather than compressed or discarded as other network does.
Our experiments suggest that CNN-style RevCol models can achieve very competitive performances on multiple computer vision tasks such as image classification, object detection and semantic segmentation, especially with large parameter budget and large dataset.
For example, after ImageNet-22K pre-training,
RevCol-XL obtains 88.2% ImageNet-1K accuracy.
Given more pre-training data, our largest model RevCol-H reaches 90.0% on ImageNet-1K, 63.0% AP box on
COCO detection test-dev set, 61.0% mIoU on ADE20k segmentation.
To our knowledge, it is the best COCO detection and ADE20k segmentation result among pure (static) CNN models.
Moreover, as a general macro architecture fashion,
RevCol can also be introduced into transformers or other neural networks, which is demonstrated to improve the performances in both computer vision and NLP tasks.
I NTRODUCTIONS
Information Bottleneck principle (IB) (Tishby et al., 2000; Tishby & Zaslavsky, 2015) rules the deep learning world.
Consider a typical supervised learning network as in Fig.
1 (a): layers close to the input contain more low-level information, while features close to the output are rich in semantic meanings.
In other words, information unrelated to the target is gradually compressed during the layer-by-layer propagation.
Although such learning paradigm achieves great success in many practical applications, it might not be the optimal choice in the view of feature learning – down-stream tasks may suffer from inferior performances if the learned features are over compressed, or the learned semantic information is irrelevant to the target tasks, especially if a significant domain gap exists between the source and the target tasks (Zamir et al., 2018).
Researchers have devoted great efforts to make the learned features to be more universally applicable, e.g . via self-supervised pre-training (Oord et al., 2018;
Devlin et al., 2018;
He et al., 2022; Xie et al., 2022) or multi-task learning (Ruder, 2017; Caruana, 1997; Sener & Koltun, 2018).
In this paper, we mainly focus on an alternative approach: building a network to learn disentangled representations .
Unlike IB learning, disentangled feature learning (Desjardins et al., 2012; Bengio et al., 2013; Hinton, 2021) does not intend to extract the most related information while discard the less related; instead, it aims to embed the task-relevant concepts or semantic words into a few decoupled dimensions respectively.
Meanwhile the whole feature vector roughly maintains as much information as the input.
It is quite analogous to the mechanism in biological cells (Hinton, 2021;
Lillicrap et al., 2020) – each cell shares an identical copy of the whole genome but has different expression intensities.
Accordingly in computer vision tasks, learning disentangled features is also reasonable: for instance, high-level semantic representations are tuned during ImageNet pre-training, meanwhile the low-level information ( e.g . locations of the edges) should also be maintained in other feature dimensions in case of the demand of down-stream tasks like object detection.
Fig.
1 (b) sketches our main idea: Reversible Column Networks (RevCol) , which is greatly inspired by the big picture of GLOM (Hinton, 2021).
Our network is composed of N subnetworks (named
[Figure]
columns ) of identical structure (however whose weights are not necessarily the same), each of which receives a copy of the input and generates a prediction.
Hence multi-level embeddings, i.e . from low-level to highly semantic representations, are stored in each column.
Moreover, reversible transformations are introduced to propagate the multi-level features from i -th column to ( i + 1) -th column without information loss.
During the propagation, since the complexity and nonlinearity increases, the quality of all feature levels is expected to gradually improve.
Hence the last column (Col N in Fig.
1 (b)) predicts the final disentangled representations of the input.
In RevCol, one of our key contributions is the design of the reversible transformations between adjacent columns.
The concept is borrowed from the family of Reversible Networks (Chang et al., 2018; Gomez et al., 2017; Jacobsen et al., 2018; Mangalam et al., 2022); however, conventional reversible structures such as RevNets (Gomez et al., 2017) (Fig. 2 (a)) usually have two drawbacks: first, feature maps within a reversible block are restricted to have the same shape * ; second, the last two feature maps in RevNets have to contain both low-level and high-level information due to the reversible nature, which may be difficult to optimize as in conflict with IB principle.
In this paper, we overcome the drawbacks by introducing a novel reversible multi-level fusion module.
The details are discussed in Sec. 2.
We build a series of CNN-based RevCol models under different complexity budgets and evaluate them in mainstream computer vision tasks, such as ImageNet classification, COCO object detection and instance segmentation, as well as ADE20K semantic segmentation.
Our models achieve comparable or better results than sophisticated CNNs or vision transformers like ConvNeXt (Liu et al., 2022b) and Swin (Liu et al., 2021).
For example, after ImageNet-22K pre-training, our RevCol-XL model obtains 88.2% accuracy on ImageNet-1K without using transformers or large convolutional kernels (Ding et al., 2022b; Liu et al., 2022b; Han et al., 2021).
More importantly, we find RevCol can scale up well to large models and large datasets.
Given a larger private pre-training dataset, our biggest model RevCol-H obtains 90.0% accuracy on ImageNet-1K classification, 63.0% AP box on COCO detection test-dev set, and 61.0% mIoU on ADE20K segmentation, respectively.
To our knowledge, it is the best reversible model on those tasks, as well as the best pure CNN model on COCO and
ADE20K which only involves static kernels without dynamic convolutions (Dai et al., 2017; Ma et al., 2020).
In the appendix, we further demonstrate RevCol can work with transformers (Dosovitskiy et al., 2020; Devlin et al., 2018) and get improved results on both computer vision and NLP tasks.
Finally, similar to RevNets (Gomez et al., 2017), RevCol also shares the bonus of memory saving from the reversible nature, which is particularly important for large model training.
Relation to previous works.
Although our initial idea on feature disentangling is derived from
GLOM (Hinton, 2021), in RevCol there are a lot of simplifications and modifications.
For example,
GLOM suggests contrastive auxiliary loss to avoid feature collapse.
Contrastive training methods need extra pairs of positive and negative samples, which is complicated and unstable.
In RevCol, reversible transformations between columns provides lossless information propagation by nature.
As for other multi-scale grid-like architectures such as HRNets (Wang et al., 2020), DEQ models (Bai et al., 2020) and FPNs (Lin et al., 2017; Tan et al., 2020), the design purpose of those models is to fuse multi-scale features rather than learn disentangled representations; therefore, in general they still follow the paradigm in Fig.
1 (a) – neither multiple entrances/exits nor reversible structures are employed.
Based on those grid-like network topology, NAS based works (Ding et al., 2021;
Wu et al., 2021; Liu et al., 2019; Ghiasi et al., 2019) search the optimized topology of network architectures for specific dataset.
However, the RevCol architecture is not limit to specific tasks or datasets.
With the reversible nature, our method maintains lossless information propagation and benefits for not only pre-training but also other down-stream tasks.
Very recently, RevBiFPN (Chiley et al., 2022) comes up with an reversible variant of FPN, which is further employed in an HRNet-like architecture.
Though our RevCol shares the similar idea of multi-scale reversible transformations with RevBiFPN, our work is done independently, which is derived from a different motivation of feature disentangling, and has much simpler architectures ( e.g . free of reversible upsampling tower) and higher performances.
We compare some of those models in Sec. 3.
M ETHOD
In this section, we introduce the design details of our Reversible Column Networks (RevCol) .
Fig.
1 (b) illustrates the top-level architecture.
Notice that for each column in RevCol, for simplicity we directly reuse existing structures such as ConvNeXt (Liu et al., 2022b), hence in the following subsections, we mainly focus on how to build the reversible connections between columns.
In addition, we introduce an plug-and-play intermediate supervision on top of each column, which further improves the training convergence and feature quality.
2. M ULTI -L EVEL R EVERSIBLE U NIT
In our network, reversible transformations plays a key role in feature disentangling without information loss, whose insight comes from Reversible Neural Networks (Dinh et al., 2014; Chang et al., 2018; Gomez et al., 2017; Jacobsen et al., 2018; Mangalam et al., 2022).
Among them, we first take a review of one representative work RevNet (Gomez et al., 2017).
As shown in Fig.
2 (a), RevNet first partitions the input x into two groups, x 0 and x 1 .
Then for later blocks, for example, block t , it takes two anterior blocks’ outputs x t − 1 and x t − 2 as input and generates the output x t .
The mapping of block t is reversible , i.e. x t − 2 can be reconstructed by two posterior blocks x t − 1 and x t .
Formally, the forward and inverse computation follow the equations † :
[Equation]
where F t denotes an arbitrary non-linear operation analogous to those residual functions in standard
ResNets ; γ is a simple reversible operation ( e.g . channel-wise scaling), whose inverse is denoted by γ − 1 .
As mentioned in the introduction, the above formulation involves too strong constraint on the feature dimensions, i.e. x t , x t +2 , x t +4 , ... have to be equal sized, which is not flexible in architecture design.
That is why RevNets (Gomez et al., 2017) introduce some non-reversible down-sampling blocks between reversible units, hence the whole network is not fully reversible.
More importantly, we find there is no clear way directly employing Eq.
1 to bridge the columns in Fig.
1 (b).
To address the issue, we generalize Eq. 1 into the following form:
[Equation]
where m is the order of the recursion ( m ≥ 2) .
Clearly, the extension is still reversible.
Then we partition every m feature maps into a group: ( x 1 , x 2 , . . .
, x m ) , ( x m +1 , x m +2 , . . .
, x 2 m ) , . . . .
Given the features within any of the group, we can easily compute the features in other groups recursively according to Eq. 2.
Compared with the original form, Eq. 2 has the following two nice properties:
• The constraint on the feature map sizes is greatly relaxed if m is relatively large.
Notice that Eq.
1 does not require feature maps within each group to be equal sized; such constraint
FusionLevel
Level
FusionLevel FusionLevel
IntermediateSupervision IntermediateSupervision
Supervision
FusionLevel
Level
FusionLevel FusionLevel
MoreColumns MoreColumns
MoreColumns MoreColumns (
[Figure]
) (c) (b)
X 1 Xt-m+1 X t-m X t X
t-
X t-2 X t X t-2 X
Level
FusionLevel FusionLevel
FusionLevel
STEM Image Input
F tF tF tF tF t t
[Figure]
X
[Figure]
only exist between groups.
Therefore, we can use tensors of different shape to represent features of different semantic levels or different resolutions.
• Eq.
2 can easily cooperate with existing network architectures, even though the latter is not reversible.
For example, we can assign m feature maps in a standard ResNet to represent the feature maps within a group ( x t , x t +1 , . . .
, x t + m − 1 ) , which is still compatible with Eq.since ResNet can be viewed as a part of ( F t , F t +1 , . .
, F t + m − 1 ) respectively.
Thus the whole network is still reversible.
Therefore, we can reorganize Eq.
2 into a multi-column fashion, as shown in Fig.
2 (b).
Each column is composed of m feature maps within a group, as well as their mother network.
We name it multi-level reversible unit , which is the basic component of our RevCol as in Fig.
1 (b).
2. R EVERSIBLE C OLUMN A RCHITECTURE 2.2. M ACRO D ESIGN
As discussed in the introduction (see Fig.
1 (b)), our network RevCol is composed of multiple subnetworks with reversible connections to perform feature disentangling.
Fig.
2 (c) elaborates the architecture design.
Following the common practice of recent models (Dosovitskiy et al., 2020; Liu et al., 2022b), first the input image is split into non-overlapping patches by a patch embedding module.
After that, patches are fed into each subnetwork (column).
Columns can be implemented with any conventional single-column architectures, e.g . ViT (Dosovitskiy et al., 2020) or ConvNeXt (Liu et al., 2022b).
We extract four-level feature maps from each column to propagate information between columns; for example, if the columns are implemented with widely-used hierarchical networks (Liu et al., 2021; He et al., 2016; Liu et al., 2022b), we can simply extract multi-resolution features from the output of each stage.
For classification tasks, we only use feature map of the last level (Level 4) in the last column for rich semantic information.
For other down-stream tasks like object detection and semantic segmentation, we use feature maps of all the four levels in the last column as they contain both low-level and semantic information.
To implement the reversible connections between columns, we adopt the multi-level reversible unit proposed in Eq. 2, but in a simplified fashion: rather than take ( m − 1) inputs for each non-linear operation F t ( · ) , we use only one low-level feature x t − 1 at the current column and one high-level feature x t − m +1 at the previous column as the input.
The simplification does not break the reversible property.
We find more inputs bring minor accuracy gain but consume much more GPU resources.
Thus Eq. 2 is simplified as:
[Equation]
Compared with conventional architectures, the macro design of our RevCol has the following three properties or advantages:
Feature disentangling.
In RevCol, the lowest level of each column maintains low-level features as it is close to the input, while the highest level in the last column is highly semantic because it is directly connected to the supervision.
Therefore, information in different levels is gradually disentangled during the (lossless) propagation between columns – some feature maps are more and more semantic and some maintain to be low-level.
Detailed analyses are presented in Appendix
F. The property brings many potential advantages, for instance, more flexible to downstream tasks which rely on both high-level and low-level features.
We argue that reversible connection plays a key role in the disentangling mechanism – some previous works like HRNet (Wang et al., 2020) involve multi-level feature fusion but without reversible connection, which may suffer from information loss and lead to inferior performances in our experiments (see Appendix C.2).
Memory saving.
The training of conventional networks takes a lot of memory footprint to store the activations during forward propagation as the demand of gradient computation.
While in our RevCol, since the connections between columns are explicitly reversible, during 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.
In Appendix C.4, we demonstrate RevCol costs roughly O (1) additional memory with the increase of column numbers.
New scaling factor for big models.
In RevCol architecture, column number serves as a new dimension in addition to depth (number of blocks), and width (channels of each block) in vanilla single-column CNNs or ViTs.
Increasing column numbers has similar income as increasing both width and depth in certain range.
2.2. M ICRO D ESIGN
We employ ConvNeXt blocks (Liu et al., 2022b) to implement each column in our network by default; other architectures, such as transformers , are also applicable (see Appendix B for details).
We make a few modifications to make ConvNeXt compatible with our macro architecture:
Fusion module.
As shown in Fig.
3, in each level of original ConvNeXt, the inputs are first downsampled in a patch-merging block.
Then the outputs are passed through a bunch of residual blocks.
In RevCol, we introduce a fusion module to fuse the feature maps from the current and previous columns (refer to Fig.
2 (c), green and blue connections).
We modify the original patch-merging block in ConvNeXt by putting the LayerNorm after the patch-merging convolution rather than before.
Channel numbers are doubled in patch-merging convolution.
We also introduce an up-sampling block, which is composed of a linear channel mapping layer, a LayerNorm normalization and a feature map interpolation layer.
We halve the channel numbers in linear channel mapping layer.
The outputs of the two blocks are summed up and then passed to the residual blocks followed by.
Kernel size.
In RevCol we revise the 7 × 7 convolutions in original ConvNeXt (Liu et al., 2022b) to 3 × 3 by default, mainly to speed up the training.
Increasing kernel size further obtains more accuracy, but not very much, partly because the our multi-column design enlarges the effective receptive field.
Please refer to Appendix C.5 for more details.
Reversible operation γ .
We adopt a learnable reversible channel-wise scaling as reversible operation γ to keep the network stable.
Each time the features are summed up in forward of Eq.
3, the magnitude grows larger, which makes the training process unstable.
Using a learnable scaling can suppress the magnitude of features.
During training, we truncate the absolute value of γ so that it will never be smaller than 1 e − 3 , because the numerical error could become large in the reverse computation when γ is too small.
2. I NTERMEDIATE S UPERVISION
Though multi-level reversible unit is able to maintain information during column iteration, the down-sample block still can discard information inside column.
Features at the end of the front columns is too close to the final output, for reversible connections simply do scaling and summation.
Such information loss leads to inferior performance.
Similar problem also happens when using deeply-supervised method (Lee et al., 2015; Szegedy et al., 2015).
To mitigate the problem of information collapse, we propose an intermediate supervision method which adds additional supervision into front columns.
For features in front columns, we hope to keep the mutual information between features and the input image as much as possible, so that the network discard less information within columns.
Consider RevCol gradually disentangle semantic and low-level information, extracting and leveraging the task-relevant information can further boost the performance.
Therefore, we need to maximize the lower bound of mutual information between features and the prediction.
Inspired by Wang et al.
(2021), we add two auxiliary heads to last level features (Level 4).
One is a decoder (He et al., 2022) which reconstructs the input image, the other is a linear classifier.
The linear classifier can be trained in a regular classification fashion with the cross-entropy (CE) loss.
The parameters of decoder are optimized by minimizing the binary cross-entropy (BCE) reconstruction loss.
Compared with commonly used L1 and L2 loss, interpreting the distribution of reconstructed logits and input image as bit probabilities (Bernoullis) outputs smoother value, which makes it more compatible with CE Loss.
For intermediate supervision at one column, the compound loss is the weighted summation of the above two loss.
Note that supervision heads may not be added to all columns.
For all the variants of RevCol, we set the number of compound loss to 4 empirically (eg.
for a 8 column RevCol, the supervision heads are added to column 2, 4, and 6, and 8).
The total loss L in is the summation of all compound loss:
[Equation]
n denotes the total number of compound loss.
L BCE and L CE denotes BCE loss and CE loss correspondingly.
α i and β i are changed linearly with the compound loss number.
When the compound loss is added in earlier columns, we use larger value α i and smaller value β i to keep
I ( h , x ) .
In later columns, value α i decreases and β i increases, which helps boost the performance.
E XPERIMENTS
We construct different RevCol variants, RevCol-T/S/B/L, to be of similar complexities to Swin transformers and ConvNeXts.
We also build a larger RevCol-XL and RevCol-H to test the scaling up capability.
These variants adopt different channel dimension C , blocks in each column B and column numbers COL .
The configuration hyper-parameters of these model variants are:
• RevCol-T: C = (64 , 128 , 256 , 512) , B = (2 , 2 , 4 , 2) , COL =• RevCol-S: C = (64 , 128 , 256 , 512) , B = (2 , 2 , 4 , 2) , COL =• RevCol-B: C = (72 , 144 , 288 , 576) , B = (1 , 1 , 3 , 2) , COL =• RevCol-L: C = (128 , 256 , 512 , 1024) , B = (1 , 2 , 6 , 2) , COL =• RevCol-XL: C = (224 , 448 , 896 , 1792) , B = (1 , 2 , 6 , 2) , COL =• RevCol-H: C = (360 , 720 , 1440 , 2880) ,
B = (1 , 2 , 6 , 2) , COL =
We conduct image classification on ImageNet dataset (Deng et al., 2009; Ridnik et al., 2021).
We also test our models on the downstream object detection task and semantic segmentation task on commonly used MS-COCO (Lin et al., 2014) and ADE20k (Zhou et al., 2017b) dataset.
Training and fine-tuning settings please refer to Appendix E. Furthermore, we show the performance of
RevCol with transformer on vision and language tasks (shown in Appendix B).
[Table]
3. I MAGE C LASSIFICATION
On ImageNet (1.28M images) (Deng et al., 2009) dataset, we train RevCol for 300 epochs with intermediate supervision.
Hyperparameters, augmentation and regularization strategies follows Liu et al.
(2022b)
We also pre-train our models on the larger ImageNet-22K dataset (Ridnik et al., 2021), which contains 14.2 million images.
In Tab. 1, we compare our RevCol variants with commonly used recent Transformers and CNNs on ImageNet-1k validation set.
Our models outperforms a large number of vanilla single-column
CNNs and Transformers with similar complexities.
For example, RevCol-S achieve 83.5% Topaccuracy, outperform ConvNeXt-S by 0.4 points.
When pre-trained with larger ImageNet-22K dataset,
RevCol-XL achieves 88.2% Top-1 accuracy.
As RevCol maintains some task-irrelevant low-level information in classification pre-training, relaxing the constraint of params and FLOPs and enlarging dataset size can further boost our models’ performance.
To further test the scaling up effectiveness of large dataset, we build a 168-million-image semi-labeled dataset (see Appendix D).
With extra data pre-training and ImageNet-1k fine-tuning, our RevCol-H achieves 90.0% top-1 accuracy .
Our results further demonstrate with RevCol, CNN models can also share the dividends of large model and massive data pre-training.
3. O BJECT D ETECTION
We evaluate our proposed RevCol on object detection task.
Experiments are conducted on the
MS-COCO dataset using the Cascade Mask R-CNN (Cai & Vasconcelos, 2019) framework.
We also finetune our largest model RevCol-H with HTC++ (Chen et al., 2019) and DINO (Zhang et al., 2022a) Framework.
In Tab. 2, we compare the AP box and AP mask with Swin/ConvNeXt in variant sizes on COCO validation set.
We find RevCol models surpass other counterparts with similar computation complexities.
Information retained in pre-training helps RevCol models acheieve better results in down-stream tasks.
When the model size grows larger, this advantage becomes more remarkable.
After finetuning under Objects365(Shao et al., 2019) dataset and DINO (Zhang et al., 2022a) framework, our largest model RevCol-H achieves 63.1% AP box on COCO detection val set.
[Table]
ConvNeXt-
[Table] [Table]
• RepLKNet-
[Table]
• RevCol-B
[Table]
RevCol
[Table]
L
[Table]
• RevCol-H (HTC++)
[Table]
• RevCol-H
[Table]
[Table] [Table]
• ConvNeXt-T (Liu et al.)
[Table]
• RevCol-T
[Table]
◦ Swin-S (Liu et al.)
[Table]
• ConvNeXt-S (Liu et al.)
[Table]
• RevCol-S
[Table]
◦ Swin-B (Liu et al.)
[Table]
• RepLKNet-B (Ding et al.)
[Table]
• ConvNeXt-B
(Liu et al.)
[Table]
• RevCol-B
[Table]
◦ Swin-B (Liu et al.)
[Table]
• RepLKNet-B (Ding et al.)
[Table]
• ConvNeXt-B
(Liu et al.)
[Table]
• RevCol-B
[Table]
◦ Swin-L (Liu et al.)
[Table]
• RepLKNet-L (Ding et al.)
[Table]
• ConvNeXt-L (Liu et al.)
[Table]
• RevCol-L
[Table]
• RevCol-H
[Table]
• RevCol-H + Mask2Former
[Table]
3. S EMANTIC S EGMENTATION
We also evaluate RevCol backbones on the ADE20K semantic segmentation task with UperNet (Xiao et al., 2018) framework.
We do not use intermediate-supervision in down-stream fine-tune process.
To further explore our model’s capacity and reach the leading performance, we utilize recent segmentation framework Mask2Former (Cheng et al., 2022), and adopt the same training settings.
[Table]
In Tab.
[Table]
we report validation mIoU with single-scale and multi-scale flip testing.
RevCol models can achieve competitive performance across different model capacities, further validating the effectiveness of our architecture design.
It’s worth noting that when use Mask2Former detector and extra pretraining data, RevCol-H achieves an mIoU of 61.0%, which shows feasible scalability towards large-scale vision applications
[Table]
3. S YSTEM - LEVEL C OMPARISON WITH SOTA F OUNDATION M ODELS
Foundation models (Kolesnikov et al., 2020; Radford et al., 2021; Yuan et al., 2021b) are generalpurpose backbones pre-trained on massive and divergent data source.
They can adapt to various down-stream tasks with limited domain-specific data.
We show comparison among various public state-of-the-art (SOTA) foundation models including Vision Transformers and Vision-Language models, namely, SwinV2 (Liu et al., 2022a), BEiT3 (Wang et al., 2022), and Florence (Yuan et al., 2021b).
As shown in Tab. 4, though our RevCol-H is purely convolutional and pre-trained on single modality dataset , the results on different tasks demonstrate remarkable generalization ability of RevCol with large scale parameters.
R ELATED W ORKS 4. D ISENTANGLE R EPRESENTATION L EARNING AND P ART -W HOLE H IERARCHY
A disentangled representation is generally described as one which separates the factors of variation, explicitly representing the important attributes of the data (Desjardins et al., 2012; Bengio et al., 2013).
Desjardins et al.
; Kulkarni et al.
; Higgins et al.
; Chen et al.
Karras et al.
(2019) seek to learn disentangled representations through generative models.
Locatello et al.
(2019) points out that unsupervised learning of disentangled representations is fundamentally impossible without inductive biases both on the considered learning approaches and the datasets.
The recent proposal of GLOM (Hinton, 2021) gives an idea of representing a part-whole hierarchy by a weight-sharing columns.
The GLOM architecture provides an interpretable part-whole hierarchies for deep neural network (Garau et al., 2022).
In RevCol, we adopt the design of using columns, but not modeling the process of formulating islands.
On the contrary, our column iteration process maintains both low-level and high-level information and gradually disentangle them.
Rather than using self-supervised methods, RevCol can be trained with supervision end-to-end.
4. R EVERSIBLE N ETWORKS
Gomez et al.
(2017) firstly propose RevNet that allow back propagation without saving intermediate activations.
The reversible design remarkably saves the training cost, since it keep O (1) GPU memory consumption as model depth scaling up.
Jacobsen et al.
(2018) propose a fully reversible network that can reverse back to the input without any information loss.
Chang et al.
(2018) develop a theoretical framework on stability and reversibility of deep neural network and derive reversible networks that can go arbitrarily deep.
Mangalam et al.
(2022) expand the reversible network scope from CNNs to Transformers.
RevBiFPN (Chiley et al., 2022), a concurrent work of ours, add the reversible connections to BiFPN (Tan et al., 2020) network.
Our RevCol maintains the information without loss inside each column rather than the whole BiFPN network in RevBiFPN.
B G ENERALIZATION TO T RANSFORMERS B. V ISION T RANSFORMER M ODELS
RevCol contains multiple light-weight sub-networks with reversible connections.
In this paper, we adopt the ConvNext micro design by default except for multi-columns fusion and smaller convolution kernel as described in Sec. 2.2.2.
However, the micro design of our RevCol is not limited to convolutional networks, but is also compatible with isotropic designing, such as the vanilla vision transformer (ViT)
(Dosovitskiy et al., 2020).
In this section, we show the micro design of RevCol can generalized to vanilla ViT, namely RevCol-ViT, with promising experimental results.
net-ViT maintains the feature resolution in the reversible columns.
Thus the patch merging blocks and up-sample blocks in the fusion modules are replaced with a simple linear projection with a post
LayerNorm.
We use the vanilla ViT building block instead of the ConvNext building block variant.
The post LayerNorms and normalized dot-product attention are used in ViT blocks to stabilize training convergence, similar to Liu et al.
(2022a).
With the properties of isotropy, we evenly arrange the building blocks in each column.
The configuration details of RevCol-ViT are:
• RevCol-ViT-S: C = (224 , 224 , 224 , 224) , B = (2 , 2 , 2 , 2) , HEAD = 4 , COL =• RevCol-ViT-B: C = (384 , 384 , 384 , 384) , B = (3 , 3 , 3 , 3) , HEAD = 6 , COL =
[Table]
• ConvNext-S ( iso. )
(Liu et al., 2022b)
[Table]
◦ RevCol-ViT-S
[Table]
◦ ViT-B
(Dosovitskiy et al.,
[Table]
◦
[Table]
-B (Touvron et al.,
[Table]
◦ Rev-ViT-B (Mangalam
[Table]
ConvNext-B ( iso. ) (
[Table]
RevCol-ViT-
[Table]
We use the same training setting with the anisotropic RevCol as described in Sec. 3.1, except that the intermediate supervision is discarded for simplicity and the stochastic depth rate is set as 0.2 for
RevCol-B. We scale down the value of last linear projection layers in each FFN accroding to the network depth in initialization, same as BEiT (Bao et al., 2021).
In Tab. 5, we compare the RevColViT with vanilla ViT and other concurrent isotropic designs.
Our RevCol-ViT surpasses vanilla vision transformer (77.9% for ViT and 81.7% for DeiT) and convolutional network ConvNeXt (82.0%) that have comparable model parameters and computational overhead on ImageNet-1k classification w.r.t.
the top-1 accuracy.
B. L ANGUAGE M ODELS
Considering the great success of applying transformer to computer vision, i.e., ViT (Dosovitskiy et al., 2020), we also made some exploration to generalize RevCol to natural language processing (NLP).
Based on the design in Appendix B.1, we can easily apply the isotropic RevCol to language models with minor modification.
To be specific, we replace the stem module in our RevCol with word embedding and positional encoding in transformer.
Then, the RevCol can be plugged into the original transformer as an encoder.
The output of the last column in RevCol will be used as the memory keys and values for the attention layers in decoder, just exactly the same as the original transformer.
We select the translation task to evaluate the potential of the RevCol in NLP.
We run experiments on the WMT’16 English-German (En-De) dataset with 4.5M sentences and larger WMT’14 EnglishFrench dataset with 36M sentences.
Each sentence is encoded by joint source and target byte pair encoding following Sennrich et al.
The details of model architecture and the BLEU score are shown in Tab. 6.
All the dataset preparation and the training configurations follows Ott et al.
(2018) and the open source project fairseq .
The models were trained for 300K steps with batch-size of 28,672 tokens on En-De and 200K steps with batch-size of 86,016 on En-Fr.
We discard the intermediate supervision for simplicity.
As shown in Tab. 6, our RevCol outperforms vanilla transformer with comparable parameters on En-De (28.67 vs . 28.43) and En-Fr (43.40 vs . 43.07), which demonstrates the RevCol’s applicability to NLP.
[Figure]
[Table]
C M ORE A NALYSIS E XPERIMENTS C. P ERFORMANCE G AIN OF R EVERSIBLE C OLUMNS A RCHITECTURE
[Figure]
[Table]
In this section
[Figure]
we evaluate the performance gain of using reversible columns.
In the first experiment, we fix a single column’s structure and FLOPs then simply add more columns to scale large and test the performance.
At the same time, we plot the vanilla single-column models with according model sizes.
As depicted in Fig. 4, compared to single-column models, using multi-column reversible architecture always gets better performance under same FLOPs.
Besides, within a certain range, scaling up
RevCol in terms of increasing column numbers can have similar performance gains compared to scaling up with both depth and channel numbers(width) in single-column models.
In the second experiment, we constrain the total FLOPs to about 4.5G and test models with different column numbers.
In other words, we add more columns but scale down the size of a single column.
In Tab. 7, we notice that adopt column number at the range of 4 to 12 can maintain the model’s performance, while more column models suffer from performance degradation.
We believe the reason is the width and depth in a single column are too low to keep representation ability.
C. R EVERSIBLE N ETWORKS VS . N ON -R EVERSIBLE N ETWORKS
In this section, we ablate different design patterns of reversible connections.
First, we build a nonreversible multi-column network using the fusion module of HRNet.
Second, we build another single column reversible ConvNeXt using the design of RevNet as shown in Fig. 2(a).
We compare the two designs with our RevCols.
The evaluation result is shown in Tab. 8.
The non-reversible multi-column network suffers from information loss during propagation, which could result in lower accuracy.
The reversible single-column network maintains information during propagation, but lack the superiority of multi-level fusion.
This experiment further indicates the effectiveness of combining the reversible design with multi-column networks.
C. P ERFORMANCE GAIN OF USING I NTERMEDIATE S UPERVISION
In this section, we evaluate the performance of RevCol-T/S/B with and without intermediate supervision on ImageNet-1K. We also evaluate the object detection task performance using 1 × training
[Table] [Table]
schedule on MS-COCO dataset.
Other settings remain the same.
From the validation results in
Tab. 9, models trained with intermediate supervision achieves 0.5% to 0.9% better top-1 accuracy.
Besides
[Table]
intermediate supervision also benefits down-stream tasks, which further demonstrates its effectiveness.
[Equation]
[Table]
[Figure]
Fig.
5 plots the GPU memory consumption with the scaling of model size.
We fix the computation complexity of a single column to 1G FLOPs and only increase column number.
Meanwhile, we measure the memory consumption in training process which includes the forward and backward propagation.
Our experiments are conducted on Nvidia Tesla V100 GPU under batch-size 64, fp16 precision and PyTorch implementation.
With the increase of model size, we can see our
RevCol keeps an O (1) GPU memory consumption, while non-reversible architecture’s memory consumption increase linearly with column number.
Note that our RevCol does not keep strictly the same GPU memory consumption as column number increase, as reversible networks need to save the operation weights in need for calculating gradients and the re-construction of feature maps in backward propagation.
C. A BLATION OF KERNEL SIZE IN CONVOLUTIONS
In original ConvNeXt, large kernel convolution achieves in better performance.
We conduct same experiments in RevCol-T. As shown in Tab. 10, for 4 column tiny models, using 5 × 5 convolution increase the ImageNet-1k Top-1 accuracy by 0.3% and the COCO AP by 0.6 for RevCol-T model.
[Table]
Larger kernel convolution has no more accuracy gain.
We consider the multi-level fusion structure already generate large receptive field and this limit the accuracy gain of using large kernel convolution.
Therefore, we fix the convolution kernel size to 3 in all variants of our model.
C. R OBUSTNESS OF THE NUMBER OF COLUMNS
In the ablation analysis of the paper, we show that when fix the total FLOPs and add more columns of RevCol, the performance first increases and then get saturated.
When the number of columns is extreme large, such as 20, the performance drop because of the representation ability of single column is limited.
When the number of columns is usual, such as 4 ˜ 12, the performances are similar, which verifies the setting robustness of the number of columns.
[Figure]
To further analyze the robustness of the number of columns, in this section, we build some RevColViT-B variants (see Appendix B for more details).
Each variant has the same number of residual blocks with the same channel dimension, but different number of columns.
In other worlds, these variants have the same channel dimension and different depth of each columns and different number of columns.
We use 32 residual blocks totally and maintain the FLOPs about 18G. Fig.
6 show the performance on ImageNet-1K of different variants.
The number of columns are 1, 2, 4, and 8, accordingly the depth of each column are 32, 16, 8, and 4.
The performance of single column variant is lower (similar to DeiT-B (Touvron et al., 2020))
because of the single column ViT can not maintain the information as multi reversible columns.
The performance is decreasing when the number of columns became larger, because of the depth of each columns is not enough.
This phenomenon indicates us that given a target FLOPs, the setting of the number of columns is robust unless the depth of each columns or channel dimension is too small.
D S EMI - LABELED P RIVATELY C OLLECTED D ATASET FOR L ARGE M ODELS D. D ATA C OLLECTION AND P SEUDO L ABEL S YSTEM
The dataset consists of around 168 million(M) images, 50M of which labeled and the remaining 118M unlabeled.
The majority of labeled images come from public datasets, e.g . ImageNet, Places365 (Zhou et al., 2017a), and Bamboo (Zhang et al., 2022b).
The others are web-crawled images annotated by in-door employees.
Unlabeled images come from weakly-annotated image-text datasets like
YFCC-100M (Thomee et al., 2016).
We do not use text annotations.
In order to utilize images of different label domains and the massive unlabeled images, we employ a multi-target label system similar to Ding et al.
(2022a) and Ghiasi et al.
We adopt a semisupervised learning strategy with ViTs, thus generating pseudo labels with continuously increased quality.
We only store soft predictions with confidence higher than 1% to save storage.
The final version of pseudo label we use are generated by a multi-head ViT-Huge teacher, which has an 89.0%
ImageNet-1k accuracy.
D. I MAGE D EDUPLICATION
Since the dataset contains large amount of unverified web-crawled images, there are probably validation or test images sneaking into our training dataset.
Works like Mahajan et al.
and
Yalniz et al.
(2019) all regard image deduplication an important procedure for fair experiments.
We first iterate over the entire dataset to filter out suspicious duplicates together with the corresponding test images based on their pseudo label distance.
This brings more than 10,000 images with high suspicion.
We look at these image pairs and finally find about 1,200 exact-duplicates and nearduplicates.
Fig.
7 shows some examples of the near-duplicates, which are difficult to detect.
Never the less, training a model without removing these duplicates gives less than 0.1% accuracy gain on
ImageNet-1k in our experiments.
We attribute this to the absence of true labels from these duplicates.
E M ORE T RAINING D ETAILS
This section gives more training details on ImageNet classification, COCO detection, and ADE20K segmentation.
E. I NTERMEDIATE S UPERVISION S ETTINGS
We add intermediate supervision in ImageNet-1k training, ImageNet-22k and extra data pre-training.
We used a 3-blocks decoder with gradually up-sampled feature maps in ImageNet-1k training.
The block setting remains the same as Sec.
2.2 We use a single layer decoder in ImageNet-22k and extra data pre-training.
For all the variants of RevCol, we set the number of compound loss n toempirically (eg. for a 8 column RevCol, the intermediate supervision is added to column 2, 4, and 6, and the original classification CE loss is also added to column 8).
α
i is set to 3, 2, 1, 0 and β i is set to 0.25, 0.5, 0.75, 1.
E. H YPERPARAMETERS U SED FOR T RAINING AND P RE - TRAINING
This section introduces the training details for main experiments, the supervised training on ImageNet and extra data.
We show this setting in Tab. 11.
All experiments in ablation studies are superivised trained on ImageNet-1K except additional descriptions and also follow settings described in this section
[Table]
E. H YPERPARAMETERS U SED FOR F INE - TUNING
This section gives the hyperparameters used for fine-tuning on ImageNet-1K and downstrea COCO object detection and instance segmentation, ADE20K semantic segmentation tasks, as shown in Tab.
12, Tab. 13 and Tab. 14.
For our largest model RevCol-H
[Table]
we pad the 3 × 3 kernel pre-trained model to 7 × 7 and 13 ×kernel size with Net2net (Chen et al., 2015) methods, then fine-tune on the COCO and ADE20K dataset.
[Table]
[Table] [Table]
[Figure]
Column Column
Column Column
Column Column
Column Column
[Figure]
Level
[Figure]
Level
Column Column
Column Column
Column Column
Column Column
F V ISUALIZATIONS OF F EATURE D ISENTANGLEMENT
In this section, we show our RevCol can disentangle features with stacked columns, which is different from the conventional sequential networks.
We use RevCol-S pre-trained on ImageNet-1K for analysis.
First, we visualize the class activation maps (CAMs) for outputs of each last layer of a level.
We adopt LayerCAM (Jiang et al., 2021) technology to generate the CAMs with the predicted classes.
Fig.
8 show the heatmaps of activation.
With the levels and columns going deeper, the features focus on the regions with more semantics.
The outputs of RevCol-S are the different levels of last column.
These features with high level semantics focus on different parts of the image and the whole part of the object, achieving disentanglement of features for task-relevant and task-irrelevant.
To quantify the disentanglement, we use Centered Kernel Alignment (CKA) similarity metric (Kornblith et al., 2019) to measure the similarity between representations in RevCol-S.
We calculate the CKA similarities between intermediate features in different levels and columns and images or labels of each category in the ImageNet val set.
Then we plot the similarities of the category with the highest label similarity in Fig.
As shown in the figure, the similarities between images and intermediate features are not clearly distinguished at different levels in Column 2-5, while the features
[Figure]
with higher levels have lower similarity to the images in Column 6-8.
The similarities between labels and intermediate features are also more distinct in columns.