PersonMAE: Person Re-Identification Pre-Training with Masked AutoEncoders

Person re-identification aims to match a specific person across different camera views. It is a kind of retrieval problem, whose challenges come from unconstrained recording conditions, occlusion, cropping errors, etc. Generally, ReID contains supervised and unsupervised task settings, which are introduced as follows.

Different from holistic ReID, the occluded counterpart is a harder task and its query set is constructed by occluded pedestrian images [35, 36, 37, 38, 39]. Zhuo et al. [36] first study this problem via the attention mechanism and occlusion simulation. The following works further disentangle discriminative cues from the occlusion with the assistance of pose landmarks [37, 40] or human parsing [38, 41].

Self-supervised pre-training aims to learn more generic feature representations from a large amount of unlabeled data, which benefits the downstream tasks. Among them, contrastive learning has been widely studied for pre-training [50, 51, 52, 53, 54], which aims to pull the representation of similar instances closer, while pushing away negative instances. In order to obtain informative negative instances, some works resort to the techniques like memory banks [50] and large batch sizes [51]. There also exist some other works [55, 56] achieving great performance without the requirement of negative instance. DINO [52] achieves overwhelming performance by emphasizing the distillation between global and local views. Recently, motivated by the success of BERT [57] in NLP, some works [12, 58, 59, 60, 61, 62, 63, 64] works with masked autoencoders (MAE) and achieve promising performance. Typically, MAE adopts the ViT backbone and aims to predict the masked patches from the remaining visible ones.

In person ReID, there exist works leveraging the success of self-supervised pre-training [9, 65, 66, 11]. Fu et al. [9] proposes a large-scale LUPerson dataset and pioneers the exploration of contrastive-based pre-training based on this dataset. Luo et al. [10] investigate contrastive-based pre-training based on ViT and improve pre-training via reducing the domain gap between pre-training and fine-tuning data. PASS [11] further proposes part-aware pre-training based on DINO. These methods demonstrate their effectiveness on the downstream ReID tasks and are all based on contrastive learning. Different from them, we aim to make the first attempt to leverage the advance of masked autoencoders into ReID pre-training.

As we briefly introduced in Section 1, a good person-ReID model should maintain three important properties: multi-level awareness, occlusion robustness, and cross-region invariance. To maintain these properties, we propose our PersonMAE framework. As shown in Figure 2, with a given image, we first generate the input image RegionA and prediction target RegionB with a cross-region generation module. Then we mask a large portion of RegionA and use the unmasked part as the online encoder input. The decoders use both the encoder output features and mask tokens to predict the whole RegionB. As is well known, both low-level clues (such as the color of clothes and hair) and high-level semantics (such as global identity information) matter in ReID. Therefore, we utilize a pixel decoder to predict the raw pixels of RegionB for low-level reasoning, and a feature decoder to predict the feature of RegionB for high-level semantic learning.

Previous general MAE-based pre-training methods always mask part of the input and then predict itself. It is an inner prediction manner that aims to model and understand the semantic content within the input image. However, this prediction manner may hardly handle the disturbance from unreliable cropping of pedestrian images. To help the model learn robust representation regardless of such misalignment, we propose our cross-region generation module.

In detail, with a given pedestrian image, we first perform random resizing as follows,

where $Ip(\cdot,\cdot)$ denotes the interpolation function, whose inputs are the original image and desired output image size. $m$ represents the maximum shift size. With the resized image $\mathbf{\tilde{X}}$, we extract two regions at the same scale with the size of $[H,W]$ as follows,

where $s_{h}^{a}=s_{w}^{a}=0$, $s_{h}^{b}\in[0,p]$ and $\ s_{w}^{b}\in[0,p/2]$ and these two tuples represent the coordinate of the upper-left corner and cropped image size, respectively. Then these two regions are separately split into non-overlapping patches, whose total number is $N=H\times W/T^{2}$ and $T$ represents the patch size. With this step, a certain region is transformed into the visual token sequence, e.g. $\mathbf{\tilde{X}_{a}}={\{\bm{x}^{i}_{a}\}}^{N}_{i=1},\bm{x}_{a}^{i}\in\mathbf{R}^{CT^{2}}$. $\bm{x}_{a}^{i}$ is obtained via squeezing the visual token and $C$ represents the image channel dimension. The same can be obtained for $\mathbf{\tilde{X}_{b}}$.

It is an outer prediction problem that the model should not only understand the content of the input RegionA, but also the invariance between RegionA and RegionB. It is much more challenging than previous inner prediction and better improves the cropping-error robustness of the model.

Occlusion is one of the most challenging problems in the ReID task, since the pedestrian may be occluded by some obstacles, e.g. trees, cars, walls, and other passengers. So it requires the model not only focus on the most salient feature of the input but all the reasonable features within it. To mimic various occlusions during training, we utilize the block-wise masking strategy introduced in BEiT [58] for the input RegionA and only use the unmasked part as the encoder input. Formally:

where $\mathcal{M}$ is the sampling function without replacement and the number of the masked positions is $r\times N$. $E_{o}(\cdot)$ denotes the online encoders with parameters $\theta_{o}$.

For previous inner prediction methods, the input and target images share the region, so they could use the same position information directly. On the contrary, our outer prediction design leads to a different region for input and target, the decoder needs to predict the target RegionB based on the feature of RegionA and their cross-region relationship.

In practice, we use absolute position embedding to model such a relationship. We build a coordinate system and set the upper-left corner of the RegionA as the origin. Denote $x$ and $y$ are the horizontal and vertical index of the predicted region. Formally:

where $x\in[0,H/T-1]$ and $y\in[0,W/T-1]$. We follow common practices to convert $r(x,y)$ into the fixed 2D position embedding $\mathbf{P}_{r}$ based on $sin(\cdot)$ and $cos(\cdot)$ operators [12].

Then the formulated shifting relation and the output $\mathbf{\tilde{F}_{a}}$ of the online encoder are jointly fed into the following two decoders for RegionB prediction, i.e., pixel regression decoder and feature prediction decoder.

As illustrated above, low-level clues are important for the ReID task. Therefore, a pixel decoder is adopted to conduct prediction at the low-level aspect via recovering all RegionB pixels. Specifically, we organize $N$ learnable tokens as the mask tokens $\mathbf{M}_{p}$. To inform the framework where the corresponding output of each mask token should predict, we further add the position embedding $\mathbf{P}_{r}$ on $\mathbf{M}_{p}$, denoted as $\mathbf{\tilde{M}}_{p}$. Then the PE-informed $\mathbf{\tilde{M}}_{p}$ is concatenated with the normalized $\mathbf{\tilde{F}_{a}}$ as the input of the pixel decoder. The pixel decoder contains two ViT blocks, followed by an MLP layer to regress the missing pixel value. Denote the output of the pixel decoder as $\mathbf{Z}_{b}=\{\bm{z}^{i}_{b}|i\in[0,N-1]\}$. During pre-training, the low-level region-consistency constraint is calculated as follows,

where $m^{i}_{b}$ and $s^{i}_{b}$ are the mean and standard deviation of the $i$-th patch and we utilize the normalized patch as the target. This constraint serves as a strong regularization term during optimization and makes the framework model human statistics via reasoning about the low-level textures.

Meanwhile, high-level semantics are the core for person re-identification, so we use the feature decoder to conduct feature prediction. Similar to the pixel regression branch, we adopt a new set of $N$ learnable tokens as the mask token $\mathbf{M}_{f}$ and associate it with the same position embedding $\mathbf{P}_{r}$, denoted as $\mathbf{\tilde{M}}_{f}$. It is concatenated with the normalized $\mathbf{\tilde{F}_{a}}$ as the input of the feature decoder. The feature decoder contains two ViT blocks, followed by a fully-connected layer to project the features on the desired semantic space. We argue that dense local supervision forces the model to capture more fine-grained details, so the prediction target is all the local features of RegionB encoded by an EMA model. Formally,

where $E_{m}(\cdot)$ is the momentum encoder, $\theta_{o}$ and $\theta_{m}$ represent the parameters of the online encoder and momentum encoder, respectively. Compared with the online encoder $E_{o}$, the parameters of $E_{m}$ are updated much slower and provide a more stable prediction target, which is beneficial for modeling the pedestrian statistics in the semantic space.

With the prediction target $\mathbf{\tilde{F}_{b}}=\{\bm{f}^{i}_{b}|i\in[0,N-1]\}$, we denote the output of feature prediction decoder as $\mathbf{Y}_{b}=\{\bm{y}^{i}_{b}|i\in[0,N-1]\}$ and the feature-level region consistency constraint is formulated as follows,

where we utilize smooth L1 loss for feature matching. $\texttt{SG}(\cdot)$ represents the stop gradient function and $\beta$ is the smoothing factor. Overall, during pre-training, the objective function is formulated as follows,

where $\lambda$ is the weighting factor. $\mathcal{L}_{p}$ and $\mathcal{L}_{f}$ are the low-level and semantic constraint, respectively.

Pre-training setting. Pre-training is conducted on 8$\times$V100 GPUs for 100 epochs. Adam is adopted for optimization. The learning rate is set to 1.2e-3, with a warmup of 20 epochs, and single-cycle cosine learning rate decay. The weight decay is set as 0.05. We only utilize horizontal flipping and random cropping on the raw data to generate the global pedestrian image. The input resolution is 256$\times$128. $\gamma$ represents momentum coefficient. It is initially set as 0.999, linearly changed to 0.9999 in the first 20 epochs, and kept as 0.9999 for the remaining epochs. $\beta$ represents the smoothing factor and is set to 2.

Supervised ReID setting. During fine-tuning, we feed the unmasked image into the pre-trained online encoder for further optimization. We follow the common practice in [9] and utilize the MGN [3] head for downstream ReID tasks. During supervised fine-tuning, we utilize the layer-wise lr decay strategy following [12, 58]. The batch size is set as 256, with 16 pedestrians in each batch. Adam optimizer is adopted with the learning rate and weight decay set as 5e-3 and 0.05, respectively. The layer decay is set as 0.65. The training lasts 60 epochs, and we utilize the cosine learning rate decay scheduler. We keep the same setting for supervised holistic and occluded ReID.

Unsupervised ReID setting. For unsupervised ReID tasks, we build upon C-Contrast [46, 10] with the ViT backbone and follow most of its settings. UDA and USL ReID only differ in model initialization. UDA utilizes the pre-trained model from the source dataset, while USL directly utilizes our LUPerson pre-trained parameters. Then the training is conducted in an unsupervised manner on the target dataset. The visual features corresponding to images in the training set are first extracted, followed by clustering for pseudo labels and corresponding clustering centers. Then the framework leverages the contrastive objective between output features and clustering centers during optimization. More details can be found in [46]. SGD optimizer is adopted. The initial learning rate is set as 0.007 and is decayed with 0.1 every 20 epochs. The batch size is set as 256, with 32 pedestrians in each batch.

We compare our method with previous state-of-the-art methods on four different downstream ReID tasks, i.e., supervised (holistic and occluded ReID), and unsupervised (UDA and USL ReID). Note that we do not adopt any post-processing method like Re-Rank [81] in our method.

Methods in the bottom part conduct self-supervised pre-training on LUPerson to initialize the backbone. Regardless of the adopted backbones, their pre-training methods are dominated by the variants of contrastive learning. UPReID [65] and PASS [11] design subtle techniques to further mine the information between global image and local patches. Our method achieves new state-of-the-art performance, i.e., 93.6 and 79.8 mAP on Martket1501 and MSMT17 datasets, respectively. It is worth mentioning that a more significant performance gain is achieved on more challenging MSMT17, i.e., 6.1% and 8% mAP gain on ViT-S and ViT-B over previous SOTA, respectively.

Supervised occluded ReID is evaluated on OccDuke benchmark as shown in Table II. It is a harder setting than the holistic counterpart since the given query pedestrian contains diverse occlusion conditions. Previous methods for this setting mainly leverage external cues or design novel architectures for feature disentanglement. ISP [71] iteratively learns feature maps at the pixel level to provide identity-guided parsing cues. FED [39] builds on the ViT-B backbone and designs two modules to eliminate the disturbance from occlusion. Compared with previous works, our method demonstrates superiority without utilizing external cues, i.e., achieving 69.5% mAP on OccDuke. It can be attributed to our explicit occlusion-aware design in the pre-training strategy, which makes our framework model pedestrian identity statistics well even under severe occlusion.

In this section, we perform ablation studies on supervised ReID to demonstrate the effectiveness of the key components. Unless stated, we adopt ViT-B as our backbone and conduct pre-training with 50 epochs.

We perform qualitative visualization of our prediction results on unseen person images in Figure 3. In each group, we show RegionA, distorted RegionA, ground-truth RegionB and our predicted RegionB in sequence. Our framework can infer the whole cloth texture or fill the whole personal belonging well (e.g. bag), even with a glance at the partial region. This strong hallucination capability may be largely attributed to the well-model pedestrian statistics during pre-training.

Furthermore, we demonstrate some retrieved pedestrian images of previous SOTA (PASS) and our method in Figure 5. We list top-5 ranking results for each query. In the upper-left part, the query pedestrian wears the white schoolbag and only shows the bag straps in the query image. However, PASS ignores this fine-grained cue and focuses on the red sweater and black coat, which leads to false retrieval. In contrast, our framework is able to identify this detail and retrieves the pedestrian correctly. For the bottom-right part, the query person is heavily occluded by a car, which is a highly challenging scenario. Our framework does not rely too much on the blue coat and retrieves correctly by capturing the cues from the bag color and hat type. These results qualitatively suggest that our method is able to satisfy the aforementioned three main properties for accurate ReID.