Pose-Oriented Transformer with Uncertainty-Guided Refinement
for 2D-to-3D Human Pose Estimation

The basic self-attention mechanism transfers the inputs $Z\in R^{N\times C}$ into corresponding $query~{}Q$, $key~{}K$ and $value~{}V$ with the same dimensions $N\times C$ by projection matrices $P^{Q},P^{K},P^{V}\in R^{C\times C}$ respectively, where $N$ denotes the sequence length, and $C$ is the number of hidden dimension.

$$Q=ZP^{Q},~{}~{}K=ZP^{K},~{}~{}V=ZP^{V},$$

(1)

Then we can calculate self-attention by:

$$A=QK^{T}/\sqrt{d},~{}~{}\textit{MH-SA}(X)={softmax}(A)V,$$

(2)

where $A\in R^{N\times N}$ denotes the attention weight matrix. Based on the basic self-attention, MH-SA further splits the $Q,K,V$ for $h$ times to perform attention in parallel and then the outputs of all the heads are concatenated.

position-wise FFN is used for non-linear feature transformation and it contains two Multilayer Perceptron (MLP) and an GELU activation layer. This procedure can be formulated as follows:

$$FFN(X)=MLP(GELU(MLP(X)))+X.$$

(3)

As MH-SA and FFN in transformer are permutation equivariant operation, additional mechanisms are required to encode the structure of input data into model. In particular, we can utilize sine and cosine functions or learnable vectors as the position embeddings, which can be formulated as

$$P_{t}=PE(t)\in R^{C},$$

(4)

where $t$ denotes the position index.

Following previous design (Zheng et al. 2021; Zhang et al. 2022), we first introduce a learnable keypoint position embeddings $K\in R^{J\times C}$ to represent the absolute position of each body joint. In addition, as shown in Figure 3, according to the distance between each joint and the root joint (Pelvis), we split body joints into five groups and design another learnable embeddings called group position embeddings, i.e. , $G\in R^{5\times C}$. Therefore, additional distance-related knowleage can be encoded into model, helping transformer better model the difficult body joints that are far from the root. In this way, the input of pose-oriented transformer encoder, $Z^{(0)}$, can be obtained by:

$$Z_{i}^{(0)}=Z_{i}+K_{i}+G_{\varphi(i)},~{}~{}for~{}~{}i\in[1,\cdots,J],$$

(5)

where $i$ is the joint index and $\varphi(i)=\mathcal{D}(i,1)$ represents the shortest path distance between $i$-th joint and the root joint.

We also propose our pose-oriented self-attention (PO-SA) that explicitly modeling the topological connections of body joints. Specifically, we compute the relative distance for each joints pair $(i,j)$, and encode it as the attention bias for the self-attention mechanism. In this way, we rewrite the self-attention in Eq (2), in which the $(i,j)$-th element of attention matrix $A$ can be computed by:

$$A_{i,j}=(Z_{i}P^{Q})(Z_{j}P^{K})^{T}/\sqrt{d}+\Phi(\mathcal{D}(i,j)),$$

(6)

where $\Phi$ is a MLP network which projects the relative distance (1-dimension) to an H-dimension vector where H is the number of heads in the SA mechanism, it makes each PO-SA have the ability to adjust the desired distance-related receptive field and the additional parameters can be ignored.

Based on the PO-SA, we can obtain output features by sending $Z^{(0)}$ to a cascaded transformer with $L_{1}$ layers. These procedure can be formulated as :

	$\displaystyle Z^{\prime l}$	$\displaystyle=\textit{PO-SA}(\textit{LN}(Z^{l-1}))+Z^{l-1},$		(7)
	$\displaystyle Z^{l}$	$\displaystyle=\textit{FFN}(\textit{LN}(Z^{\prime l}))+Z^{\prime l},$		(8)

where $LN(\cdot)$ represents the layer normalization and $l\in\left[1,2,\cdots,L_{1}\right]$ is the index of POT encoder layers.

In the regression head, we apply a MLP on the output feature $Z^{L_{1}}$ to perform pose regression, generating the first-stage 3D pose $\widetilde{Y}\in R^{J\times 3}$.

We first model the uncertainty for each joint. Specifically, features of POT encoder $Z^{L_{1}}$ are sent to another uncertainty estimation head, producing the uncertainty $\sigma\in R^{J\times 3}$ of the first-stage 3D poses by using an uncertainty estimation loss $\mathcal{L}_{\sigma}$ (Kendall and Gal 2017).

Instead of directly utilizing the first-stage 3D predictions $\widetilde{Y}$, we randomly sample 3D coordinates $\bar{Y}$ around $\widetilde{Y}$ according to a Gaussian distribution $\mathcal{N}(\widetilde{Y},\sigma)$ with the predicted uncertainty $\sigma$ as variance, and send the sampled coordinates to UGRN. This uncertainty-guided sampling strategy ensures that the sampled coordinates have large variance on difficult joints, which requires the model to focus more on making use of context from other joints to compensate for the difficult joint predictions, thus further enhancing the model robustness.

To enable correct back-propagation, we employ a re-parameterization trick to draw a sample $\epsilon$ from the standard Gaussian distribution $\mathcal{N}(\textbf{0},\textbf{1})$ randomly, $i.e.$, $\epsilon\sim\mathcal{N}(\textbf{0},\textbf{1})$. In this way, we can obtain the sampled 3D coordinates by:

$$\bar{Y}=\widetilde{Y}+\sigma\cdot\epsilon.$$

(9)

Note that this sample strategy is only implemented in the training stage. In the inference stage, we set $\bar{Y}=\widetilde{Y}$ directly.

After obtaining the sampled 3D pose $\bar{Y}$, we first concatenate it with the input 2D pose $X$ and obtain $\widetilde{X}$, i.e., $\widetilde{X}=\textit{Concat}(\bar{Y},X)$. Then we project $\widetilde{X}$ to feature embeddings $\widetilde{Z}$ and equip them with keypoint position embeddings $K$ and group position embedding $G$:

$$\widetilde{Z}_{i}^{(0)}=\widetilde{Z}_{i}+K_{i}+G_{\varphi(i)},~{}~{}for~{}~{}i\in[1,J].$$

(10)

Next, $\tilde{Z}_{i}^{(0)}$ is sent to the following $L_{2}$ transformer layers of UGRN to perform uncertainty-guided refinement. The transform layers of UGRN is similar to those of POT, but we replace the distance-related term of Eq. 6 with uncertainty guildance to dynamically adjust the attention weights:

$$A_{i,j}=(Z_{i}P^{Q})(Z_{j}P^{K})^{T}/\left(\sqrt{d}\cdot\textit{Sum}(\sigma_{j})\right),$$

(11)

where $\sigma_{j}\in R^{3}$ is the predicted uncertainty of $j$-th joint. The above uncertainty-guided self-attention (UG-SA) ensures that the body joints with high uncertainty will contribute less in the self-attention mechanism, which can not only alleviate the error propagation, but also enhance the context understanding ability of the model.

Finally, we apply another regression head to $\tilde{Z}^{L2}$ and generate our second-stage refined 3D pose $\hat{Y}\in R^{J\times 3}$.

We first train our POT for the first-stage 3D pose regressing. The objective function can be formulated as :

$$\mathcal{L}_{\mathrm{stageI}}=\frac{1}{J}\sum_{i=1}^{J}{\left(\left\|\widetilde{Y}_{i}-Y_{i}\right\|^{2}\right)},$$

(12)

where $\widetilde{Y}_{i}$ and $Y_{i}$ are the estimated first-stage 3D positions and the ground truth of $i$-th joint respectively.

We aim to predict the uncertainty correctly as well as estimate an accurate refined 3D pose in Stage II. During this stage, we freeze the model parameters of POT and only train the UGRN for stable results. Following (Kendall and Gal 2017), we set our uncertainty estimation loss as:

$$\mathcal{L}_{\sigma}=\frac{1}{J}\sum_{i=1}^{J}{\left(\left\|\frac{\widetilde{Y}_{i}-Y_{i}}{\sigma_{i}}\right\|^{2}+\log(\|\sigma_{i}\|^{2})\right)}.$$

(13)

In addition, we also apply L2 loss to minimize the errors between the refined 3D poses and ground truths:

$$\mathcal{L}_{\mathrm{refine}}=\frac{1}{J}\sum_{i=1}^{J}{\left(\left\|\hat{Y}_{i}-Y_{i}\right\|^{2}\right)},$$

(14)

The final loss function of Stage II is computed by $\mathcal{L}_{\mathrm{stageII}}=\mathcal{L}_{\mathrm{refine}}+\lambda\mathcal{L}_{\sigma}$, where $\lambda$ is the trade-off factor. We set $\lambda$ to 0.001 such that the two loss terms are of the same order of magnitude.

Human3.6M dataset (Ionescu et al. 2013) is widely used in the 3D HPE task which provides 3.6 million indoor RGB images, including 11 subjects actors performing 15 different actions. For fairness, we follow previous works (Martinez et al. 2017; Zhao et al. 2019; Xu and Takano 2021) and take 5 subjects (S1, S5, S6, S7, S8) for training and the other 2 subjects (S9, S11) for testing. In our work, We evaluate our proposed method and conduct ablation study on the Human3.6M dataset. Besides, the MPI-INF-3DHP (Mehta et al. 2017) test set provides images in three different scenarios: studio with a green screen (GS), studio without green screen (noGS) and outdoor scene (Outdoor). We also apply our method to it to demonstrate the generalization capabilities of our proposed method.

For Human3.6M, we follow previous works (Martinez et al. 2017; Zhao et al. 2019) to use the mean per-joint position error (MPJPE) as evaluation metric. MPJPE computes the per-joints mean Euclidean distance between the predicted 3D joints and the ground truth after the origin (pelvis) alignment. For MPI-INF-3DHP, we employ 3D-PCK and AUC as evaluation metrics.

In our experiment, we set the dimension of embeddings to 96 and adopts 6 heads for self-attention with a dropout rate of 0.25. The MLP ratio of FFN is set to 1.5 to reduce the model parameters. We implement our method within the PyTorch framework. During the training stage, we adopt the Adam  (Kingma and Ba 2014) optimizer. For both Stage I and Stage II, the learning rate is initialized to 0.001 and decayed by 0.96 per 4 epochs, and we train each stage for 25 epochs using a mini-batch size of 256. We initialize weights of the our model using the initialization method described in (Glorot and Bengio 2010). We also adopt Max-norm regularization to avoid overfitting.

We first diagnose how each pose-oriented design in the POT affects the performance. In this section, the UGRN is excluded and the first stage 3D pose $\widetilde{Y}$ is used for evaluated. As shown in Table 3, our method achieves the best performance when all the pose-oriented designs are included. Compared with only using keypoint position embeddings , we achieve 0.88mm (37.57mm to 36.69mm) improvement by adding the distance-related group position embeddings, proving that the representation of difficult joints is effectively facilitated. In addition, by replacing the standard self-attention with our PO-SA, we also achieve 1.10mm (36.69mm to 35.59mm) improvement with only 0.01M model parameters increase, which reflects the benefits of enhancing the ability of modeling the topological interactions.

We then inspect how uncertainty-guided refinement benefits performance. It can be seen from Table 4 that our first-stage prediction obtained directly by POT can achieve 35.59 mm in MPJPE, while adding UGRN for refinement can bring 0.83mm (35.59mm to 34.72mm) performance improvement, and UG-sampling can facilitate the learning procedure and further bring 0.9 mm (34.72mm to 33.82mm) gains. To demonstrate that the performance improvement is not brought by the increased model parameters, we also test other refinement model design using other kinds of self-attention, and the results are shown in Table 5. When we replacing UG-RA with standard MH-SA, the performance degrades from 34.72mm to 35.22mm. In addition, when using the proposed PO-SA in the UGRN, the performance also degrades (34.72mm to 35.07mm), which reflects that the uncertainty-related information is more important than distance-related information in the second refinement stage.

Table 6 reports how different parameters impact the performance and the complexity of our model. The results show that, enlarging the embedding dimension from 48 to 96 can boost the performance, but using dimensions larger than 96 cannot bring further benefits. In addition, we observe the best performance when using 12 and 3 transformer layers in POT encoder and UGRN, respectively, and no more gains can be obtained by stacking more layers. Therefore, we set the basic setting to $L_{1}=12$, $L_{2}=3$, and $C=96$.

In Table 7, We compare both the accuracy and the model complexity with other benchmarks on the Human3.6M dataset. We provide two configurations of our method, in which the embedding dimension of Our-S is 48 while that of Our-L is set to 96. Results show that our method can achieve better results with even much fewer parameters.

In Figure 4, we present the average estimation errors of different body joints according to its group index. It can be seen that, both our group position embedding and UGRN bring more performance improvement for group 4 and 5, in which joints are far from the root joint. The results confirm that our benefit mainly comes form the difficult joints.

Figure 5 demonstrates some qualitative results on the Human3.6M dataset compared with Graformer (Zhao, Wang, and Tian 2022). It can be seen that our method can make accurate pose prediction, especially for the difficult joints that are far from the root.