Not All Tokens Are Equal: Human-centric Visual Analysis via
Token Clustering Transformer

Modeling the interaction between human and environment, and the relationships between body parts is a key point in human-centric vision tasks. With global receptive fields, transformers have achieved great success recently in 2d pose estimation [33, 74, 41, 69], 3d pose estimation [80, 15, 19] and 3d mesh reconstruction [35, 34]. Prior works can be divided into two categories, i.e. refining image features and aggregating part-wise features.

The first kind of method applies transformers to extract better image features. For example, TransPose [69] refines the image features extracted by CNN with extra transformer encoder layers. HRFormer [74] applies transformer blocks in HRNet [50] structure. Our proposed TCFormer also belongs to the first kind.

The second kind of method applies transformers to aggregate features for different parts. For example, TFPose [41] and TokenPose [33] use stacked transformers to generate keypoint-wise features, from which the keypoint location is predicted by regressing coordinates and heatmaps respectively. Mesh Graphormer [35] designs a head with transformers and graph convolution layers. The head aggregates joint-wise and vertex-wise features from image feature maps and predicts the 3D location of human joints and mesh vertices.

Most prior works in both categories generate vision tokens with fixed grid token region, which is sub-optimal for human-centric tasks. In contrast, the token regions of TCFormer are learned automatically rather than using a hand-crafted design. The learned token regions are flexible in location, shape and size according to the semantic meanings.

Recently, there are increasing exploration about dynamic token generation in vision transformers [48, 61, 75, 59]. DynamicViT [48] and PnP-DETR [59] pick up important tokens by predicting token-level scores. For unimportant tokens, DynamicViT simply discards them, while PnP-DETR represents them with sparse feature vectors. DVT [61] builds vision tokens with grids in different resolutions according to the difficulty of classification. PS-ViT [75] builds tokens with a fixed grid size and progressively adjusts the grid centers during processing.

The methods mentioned above are all variants of grid-based token generation. They modify the number, resolution, and centers of grids specifically. In contrast, the token regions of our proposed TCFormer are not restricted by grid structure and are more flexible in three aspects, i.e. location, shape, and size. Firstly, we assign image regions to certain vision tokens based on the semantic similarity instead of the spatial proximity. Secondly, the image region of a token is not restricted to a rectangular shape. Regions with the same semantic meaning, even if they are non-adjacent, can be represented by a single token. Thirdly, vision tokens in the same stage may have different sizes. For example, in Figure 1 (b), the token for the background presents a large region, while the token for the human face only presents a small region. Such property is helpful for retaining important details among all the stages of TCFormer.

Clustering-based feature aggregation methods are wildly explored in research areas of point clouds [47] and graph representations [71]. PointNet++ [47] downsamples point clouds by farthest point sampling and then aggregates features of the k-nearest neighbors. DIFFPOOL [71] predicts soft cluster assignment matrices for an input graph, which are used to hierarchically coarsen the graph.

These methods are specially designed for point cloud and graph data, and cannot be directly applied to the image-based vision transformers. In contrast, we cluster the tokens in the different stages for image-based human-centric vision tasks. To the best of our knowledge, it is the first time that clustering is utilized for vision token generation.

As shown in Figure 2, our proposed Token Clustering Transformer (TCFormer) follows the popular multi-stage architecture. TCFormer consists of 4 hierarchical stages and a Multi-stage Token Aggregation (MTA) head. Each stage contains several stacked transformer blocks. Between two adjacent stages, a Clustering-based Token Merge (CTM) block is inserted to merge tokens and generate tokens for the next stage. MTA head aggregates token features from all stages and outputs the final heatmaps.

Figure 3 shows the transformer block used in TCFormer. Following [60], a spatial reduction layer is used to reduce the computational complexity. The spatial reduction layer first transforms the vision tokens to feature maps and then reduces the feature map resolution with a strided convolutional layer. The pixels in the processed feature maps, which are much fewer than the vision tokens, are fed into the multi-head attention module as keys and values. The multi-head attention module aggregates feature between tokens. Inspired by [32, 62, 74], we utilize a depth-wise convolutional layer to capture local feature and positional information, and remove explicit positional embedding.

As shown in Figure 2, our Clustering-based Token Merge (CTM) block has two processes, i.e. token clustering and feature merging. We apply token clustering to group vision tokens into a certain number of clusters using token features, and then apply feature merging to merge the tokens in the same cluster into a single token.

Token Clustering. In the token clustering process, we utilize a variant of k-nearest neighbor based density peaks clustering algorithm (DPC-KNN) [11].

Given a set of tokens $X$, we compute the local density $\rho$ of each token according to its k-nearest neighbors:

where $\mathrm{KNN}\left({x}_{i}\right)$ denotes the k-nearest neighbors of a token $i$. $x_{i}$ and $x_{j}$ are their corresponding token features.

Then, for each token, we compute the distance indicator as the minimal distance between it and any other token with higher local density. For the token with the highest local density, its indicator is set as the maximal distance between it and any other tokens.

where $\delta_{i}$ is the distance indicator and $\rho_{i}$ is the local density.

We combine the local density and the distance indicator to get the score of each token as $\rho_{i}\times\delta_{i}$. Higher scores mean higher potential to be cluster centers. We determine cluster centers by selecting the tokens with the highest scores, and then assign other tokens to the nearest cluster center according to the feature distances.

Feature Merging. For token features merging, an intuitive method is to average the token features in the cluster directly. However, even for the tokens with similar semantic meanings, the importance is not totally the same.

Inspired by [48], we introduce an importance score $P$ to explicitly represent the importance of each token, which is estimated from the token features. The token features are averaged with the guidance of the token importance:

where $C_{i}$ means the set of the $i$-th cluster, $x_{j}$ and $p_{j}$ are the original token features and the corresponding importance score respectively, and $y_{i}$ is the features of the merged token. The token region of the merged token is the union of the original token regions.

As shown in Figure 2, the merged tokens are fed into a transformer block as queries $Q$, and the original tokens are used as keys $K$ and values $V$. To ensure that important tokens contribute more to the output, the importance score $P$ is added to the attention weight as follows:

where $d_{k}$ is the channel number of the queries. We omit the multi-head setting and the spatial reduction layer here for clarity. Introducing the token importance score equips our CTM block with the capability to focus on the critical image features when merging vision tokens.

We adopt an efficient GPU implementation of CTM block. The clustering and feature merging parts only cost $9.4\%$ of the forward time of TCFormer.

Prior works [50, 74] prove the benefits of feature aggregation in multiple stages for human-centric vision tasks. In order to aggregate features, we propose a transformer-based Multi-stage Token Aggregation (MTA) head, which is able to maintain details in all the stages.

Figure 4 (a) shows the token upsampling process. During the token merging process (Section 3.3), each token is assigned to a cluster and each cluster is represented by a single merged token. We record the relationship between the original tokens and merged tokens. In the token upsampling process, we use the recorded information to copy the merged token features to the corresponding upsampled tokens. As shown in Figure 4 (b), after the token upsampling, MTA head adds the token features in the previous stage to the upsampled vision tokens. The vision tokens are then processed by a transformer block. Such processing is executed progressively until all vision tokens are aggregated. The final tokens, whose token region is a single pixel in the high-resolution feature map, can be easily reshaped to feature maps for further processing.

Why not CNN? Most prior works [60, 74, 38] transform vision tokens to feature maps first and aggregate multi-stage features with convolutional layers. However, as shown in Figure 5, vision tokens for the regions with important details, such as the human face, are in small size. Transforming vision tokens to low-resolution feature maps involves feature averaging of these tokens, which leads to the loss of details. Such loss can be avoided by transforming vision tokens in all stages to high-resolution feature maps, but it will lead to unacceptable complexity and memory cost. Our MTA head, which aggregates features in the form of tokens and transforms the final tokens to high-resolution feature maps at the end, preserves image details in all stages with relatively low complexity.

Transform between Vision Tokens and Feature Maps. The convolutional process is non-trivial for irregular tokens. We need to transform tokens to feature maps before this process and perform the inverse transform after that. We regard every pixel in feature maps as a representation of a rectangular grid region. When transforming vision tokens to feature maps, we first find the related vision tokens for every pixel, i.e. the vision tokens whose token region overlaps with the grid region of the pixel. Then we average the token features according to the overlapped areas to get the image feature of every pixel. When transforming feature maps to vision tokens, we directly set the token features as the average image features in the token region.

Cluster Numbers. In the CTM block, the cluster number can be any value smaller than the input token number. In this work, we set the cluster number to a quarter of the input token number to keep the token numbers in every stage the same as those of the prior works [60, 38].

Whole-body pose estimation targets at localizing fine-grained keypoints on the entire human body including the face, the hands, and the feet, which requires the ability to capture detailed information.

Settings. We conduct experiments on COCO-WholeBody V1.0 dataset [23]. COCO-WholeBody dataset is built upon the popular COCO dataset [36] with additional whole-body pose annotations. The full pose contains 133 keypoints, including 17 for the body, 6 for the feet, 68 for the face, and 42 for the hands. Following [36, 23], we use OKS-based Average Precision (AP) and Average Recall (AR) for evaluation.

We follow most of the default training and evaluation settings of mmpose [9] and replace Adam [28] with AdamW [40] with momentum of $0.9$ and weight decay of $1\times 10^{-2}$.

Results. Table 1 shows the experimental results comparing TCFormer with the state-of-the-art models on COCO-WholeBody V1.0 dataset [23]. The results show that the whole-body pose estimation accuracy of TCFormer (57.2% AP and 67.8% AR) is higher than those of the state-of-the-art top-down methods, e.g. HRNet [50], by a large margin.

The size of the hand in the input image is relatively small, which makes the estimation of hand keypoint extremely difficult and heavily reliant on the capability of the model in capturing details. As shown in Table 1, most models achieve much lower performance on the hand than the other parts. Our TCFormer achieves a large gain in hand keypoint estimation, i.e. $6.2\%$ AP higher than HRNet-w32 [50] and $5.3\%$ AP higher than SBL-Res152 [66], which demonstrates the excellent capability of TCFormer in capturing critical image details with small sizes.

Prior works of 3D human mesh reconstruction can be divided into model-based [27, 30, 29, 29, 77, 26] and model-free methods [34, 35, 7, 43, 78]. We build a model-based method by combining TCFormer and HMR head [27].

Dataset. We evaluate our model on 3DPW dataset [56] and Human3.6M dataset [20]. 3DPW is composed of more than 51K frames with accurate 3D mesh in the wild, and Human3.6M is a large-scale indoor dataset for 3D human pose estimation. Mean Per Joint Position Error (MPJPE) and the error after Procrustes alignment (PA-MPJPE) are reported. Following the setting of [27], we train our model with a mixture of datasets, including Human3.6M [20], MPI-INF-3DHP [42], LSP [24], LSP-Extended [25], MPII [1] and COCO [36]. For fair comparisons with the recent methods [26, 13], we also use the training set of 3DPW and the pseudo-label provided by  [26].

Settings. We crop the image with the ground truth bounding box and resize it to the resolution of $224\times 224$. During training, we augment the data with random scaling, random rotation, random flipping, and random color jittering. The model is trained with 8 GPUs with a batch size of 32 in each GPU for 80 epochs. We use Adam optimizer [28] with the learning rate of $2.5\times 10^{-4}$ and do not use learning rate decay. The settings are the same as that in [26, 13].

Results. We compare TCFormer with the state-of-the-art SMPL-based methods and show the results in Table 2. TCFormer outperforms most of the prior works [27, 29, 77, 26] with similar structures and model complexity. TCFormer even obtains competitive results compared with the DSR method that uses extra dense supervision [13]. The results of human mesh estimation further validate the effectiveness of TCFormer in capturing important image features. It works well not only on the dense prediction task, but also on the global feature based regression task.

Settings. We perform experiments on WFLW [64] dataset, which consists of $7,500$ training and $2,500$ testing images with $98$ landmarks. The evaluation is conducted on the test set and several subsets: large pose, expression, illumination, make-up, occlusion, and blur. We use the normalized mean error (NME) as the evaluation metric and the inter-ocular distance for normalization. We apply the same training and evaluation settings as that of [58]. For fair comparisons, we use a lightweight version of TCFormer (TCFormer-Light) with similar model complexity as [58].

Results. As shown in Table 3, TCFormer achieves superior performance ($4.28\%$ NME), compared to the other state-of-the-art methods on the full test set and all the subsets. TCFormer even has lower error than the methods with extra information, such as PDB [14] which uses strong data augmentation and LAB [64] which uses extra boundary information. The performance on face alignment validates the versatility of TCFormer beyond human body estimation.

In order to evaluate the versatility of TCFormer on general vision tasks, we also extend it to image classification.

Settings. We perform experiments on the ImageNet-1K dataset [49]. We apply the setting totally the same as [60]. TCFormer without the MTA head is trained from scratch for 300 epochs with batch size 128 and evaluated on the validation set with a center crop of $224\times 224$ patch. For more details, please refer to the supplementary.

Results. Although the target of our model is not image classification, experimental results on ImageNet-1K show that TCFormer achieves competitive performance (82.4% Top-1 Acc.) compared with the state-of-the-art architectures, which indicates that our CTM block also works well in extracting general image features.

We conduct ablative analysis on the task of whole-body pose estimation as shown in Table 1.

Effect of CTM. To validate the effect of CTM (Section 3.3), we build a baseline transformer network by replacing CTM with a strided convolutional layer, and find that the performance drops significantly ($-3.7\%$ AP and $-3.9\%$ AR). The performance drop for the parts relying on detailed information, such as the foot and the hand, is more significant than the others. The estimation AP for the foot and the hand decrease by $13.6\%$ and $5.6\%$ respectively, while the AP for the human body only drops $2.4\%$. It demonstrates that the performance drop is mainly caused by the loss of image details and validates the effectiveness of our CTM block in capturing image details in small size.

Effect of MTA Head. To validate the effect of MTA Head (Section 3.4), we replace it with a deconvolutional head [66] and notice a performance drop of $-1.9\%$ AP and $-1.6\%$ AR. Especially, the performance drop on foot ($-4.0\%$ AP) and hand ($-3.6\%$ AP) is more obvious than that of the body ($-1.2\%$ AP). The results show that MTA Head is beneficial to preserving image details of body parts.

Figure 6 shows some qualitative results for human whole-body pose estimation, 3D mesh reconstruction, and face alignment. Figure 7 shows some examples of the vision token distribution in the above tasks. As shown in Figure 7, the vision tokens focus on the foreground regions and represent background areas with only very few tokens, even when the background is complex. The vision tokens with fine spatial size are used for the area with important details, for example, the face and hand regions in whole-body estimation and 3D mesh reconstruction. For face alignment, TCFormer allocates fine tokens to the face edge regions. The dynamic vision token allows TCFormer to capture image details more efficiently and achieve better performance.