PAFormer : Part Aware Transformer for Person Re-identification

Prior to the emergence of Vision Transformer, almost all ReID approaches were based on CNN. BoT [20] has proposed to use a CNN backbone to extract global features and to apply cross-entropy loss and triplet loss as supervision. Various subsequent researches built upon this ReID loss have also been proposed. However, because of the challenges faced by those relying soley on global features, such as occlusion, part-based approaches have emerged.

In PCB [1], stripe-based methods are proposed to measure the distances between samples within the same region. Zhao et al.[5] introduces CNN-based attention networks, allowing the model to autonomously find important regions in the image for ReID. ABDNet [4] builds upon this attention network concept and applies orthogonality regularization to each channel, enabling the extraction of information from various body parts. SPReID [21] employs a human semantic parsing model to aggregate features only from the regions corresponding to each body part. AACN [22] and BPBReID [16] directly supervise formation of attention maps by utilizing pose estimation models. These approaches suggeset explicit answers on how attention maps should be formed.

After the introduction of TransReID [9], which utilizes a jigsaw patch module to capture local features, many transformer-based ReID methods have emerged. LA Transformer [3] implements stripe-based methods on the ViT architecture. PAT [11] adopts an encoder-decoder structure in Transformer and uses cosine dissimilarity loss to encourage each part token to gather information from different body parts. AAFormer [12] enhances distinctiveness of each token by physically limiting the number of patch tokens that can interact with each part token using Optimal Transport algorithm [23]. NFormer [8] effectively removes noise by performing self-attention only between reciprocal neighbor patch tokens. HAT [10] utilizes multi-scale information from ResNet [24] backbone to aggregate hierarchical information. Rest-ReID [25] combines attention-guided graph convolutional networks with Transformer to find correlations between body parts. PFD [13] leverages an embedding of pose heatmaps as additional information to Transformer-based models.

Our model employs a new learnable prototype named ‘pose token’. The key advantage of using a pose token is to learn to extract features of the same body part across diverse samples. This allows us to shift the main focus from aligning IDs in the training set to estimating correlations between patch tokens and body parts, which significantly aids in facilitating part-to-part comparison. Furthermore, the use of a clear guidance about human body part ensures that each pose token is trained to specifically represent a distinct body part. This attribute also enables the pose token to function as a valuable indicator for detecting occlusions in specific body parts. Detailed explanations regarding the training and application of the pose token will be provided in following subsections.

We expect the pose token to demonstrate a high degree of similarity with their associated patch tokens only. We aim to leverage this similarity in the forthcoming cross-attention and aggregation layers, as detailed in subsequent subsections. To achieve this, a prerequisite step is required, which enhances the similarity among patch tokens that share the same anatomical semantics. To accomplish this, we preprocess the patch tokens using self-attention (SA), a mechanism that has already been demonstrated in various studies for effectively performing this task. In Figure 5, it is evident that patch tokens that share anatomical semantics converge in the vanilla ViT-based ReID model.

Additionally, by utilizing the self-attention mechanism, we can harness the robust global feature extraction capabilities of ViT. Similar to the vanilla ViT, we employ a class token to extract global features and apply ReID loss on it:

where $L_{ID}$ and $L_{tri}$ are cross-entropy loss and triplet loss, respectively. Also, we incorporate SA layer among pose tokens, with the objective of empowering PAFormer to learn associations among prototypes, akin to DETR [27] or PAT [11].

In order to imbue PAFormer with part awareness, during training, we create a ground truth heatmaps generated by post-processing initial keypoint heatmaps obtained from a pose estimation model, PifPaf [28]. The creation process of this ground truth heatmaps is detailed in section 5.3. The ground truth heatmaps directly guide how the attention maps between the pose token and patch token in the cross-attention (CA) layer should be formed. This guidance ensures that the pose token accurately perform localization of body parts. The loss function required for this process is as follows:

where $P$ is the number of pose tokens, $MSE$ denotes mean square error, and $GT$ represents the ground truth heatmap. $\overline{Attn}_{pose}$ is the averaged value of all $Attn_{pose}$ (12 layers$\times$12 heads) present in CA layers of PAFormer. In contrast to existing methods, PAFormer allows prototypes to concentrate on identifying patch tokens corresponding to specific body parts.

While some methods which utilize a pose estimation model have already been proposed, they often use a separate body part localization module within their models or continue to rely on an external pose estimation model during inference. However, PAFormer takes a different approach by learning how attention maps should be formed during the cross-attention process. As a result, it does not require pose heatmaps during inference.

The underlying principle of learning part awareness through a simple mean square error is as follows:

where $W_{Q}$ and $W_{K}$ refer to fully connected layers performing query and key transforms, respectively. $X_{pose}$ and $X_{patch}$ represents the pose tokens and the patch tokens, respectively. Expanding upon the expression, the product of $W_{Q}$ and $W_{K}$ can be expressed as a single weight matrix, denoted as $W$. The application of MSE loss offers supervision to PAFormer, allowing $W$ to accentuate channels associated with specific body parts. Consequently, the dot product of paired patch tokens and pose tokens increases, causing the attention score to selectively rise only for patches belonging to the specific body part, as intended.

Once the association of each pose token with specific patch tokens is estimated in the cross-attention layer, we leverage these probabilities to aggregate the refined values of the patch tokens, thereby extracting partial features. This stage is denoted as an aggregation layer. Aggregated partial features $z_{p}$ can be expressed as:

where $\overline{Attn}_{pose}^{p}$ is $p$-th pose token’s averaged attention map across all heads from a CA layer. In typical ViT-based methods, a value transformation by a fully-connected network are applied to patch tokens when creating partial features. However, in the aggregation layer, no separate transformation is applied because the information from patch tokens should be directly accepted based on the association probabilities. Similar to the conventional Transformer architecture, as each layer is traversed, the partial features $z_{p}$ are also updated by layer normalization, feed-forward networks, and residual connections.

Relying solely on the pose heatmap for model guidance can be vulnerable if the heatmap is inaccurately generated. To mitigate this vulnerability, we supplement the guidance by incorporating the ReID loss for partial features:

where $z_{i}^{L}$ denotes the aggregated partial features from $i$-th pose token at the last PAFormer block. We minimize attention to regions that do not correspond to body parts by $L_{pose}$. Subsequently, with the assistance of $L_{ReID}^{part}$, we perform an additional filtering, addressing areas that may have been overlooked by $L_{pose}$. However, to prevent potential issues stemming from the dominant influence of the ReID loss, we assign a substantially higher weight to $L_{pose}$ than $L_{ReID}^{part}$. This strategic weighting ensures that the ReID loss functions merely as a complementary element.

While many other methods that consider occlusion rely on the confidence score of the output of a module responsible for performing body part localization to address this issue, PAFormer employs a learning-based prediction network for a more sophisticated treatment of the problem, as in Figure 4. This is attributed to the design of the model in preceding stages, which ensures that the pose token serves as a representative for various body parts. For the ground truth visibility score, we assume that a body part is invisible if the maximum value in the pose heatmap does not exceed $\theta_{p}$. As the pose estimation model yields varying confidence scores for different body parts, $\theta_{p}$ is set differently for each body part. To train the visibility predictor, we employ $L_{vis}$ which uses binary cross-entropy loss between the obtained output and the ground truth:

where $v$ denotes a predicted visibility score, and $v_{GT}$ denotes ground truth visibility score from the pose estimation model.

It is essential to refrain from calculating ID loss or triplet loss for occluded body parts. To prevent this, we adapt the ReID loss by incorporating the ground truth visibility score as a form of teacher forcing to guarantee the exclusion of occluded parts from the learning process. Particularly, during hard triplet mining [29], to deter the selection of occluded regions, a value of 0 is multiplied to the distance for positive cases, while an extremely large value is multiplied for negative cases.

The overall loss $L$ is calculated as follow:

$\lambda$ is a weight for $L_{pose}$ and is heuristically set to 10. All modules in PAFormer are learned together.

During inference phase, PAFormer calculate distances as follows:

where $d^{i,j}$ is a distance between $i$-th and $j$-th samples, and $v$ denotes the predicted visibility score. Among the subscripts in $d$, ‘CLS’ refers to the class token, and $p$ indicates the specific pose token’s identifier.

Generally, time complexity of Transformer is known to be $O(N^{2}d)$, where $N$ and $d$ are the number of tokens and channels, respectively. Since PAFormer performs self-attention similarly, it maintains the same $O(N^{2}d)$ time complexity. Additionally, the pose token self-attention process contributes $O(P^{2}d)$, and the cross-attention process adds $O(PNd)$, resulting in an overall time complexity of $O((N^{2}+P^{2}+PN)d)$. Given that $N$ is usually much larger than $P$ ($N$: three-digit, $P$: one-digit), the additional burden introduced by PAFormer is considered to be negligible.

Three widely used ReID datasets are chosen for our experiments and their statistics are shown in Table 1. Images in every dataset are resized to $256\times 128$ or $384\times 128$. Then, the training images are augmented with random horizontal flipping, cropping, padding, grayscale [30] and erasing [31]. Among these methods, we apply all except random erasing and grayscale augmentation to the pose mask. Additionally, for regions affected by random erasing in the original data, we set the mask value to 0.

We adopt ImageNet pre-trained ViT-B/16 as the backbone. The model is trained with SGD optimizer with a momentum of 0.9 and a weight decay of 1e-4. The learning rate is set to 0.008 with a cosine decay and the batch size is determined to be 64. We also use sliding-window ($S=12$) and side information embedding [9]. Both training and testing processes are implemented with two NVIDIA GeForce RTX 4080 GPUs. We conducted training for 320 epochs and evaluated the performance every 20 epochs. Among these evaluations, we choose the one with the highest performances.

We use PifPaf [28] to obtain pose heatmaps from the input image. However, instead of using the heatmaps predicted by PifPaf directly, we employ several modification steps to adapt them as ground truth heatmaps.

Firstly, we apply random erasing augmentation [31] to pose heatmap as well. The probabilities corresponding to the erased areas in original image are all set to 0.

Secondly, we combine heatmaps of small fragments such as keypoints and joints to form masks for various body parts. During this process, we regard heatmaps from PifPaf with a maximum confidence score below $\theta_{p}$ as noise and ignore them. The value of $\theta_{p}$ is set differently for each body part and for each dataset, as in Table. 2.

Thirdly, we normalize the combined heatmaps. Given our objective is to avoid focusing solely on discriminative parts, we assume that the influence of all patch tokens belonging to a specific body part is uniform. Consequently, value in the ground truth mask is either 0 or ${1\over K}$, where $K$ represents the count of non-zero values in the ground truth mask. Following this transformation, we normalize the mask by dividing it by the sum of its constituent values, ensuring a total sum of 1.

The ground truth visibility score is also determined by $\theta_{p}$. Pseudo labels are assigned as ‘visible’ only for body parts with a maximum confidence score greater than or equal to $\theta_{p}$.

Table 3 shows the performances of our PAFormer and other previous ReID methods. We adopt mean Average Precision (mAP) and Rank-1 score (R1) for evaluations. PAFormer demonstrates its strong capabilities by achieving state-of-the-art performance on various datasets, outperforming competitors and securing at least the second-place position. The performances using the holistic datasets and the occluded dataset have been achieved with $P=5$ and $P=6$, respectively.

PAFormer exhibits strong performance even for 384$\times$128 images, achieving 90.9 in mAP and 96.1 in Rank-1 score. It records the highest mAP and the second-highest Rank-1 score among the comparison groups, demonstrating excellent performance.

When applied to images resized to $384\times 128$, PAFormer achieves a performance of 84.2 in mAP and 92.5 in Rank-1 score, which are both the highest scores among the comparisons. The high performance on both Market-1501 and DukeMTMC-reID datasets demonstrates the ReID ability of PAFormer on holistic datasets.

PAFormer’s training strategy hinges on the output pose heatmap from a pose estimation model, revealing vulnerabilities in scenarios with two or more individuals, as depicted in Table 8. This challenge is not so much a flaw in our methodology but rather stems from inaccuracies in generating ground truth masks. By using models adept at multi-person pose estimation and leveraging precisely annotated data, PAFormer has the potential to achieve enhanced performance in such situations.