Scalable Vision Transformers with Hierarchical Pooling

Introduction

Equipped with the self-attention mechanism that has strong capability of capturing long-range dependencies, Transformer based models have achieved significant breakthroughs in many computer vision (CV) and natural language processing (NLP) tasks, such as machine translation , image classification , segmentation and object detection . However, the good performance of Transformers comes at a high computational cost. For example, a single Transformer model requires more than 10G Mult-Adds to translate a sentence of only 30 words. Such a huge computational complexity hinders the widespread adoption of Transformers, especially on resource-constrained devices, such as smart phones.

To improve the efficiency, there are emerging efforts to design efficient and scalable Transformers. On the one hand, some methods follow the idea of model compression to reduce the number of parameters and computational overhead. Typical methods include knowledge distillation , low-bit quantization and pruning . On the other hand, the self-attention mechanism has quadratic memory and computational complexity, which is the key efficiency bottleneck of Transformer models. The dominant solutions include kernelization , low-rank decomposition , memory , sparsity mechanisms, etc.

Despite much effort has been made, there still lacks specific efficient designs for Visual Transformers considering taking advantage of characteristics of visual patterns. In particular, ViT models maintain a full-length sequence in the forward pass across all layers. Such a design can suffer from two limitations. Firstly, different layers should have different redundancy and contribute differently to the accuracy and efficiency of the network. This statement can be supported by existing compression methods , where each layer has its optimal spatial resolution, width and bit-width. As a result, the full-length sequence may contain huge redundancy. Secondly, it lacks multi-level hierarchical representations, which is well known to be essential for the success of image recognition tasks.

To solve the above limitations, we propose to gradually downsample the sequence length as the model goes deeper. Specifically, inspired by the design of VGG-style and ResNet-style networks, we partition the ViT blocks into several stages and apply the pooling operation (e.g., average/max pooling) in each stage to shrink the sequence length. Such a hierarchical design is reasonable since a recent study shows that a multi-head self-attention layer with a sufficient number of heads can express any convolution layers. Moreover, the sequence of visual tokens in ViT can be analogous to the flattened feature maps of CNNs along the spatial dimension, where the embedding of each token can be seen as feature channels. Hence, our design shares similarities with the spatial downsampling of feature maps in CNNs. To be emphasized, the proposed hierarchical pooling has several advantages. (1) It brings considerable computational savings and improves the scalability of current ViT models. With comparable floating-point operations (FLOPs), we can scale up our HVT by expanding the dimensions of width/depth/resolution. In addition, the reduced sequential resolution also empowers the partition of the input image into smaller patch sizes for high-resolution representations, which is needed for low-level vision and dense prediction tasks. (2) It naturally leads to the generic pyramidal hierarchy, similar to the feature pyramid network (FPN) , which extracts the essential multi-scale hidden representations for many image recognition tasks.

In addition to hierarchical pooling, we further propose to perform predictions without the class token. Inherited from NLP, conventional ViT models equip with a trainable class token, which is appended to the input patch tokens, then refined by the self-attention layers, and is finally used for prediction. However, we argue that it is not necessary to rely on the extra class token for image classification. To this end, we instead directly apply average pooling over patch tokens and use the resultant vector for prediction, which achieves improved performance. We are aware of a concurrent work that also observes the similar phenomenon.

Our contributions can be summarized as follows:

We propose a hierarchical pooling regime that gradually reduces the sequence length as the layer goes deeper, which significantly improves the scalability and the pyramidal feature hierarchy of Visual Transformers. The saved FLOPs can be utilized to improve the model capacity and hence the performance.

Empirically, we observe that the average pooled visual tokens contain richer discriminative patterns than the class token for classification.

Extensive experiments show that, with comparable FLOPs, our HVT outperforms the competitive baseline DeiT on image classification benchmarks, including ImageNet and CIFAR-100.

Related Work

The powerful multi-head self-attention mechanism has motivated the studies of applying Transformers on a variety of CV tasks. In general, current Visual Transformers can be mainly divided into two categories. The first category seeks to combine convolution with self-attention. For example, Carion et al. propose DETR for object detection, which firstly extracts visual features with CNN backbone, followed by the feature refinement with Transformer blocks. BotNet is a recent study that replaces the convolution layers with multiheaded self-attention layers at the last stage of ResNet. Other works also present promising results with this hybrid architecture. The second category aims to design a pure attention-based architecture without convolutions. Recently, Ramachandran et al. propose a model which replaces all instances of spatial convolutions with a form of self-attention applied to ResNet. Hu et al. propose LR-Net that replaces convolution layers with local relation layers, which adaptively determines aggregation weights based on the compositional relationship of local pixel pairs. Axial-DeepLab is also proposed to use Axial-Attention , a generalization form of self-attention, for Panoptic Segmentation. Dosovitskiy et al. first transfers Transformer to image classification. The model inherits a similar architecture from standard Transformer in NLP and achieves promising results on ImageNet, whereas it suffers from prohibitively expensive training complexity. To solve this, the following work DeiT propose a more advanced optimization strategy and a distillation token, with improved accuracy and training efficiency. Moreover, T2T-ViT aims to overcome the limitations of simple tokenization of input images in ViT and propose to progressively structurize the image to tokens to capture rich local structural patterns. Nevertheless, the previous literature all assumes the same architecture to the NLP task, without the adaptation to the image recognition tasks. In this paper, we propose several simple yet effective modifications to improve the scalability of current ViT models.

Transformer-based models are resource-hungry and compute-intensive despite their state-of-the-art performance. We roughly summarize the efficient Transformers into two categories. The first category focuses on applying generic compression techniques to speed up the inference, either based on quantization , pruning , and distillation or seeking to use Neural Architecture Search (NAS) to explore better configurations. Another category aims to solve the quadratic complexity issue of the self-attention mechanism. A representative approach is to express the self-attention weights as a linear dot-product of kernel functions and make use of the associative property of matrix products to reduce the overall self-attention complexity from $\mathcal{O}(n^{2})$ to $\mathcal{O}(n)$. Moreover, some works alternatively study diverse sparse patterns of self-attention , or consider the low-rank structure of the attention matrix , leading to linear time and memory complexity with respect to the sequence length. There are also some NLP literatures that tend to reduce the sequence length during processing. For example, Goyal et al. propose PoWER-BERT, which progressively eliminates word tokens during the forward pass. Funnel-Transformer presents a pool-query-only strategy, pooling the query vector within each self-attention layer. However, there are few literatures targeting improving the efficiency of the ViT models.

To compromise FLOPs, current ViT models divide the input image into coarse patches (i.e., large patch size), hindering their generalization to dense predictions. In order to bridge this gap, we propose a general hierarchical pooling strategy that significantly reduces the computational cost while enhancing the scalability of important dimensions of the ViT architectures, i.e., depth, width, resolution and patch size. Moreover, our generic encoder also inherits the pyramidal feature hierarchy from classic CNNs, potentially benefiting many downstream recognition tasks. Also note that different from a concurrent work which applies 2D patch merging, this paper introduces the feature hierarchy with 1D pooling. We discuss the impact of 2D pooling in Section 5.2.

Proposed Method

In this section, we first briefly revisit the preliminaries of Visual Transformers and then introduce our proposed Hierarchical Visual Transformer.

Suppose that the encoder in a Transformer consists of $L$ blocks. Each block contains a multi-head self-attention (MSA) layer and a position-wise multi-layer perceptron (MLP). For each layer, layer normalization (LN) and residual connections are employed, which can be formulated as follows

where $l\in[1,...,L]$ is the index of Transformer blocks. Here, a MLP contains two fully-connected layers with a GELU non-linearity . In order to perform classification, ViT applies a layer normalization layer and a fully-connected (FC) layer to the first token of the Transformer encoder’s output ${\bf X}_{L}^{0}$. In this way, the output prediction ${\bf y}$ can be computed by

2 Hierarchical Visual Transformer

In this paper, we propose a Hierarchical Visual Transformer (HVT) to reduce the redundancy in the full-length patch sequence and construct a hierarchical representation. In the following, we first propose a hierarchical pooling to gradually shrink the sequence length and hence reduce the computational cost. Then, we propose to perform predictions without the class token. The overview of the proposed HVT is shown in Figure 2.

We propose to apply hierarchical pooling in ViT for two reasons: (1) Recent studies on Transformers show that tokens tend to carry redundant information as it goes deeper. Therefore, it would be beneficial to reduce these redundancies through the pooling approaches. (2) The input sequence projected from image patches in ViT can be seen as flattened CNN feature maps with encoded spatial information, hence pooling from the nearby tokens can be analogous to the spatial pooling methods in CNNs.

Motivated by the hierarchical pipeline of VGG-style and ResNet-style networks, we partition the Transformer blocks into $M$ stages and apply downsampling operation to each stage to shrink the sequence length. Let $\{b_{1},b_{2},\dots,b_{M}\}$ be the indexes of the first block in each stage. At the $m$-th stage, we apply a 1D max pooling operation with a kernel size of $k$ and stride of $s$ to the output of the Transformer block $b_{m}\in\{b_{1},b_{2},\dots,b_{M}\}$ to shrink the sequence length.

Note that the positional encoding is important for a Transformer since the positional encoding is able to capture information about the relative and absolute position of the token in the sequence . In Eq. (1) of ViT, each patch is equipped with positional embedding ${\bf E}$ at the beginning. However, in our HVT, the original positional embedding ${\bf E}$ may no longer be meaningful after pooling since the sequence length is reduced after each pooling operation. In this case, positional embedding in the pooled sequence needs to be updated. Moreover, previous work in NLP also find it important to complement positional information after changing the sequence length. Therefore, at the $m$-th stage, we introduce an additional learnable positional embedding ${\bf E}_{b_{m}}$ to capture the positional information, which can be formulated as

where ${\bf X}_{b_{m}}$ is the output of the Transformer block $b_{m}$. We then forward the resulting embeddings $\hat{{\bf X}}_{b_{m}}$ into the next Transformer block $b_{m}+1$.

2.2 Prediction without the Class Token

Previous works make predictions by taking the class token as input in classification tasks as described in Eq. (4). However, such structure relies solely on the single class token with limited capacity while discarding the remaining sequence that is capable of storing more discriminative information. To this end, we propose to remove the class token in the first place and predict with the remaining output sequence on the last stage.

Specifically, given the output sequence without the class token on the last stage ${\bf X}_{L}$, we first apply average pooling, then directly apply an FC layer on the top of the pooled embeddings and make predictions. The process can be formulated as

3 Complexity Analysis

In this section, we analyse the block-wise compression ratio with hierarchical pooling. Following ViT , we use FLOPs to measure the computational cost of a Transformer. Let $n$ be the number of tokens in a sequence and $d$ is the dimension of each token. The FLOPs of a Transformer block $\phi_{BLK}(n,d)$ can be computed by

where $\phi_{MSA}(n,d)$ and $\phi_{MLP}(n,d)$ are the FLOPs of the MSA and MLP, respectively. Details about Eq. (7) can be found in the supplementary material.

Without loss of generality, suppose that the sequence length $n$ is reduced by half after performing hierarchical pooling. In this case, the block-wise compression ratio $\alpha$ can be computed by

Clearly, Eq. (8) is monotonic, thus the block-wise compression ratio $\alpha$ is bounded by $(2,4)$, i.e., $\alpha\in(2,4)$.

Discussions

In CNNs, feature maps are usually downsampled to smaller sizes in a hierarchical way . In this paper, we show that this principle can be applied to ViT models by comparing the visualized feature maps from ResNet conv4_2, DeiT-S block1 and HVT-S-1 block1 in Figure 3. From the figure, in ResNet, the initial feature maps after the first convolutional layer contain rich edge information. After feeding the features to consecutive convolutional layers and a pooling layer, the output feature maps tend to preserve more high-level discriminative information. In DeiT-S, following the ViT structure, although the image resolution for the feature maps has been reduced to 14 $\times$ 14 by the initial linear projection layer, we can still observe clear edges and patterns. Then, the features get refined in the first block to obtain sharper edge information. In contrast to DeiT-S that refines features at the same resolution level, after the first block, the proposed HVT downsamples the hidden sequence through a pooling layer and reduces the sequence length by half. We then interpolate the sequence back to 196 and reshape it to 2D feature maps. We can find that the hidden representations contain more abstract information with high discriminative power, which is similar to ResNet.

2 Scalability of HVT

The computational complexity reduction equips HVT with strong scalability in terms of width/depth/patch size/resolution. Take DeiT-S for an example, the model consists of 12 blocks and 6 heads. Given a 224$\times$224 image with a patch size of 16, the computational cost of DeiT-S is around 4.6G FLOPs. By applying four pooling operations, our method is able to achieve nearly 3.3$\times$ FLOPs reduction. Furthermore, to re-allocate the reduced FLOPs, we may construct wider or deeper HVT-S, with 11 heads or 48 blocks, then the overall FLOPs would be around 4.51G and 4.33G, respectively. Moreover, we may consider a longer sequence by setting a smaller patch size or using a larger resolution. For example, with a patch size of 8 and an image resolution of 192$\times$192, the FLOPs for HVT-S is around 4.35G. Alternatively, enlarging the image resolution into 384$\times$384 will lead to 4.48G FLOPs. In all of the above mentioned cases, the computational costs are still lower than that of DeiT-S while the model capacity is enhanced.

It is worth noting that finding a principled way to scale up HVT to obtain the optimal efficiency-vs-accuracy tradeoff remains an open question. At the current stage, we take an early exploration by evenly partitioning blocks and following model settings in DeiT for a fair comparison. In fact, the improved scalability of HVT makes it possible for using Neural Architecture Search (NAS) to automatically find optimal configurations, such as EfficientNet . We leave for more potential studies for future work.

Experiments

To investigate the effectiveness of HVT, we compare our method with DeiT and a BERT-based pruning method PoWER-BERT . DeiT is a representative Vision Transformer and PoWER progressively prunes unimportant tokens in pretrained BERT models for inference acceleration. Moreover, we consider two architectures in DeiT for comparisons: HVT-Ti: HVT with the tiny setting. HVT-S: HVT with the small setting. For convenience, we use “Architecture-$M$” to represent our model with $M$ pooling stages, e.g., HVT-S-1.

We evaluate our proposed HVT on two image classification benchmark datasets: CIFAR-100 and ImageNet . We measure the performance of different methods in terms of the Top-1 and Top-5 accuracy. Following DeiT , we measure the computational cost by FLOPs. Moreover, we also measure the model size by the number of parameters (Params).

For experiments on ImageNet, we train our models for 300 epochs with a total batch size of 1024. The initial learning rate is 0.0005. We use AdamW optimizer with a momentum of 0.9 for optimization. We set the weight decay to 0.025. For fair comparisons, we keep the same data augmentation strategy as DeiT . For the downsampling operation, we use max pooling by default. The kernel size $k$ and stride $s$ are set to 3 and 2, respectively, chosen by a simple grid search on CIFAR100. Besides, all learnable positional embeddings are initialized in the same way as DeiT. More detailed settings on the other hyper-parameters can be found in DeiT. For experiments on CIFAR-100, we train our models with a total batch size of 128. The initial learning rate is set to 0.000125. Other hyper-parameters are kept the same as those on ImageNet.

1 Main Results

We compare the proposed HVT with DeiT and PoWER, and report the results in Table 1. First, compared to DeiT, our HVT achieves nearly 2$\times$ FLOPs reduction with a hierarchical pooling. However, the significant FLOPs reduction also leads to performance degradation in both the tiny and small settings. Additionally, the performance drop of HVT-S-1 is smaller than that of HVT-Ti-1. For example, for HVT-S-1, it only leads to $1.80\%$ drop in the Top-1 accuracy. In contrast, it results in $2.56\%$ drop in the Top-1 accuracy for HVT-Ti-1. It can be attributed to that, compared with HVT-Ti-1, HVT-S-1 is more redundant with more parameters. Therefore, applying hierarchical pooling to HVT-S-1 can significantly reduce redundancy while maintaining performance. Second, compared to PoWER, HVT-Ti-1 uses less FLOPs while achieving better performance. Besides, HVT-S-1 reduces more FLOPs than PoWER, while achieving slightly lower performance than PoWER. Also note that PoWER involves three training steps, while ours is a simpler one-stage training scheme.

Moreover, we also compare the scaled HVT with DeiT under similar FLOPs. Specifically, we enlarge the embedding dimensions and add extra heads in HVT-Ti. From Table 1 and Figure 4, by re-allocating the saved FLOPs to scale up the model, HVT can converge to a better solution and yield improved performance. For example, the Top-1 accuracy on ImageNet can be improved considerably by 3.03% in the tiny setting. More empirical studies on the effect of model scaling can be found in Section 5.2.

2 Ablation Study

To investigate the effect of the prediction without the class token, we train DeiT-Ti with and without the class token and show the results in Table 2. From the results, the models without the class token outperform the ones with the class token. The performance gains mainly come from the extra discriminative information stored in the entire sequence without the class token. Note that the performance improvement on CIFAR-100 is much larger than that on ImageNet. It may be attributed that CIFAR-100 is a small dataset, which lacks varieties compared with ImageNet. Therefore, the model trained on CIFAR-100 benefits more from the increase of model’s discriminative power.

We train HVT-S with different pooling stages $M\in\{0,1,2,3,4\}$ and show the results in Table 4. Note that HVT-S-0 is equivalent to the DeiT-S without the class token. With the increase of $M$, HVT-S achieves better performance with decreasing FLOPs on CIFAR-100, while on ImageNet we observe the accuracy degrades. One possible reason is that HVT-S is very redundant on CIFAR-100, such that pooling acts as a regularizer to avoid the overfitting problem and improves the generalization of HVT on CIFAR-100. On ImageNet, we assume HVT is less redundant and a better scaling strategy is required to improve the performance.

To investigate the effect of different downsampling operations, we train HVT-S-4 with three downsampling strategies: convolution, average pooling and max pooling. As Table 3 shows, downsampling with convolution performs the worst even it introduces additional FLOPs and parameters. Besides, average pooling performs slightly better than convolution in terms of the Top-1 accuracy. Compared with the two settings, HVT-S-4 with max pooling performs much better as it significantly surpasses average pooling by 5.05% in the Top-1 accuracy and 2.17% in the Top-5 accuracy. The result is consistent with the common sense that max pooling performs well in a large variety of settings. To this end, we use max pooling in all other experiments by default.

One of the important advantages of the proposed hierarchical pooling is that we can re-allocate the saved computational cost for better model capacity by constructing a model with a wider, deeper, larger resolution or smaller patch size configuration. Similar to the CNNs literature , we study the effect of model scaling in the following.

Based on HVT-S-4, we first construct deeper models by increasing the number of blocks in Transformers. Specifically, we train 4 models with different number of blocks $L\in\{12,16,20,24\}$. As a result, each pooling stage for different models would have 3, 4, 5, and 6 blocks, respectively. We train 4 models on CIFAR-100 and report the results in Table 5. From the results, we observe no more gains by stacking more blocks in HVT.

Based on HVT-Ti-4, we then construct wider models by increasing the number of self-attention heads. To be specific, we train 4 models with different numbers of self-attention heads, i.e., 3, 6, 12, and 16, on CIFAR-100 and report the results in Table 6. From the results, our models achieve better performance with the increase of width. For example, the model with 16 self-attention heads outperforms those with 3 self-attention heads by 6.79% in the Top-1 accuracy and 1.38% in the Top-5 accuracy.

Based on HVT-S-4, we further construct models with larger input image resolutions. Specifically, we train 4 models with different input image resolutions, i.e., 160, 224, 320, and 384, on CIFAR-100 and report the results in Table 7. From the results, with the increase of image resolution, our models achieve better performance. For example, the model with the resolution of 384 outperforms those with the resolution of 160 by 2.47% in the Top-1 accuracy and 1.12% in the Top-5 accuracy. Nevertheless, increasing image resolutions also leads to high computational cost. To make a trade-off between computational cost and accuracy, we set the image resolution to 224 by default.

We finally train HVT-S-4 with different patch sizes $P\in\{8,16,32\}$ and show the results in Table 8. From the results, HVT-S-4 performs better with the decrease of patch size. For example, when the patch size decreases from 32 to 8, our HVT-S achieves 9.14% and 4.03% gain in terms of the Top-1 and Top-5 accuracy. Intuitively, a smaller patch size leads to fine-grained image patches and helps to learn high-resolution representations, which is able to improve the classification performance. However, with a smaller patch size, the patch sequence will be longer, which significantly increases the computational cost. To make a balance between the computational cost and accuracy, we set the patch size to 16 by default.

Compared to 1D pooling, 2D pooling brings more requirements. For example, it requires a smaller patch size to ensure a sufficient sequence length. Correspondingly, it is essential to reduce the heads at the early stages to save FLOPs and memory consumption from high-resolution feature maps. Besides, it also requires to vary the blocks at each stage to control the overall model complexity. In Table 9, we apply 2D pooling to HVT-S-2 and compare it with DeiT-S. The results show that HVT-S-2 with 2D pooling outperforms DeiT-S on CIFAR100 by a large margin with similar FLOPs. In this case, we assume that HVT can achieve promising performance with a dedicated scaling scheme for 2D pooling. We will leave this exploration for future work.

Conclusion and Future Work

In this paper, we have presented a Hierarchical Visual Transformer, termed HVT, for image classification. In particular, the proposed hierarchical pooling can significantly compress the sequential resolution to save computational cost in a simple yet effective form. More importantly, this strategy greatly improves the scalability of visual Transformers, making it possible to scale various dimensions - depth, width, resolution and patch size. By re-allocating the saved computational cost, we can scale up these dimensions for better model capacity with comparable or fewer FLOPs. Moreover, we have empirically shown that the visual tokens are more important than the single class token for class prediction. Note that the scope of this paper only targets designing our HVT as an encoder. Future works may include extending our HVT model to decoder and to solve other mainstream CV tasks, such as object detection and semantic/instance segmentation. In addition, it would be interesting to find a principled way to scale up HVT that can achieve better accuracy and efficiency.

Acknowledgements

This research is partially supported by Monash FIT Start-up Grant and Sensetime Gift Fund.

References

Appendix

We organize our supplementary material as follows.

In Section S1, we elaborate on the components of a Transformer block, including the multi-head self-attention layer (MSA) and the position-wise multi-layer perceptron (MLP).

In Section S2, we provide details for the FLOPs calculation of a Transformer block.

S1 Transformer Block

S1.2 Position-wise Multi-Layer Perceptron

Let ${\bf X}^{{}^{\prime}}$ be the output from the MSA layer. An MLP layer which contains two fully-connected layers with a GELU non-linearity can be represented by

S2 FLOPs of a Transformer Block

We denote $\phi(n,d)$ as a function of FLOPs with respect to the sequence length $n$ and the embedding dimension $d$. For an MSA layer, The FLOPs mainly comes from four parts: (1) The projection of ${\bf Q}$,${\bf K}$,${\bf V}$ matrices $\phi_{qkv}(n,d)=3nd^{2}$. (2) The calculation of the attention map $\phi_{A}(n,d)=n^{2}d$. (3) The self-attention operation $\phi_{O}(n,d)=n^{2}d$. (4) And finally, a linear projection for the concatenated self-attention outputs $\phi_{proj}(n,d)=nd^{2}$. Therefore, the overall FLOPs for an MSA layer is

By combining Eq. (14) and Eq. (15), we can get the total FLOPs of one Transformer block