Vision Pair Learning: An Efficient Training Framework for Image Classification

Transformer architecture has attracted much attention since its first application in machine translation, and it has proven effective in various tasks of natural language processing. As a result, the transformer is introduced to computer vision recently. Although researchers have successfully adapted the transformer for CV tasks, currently CNN based models are superior to transformer based ones in terms of performance and speed. Main reasons include 1) The special structure of transformer lets it pay more attention to global features but ignore detailed information. Because of that, pure transformer-based methods often fail in fine grained tasks such as small object detection. 2) Problems also arise in model training. Due to the relatively large number of parameters, training transformer tends to be more computational expensive than CNN models. For example, comparing with ViT-L/16 (Dosovitskiy et al. 2020), training EfficientNetV2 (Tan and Le 2021) is $5\times$-$11\times$ faster. This offers unfriendly hardware requirement and slows down prototyping experiments.

To tackle these problems, effective components of CNN are combined with transformer, for both enhancing performance and decreasing computation cost. These approaches generally fall into two categories, 1) One may utilize CNN blocks to minimize the information loss brought by image tokenization, and reduce the dimension of input feature, so as to balance the performance and speed, e.g. Visual Transformers (Wu et al. 2020), T2T-ViT (Yuan et al. 2021). 2) Modify the general design of network. TNT (Han et al. 2021) follows the idea of “deep-narrow” and makes the transformer deeper and slimmer. CvT (Wu et al. 2021) replaces linear projection with convolution projection for less computation. Other widely-used techniques, e.g. pruning (Zhu et al. 2021) and sliding window (Liu et al. 2021), are adopted as well.

The performance of above methods are satisfactory, but mixing the advantages of CNN and transformer often relies on elaborated and heuristic network design, and may hurt the transferability. In this paper, we propose a pair learning method based on CNN and transformer for image classification tasks. The method enables two different networks to learn from each other and achieve better performance together. The structure is shown in Figure 1.

The network consists of three parts, CNN branch, transformer branch and Pair Learning Module (PLM). Since the target of VPL is to build a high-level training framework for better representation learning, we directly use CNN and transformer baseline models as branches without modification. Being the core part of the model, PLM introduces a contrastive-learning-like task. The most common idea of contrastive learning (He et al. 2020; Chen et al. 2020b; Grill et al. 2020; Chen and He 2020; Caron et al. 2020) is to feed a pair of randomly transformed images to two parameter-sharing branches and predict whether they are from the same source image by the extracted embeddings. Different from that, the branches of our method read the same image as input but are distinct from each other in terms of basic network block. Due to the natural difference of the two branches’ network architecture and capacity, the primary model collapse problem in contrastive learning is easily avoided. Meanwhile due to the different inductive biases of global attention and local convolution, the transformer and CNN branches produce representations of diverse views. And the task forces the high-level representation learned by two branches to interact, thus providing extra supervision signals and accelerating the convergence. A distillation-like loss is utilized as well.

Considering that the two branches are of different data efficiency, we additional adjust the training strategy so that the branches can choose to learn from partners during the appropriate stage of the training process. The heterogeneous model with PLM and modified training protocol together constitute vision pair learning. After training, PLM is removed and branches can be used for inference without extra computation.

To summarize, our contributions are three-fold:

• •

  We propose a new training framework named vision pair learning. By building a heterogeneous model composed of CNN and transformer, and introducing pair learning module (PLM) and multi-stage training strategy, the framework offers both branches faster convergence speed and superior performance in image classification task. After training, PLM is removed and CNN and transformer work without additional component, therefore vision pair learning brings no extra computation burden for inference.

• •

  With the help of pair learning, ViT-Base obtains $81\%$ top-1 accuracy on ImageNet validation set (Deng et al. 2009) in 100-epoch training, which surpasses our reimplemented $77.1\%$ baseline (Touvron et al. 2020). Furthermore, with 300-epoch training, ViT-Base can reach $83.47\%$ and exceeds the official experimental result $81.8\%$ by $1.67$ percent. And ResNet-50 achieve competitive result $79.61\%$. In addition, RegNet-16GF (Radosavovic et al. 2020) and ViT-Base are trained jointly as well for further confirming the effectiveness of the proposed method. The accuracy is increased by $2.36\%$ and $2\%$ respectively compared to independent training.

• •

  Experiments on more image classification datasets suggest that vision pair learning is a robust and efficient training framework. And ablation studies prove the effectivity of proposed component including simultaneous training, restricted gradient flow, the necessity of different two branches, and so on.

Transformer (Vaswani et al. 2017) and its applications (Devlin et al. 2018; Brown et al. 2020) have achieved great success in various natural language processing tasks in recent years. Different from traditional RNN model, its attention-only building block is effective in modeling long-range correlation of input sequences. And large number of parameters introduced by transformer architecture increase the model capacity, as well as the difficulty of optimization. When optimizing a big model, abundant training data is usually necessary for a satisfactory performance. To better apply transformer to data-scarce tasks, NLP researchers propose to utilize the pretraining technique. A transformer is typically pretrained in unsupervised manner with large-scale corpus, and then fine-tuned in downstream tasks. For examples, BERT (Devlin et al. 2018) exploits the encoder part of transformer for sequence modeling and propose two unsupervised pretraining tasks, masked language model and next sentence prediction. BERT is proven powerful and reaches the state-of-the-art in $11$ downstream NLP tasks. Transformer decoder based GPT series models (Radford et al. 2018, 2019; Brown et al. 2020), which lead the trend of expanding model size, suggest that pretrained casual language model is good at few-shot/zero-shot learning and natural language generation.

Due to the amazing performance transformer presents in NLP, researchers attempt to introduce it to various CV tasks. (Girdhar et al. 2019) proposes an action transformer model for recognizing and localizing human actions in video clips. The transformer encoder based model combines I3D (Carreira and Zisserman 2017) and RPN (Ren et al. 2016) simultaneously, and makes a significant margin in Atomic Visual Actions dataset (Gu et al. 2018). DETR (Carion et al. 2020) transfers transformer to object detection task. The main idea is to learn a 2D representation of an image using CNN backbone and feed it to the transformer to make the final detection prediction. DETR demonstrates significantly better performance on large objects but fails on small objects. Since the results are not particularly satisfactory, its follow-up UP-DETR (Dai et al. 2020) puts forward a random query patch detection method and boosts the performance of DETR with faster convergence and higher precision. IPT (Chen et al. 2020a) generates corrupted image pairs from ImageNet (Deng et al. 2009) and pretrains transformer on them. By fine-tuning the model in low-level CV tasks such as denoising, super-resolution and deraining, IPT outperforms contemporaneous approaches.

Though much progress transformer has made in CV tasks, the majority of proposed approaches just use transformer as a component. And researchers continue to explore transformer-only solutions. ViT (Dosovitskiy et al. 2020) firstly shows that a pure transformer is promising in image classification tasks. Even though its performance is weaker than some CNN-based methods, it can overtake them when pretrained with huge datasets. DeiT (Touvron et al. 2020) based on ViT further proves transformer can outperform CNN-based methods trained on the same dataset by optimizing the training strategies. Furthermore, the authors presents a distillation strategy for better performance but same inference speed.

Despite that transformer is promising in CV tasks, it is rarely used in production due to the huge computation resources it requires. VPL aims at reducing the cost by combining both advantages of transformer and CNN. It can accelerate the training speed and achieve comparable results with less training iterations. A similar approach (Zhang et al. 2018) has been proposed before, which pays more attention to promoting CNN series by using distillation iteratively. We will introduce the detailed differences between DML and our approach in next section.

Given a batch of input images $\bm{X}=\{\bm{x}_{1},\dots,\bm{x}_{N}\}$, image classification task aims to predict the ground-truth label $\bm{Y}=\{\bm{y}_{1},\dots,\bm{y}_{N}\}$, where $\bm{y}_{i}\in\mathbb{N}^{C}$ is one-hot vector, $N$ is the batch size and $C$ is the class number. The proposed vision pair learning is composed of two branches, pair learning module and multi-stage training. The two branches are original CNN and transformer baseline models, and the pair learning module provide extra supervision signals for them. And the multi-stage training helps branches learn from their partners.

In this section, we first introduce the implementation details. And then we evaluate the performance of VPL on ImageNet classification dataset (Deng et al. 2009). More transfer learning experiments are conducted in various datasets as well, including CIFAR-10 (Krizhevsky, Hinton et al. 2009), CIFAR-100 (Krizhevsky, Hinton et al. 2009), Oxford Flowers-102 (Nilsback and Zisserman 2008), and Stanford Cars (Krause et al. 2013). The statistics of the datasets are in Table 1. Last, we analyze the importance of different components of the VPL in ablation study.

In this paper, we have proposed a new training framework named vision pair learning (VPL) for image classification task. The VPL makes use of the advantages of CNN and transformer and promote theirs performances simultaneously. We prove the effectivity of VPL in ImageNet series datasets and four downstream datasets. And the experiments further verify its ability in building strong representation. Future work will be on improving the model’s capability and flexibility to deal with more different structures and verify its ability in other vision fields. We hope the viewpoints and experimental results described in this work will be helpful to the follow-up work.