Vision Permutator: A Permutable MLP-Like Architecture for Visual Recognition

Introduction

Recent studies (Tolstikhin et al., 2021; Touvron et al., 2021a) have shown that pure multi-layer perceptron based networks perform well in ImageNet classification (Deng et al., 2009). Compared to convolutional neural networks (CNNs) and vision transformers that employ spatial convolutions or self-attention layers to encode spatial information, MLP-like networks (a.k.a., MLPs) make use of pure fully-connected layers (or called $1\times 1$ convolutions) and hence are more efficient in both training and inference (Tolstikhin et al., 2021). However, the good performance of MLPs in image classification largely benefits from training on large-scale datasets (e.g., ImageNet-22K and JFT-300M). Without the support of sufficiently large amount of training data, their performance still lags largely behind CNNs (Tan & Le, 2019; Brock et al., 2021; Zhang et al., 2020) and vision transformers (Jiang et al., 2021; Touvron et al., 2021b; Liu et al., 2021b).

In this work, we are interested in exploiting the potential of MLPs with using merely the ImageNet-1k data for training and target data-efficient MLPs. To this end, we propose the Vision Permutator architecture. Specially, Vision Permutator innovates the existing MLP architectures by presenting a new layer structure that can more effectively encode spatial information based on the basic matrix multiplication routine. Unlike current MLP-like models, such as Mixer (Tolstikhin et al., 2021) and ResMLP (Touvron et al., 2021a), that encode spatial information by flattening the spatial dimensions first and then conducting linear projection along the spatial dimension (i.e., operating on tokens with shape “tokens$\times$channels”), leading to the loss of positional information carried by 2D feature representations, Vision Permutator maintains the original spatial dimensions of the input tokens and separately encode spatial information along the height and width dimensions to preserve positional information.

To be specific, our Vision Permutator begins with a similar tokenization operation to vision transformers, which uniformly splits the input image into small patches and then maps them to token embeddings with linear projections, as depicted in Figure 1. The resulting token embeddings with shape “height$\times$width$\times$channels” are then fed into a sequence of Permutator blocks, each of which consists of a Permute-MLP for spatial information encoding and a Channel-MLP for channel information mixing. The Permute-MLP layer, as depicted in Figure 2, consists of three independent branches, each of which encodes features along a specific dimension, i.e., the height, width or channel dimension. Compared to existing MLP-like models that mix the two spatial dimensions into one, our Vision Permutator separately processes the token representations along these dimensions, resulting in tokens with direction-specific information, which has been demonstrated essential for visual recognition (Hou et al., 2021; Wang et al., 2020).

Experiments show that our Vision Permutator can largely improve the classification performance of existing MLP-like models. Taking the small-sized Vision Permutator as an example, it attains 81.5% top-1 accuracy on ImageNet without any extra training data. Scaling up the model to 55M and 88M, we can further achieve 82.7% and 83.2% accuracy, respectively.

Related Work

Modern deep neural networks for image classification can be mainly categorized into three different classes: convolutional neural networks (CNNs), vision transformers (ViTs), and multi-layer perceptron based models (MLPs). In the following, we will briefly describe the development trend of each type of networks and state the differences of the proposed Vision Permutator from previous work.

CNNs, as the de-facto standard networks in computer vision for years, have been deeply studied. Early CNN models, such as AlexNet (Krizhevsky et al., 2012) and VGGNet (Simonyan & Zisserman, 2014), mostly adopt structures with a stack of spatial convolutions (with kernel size $\geq 3$) and pooling operations. Later, ResNets and their variants (He et al., 2016; Xie et al., 2017; Zagoruyko & Komodakis, 2016) introduce skip connection and building blocks with bottleneck structure into CNNs, enabling training very deep networks possible. Inceptions (Szegedy et al., 2015; 2016) renovate the design of traditional building block structure and utilize multiple parallel paths of sets of specialized filters. Attention mechanisms (Hu et al., 2018; 2019; Wang et al., 2018; Bello et al., 2019; Liu et al., 2020; Chen et al., 2018) break through the limitations of convolutions in capturing local features and further promote the development of CNNs. Our work can also be regarded as a special CNN. Different from previous CNNs that globally aggregate the locally captured features with spatial convolutions, our Vision Permutator is composed of pure $1\times 1$ convolutions but can encode global information.

Our work is also related to vision transformers (Dosovitskiy et al., 2020). Unlike CNNs that exploit local convolutions to encode spatial information, vision transformers takes advantage of the self-attention mechanism to capture global information and have been the prevailing research direction in image classification recently. Since then, a great number of transformer-based classification models appear, aiming at advancing the original vision transformer by either introducing locality (Zhou et al., 2021b; Vaswani et al., 2021; Wu et al., 2021; Liu et al., 2021b; Han et al., 2021; Yuan et al., 2021), or scaling the depth (Zhou et al., 2021a; Touvron et al., 2021b), or tailoring powerful optimization strategies (Jiang et al., 2021). Different from the aforementioned methods, our Vision Permutator eliminates the dependence on self-attention and hence is more efficient.

Very recently, there are also some work (Tolstikhin et al., 2021; Touvron et al., 2021a; Liu et al., 2021a; Guo et al., 2021) targeting at developing pure MLP-like models for ImageNet classification. To encode rich spatial information with MLPs, these methods flatten the spatial dimensions and treat the three-dimensional (height, width, and channel) token representations as a two-dimensional input table. Differently, our Vision Permutator operates on three-dimensional feature representations and encodes spatial information separately along the height and width dimensions. We will show the advantages of the proposed Vision Permutator over existing MLP-like models in our experiment section.

Vision Permutator

The basic architecture of the proposed Vision Permutator can be found in Figure 1. Our network takes an image of size $224\times 224$ as input and uniformly splits it into a sequence of image patches ($14\times 14$ or $7\times 7$). All the patches are then mapped into linear embeddings (or called tokens) using a shared linear layer as (Tolstikhin et al., 2021). We next feed all the tokens into a sequence of Permutators to encode both spatial and channel information. The resulting tokens are finally averaged along the spatial dimensions, followed by a fully-connected layer for class prediction. In the following, we will detail the proposed Permutator block and the network settings.

where, $\operatorname{LN}$ refers to LayerNorm. The output $\mathbf{Z}$ will serve as the input to the next Permutator block until the last one.

Weighted Permute-MLP: In Eqn. 3, we simply fuse the outputs from all three branches with element-wise addition. Here, we further improve the above Permute-MLP by recalibrating the importance of different branches and present Weighted Permute-MLP. This can be easily implemented by exploiting the split attention (Zhang et al., 2020). What is different is that the split attention is applied to $\mathbf{X}_{H}$, $\mathbf{X}_{W}$, and $\mathbf{X}_{C}$ instead of a group of tensors generated by a grouped convolution. In the following, we use the weighted Permute-MLP in Permutator by default.

2 Various Configurations of Vision Permutator

We summarize various configurations of the proposed Vision Permutator in Table 1. We present three different versions of Vision Permutator (ViP), denoted as ‘ViP-Small’, ‘ViP-Medium’, and ‘ViP-Large’ respectively, according to their model size. Notation ‘ViP-Small/14’ denotes the small-sized model with patch size $14\times 14$ in the starting patch embedding module. In ‘ViP-Small/16’ and ‘ViP-Small/14’, there is only one patch embedding module, which is then folloed by a sequence of Permutators. The total number of Permutators for them are 16.

Our ‘ViP-Small/7,’ ‘ViP-Medium/7,’ and ‘ViP-Large/7’ have two stages, each of which starts with a patch embedding module. For these models, we add a few Permutators targeting at encoding fine-level token representations which we found beneficial to the model performance. In our experiment section, we will show the advantage of encoding fine-level token representations.

Experiments

We report of the results of our proposed Vision Permutator on the widely-used ImageNet-1k (Deng et al., 2009) dataset. The code is implemented based on PyTorch (Paszke et al., 2019) and the timm (Wightman, 2019) toolbox. Note that in training, we do not use any extra training data.

We adopt the AdamW optimizer (Loshchilov & Hutter, 2017) with a linear learning rate scaling strategy $lr=10^{-3}\times\frac{batch\_size}{1024}$ and $5\times 10^{-2}$ weight decay rate to optimize all the models as suggested by previous work (Touvron et al., 2020; Jiang et al., 2021). The batch size is set to 2048 which we found works better than 1024 in our Vision Permutator. Stochastic Depth (Huang et al., 2016) is used. Detailed drop rates can be found in Table 1. We train our models on the ImageNet dataset for 300 epochs. For data augmentation methods, we use CutOut (Zhong et al., 2020), RandAug (Cubuk et al., 2020), MixUp (Zhang et al., 2017), and CutMix (Yun et al., 2019). Note that we do not use positional encoding in our Vision Permutator as we found it hurts the performance. Training small-sized Vision Permutator models requires a machine node with 8 NVIDIA V100 GPUs (32G memory). Two nodes are needed for medium-sized and large-sized Vision Permutator models.

2 Main Results on ImageNet

In this subsection, we compare our proposed Vision Permutator with previous CNN-based, Transformer-based, and MLP-like models. We first compare our proposed Vision Permutator with recent MLP-like models in Table 2. The ‘Train size’ and ‘Test Size’ refer to the training resolution and test resolution, respectively. Our ViP-Small/7 model with 25M parameters achieves top-1 accuracy of 81.5%. This result is already better than most of the existing MLP-like models and comparable to the best one gMLP-B (Liu et al., 2021a) with 73M parameters. Scaling up the model to 55M allows our ViP-Medium/7 to attain 82.7% accuracy, which is better than all other MLP-like models as shown in Table 2. Further increasing the model size to 88M leads to a better result 83.2%.

We argue that the main factor leading to the improvement for our Vision Permutator is the way of encoding spatial information as described in Sec. 3. Different from concurrent popular MLP-like models listed in Table 2, we separately encoding the token representations along the height and width dimensions, generating position-sensitive outputs that are crucial for locating and identifying objects of interest (Hou et al., 2021; Wang et al., 2020). In addition, our Vision Permutator encodes not only coarse-level token representations (with $16\times 16$ tokens) but also features at fine-level (with $32\times 32$ tokens). We will detail this in next subsection.

In Table 3, we show the comparison with classic CNN-based and transformer-based models. Compared with classic CNNs, like ResNets (He et al., 2016), SE-ResNeXt (Xie et al., 2017; Hu et al., 2018), and RegNet (Radosavovic et al., 2020), our Vision Permutator with similar model size constraint receives better results. Taking the ViP-Small/7 model as an example, the performance is 81.5%, which is even better than ResNeSt-50 (81.5% v.s. 81.1%). Compared to some transformer-based models, such as DeiT (Touvron et al., 2020), T2T-ViT (Yuan et al., 2021), and Swin Transformers (Liu et al., 2021b), our results are also better. However, there is still a large gap between our Vision Permutator and recent state-of-the-art CNN- and transformer-based models, such as NFNet (Brock et al., 2021) (86.5%), LV-ViT (Jiang et al., 2021) (86.4%), and CaiT (Touvron et al., 2021b) (86.5%). We believe there is still a large room for improving MLP-like models, just like what happened in the research field of vision transformers.

3 Ablation Analysis

In this subsection, we conduct a series of ablation experiments on fine-level information encoding, model scaling, data augmentation, and the proposed Permutator. We take the ViP-Small/14 model as baseline.

Importance of Fine-level Token Representation Encoding: We first show that encoding finer-level token representations is important for MLP-like models. We demonstrate this argument in two ways: I) Adjusting the patch size in the initial patch embedding layer and keep the backbone unchanged; II) Halving the patch size for each patch side and introducing a few Permutators to encode fine-level token representations. Table 4 summaries the performance for ViP-Small/16, ViP-Small/14, and ViP-Small/7. Compared to ViP-Small/16, ViP-Small/14 has smaller initial patch size and more input tokens to the Permutators. According to the results, ViP-Small/14 yields better performance than ViP-Small/16 (80.5% v.s. 79.8%). Despite more tokens and more parameters used in ViP-Small/14, the efficiency (throughput) does not change much. This indicates that we can appropriately use smaller initial patch size to improve the model performance.

We further reduce the initial patch size from $14\times 14$ to $7\times 7$. Compared to ViP-Small/14, ViP-Small/7 adopts 4 Permutators to encode fine-level token representations (with $32\times 32$ tokens). As shown in Table 4, such a slight modification can largely boost the performance and reduce the number of learnable parameters. The top-1 accuracy is improved from 80.5% to 81.5%. This demonstrates that encoding fine-level token representations does help in improving our model performance but a disadvantage is that the efficiency goes down a little.

Role of the model scale: Scaling up models for deep neural networks is always an effective way to improve model performance. Here, we show the influence of model scaling on the proposed Vision Permutator by increasing the number of layers and hidden dimension. Table 5 lists the results for three different versions of the proposed Vision Permutator: ViP-Small/7, ViP-Medium/7, and ViP-Large/7. We can see that increasing the number of layers and hidden dimension yields better results for our Vision Permutator. The ViP-Medium/7 can raise the performance of ViP-Small/7 to 82.7% with a performance gain of more than 1%. Further increasing the model size results in better performance 83.2%.

Effect of Data Augmentations: Data augmentation has been demonstrated an effective and efficient way to lift the model performance in deep learning (He et al., 2019; Touvron et al., 2020; Jiang et al., 2021). Four commonly-used data augmentation methods should be Random Augmentation (Cubuk et al., 2020), CutOut (Zhong et al., 2020), MixUp (Zhang et al., 2017), and CutMix (Yun et al., 2019). Here, we show how each method influences the model performance. The results have been shown in Table 6. Without any data augmentation, we achieve 75.3% top-1 accuracy for our ViP-Small/14 model. Using Random Augmentation improves the performance to 77.7% (+2.4%). Adding CutOut lifts the result to 78.0% (+2.7%). Adding MixUp yields 80.2% top-1 accuracy (+4.9%) and the result is further improved to 80.6% (+5.3%) by using CutMix. These experiments indicate that data augmentation is extremely important in training Vision Permutator as happened in training CNNs (He et al., 2019) and vision transformers (Touvron et al., 2020; Jiang et al., 2021).

Ablation on Permutator: In this paragraph, we demonstrate the importance of encoding spatial information along the height and width dimensions separately and show how weighted Permutator helps in improving model performance. In Table 7, we summarize the results under different Permutator settings. Detailed description on each setting can be found in the caption. We can see that discarding either height information encoding or width information encoding leads to worse performance (80.2% v.s. 72.8% or 72.7%). This demonstrates that encoding both height and width information is important. In addition, we can also observe that replacing the vanilla Permute-MLP with the Weighted Permute-MLP can further improve the performance from 80.2% to 80.6%.

Conclusions and Future Work

In this paper, we present a novel MLP-like network architecture for visual recognition, termed Vision Permutator. We demonstrate that separately encoding the height and width information can largely improve the model performance compared to recent MLP-like models that deem the two spatial dimensions as one. Our experiments also give full support of this.

Despite the large improvement over concurrent popular MLP-like models, a clear downside of the proposed Permutator is the scaling problem in spatial dimensions, which also exists in other MLP-like models. As the shapes of the parameters in fully-connected layers are fixed, it is impossible to process input images with arbitrary shapes. This makes MLP-like models difficult to be used in down-stream tasks with various-sized input images.

Our future work will be continuously put on the development of MLP-like models considering the high efficacy in parallelization. Specifically, we will continue to conquer the limitations of MLP-like models in processing input images with arbitrary shapes and their applications in down-stream tasks, such as object detection and semantic segmentation.

References