Exploring and Improving Mobile Level Vision Transformers

The most representative traditional vision transformer is DeiT, which consists of three components, patch embedding, transformer blocks and class token. We use DeiT to represent traditional vision transformer and analyze its components.

The first part of DeiT is the patch embedding module, which splits the input image into $14\times 14$ patches evenly. Each patch contains $16\times 16$ pixels, and it can be seen as a vector of size $16\times 16\times 3=768$. Then every patch is projected into a certain dimension using a fully-connected(FC) layer. These projected vectors are then forwarded as input to the following transformer blocks.

The main components are several transformer blocks. Each block consists of a multi-head self-attention (MSA) module and a fast-forward network (FFN) module. The MSA module adopts self-attention to realize the communication between patches, which captures the global information of the whole image. The FFN module is a multi-layer perceptron (MLP), which extracts the information inside every patch. We shall notice that, both the MSA module and FFN module are not experts at extracting local low-level information.

The final part is the class token. The class token is processed as a special patch, which is added to the $14\times 14$ patches and these 197 patches will go through all transformer blocks together. The class token is regarded as an indicator to represent the class of the input image. After the transformer blocks, we pass the class token to a FC head to get the final prediction.

While the first and last part of DeiT is contributing little to the overall FLOPs, the majority of DeiT parameters are in the second part. When compressing the DeiT into the mobile level, it’s a good plan to reduce the parameters of the second part. However, We show that, the existing patch embedding and class token are not well designed if we simply compress the transformer blocks.

Following the official compression process from DeiT-Base to DeiT-Tiny, the compression of DeiT lies in two aspects, the channel of every patch and the number of transformer blocks. When we reduce the channel of patches, the information loss in patch embedding becomes very large. The channel number of DeiT-Base is 768, which equals the dimension of every patch in original images($16\times 16\times 3=768$). DeiT-Small and DeiT-Tiny reducing the channel number to 384 and 192, it will lose almost 1/2 and 3/4 information in the patch embedding blocks respectively. When we further compress the model to mobile-level, the information loss becomes quite severe and greatly affects the final accuracy. Therefore, a more advanced patch embedding module that can extract the information from the original image to a low-dimension patch is urgently needed. What’s more, when we reduce the number of transformer blocks, the drawbacks that MSA and FFN cannot extract low-level information effectively will be exposed. We can also use a carefully designed patch embedding module to overcome these drawbacks by extracting more high-level information into every patch.

In addition, when we reduce the number of transformer blocks, the information communication between the class token and other patches is decreased to some extent. Thus, the information in the class token is not good enough for the final prediction, which also highly affects the accuracy. A more effective module to replace the class token is also essential for mobile level vision transformers.

We propose a novel irregular patch embedding (IPE) module for mobile level vision transformers. We introduce our design in three steps progressively.

Firstly, we propose to use convolution layers in the patch embedding module. There are many advantages of using convolution layers. Many researches show that the convolution layer is skilled at extracting low-level local information, which is exactly a complement for transformer blocks. We can adopt a very elegant convolutional module to extract patches with rich high-level information for the following transformer blocks. Thus the information can be concentrated in low-dimensional patches, which effectively relieves the information loss in the traditional patch embedding. Therefore, using convolution layers as a patch embedding module is very suitable for mobile level vision transformers. More specifically, we use the combination of depth-wise convolution and point-wise convolution to further reduce resource consumption.

Secondly, we find that $14\times 14$ high-dimensional patches are too overwhelmed for mobile level vision transformers to afford. On the one hand, if we reduce the channels of patches and number of blocks to solve the above issue, the remaining transformer cannot sufficiently process that many patches. On the other hand, we are capable of reducing the number of patches while increasing the channels and numbers of blocks to find a better balance in the mobile level models. What’s more, according to the first step, we can manage to generate quite high-quality patches, which contain more high-level information. The quality of patches helps relieve the influence of decreased patches. To some extent, the quality of patches is even more crucial than the number of patches in mobile level vision transformers.

Thirdly, we produce irregular patches in mobile level vision transformers. DeiT splits the input image into $14\times 14$ parts evenly and obtains regular patches from them, making these patches have the same-level receptive field. However, the type of classes in ImageNet dataset varies and these classes exhibit quite dissimilar characteristics. Some object classes like dog and cat are more likely to be discriminated based on small-region high-level information, while the whole-image texture-level information turns out to be more important for other scene classes like cliff and lakeside. Therefore, using regular patches is not efficient. We tend to generate irregular patches which contain different receptive fields to further improve mobile level vision transformers’ performance. The detailed architecture is shown in Figure 3. Concretely, instead of $8\times 8$ patches, we use $7\times 7+4\times 4+1\times 1$ patches. These patches are generated with three parallel modules which consist of inverted residual bottleneck blocks with SE layer. We use different numbers of strides to generate patches with different receptive fields. The vision transformer blocks own the great power to effectively communicate between these patches and excavate the most useful information for final prediction.

In DeiT, there is a learnable class token to predict the classification results. The class token is randomly initialized and optimized via gradient descent during training. The class token is concatenated with patches together to be processed by vision transformer blocks. In order to extract useful information for recognition, the class token communicates with other patches in the MSA module. When it comes to the final prediction, only the class token is used, and the patches are abandoned. This design is suitable for large models, where the class token naturally gathers enough information from adequate transformer blocks. In that case, other patches’ information is not necessary, and the participation of other patches may even harm the prediction.

However, the situation in the mobile level vision transformer is very different. We have to carefully utilize every patches’ information under small computation limitations. We find that adopting the class token is much worse than integrating all patches’ information in mobile level vision transformers. We conjecture that the class token cannot gather enough information with a limited number of transformer blocks. Therefore, how to effectively using all patches’ information is a key problem in mobile level vision transformers.

To tackle the above problem, we replace the class token with a designed simple patch merging module that can efficiently gather all patches’ information. In Section 3.2, we have mentioned that we are using irregular patches, and different images may focus on different patches. Based on this, we propose to adaptively fuse patches, which means we assign the patches different weights according to each image. Then we gather all patches together. The adaptive patch fusion module has a quite simple mechanism as shown in Figure 4. We generate an adaptive weight for each image and maintain a global weight for all images. The final weights is obtained by multiplying these two weights together. The adaptive weights are calculated using a multi-layer perceptron. The global weights are updated by gradients during training and fixed during testing, which is computed by a linear layer. This design is much more advantageous than the class token to mobile level vision transformers. More experimental results and visualized analyses are available in Section 5.2.

Overall, we propose IPE and APM to boost the performance of traditional uniform vision transformers like DeiT. We set up different sizes of transformers, the details are summarized in Figure 1.

We also evaluate our proposed method on other datasets. CIFAR-10 and CIFAR-100 are relatively small datasets that contain 50,000 images for training and 10,000 images for validation. The size of images in the CIFAR datasets is only $32\times 32$. We upsample them to $224\times 224$ following DeiT [34]. But the number of training epochs in DeiT is too large (7200 epochs), we use extremely fewer epochs (300 epochs) in order to save tremendous training costs. Besides, we evaluate our method on iNautural-18 and iNatural-19 datasets. We also resize the images to $224\times 224$ and train models for 300 epochs. The result is summarized in Table 3. We can find that, our methods’ performance is much better than the baseline DeiT method on downstream tasks.

We experiment with different types of patch embedding and summarize the results in Table 4. We report different models’ top-1 accuracy on ImageNet. We only train these models for 100 epochs to save resources. Please note that, we only change the patch embedding in this table. (The channel of patches and number of transformer blocks is changed in order to meet the FLOPs constraint.) And we have the class token kept and do not adopt our proposed adaptive patch merging module.

The model (a) is the DeiT-tiny baseline whose FLOPs is around 1200M. The model (b)’s FLOPs is also around 1200M. Other models’ are around 900M. As is shown in this table, when we introduce convolution layers to the patch embedding (PE) module, the performance is significantly improved. This is because the convolution layers are good at processing low-level features and extracting more useful information into the low-dimensional features for every patch. Comparing (b) and (c), the performance is improved with fewer FLOPs. This phenomenon verifies our hypothesis that the quality of patches is more critical than the quantity. Fewer high-quality patches is more practical than a large amount of low-dimensional patches. After we split the patches into different irregular patches, the performance is further improved as shown in (d).

We experiment with different patch merging methods in this section. We also train the models on ImageNet for 100 epochs and report their top-1 accuracy. In Table 5, we only manipulate the patch merging method, and we always use the irregular patch embedding.

Comparing (a) and (b), we surprisingly find that the simplest average pooling is much better than the class token, which is contrary to other large vision transformers. We conjecture in mobile level vision transformers, the class token cannot effectively gather adequate information with fewer transformer blocks. And the average pooling can also integrate all patches together to achieve better performance.

Based on the observation of (b), we propose our adaptive patch merging module, which can gather patches information effectively. As shown in Table 5 (c), our proposed module can further improve the performance of average pooling.

We also visualize some images’ patch merging weights in Figure 5. The input image and patch weight for the input are listed. The lighter means a larger value in the patch weight. We can easily conclude that for small objects like birds, the patch weight will focus on the target objects’ area. For objects with obvious characteristics like a rooster, the patch weight can correctly focus on the most discriminative parts. Please note that, for most images, the patch weights of corners are relatively small, which is consistent with our intuition that corners are not important to classify these images.

We also investigate the usage of positional encoding in our framework. Since we didn’t change the transformer blocks of DeiT, we continue to use the same positional encoding as DeiT. And we show the results of our model with and without positional encoding in Table 6. We find that, the positional encoding is not essential for our model. The results in Table 6 are quite similar. We conjecture that, the irregular patch embedding module already encodes enough positional information inside the patch, and this kind of positional encoding doesn’t provide extra valuable information for the transformer blocks. A more suitable positional encoding may further improve the performance, we leave this part for future study.