SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jin Li $^{1,\dagger}$  Yaoming Wang $^{1,\dagger}$  Xiaopeng Zhang $^{2}$  Bowen Shi $^{1}$  Dongsheng Jiang $^{2}$  Chenglin Li $^{1}$  Wenrui Dai $^{1}$  Hongkai Xiong $^{1}$  Qi Tian $^{2*}$  Shanghai Jiao Tong University $^{2}$  Huawei Cloud {deserve_lj, wang_yaoming, sjtu_shibowen, lcl1985, daiwenrui, xionghongkai}@sjtu.edu.cn; {zhangxiaopeng12, jiangdongsheng1, tian.qi1}@huawei.com
Vision transformers (ViTs) have emerged as a prevalent architecture for vision tasks owing to their impressive performance. However, when it comes to handling long token sequences, especially in dense prediction tasks that require high-resolution input, the complexity of ViTs increases significantly. Notably, dense prediction tasks, such as semantic segmentation or object detection, emphasize more on the contours or shapes of objects, while the texture inside objects is less informative. Motivated by this observation, we propose to apply adaptive resolution for different regions in the image according to their importance. Specifically, at the intermediate layer of the ViT, we utilize a spatial-aware density-based clustering algorithm to select representative tokens from the token sequence. Once the representative tokens are determined, we proceed to merge other tokens into their closest representative token. Consequently, semantic similar tokens are merged together to form low-resolution regions, while semantic irrelevant tokens are preserved independently as high-resolution regions. This strategy effectively reduces the number of tokens, allowing subsequent layers to handle a reduced token sequence and achieve acceleration. At the output layers, the resolution of the feature map is restored by unfolding the merged tokens for task prediction. As a result, our method significantly accelerates ViTs for dense prediction tasks. We evaluate our proposed method on three different datasets and observe promising performance. For example, the "Segmenter ViT-L" model can be accelerated by  $48\%$  FPS without fine-tuning, while maintaining the performance. Additionally, our method can be applied to accelerate fine-tuning as well. Experimental results demonstrate that we can save  $52\%$  training time while accelerating  $2.46 \times$  FPS with only a  $0.09\%$  performance drop. The code is available at https://github.com/caddyless/ailurus/tree/main.
Transformers have shown significant advancements in various vision tasks such as image classification [9, 27, 33, 18], object detection [38, 34], and semantic segmentation[26, 4]. Despite their impressive performance across various visual tasks, the complexity of these models poses challenges for fine-tuning and deployment, particularly as their capacity continues to grow [11, 13, 39]. This complexity issue is particularly relevant for dense prediction tasks that require high-resolution input. Efforts have been made to address this challenge by designing efficient ViT models [36, 6, 20, 24, 19, 16, 23, 31, 14, 1, 17]. However, most of these works are primarily focused on classification tasks and are not applicable to dense prediction tasks. Recently, [15] proposed to expedite well-trained ViTs for dense
*: Corresponding author. †: Equal contribution. This work was done when Jin Li interned at Huawei Cloud.
prediction tasks through super-pixel clustering. Nevertheless, the clustering operation can be only conducted in the relatively deep layers, resulting in limited acceleration ratio and scalability.
In this paper, we introduce a novel approach, namely Adaptive resolution with spatial-awaRe clustering (AiluRus), to accelerate ViTs for dense prediction tasks. We find that dense prediction tasks focus more on the shape or contour of objects rather than the texture. For instance, in segmentation maps, the contour of objects carries crucial information, while the interior regions are filled with the same prediction values, indicating their lower informativeness compared to the boundaries. Motivated by this observation, we propose a token pruning technique that incorporates adaptive resolution to represent different regions in an image. Specifically, we allocate more tokens to critical regions that contribute to decision-making and fewer tokens to less informative regions. The main challenge is to determine a reasonable assignment of resolutions. To address this issue, we utilize density-based clustering algorithms[10, 21] to generate the assignment of each token, where spatial information is incorporated to encourage neighboring tokens to have the same assignment. Tokens that have the same assignments are averaged to produce the representative tokens. These representative tokens could be the original ones, which correspond to informative regions, or the average of several tokens, which correspond to less informative regions. This approach enables us to reduce the length of the token sequence at intermediate layers, thereby accelerating the model. At the output stage, we restore the original resolution for prediction tasks by assigning the value of the representative token to its corresponding regions.
We provide compelling visualizations in Figure Fig. 2 to support the reasonableness of the generated assignments and their limited impact on visual perception. To assess the effectiveness of our proposed method, we adopt the benchmark of [15] and integrate our method into well-trained models without fine-tuning. Our experimental results demonstrate that AiluRus effectively accelerates ViTs and outperforms previous methods, particularly in scenarios with high acceleration ratios. Specifically, AiluRus achieves a  $48\%$  increase in FPS for Segmenter ViT-L while maintaining the performance. Moreover, we further apply AiluRus to accelerate the fine-tuning process. Experiments show that AiluRus reduces training time by  $52\%$  while achieving a  $2.46\times$  increase in FPS with only a  $0.09\%$  drop in performance. These findings demonstrate the effectiveness of AiluRus in accelerating ViTs.
In summary, we list our contributions as follows:
- We propose to apply adaptive resolution on the feature map of ViT-based dense prediction tasks for acceleration without fine-tuning.
- We propose to generate the resolution assignments through the proposed spatial-aware DPC algorithm. Visualizations demonstrate that the produced assignments have little influence on visual perception and thus could expedite models without fine-tuning.
- Our proposed AiluRus can be used to accelerate well-trained models or pre-trained models for inference or fine-tuning. Experiments show that AiluRus could significantly accelerate models with a negligible performance drop.
Vision transformer for dense prediction tasks. Transformers have gained immense popularity in Natural Language Processing (NLP) tasks, and there have been considerable efforts to extend their success to computer vision tasks. DETR [38] introduced transformers as the detection head in a convolutional neural network, opening up new avenues for utilizing transformers in dense prediction tasks. This work has inspired the development of hybrid-transformer architectures aimed at facilitating dense prediction tasks [3, 30]. Other works have proposed pure transformer architectures, which have achieved significant progress in recent advances [26, 4, 29]. In this paper, instead of proposing a new architecture or framework for dense prediction tasks, we focus on accelerating existing dense prediction methods.
Efficient vision transformers. One of the primary strategies for improving the efficiency of Vision Transformers (ViTs) is to reduce the complexity of the self-attention operation. The conventional self-attention operation involves establishing interactions between any two tokens, resulting in quadratic complexity with respect to the number of tokens. To address this challenge, recent approaches aim to approximate the self-attention results through clustering based on the sparse and low-rank properties of self-attention [36, 6, 28].
Figure 1: The framework of AiluRus. we focus on a specific intermediate layer of the ViT and apply our AiluRus using the spatial-aware DPC algorithm. This algorithm searches for cluster centers based on local density and distance indicators, and then averages the tokens in each cluster to obtain a representative token for the following layers. To ensure that each representative token is weighted appropriately, we re-weight them based on the number of tokens they represent. Finally, we unfold the representative tokens at the output end to recover the original resolution. For a more detailed explanation of our method, please refer to Section 3.
Another line of works focuses on token pruning [20, 19, 16, 23, 31, 14, 1, 8]. These methods aim to gradually discard less informative tokens at different layers and retain only a subset of tokens at the output end to accelerate the model. However, most of these approaches are designed for classification tasks and are less practical for dense prediction tasks. For instance, EViT [16], Evo-ViT [31], and PCAE [14] select informative tokens based on a single criterion, such as the attention weight to the class token or similarity to the mean token, which is not suitable for dense prediction tasks where there are many objects belonging to various categories in the image. DynamicViT [20] and Ada-ViT [19] rely on specific architectures and require re-training, which may lead to additional computational overhead in practical applications. ToMe [1] progressively merges a percentage of the most similar tokens between bipartite tokens but has only been verified in classification tasks.
Some works aim to design more efficient ViTs by introducing learnable token merging or token pruning modules [2, 12, 32, 25]. For example, STViT [2] replaces redundant tokens with a few semantic tokens, which can be regarded as cluster centers for acceleration. PaCa-ViT [12] proposes patch-to-cluster attention to address the semantic gap between patch-to-patch attention in visual tasks and its NLP counterpart. TCFoMa [32] merges tokens from less informative regions via clustering to emphasize critical regions for human-centric tasks. Although these methods achieve promising results for efficient ViTs, some of them rely on specific architectures [32, 12] and all require fine-tuning. Recently, Liang et al. [15] proposed a method to expedite well-trained large ViTs for dense prediction tasks using superpixel clustering, which does not require fine-tuning. However, this strategy can only be applied in relatively deep layers, resulting in limited improvements in efficiency.
A typical ViT requires sequential input and thus reshapes the input image  $X \in \mathbb{R}^{H \times W \times 3}$  as a token sequence  $X_{p} \in \mathbb{R}^{(H*W)/p^2 \times 3*p^2}$ , where  $p$  indicates the patch size. As the patch size is fixed for a given ViT, the number of tokens depends on the resolution of the input image, and thus high-resolution images, which are usually required for dense prediction tasks, inevitably suffer from high computational complexity. One intuitive way toward efficient ViTs is to reduce the number of tokens. However, reducing tokens inevitably accompanies information loss and results in performance degradation. We notice that dense prediction tasks such as detection and segmentation mainly focus on the shape and contour of objects while less caring about the texture inside objects, or irrelevant background. Based on this observation, we propose an adaptive resolution strategy for accelerating dense prediction tasks. Our framework is illustrated in Fig. 1.
3.2 Adaptive Resolution
For an input token sequence  $Z \in \mathbb{R}^{N \times D}$ , the target is to generate  $M$  representative tokens according to  $Z$  where  $M \leq N$ . These representative tokens can be either the original tokens in  $Z$  that correspond to informative regions in the image or the average of several tokens in  $Z$  that correspond to less important regions for decision-making. In this way, different regions are represented by different numbers of tokens, i.e., resolutions based on their importance. The main challenge is to generate a proper assignment of  $Z$  that minimizes information loss.
To generate this assignment, we propose to apply density-based clustering algorithms, specifically DPC [21, 10]. We are motivated by two reasons. For one thing, DPC does not rely on iterative updates like traditional clustering methods such as K-means, making it more suitable for latency-sensitive scenarios. For another thing, DPC searches for cluster centers among the input data, and thus specific tokens can be independently preserved, which enables it to preserve details for informative regions. We find that over  $20\%$  cluster centers are independently preserved when selecting 400 cluster centers from 1600 tokens (Please refer to supplementary for details). In contrast, the cluster centers in K-means are linearly weighted by input, which will lead to information distortion and affect decision-making for fine-grained regions. However, conventional DPC algorithms only consider the relations of data points in the feature space but ignore their intrinsic spatial structure. Since image patches have clear spatial relationships that are critical for dense prediction, directly applying DPC may lead to unreasonable assignments and performance degradation. To address this issue, we propose to incorporate spatial information into clustering.
Spatial-aware DPC. DPC selects cluster centers from input data points based on their product of local density  $\rho$  and distance indicator  $\delta$ , where a large  $\rho * \sigma$  indicates the high potential to be the cluster center. We will explain how we calculate  $\rho, \delta$ , and incorporate the spatial information in the following. Specifically, we calculate the local density of each token by:
$$
\rho_ {i} = \exp \left(- \frac {1}{k} \sum_ {z _ {j} \in \mathcal {K}} \sigma \left(z _ {i}, z _ {j}\right) * s (i, j)\right) \tag {1}
$$
$$
s (i, j) = \left\{ \begin{array}{l l} (1 - \alpha) r a n k (j) / \lambda + \alpha & r a n k (j) \leq \lambda \\ i n f & r a n k (j) \geq \lambda \end{array} \right. \tag {2}
$$
where  $\sigma(\cdot, \cdot)$  denotes the distance metric,  $\mathcal{K} = KNN(z_i)$  and we apply the Euclidean distance here,  $k$  indicates the number of neighbors used to calculate the local density,  $s(\cdot, \cdot)$  is the introduced spatial information where  $\alpha$  is the hyperparameter to control the strength of the spatial constraint, and  $\lambda$  is the number of spatial neighbors.  $s(\cdot, \cdot)$  assigns different weights for  $\sigma(z_i, z_j)$  according to their spatial relation. Specifically, for tokens that are not  $\lambda$  nearest,  $s(i, j)$  assigns the maximum weight for them while assigning the value of  $\alpha$  to 1 for the  $\lambda$  nearest tokens.  $s(\cdot, \cdot)$  encourages spatially adjacent tokens to be merged first. With the introduced spatial information, the local density of each token only depends on the  $\lambda$  spatial neighbors, and each cluster at most merges  $\lambda$  tokens. These properties enable the produced assignments to avoid extreme cases where tokens far away in space are merged or too many tokens are merged together.
The distance indicator  $\delta$  is calculated by:
$$
\delta_ {i} = \left\{ \begin{array}{l l} \min  _ {j: \rho_ {j} > \rho_ {i}} \sigma \left(z _ {i}, z _ {j}\right) * s (i, j), & \exists \rho_ {j} > \rho_ {i} \\ \inf , & o t h e r w i s e \end{array} \right. \tag {3}
$$
With  $\rho_{i}$  and  $\delta_{i}$ , we rank tokens according to  $\rho_{i} * \delta_{i}$  and select top  $M$  tokens as the cluster centers. The remaining tokens are assigned to the closest cluster centers, and tokens belonging to the same cluster centers are merged together as the representative token.
Token re-weight. As the produced representative tokens correspond to different numbers of original tokens (from dozens to only one), there is a distribution gap between the representative tokens and original tokens. For example, tokens corresponding to large objects may be merged into a few representative tokens, which results in inconsistent self-attention results. To minimize this gap, we assign different weights for each representative token during self-attention. In conventional ViTs, the attention of token  $i$  to token  $j$  is given as:
$$
a _ {i j} = \frac {\exp \left(q _ {i} ^ {T} k _ {j} / s\right)}{\sum_ {i} \exp \left(q _ {i} ^ {T} k _ {j} / s\right)} \tag {4}
$$
Figure 2: Visualization to clustering results. The first and fourth columns display the original image, the third and sixth columns show the produced assignments, where tokens with the same assignment are marked with the same color, and the reconstructed low-resolution images are presented in the second and fifth columns.
where  $q_{i}$  is the query of token i,  $k_{j}$  is the key of token j and  $s$  is the scale constant. To minimize the differences brought by token reduction, the same token is expected to have similar attention values in both the original token sequence and the representative token sequence. To this end, we group tokens belonging to the same representative token in self-attention. As only similar tokens are merged, it can be assumed that their value of  $q^{T}k$  is also similar, and thus Eq. (4) can be written as:
$$
a _ {i j} = \frac {\exp \left(q _ {i} ^ {T} k _ {j} / s\right)}{\sum_ {n} \sum_ {i \in m _ {n}} q _ {i} ^ {T} k _ {j}} \approx \frac {m _ {n} \exp \left(q _ {n} ^ {T} k _ {j} / s\right)}{\sum_ {n} m _ {n} \exp \left(q _ {n} ^ {T} k _ {j} / s\right)} \tag {5}
$$
where  $m_{n}$  denote the  $n_{th}$  representative token. We notice that this trick is also used in [1].
3.3 Visualization
To evaluate the effectiveness of the clustering algorithm, we visualize the assignments and the reconstructed low-resolution images. Specifically, we apply our spatial-aware density-based clustering method with 400 cluster centers to the output of the second layer of Segmenter ViT-L, which consists of 1600 tokens. We also reconstruct the entire image using 400 patches corresponding to the cluster centers based on the assignments. The visualizations are shown in Fig. 2. Our results indicate that the reconstructed images have a similar appearance to the original images, suggesting that the produced assignments are reasonable and that  $1/4$  of the tokens are capable of capturing most of the shape and contour information of the original images. Please note that although some regions may appear to have the same color in the visualization of assignments due to the use of 400 different colors, they may actually have different assignments. These findings provide strong evidence for the effectiveness of our proposed method in generating high-quality representative tokens for dense prediction tasks.
Table 1: GFLOPs, FPS and mIoU of AiluRus under different acceleration ratio. The baseline results refer to Segmenter [26] ViT-L. The best results are marked in bold.
Methods	Slight	Mild	Extreme
GFLOPs	FPS	mIoU	GFLOPs	FPS	mIoU	GFLOPs	FPS	mIoU
Results on ADE20k
Baseline	659.0	6.55	51.82	659.0	6.55	51.82	659.0	6.55	51.82
ACT [36]	611.1	6.01	51.69	545.2	6.16	51.24	533.5	6.33	48.03
ToMe [1]	516.2	6.97	51.66	448.7	8.31	50.96	321.3	10.75	47.12
EViT [16]	572.0	7.58	51.52	500.2	8.50	50.37	351.8	12.03	38.89
Expedite [15]	529.8	7.92	51.93	443.8	9.51	51.56	309.4	13.51	47.96
AiluRus	478.8	8.72	52.17	427.8	9.53	51.79	300.8	14.14	50.21
Results on CityScapes
Baseline	995.6	4.20	79.14	995.6	4.20	79.14	995.6	4.20	79.14
ACT [36]	906.3	4.76	79.00	742.7	4.49	78.71	730.4	5.32	75.42
ToMe [1]	760.8	5.20	78.37	651.5	5.50	77.81	448.5	7.84	71.23
EViT [16]	822.7	5.27	79.03	707.2	5.96	78.49	506.2	8.68	68.14
Expedite [15]	840.9	4.82	78.82	691.0	5.89	78.38	529.6	8.02	76.20
AiluRus	710.9	5.88	78.83	669.8	6.65	78.73	461.5	9.36	77.38
Results on Pascal Context
Baseline	338.7	14.7	58.07	338.7	14.7	58.07	338.7	14.7	58.07
ACT [36]	306.7	11.1	58.04	299.0	11.7	57.88	298.3	11.9	56.08
ToMe [1]	269.8	13.9	57.67	236.5	17.0	57.24	172.4	18.2	54.25
EViT [16]	271.7	16.0	57.94	261.0	17.7	56.99	184.4	23.5	48.57
Expedite [15]	251.2	18.2	58.27	201.3	21.6	57.85	161.0	25.0	55.08
AiluRus	241.2	19.8	57.95	224.3	21.9	57.91	157.7	28.4	57.02
Table 2: GFLOPs, FPS and mIoU of AiluRus under different acceleration ratio. The baseline results refer to Segmenter [26] ViT-B. The best results are marked in bold.
Methods	Slight	Mild	Extreme
GFLOPs	FPS	mIoU	GFLOPs	FPS	mIoU	GFLOPs	FPS	mIoU
Results on ADE20k
Baseline	124.7	32.2	48.48	124.7	32.2	48.48	124.7	32.2	48.48
ACT [36]	105.1	26.0	48.39	105.1	25.9	47.55	105.1	26.1	44.01
ToMe [1]	100.2	31.2	47.99	88.6	32.4	46.96	66.4	35.6	40.85
EViT [16]	109.4	34.1	48.44	96.5	37.8	48.05	69.4	46.5	38.27
Expedite [15]	110.2	29.8	46.74	97.4	34.6	46.14	87.1	39.4	45.17
AiluRus	62.3	37.9	48.59	50.3	45.9	48.38	40.2	55.3	47.32
In this section, we begin by comparing AiluRus with recent SOTA methods on semantic segmentation tasks in Section 4.1 which includes a more detailed comparison with the reminiscent method [15]. Subsequently, in Section 4.2, we evaluate the performance of AiluRus on object detection and instance segmentation tasks to assess its generalization ability across various dense prediction tasks. Moving on to Section 4.3, we investigate the applicability of AiluRus in the fine-tuning process to enable acceleration. Additionally, in Section 4.4, we conduct ablation experiments to study the impact of different hyper-parameters of AiluRus on the overall performance. Furthermore, we provide supplementary experiments in the appendix, including the application of AiluRus in expediting classification tasks and text-based video generation tasks. We also delve into the reasons behind the superior performance of AiluRus compared to Expedite in these tasks. For more detailed information, please refer to our appendix.
4.1 Comparison to other methods
We follow the benchmark in [15] to adapt the proposed method to the Segmenter [26] framework built on ViT [9]. Specifically, we load the parameters of the officially released models and integrate
Figure 3: The ablation study on the number of clusters with a fixed cluster location equal to 2. All results are obtained from the officially released checkpoints, and the ablation study on ViT-L and ViT-B are shown in (a) and (b) respectively.
Figure 4: The ablation study on the cluster location with a fixed cluster number equal to 400. All results are obtained from the officially released checkpoints, and the ablation study on ViT-L and ViT-B are shown in (a) and (b) respectively.
AiluRus into the backbone. The adaptive resolution is applied to the output of a certain intermediate layer, and the produced representative tokens are processed by the following layers. With the reduced feature size, the inference speed of models can be significantly improved. We find that the decoder in Segmenter is robust to the reduced feature size, and thus the resolution is recovered after the decoder instead of the backbone.
Comparisons to SOTA methods. We conduct experiments to compare AiluRus with recent SOTA efficient ViT methods across different datasets and architectures. The results are presented in Tab. 1 and Tab. 2. Note that Segmenter [26] only provides ViT-B for the ADE20K dataset. Hence, we only compare different methods using ViT-B on the ADE20K dataset. We follow the benchmark set by [15] and report the GFLOPS, FPS, and mIoU under three different acceleration ratios. The results demonstrate that AiluRus consistently achieves the best trade-off between performance and efficiency compared to other methods across various datasets and acceleration ratios. The advantages of AiluRus are particularly evident in the extreme scenario. For instance, as presented in Tab. 1, with higher FPS, AiluRus outperforms Expedite by  $2.25 \uparrow$  mIoU on ADE20K,  $1.18 \uparrow$  mIoU on CityScapes, and  $1.94 \uparrow$  mIoU on Pascal Context, indicating that AiluRus is much more robust to high acceleration ratios. Such a property enables AiluRus achieve more significant acceleration ratios at an acceptable performance drop.