Pretrained equivariant features improve unsupervised landmark discovery

Producing models that learn with minimal human supervision is crucial for tasks that are otherwise expensive to solve with manual supervision. Discovering landmarks within natural images falls into this category since manually annotating landmark points at the pixel level scales poorly with the number of annotations and the size of the dataset. Recent advancements in Deep Learning have brought significant performance improvements to various computer vision tasks. This work focuses on using deep-learning models for unsupervised discovery of landmarks, as adopted in several recent articles (Thewlis et al., 2017b; Jakab et al., 2018; Thewlis et al., 2019; Zhang et al., 2018).

Equivariance, the fundamental concept most modern unsupervised landmark discovery approaches are based upon, formalizes the intuitive property that, when an image is deformed through the action of a geometric transformation, semantically meaningful landmarks associated with that image should also be deformed through the action of the same deformation. Most recent approaches for unsupervised landmark discovery leverage this equivariance property of landmark points.

The tasks of discovering and detecting semantically meaningful landmarks have been extensively studied. Developed before the advent of more recent deep-learning techniques and among the large number of proposed methods for solving these problems, it is worth mentioning the supervised and weakly-supervised learning approaches based on templates (Zhu and Ramanan, 2012; Pedersoli et al., 2014), active appearance (Cootes et al., 1998; Matthews and Baker, 2004; Cristinacce and Cootes, 2008), as well as the regression-based models (Valstar et al., 2010; Dantone et al., 2012; Cao et al., 2014; Ren et al., 2014). Most of the recently proposed methodologies such as Cascade CNN (Sun et al., 2013), CFAN (Zhang et al., 2014a), coordinate regression (Nibali et al., 2018) and several variants (Zhang et al., 2014b, 2015; Yu et al., 2016; Xiao et al., 2016; Wu et al., 2017; Xiao et al., 2017) leverage deep-convolutional feature extractors for solving the task of supervised landmark detection.

In contrast, the approach described in this text is fully unsupervised and yet can perform more robustly than several supervised approaches. Among other unsupervised landmark discovery pipelines, the approach of Thewlis et al. (2017b) was one of the first ones to leverage the equivariance property of landmarks. Equivariance was also used to discover 3D landmarks in the work of Suwajanakorn et al. (2018). Several other approaches are based on the creation of representations that disentangle geometric structures from appearance. Such approaches can be seen in the work of Rocco et al. (2017), Warpnet (Kanazawa et al., 2016), as well as Shu et al. (2018). In the work of Thewlis et al. (2017a), objects are mapped to labels for capturing structures in the image. Similar ideas can also be found (Thewlis et al., 2019).

Several approaches attempt to disentangle the latent representation of images into visual and structural content by comparing representation between related images. These related images can be obtained through synthetic image deformations, or by exploiting the temporal continuity of video streams. Jakab et al. (2018) reconstructs images by combining the visual and the structural content to discover landmarks. Simultaneous image generation and landmark discovery can also be seen in (Zhang et al., 2018; Reed et al., 2016, 2017). Although this set of methods has proven to be successful in some situations, the reliance on reconstructing high-resolution images is problematic in several aspects: these methods require high-capacity models, are often slow to train, and perform sub-optimally in situations where images are corrupted by a high level of noise or in the presence of artifact, as is, for example, common in medical applications. The proposed approach does not rely on image reconstruction and the numerical experiments presented in Section 5 demonstrate that it can still perform competitively on challenging tasks involving human-pose estimations, a task generally solved by leveraging the temporal continuity of adjacent video frames.

Landmarks can be defined as a finite set of points that define the structure of an object. A particular class of objects is expected to have the same set of landmarks across different instances. In many scenarios, the variations in the locations of the landmarks across different instances of the same object can be used for downstream tasks (e.g., facial recognition, shape analysis, tracking).

Consider an image $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$ represented as a tensor with height $H$, width $W$ and number of channels $C$. This image can also be thought of as (the discretization of) a mapping from $\Lambda\triangleq[0,H]\times[0,W]$ to $\mathbb{R}^{C}$. We are interested in designing a landmark extraction function $\mathbf{y}(\mathbf{x})=[\mathbf{y}_{1}(\mathbf{x}),\ldots,\mathbf{y}_{K}(\mathbf{x})]\in\Lambda^{K}$ that associates to the image $\mathbf{x}$ a set of $K\geq 1$ landmarks $\{\mathbf{y}_{k}(\mathbf{x})\}_{k=1}^{K}$ where $\mathbf{y}_{k}(\mathbf{x})\in\Lambda$ for $1\leq k\leq K$. In this work, the mapping $\mathbf{y}:\mathbb{R}^{H\times W\times C}\to\Lambda^{K}$ is parametrized by a neural network, although it is not a requirement. Occasionally, in order to avoid notational clutter, we omit the dependence of $\mathbf{y}$ on the input image $\mathbf{x}$.

Consider a bijective mapping $\mathbf{g}:\Lambda\to\Lambda$ that describes a geometric deformation within of the domain $\Lambda$. Examples include translations, rotations, elastic deformations, with a proper treatment of the boundary conditions. To any function $\varphi:\Lambda\to\mathbb{R}^{C}$ can be associated the function $\mathbf{g}^{\sharp}\circ\varphi:\Lambda\to\mathbb{R}^{C}$ defined as $\mathbf{g}^{\sharp}\circ\varphi(\lambda)=\varphi(\mathbf{g}^{-1}(\lambda))$ for all $\lambda\in\Lambda$. Since images can be thought of as discretizations of functions from $\Lambda$ to $\mathbb{R}^{C}$, the function $\mathbf{g}^{\sharp}$ also defines, with a slight abuse of notation, a deformation mapping on the set of images. With these notations, the function $\mathbf{g}$ describes the transformation of coordinates (i.e. elements of $\Lambda$), while the function $\mathbf{g}^{\sharp}$ describes the action of the transformation $\mathbf{g}$ when applied to all pixels of a given image. A landmark extraction function $\mathbf{y}_{k}:\mathbb{R}^{H\times W\times C}\to\Lambda$ is said to satisfy the equivariance property if the condition

$\displaystyle\mathbf{y}_{k}[\mathbf{g}^{\sharp}(\mathbf{x})]=\mathbf{g}[\mathbf{y}_{k}(\mathbf{x})]$

(1)

holds for all images $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$. Equation (1) translate the intuitive property that, when an image is deformed under the action of a deformation mapping $g^{\sharp}$, the associated landmarks are deformed through the action of the same deformation.

In this article, we follow the standard approach (Jakab et al., 2018; Thewlis et al., 2017b) of implementing the landmark extraction mapping $\mathbf{y}_{k}:\mathbb{R}^{H\times W\times C}\to\Lambda$ as a differentiable function by expressing it as the composition of a heatmap function $\mathbf{h}_{k}:\mathbb{R}^{H\times W\times C}\to\mathbb{R}^{H\times W}$ with the standard spatial soft-argmax (Chapelle and Wu, 2010) function $\sigma_{\mathrm{arg}}:\mathbb{R}^{H\times W}\to\Lambda$,

$\displaystyle\mathbf{x}\underset{\mathbf{h}_{k}}{\longrightarrow}\;\mathbf{h}_{k}(\mathbf{x})\underset{\sigma_{\mathrm{arg}}}{\longrightarrow}\mathbf{y}_{k}(\mathbf{x})$

so that $\mathbf{y}_{k}(\mathbf{x})=\sigma_{\mathrm{arg}}\circ\mathbf{h}_{k}(\mathbf{x})\in\Lambda$. In order to train a landmark extraction model, it is consequently a natural choice to try to minimize the equivariance loss

$\displaystyle\mathcal{L}_{\text{eqv}}(\mathbf{x},\mathbf{g})=\frac{1}{K}\sum_{k=1}^{K}\left\|\mathbf{y}_{k}[\mathbf{g}^{\sharp}(\mathbf{x})]\;-\;\mathbf{g}[\mathbf{y}_{k}(\mathbf{x})]\right\|^{2}.$

(2)

This formulation was first used by Thewlis et al. (2017b), and then subsequently adopted by several other articles (Thewlis et al., 2019; Zhang et al., 2018). Unfortunately, minimizing the equivariance loss alone generally leads to a trivial solution: all landmarks coincide. In practice, it is consequently necessary to consider an additional diversity loss

$\displaystyle\mathcal{L}_{\text{div}}(\mathbf{x})=\sum_{\lambda\in\Lambda}{\left(\sum_{k\in K}[\sigma\circ\mathbf{h}_{k}(\mathbf{x})]_{\lambda}-\max_{1\leq k\leq K}[\sigma\circ\mathbf{h}_{k}(\mathbf{x})]_{\lambda}\right)}$

(3)

where $\sigma:\mathbb{R}^{H\times W}\to(0,1)^{H\times W}$ is the usual spatial softmax function, although other choices are certainly possible. The notation $[\sigma\circ\mathbf{h}_{k}(\mathbf{x})]_{\lambda}$ denotes the spatial softmax evaluated at the coordinate $\lambda\in\Lambda$. The loss (3) was first proposed in Thewlis et al. (2017b) in conjunction with the equivariance loss to penalize the concentration of different heatmaps at the same coordinate. The diversity loss is often computed for small image patches rather than individual coordinates. Landmarks can be learned by directly minimizing the total loss $\mathcal{L}_{\text{div}}+\mathcal{L}_{\text{eqv}}$ on a dataset on unlabeled images and randomly generated deformation mappings $\mathbf{g}$.

this section defines a measure of equivariance for convolutional features. This allows us to measure the equivariance of intermediate convolutional features learned by the set of methods described in Section 2.

To this end, in our study of the equivariance properties of intermediate convolutional features, we exploit a labelled dataset $\mathcal{D}$ whose images have been annotated with $J\geq 1$ semantically meaningful landmark points: to each image $\mathbf{x}$ is associated $\mathbf{y}_{1},\ldots,\mathbf{y}_{J}\in\Lambda$ landmarks. For such an image $\mathbf{x}$ and a geometric deformation $\mathbf{g}$, consider the deformed image $\mathbf{x}^{\prime}=\mathbf{g}^{\sharp}(\mathbf{x})$ as well as the associated convolutional features $\mathbf{f}$ and $\mathbf{f}^{\prime}$. For each landmark $\mathbf{y}_{j}$ associated to $\mathbf{x}\in\mathcal{D}$, we consider the location $\widehat{\mathbf{y}}^{\prime}_{j}\in\Lambda$ in the deformed image $\mathbf{x}^{\prime}$ whose representation is the most similar to $\mathbf{f}^{\prime}(\mathbf{y}_{j})$,

$\displaystyle\widehat{\mathbf{y}}^{\prime}_{j}=\underset{\lambda\in\Lambda}{\mathbf{argmax}}\;\operatorname{sim}\big{(}\mathbf{f}^{\prime}(\lambda),\,\mathbf{f}(\mathbf{y}_{j})\big{)}.$

(4)

In (4) we have used the cosine similarity,

$\displaystyle\operatorname{sim}(\mathbf{f}_{1},\mathbf{f}_{2})\triangleq\biggl{<}\frac{\mathbf{f}_{1}}{\|\mathbf{f}_{1}\|},\,\frac{\mathbf{f}_{2}}{\|\mathbf{f}_{2}\|}\biggl{>},$

(5)

although other choices are indeed possible. For discriminative and equivariant convolutional features, we expect the Euclidean distance between $\widehat{\mathbf{y}}^{\prime}_{j}$ and $\mathbf{y}_{j}$ to be small. Consequently, for a distance threshold $d>0$, the accuracy at threshold $d$ is defined as

$\displaystyle\mathrm{Acc}(d)\triangleq\frac{1}{|\mathcal{D}|\cdot|J|}\sum_{j\in J,\mathbf{x}_{i}\in\mathcal{D}}\mathbf{I}\big{(}\|\widehat{\mathbf{y}}^{\prime}_{i,j}-\mathbf{y}^{\prime}_{i,j}\|\leq d\big{)}.$

(6)

where $\widehat{\mathbf{y}}^{\prime}_{i,j}$ and $\mathbf{y}^{\prime}_{i,j}$ denote the $j$-th landmark quantities associated to image $\mathbf{x}_{i}\in\mathcal{D}$. In the sequel, the curve $\{(d,\mathrm{Acc}(d)):d\!\geq\!0\}$ will be referred to as the accuracy curve. An example is depicted in Figure 1.

Figure 1 presents examples from the BBC Pose test dataset where random rotations and random elastic deformations were applied. The circles denote the locations $\widehat{\mathbf{y}}^{\prime}_{i}$ estimated through Equation (4) and the crosses indicate the ground truth locations $\bar{\mathbf{y}}^{\prime}_{i}$; these locations were obtained by training a landmark extractor in an end-to-end manner. The right-most plots of Figure 1 report the accuracy curves $\{(d,\mathrm{Acc}(d))\,:\,d>0)\}$ estimated from a set of $J=7$ annotated landmarks. Each of the four accuracy curves is associated to one of four intermediate convolutional features coined Layer 1 to Layer 4. Here, Layer 1 is the output of the feature extractor $\mathrm{F}(\cdot)$ while Layer (2, 3, 4) are representation situated in between Layer 1 and the final heatmaps generated through the complete network $\mathrm{T}\circ\mathrm{F}$. For comparison, the black dashed line shows the performance associated with locations $\widehat{\mathbf{y}}^{\prime}_{i}$ generated uniformly at random within $\Lambda=[0,H]\times[0,W]$. The results indicate that, even though the intermediate convolutional features do possess some degree of equivariance (i.e. better than random), the performance is relatively poor with an accuracy $\mathrm{Acc}(d)=57\%$ for a distance threshold of $d=20$ pixels.

The previous section demonstrates that the direct minimization of the combined loss $\mathcal{L}_{\text{eqv}}+\mathcal{L}_{\text{div}}$ is not efficient at inducing equivariance properties to the intermediate convolutional features. In contrast, the first step of the methodology described in this section aims at directly enforcing that the intermediate features enjoy enhanced equivariance properties. In a second step, these pre-trained features are leveraged within more standard landmark discovery pipelines. Although we can choose to pre-train any of the convolutional features of the complete landmark extractor $\mathrm{T}\circ\mathrm{F}$, our experiments indicate that it is most efficient to concentrate on the final convolutional features (i.e. layer 1 with the terminology of the previous section) generated by the feature extractor network $\mathrm{F}(\cdot)$. As depicted in the right-most plot of Figure 1, when end-to-end training is used, the layer 1 is the worst in terms of equivariance among the four layers considered in the previous section. At a heuristic level, this can be explained by the fact that the layer 1 is the further “away" from the training signal among the four layers considered.

Our proposed methodology proceeds in two steps. Step 1: train the feature extractor convolutional network $\mathrm{F}$ with the contrastive learning approach described in the remainder of this section. Step 2: freeze the parameters of the feature extractor network $\mathrm{F}(\cdot)$ and train the complete network $\mathrm{T}\circ\mathrm{F}(\cdot)$ with a standard landmark discovery approach, i.e. minimize the combined objective $\mathcal{L}_{\text{eqv}}+\mathcal{L}_{\text{div}}$.

In this section, we illustrate the performance gains in terms of equivariance of the intermediate features when adopting our proposed two-step method. Figure 2 shows some examples of deformed images from the BBC Pose dataset with the ground truth and predicted locations of $J=7$ landmarks: the setting is similar to the one of Figure 1. The predicted locations were estimated through maximizing the similarity, as described in Equation (4). The third plot shows the accuracy of the four intermediate representations layers (1,2,3,4) previously described at the end of Section 4: note that all layers benefited in terms of equivariance from the two-step approach, even though only layer 1 was pre-trained with contrastive learning. The right-most plot compares the equivariance accuracy curve associated to Layer 1 when the model is trained in an end-to-end fashion and when trained with our proposed approach: there is a considerable boost in equivariance accuracy (6). At a distance threshold of $d=20$ pixels, the end-to-end method leads to an accuracy of $\mathrm{Acc}(d=20)=57\%$ while the proposed two-step approach leads to an equivariance accuracy of $\mathrm{Acc}(d=20)=87\%$.

Our study shows that the intermediate neural representations learned by standard end-to-end approach leveraging equivariance losses enjoy poor equivariance properties. Instead, for unsupervised landmark discovery, we proposed to use contrastive learning for pre-training equivariant intermediate neural representations. The numerical experiments demonstrate that this simple strategy, which can naturally be used within several other pipelines and tasks, can lead to a significant performance boost even in challenging situations such as the one described in Section 5.