Sparse Local Patch Transformer for Robust Face Alignment and Landmarks Inherent Relation Learning

Coordinate regression methods [42, 41, 37, 12] regress the coordinates of landmarks from feature map directly via FC layers. To further improve the robustness, diverse cascaded networks [30, 17] and recurrent networks [38] are proposed to achieve face alignment with multi stages. Despite coordinate regression methods have an innate potential to learn the inherent relation, it commonly requires a huge number of samples for training. To address the problem, Qian et al. [26] and Dong et al. [9] expand the number of training samples by style transfer; Browatzki et al. [3] and Dong et al. [10] leverage the unlabeled dataset to train the model. In recent years, state-of-the-art works employ the structure information of face as the prior knowledge for better performance. Lin et al. [24] and Li et al. [22] model the interaction between landmarks by a graph convolutional network (GCN). However, the adjacency matrix of GCN is fixed during inference and cannot adjust case by case. Learning an adaptive inherent relation is crucial for robust face alignment. Unfortunately, there is no work yet on this topic, and we propose a method to fill this gap.

Heatmap regression methods [25, 29, 34, 7] output an intermediate heatmap for each landmark and consider the pixel with highest intensity as the optimal output. Therefore, it leads to quantization errors since the heatmap is commonly much smaller than the input image. To eliminate the error, Kumar et al. [18] estimate the uncertainty of predicted landmark locations; Lan et al [19] adopt an additional decimal heatmap for subpixel estimation; Huang et al. [15] further regress the coordinate from an anisotropic attention mask generated from heatmaps. Moreover, heatmap regression methods also ignore the relation between landmarks. To construct the relation between neighboring points, Wu et al. [36] and Wang et al. [35] take advantage of facial boundaries as the prior knowledge; Zou et al. [47] cluster landmarks with a graph model to provide structural constraints. However, they still cannot explicitly model an inherent relation between the landmarks with long distance.

The vision transformer [11] proposed recently enables the model to attend the area with a long distance. Besides, the attention mechanism in transformer can generate an adaptive global attention for different tasks, such as object detection [5, 46] and human pose estimation [23], and in principle, we envision that it can also learn an adaptive inherent relation for face alignment. In this paper, we demonstrate the capability of SLPT for learning the relation.

As shown in Fig.2, Sparse Local Patch Transformer (SLPT) consists of three parts, the patch embedding & structure encoding, inherent relation layers and prediction heads.

Patch embedding & structure encoding: ViT [11] divides an image or a feature map $\bm{I}\in\mathbb{R}^{H_{I}\times W_{I}\times C}$ into a grid of $\frac{H_{I}}{P_{h}}\times\frac{W_{I}}{P_{w}}$ with each patch of size $P_{h}\times P_{w}$ and maps it into a $d$-dimension vector as the input. Different from ViT, for each landmark, the SLPT crops a local patch with the fixed size $\left(P_{h},P_{w}\right)$ from the feature map as its supporting patch, whose center is located at the landmark. Then, the patches are resized to $K\times K$ by linear interpolation and mapped into a series of vectors by a CNN layer. Hence, each vector can be viewed as the representation of the corresponding landmark. Besides, to retain the relative position of landmarks in a regular face shape (structure information), we supplement the representations with a series of learnable parameters called structure encoding. As shown in Fig.3, the SLPT learns to encode the distance between landmarks within the regular facial structure in the similarity of encodings. Each encoding has high similarity with the encoding of neighboring landmark (eg. left eye and right eye).

Inherent relation layer: Inspired by Transformer [32], we propose inherent relation layers to model the relation between landmarks. Each layer consists of three blocks, multi-head self-attention (MSA) block, multi-head cross-attention (MCA) block, and multilayer perceptron (MLP) block, and an additional Layernorm (LN) is applied before every block. Based on the self-attention mechanism in MSA block, the information of queries interact adaptively for learning a $query-query$ inherent relation. Supposing the $l$-th MSA block obtains $H$ heads, the input $T^{l}$ and landmark queries $Q$ with $C_{I}$-dimension are divided into $H$ sequences equally ($T^{l}$ is a zero matrix in $1$st layer). The self-attention weight of the $h$-th head $\bm{A}_{h}$ is calculated by:

where $\bm{W}^{q}_{h}$ and $\bm{W}^{k}_{h}$ $\in\mathbb{R}^{C_{h}\times C_{h}}$ are the learnable parameters of two linear layers. $\bm{T}^{l}_{h}\in\mathbb{R}^{N\times C_{h}}$ and $\bm{Q}_{h}\in\mathbb{R}^{N\times C_{h}}$ are the input and landmark queries respectively of the $h$-th head with the dimension $C_{h}=C_{I}/H$. Then, MSA block can be formulated as:

where $\bm{W}_{h}^{v}\in\mathbb{R}^{C_{h}\times C_{h}}$ and $\bm{W}_{P}\in\mathbb{R}^{C_{I}\times C_{I}}$ are also the learnable parameters of linear layers.

The MCA block aggregates the representations of facial landmarks based on the cross-attention mechanism for learning an adaptive $representation-query$ relation. As shown in the rightmost images of Fig.2, by taking advantage of the cross attention, each landmark can employ neighboring landmarks for coherent prediction and the occluded landmark can be predicted according to the representations of visible landmarks. Similar to MSA, MCA also has $H$ heads and the attention weight in the $h$-th head $\bm{A}_{h}^{\prime}$ can be calculated by:

Where $\bm{W}^{\prime q}_{h}$ and $\bm{W}^{\prime k}_{h}\in\mathbb{R}^{C_{h}\times C_{h}}$ are learnable parameters of two linear layers in the $h$-th head. $\bm{T}_{h}^{\prime l}\in\mathbb{R}^{N\times C_{h}}$ is the input $l$-th MCA block; $\bm{P}_{h}\in\mathbb{R}^{N\times C_{h}}$ is the structure encodings; $\bm{R}_{h}\in\mathbb{R}^{N\times C_{h}}$ is the landmark representations. MCA block can be formulated as:

where $\bm{W}_{h}^{\prime v}\in\mathbb{R}^{C_{h}\times C_{h}}$ and $\bm{W}^{\prime}_{P}\in\mathbb{R}^{C_{I}\times C_{I}}$ are also the learnable parameters of linear layers in MCA block.

Supposing predicting $N$ pre-defined landmarks, the computational complexity of the MCA that employ sparse local patches $\Omega(S)$ and full feature map $\Omega(F)$ is:

Compared to using the full feature map, the number of representations decreases from $\frac{H_{I}}{P_{h}}\times\frac{W_{I}}{P_{w}}$ to $N$ (with the same input size, $\frac{H_{I}}{P_{h}}\times\frac{W_{I}}{P_{w}}$ is $16\times 16$ in the related framework [5]), which decreases the computational complexity significantly. For a 29 landmark dataset [4], $\Omega(S)$ is only $1/5$ of $\Omega(F)$ ($H=8$ and $C_{h}=32$ in the experiment).

Prediction head: the prediction head consists of a layernorm to normalize the input and a MLP layer to predict the result. The output of the inherent relation layer is the local position of the landmark with respect to its supporting patch. Based on the local position on the $i$-th patch $\left(t_{x}^{i},t_{y}^{i}\right)$, the global coordinate of the $i$-th landmark $\left(x^{i},y^{i}\right)$ can be calculated by:

where $(w^{i},h^{i})$ is the size of the supporting patch.

To further improve the performance and robustness of SLPT, we introduce a coarse-to-fine framework trained in an end-to-end method to incorporate with the SLPT. The pseudo-code in Algorithm 1 shows the training pipeline of the framework. It enables a group of initial facial landmarks $\bm{S}_{0}$ calculated from the mean face in the training set to converge to the target facial landmarks gradually with several stages. Each stage takes the previous landmarks as center to crop a series of patches. Then, the patches are resized into a fixed size $K\times K$ and fed into the SLPT to predict the local point on the supporting patches. Large patch size in the initial stage enables the SLPT to obtain a large receptive filed that prevents the patch from deviating from the target landmark. Then, the patch size in the following stages is $1/2$ of its former stage, which enables the local patches to extract fine-grained features and evolve into a pyramidal form. By taking advantage of the pyramidal form, we can observe a significant improvement for SLPT. (see Section 4.5).

We employ the normalized L2 loss to provide the supervision for stages of the coarse-to-fine framework. Moreover, similar to other works [25, 29], providing additional supervision for the intermediate output during the training is also helpful. Therefore, we feed the intermediate output of each inherent relation layer into a shared prediction head. The loss function is written as:

where $S$ and $D$ indicate the number of coarse-to-fine stage and inherent relation layer respectively. $\left(x_{gt}^{k},y_{gt}^{k}\right)$ is the labeled coordinate of the $k$-th point. $\left(x^{ijk},y^{ijk}\right)$ is the coordinate of $k$-th point predicted by $j$-th inherent relation layer in $i$-th stage. $d$ is the distance between outer eye corners that acts as a normalization factor.

Experiments are conducted on three popular benchmarks, including WFLW [36], 300W[28] and COFW[4].

WFLW dataset is a very challenging dataset that consists of 10,000 images, 7,500 for training and 2,500 for testing. It provides 98 manually annotated landmarks and rich attribute labels, such as profile face, heavy occlusion, make-up and illumination.

300W is the most commonly used dataset that includes 3,148 images for training and 689 images for testing. The training set consists of the fullset of AFW [45], the training subset of HELEN [20] and LFPW [2]. The test set is further divided into a challenging subset that includes 135 images (IBUG fullset [28]) and a common subset that consists of 554 images (test subset of HELEN and LFPW). Each image in 300W is annotated with 68 facial landmarks.

COFW mainly consists of the samples with heavy occlusion and profile face. The training set includes 1,345 images and each image is provided with 29 annotated landmarks. The test set has two variants. One variant presents 29 landmarks annotation per face image (COFW), The other is provided with 68 annotated landmarks per face image (COFW68 [14]). Both contains 507 images. We employ the COFW68 set for cross-dataset validation.

Referring to other related work [18, 35, 24], we evaluate the proposed methods with standard metrics, Normalized Mean Error (NME), Failure Rate (FR) and Area Under Curve (AUC). NME is defined as:

where $\bm{S}$ and $\bm{S}_{gt}$ denote the predicted and annotated coordinates of landmarks respectively. $\bm{p}^{i}$ and $\bm{p}^{i}_{gt}$ indicate the coordinate of $i$-th landmark in $\bm{S}$ and $\bm{S}_{gt}$. $N$ is the number of landmarks, $d$ is the reference distance to normalize the error. $d$ could be the distance between outer eye corners (inter-ocular) or the distance between pupil centers (inter-pupils). FR indicates the percentage of images in the test set whose NME is higher than a certain threshold. AUC is calculated based on Cumulative Error Distribution (CED) curve. It indicates the fraction of test images whose NME(%) is less or equal to the value on the horizontal axis. AUC is the area under CED curve, from zero to the threshold for FR.

Each input image is cropped and resized to $256\times 256$. We train the proposed framework with Adam [8], setting the initial learning rate to $1\times 10^{-3}$. Without specifications, the size of the resized patch is set to $7\times 7$ and the framework has 6 inherent relation layers and 3 coarse-to-fine stages. Besides, we augment the training set with random horizontal flipping ($50\%$), gray ($20\%$), occlusion ($33\%$), scaling ($\pm 5\%$), rotation ($\pm 30^{\circ}$), translation ($\pm 10px$). We implement our method with two different backbone: a light HRNetW18C [34] (the modularized block number in each stage is set to 1) and Resnet34[16]. For the HRNetW18C-lite, the resolution of feature map is $64\times 64$, and for the Resnet34, we extract representations from the output feature maps of stages C2 through C5. (see Appendix A.1).

WFLW: as tabulated in Table 1 (more detailed results on the subset of WFLW are in Appendix A.2), SLPT demonstrates impressive performance. With the increasing of inherent layers, the performance of SLPT can be further improved and outperforms the ADNet (see Appendix A.5). Referring to DETR, we also implement a Transformer based method that employs the full feature map for face alignment. The number of the input tokens is $16\times 16$. With the same backbone (HRNetW18C-lite), we observe an improvement of 12.10% in NME, and the number of training epoch is $8\times$ less than the DETR (see Appendix A.3). Moreover, the SLPT also outperforms the coordinate regreesion and heatmap regression methods significantly. Some qualitative results are shown in Fig. 4. It is evident that our method could localize the landmarks accurately, in particular for face images with blur (2$nd$ row in Fig.4), profile view (1$st$ row in Fig.4) and heavy occlusion (3$rd$ and 4$th$ row in Fig.4).

300W: the comparison result is shown in Table 2. Compared to the coordinate and heatmap regression methods (HRNetW18C [34]), SLPT still achieves an impressive improvement of 9.69% and 4.52% respectively in NME on the fullset. However, the improvement on 300W is not as significant as WFLW since learning an adaptive inherent relation requires a large number of annotated samples. With limited training samples, the methods with prior knowledge, such as facial boundaries (Awing and ADNet) and affined mean shape (SDL), always achieve better performance.

COFW: We conduct two experiments on COFW for comparsion, the within-dataset validation and cross-dataset validation. For the within-dataset validation, the model is trained with 1,345 images and validated with 507 images on COFW. The inter-ocular and inter-pupil NME of SLPT and the state-of-the-art methods are reported in Table 3 respectively. In this experiment, the number of training sample is quite small, which leads to the significant degradation of the coordinate regression methods, such as SDFL, LAB. Nevertheless, SLPT still maintains excellent performance and yields the second best performance. It improves the metric by 3.77% and 11.00% in NME over the heatmap regression and coordinate regression methods respectively.

For the cross-dataset validation, the training set includes the complete 300W dataset (3,837 images) and the test set is COFW68 (507 images with 68 landmark annotation). Most samples of COFW68 are under heavy occlusion. The inter-ocular NME and FR of SLPT and the state-of-the-art methods are reported in Table 4. Compared to the methods based on GCN (SDL and SDFL), the SLPT (HRNet) achieves impressive result, as low as 4.10% in NME. The result illustrates that the adaptive inherent relation of SLPT works better than the fixed adjacency matrix of GCN for robust face alignment, especially for the condition of heavy occlusion.

Evaluation on different coarse-to-fine stages: to explore the contribution of the coarse-to-fine framework, we train the SLPT with different number of coarse-to-fine stages on the WFLW dataset. The NME, AUC_0.1 and FR_0.1 of each intermediate stage and the final stage are shown in Table 5. Compared to the model with only one stage, the local patches in multi-stages model evolve into a pyramidal form, which improves the performance of intermediate stages and final stage significantly. When the stage increases from 1 to 3, the NME of the first stage decreases dramatically from 4.79% to 4.38%. When the number of stages is more than 3, the performance converges and additional stages cannot bring any improvement to the model.

Evaluation on MSA and MCA block: To explore the influence of query-query inter relation (eq.1) and representation-query inter relation (eq.3) created by MSA and MCA blocks, we implement four different models with/without MSA and MCA, ranging from 1 to 4. For the models without MCA block, we utilize the landmark representations as the queries input. The performance of the four models are tabulated in Table 6. Without MSA and MCA, each landmark is regressed merely based on the feature of the supporting patches in model 1. Nevertheless, it still outperforms other coordinate regression methods because of the coarse-to-fine framework. When self-attention or cross-attention is introduced into the model, the performance is boosted significantly, reaching at 4.20% and 4.17% respectively in terms of NME. Moreover, the self-attention and cross-attention can be combined to improve the performance of model further.

Evaluation on structure encoding: we implement two models with/without structure encoding to explore the influence of structural information. With structural information, the performance of SLPT is improved, as shown in Table 7.

Evaluation on computational complexity: the computational complexity and parameters of SLPT and other SOTA methods are shown in Table 8. The computational complexity of SLPT is only $1/8$ to $1/5$ FLOPs of the previous SOTA methods (AVS and AWing), demonstrating that learning inherent relation is more efficient than other methods. Although SLPT runs three times for coarse-to-fine localization, patch embedding and linear interpolation procedures, we do not observe a significant increasing of computational complexity, especially for 29 landmarks, because the sparse local patches lead to less tokens.

Besides, the influence of patch size and inherent layer number are shown in the Appendix A.4 and A.5.

We calculate the mean attention weight of each MCA and MSA block on the WFLW test set, as shown in Fig.5. We find out that the MCA block tends to aggregate the representation of the supporting and neighboring patches to generate the local feature, while MSA block tends to pay attention to the landmarks with a long distance to create the global feature. That is why the MCA block can incorporate with the MSA block for better performance.

As discussed in Section 4.3, we construct multi-level feature maps for ResNet34, as shown in Fig.6. Supposing the feature map size of $k$-th stage in ResNet34 is $W_{k}\times H_{k}\times d_{k}$, we firstly adopt a $1\times 1$ CNN layer to reduce the channels from $d_{k}$ to $C_{I}/4$. Then, the SLPT crops $N$ patches whose size is $P_{Wk}\times P_{Hk}$ from each level and resizes these patches to $K\times K$. Note that $P_{Wk}\times P_{Hk}$ is $W_{k}/4\times H_{k}/4$ in the initial coarse-to-fine stage and is reduced by half in each following stage. Finally, the resized patches from different levels are concatenated on the channel dimension which is $C_{I}$. As the result, the SLPT can utilize both high level and low level features for face alignment.

The comparison results on WFLW test set and its subsets are tabulated in Table 9. SLPT yields the best performance in NME and works at SOTA level on all subsets.

The convergence curves of SLPT and DETR is shown in Fig.7. The DETR achieves 4.71% NME at 391 epochs on WFLW test set. The SLPT achieves better performance with around 8$\times$ less training epochs. With the increasing of training epochs, the performance of SLPT is improved further, achieving 4.14% NME at 140 epochs.

Each local patch is resized to $K\times K$ and then projected into a vector by a CNN layer with $K\times K$ kernel size. In this section, we explore the influence of the patch size on WFLW test set, as tabulated in Table 10. Compared to $7\times 7$ patches, the $5\times 5$ patches lose more information because of the lower resolution, which leads to degradation of the performance. When the patch size is extended from $7\times 7$ to $9\times 9$, the parameters of the CNN layer is doubled, which leads to the overfitting on the training set. Therefore, we can also observe a slight degradation with $9\times 9$ patch size, from 4.14% to 4.16% in NME.

Table 11 demonstrates the influence of inherent relation layer number. The performance of SLPT relies on the inherent relation layer heavily. When the number of inherent relation layers increases from 2 to 12, We can observe a significant improvement, from 4.19% to 4.12% in NME. Nevertheless, too many inherent relation layers also increase the parameters and computational complexity dramatically. Considering the real-time capability, we choose the model with 6 inherent relation layers as the optimal model.

We visualize the predicted results and adaptive inherent relation maps for the samples of COFW, 300W and WFLW, as shown in Fig.8, Fig.9 and Fig.10 respectively. In the inherent relation maps, we connect each point to the point with highest cross-attention weight. The SLPT tends to utilize the visible landmarks to localize the landmarks with heavy occlusion for robust face alignment. For other landmark, it relies more on its neighboring landmark.