TUSK: Task-Agnostic Unsupervised Keypoints

Existing unsupervised methods for keypoint learning rely heavily on the assumption that a specific keypoint type (e.g. elbow, digit, abstract geometric shape) appears only once in an image. This greatly limits their applicability, as each instance must be isolated before applying the method-an issue that is never discussed or evaluated. We thus propose a novel method to learn Task-agnostic, UnSupervised Keypoints (TUSK) which can deal with multiple instances. To achieve this, instead of the commonly-used strategy of detecting multiple heatmaps, each dedicated to a specific keypoint type, we use a single heatmap for detection, and enable unsupervised learning of keypoint types through clustering. Specifically, we encode semantics into the keypoints by teaching them to reconstruct images from a sparse set of keypoints and their descriptors, where the descriptors are forced to form distinct clusters in feature space around learned prototypes. This makes our approach amenable to a wider range of tasks than any previous unsupervised keypoint method: we show experiments on multiple-instance detection and classification, object discovery, and landmark detection-all unsupervised-with performance on par with the state of the art, while also being able to deal with multiple instances.

Introduction
Keypoint-based methods are a popular off-the-shelf building block for a wide array of computer vision tasks, including city-scale 3D mapping and re-localization [35,48], landmark detection on human faces [49,57,21,33], human pose detection [49,57,33], object detection [34,28,12,42], scene categorization [27,14], image retrieval [24,40], robot navigation [38,11], and many others. As elsewhere in computer vision, learning-based strategies have become the de-facto standard in keypoint detection and description [49,6,40,57,21,8,33,50]. This is accomplished by tailoring a method to a specific task by means of domain-specific supervision signals-methods are often simply not applicable outside the use case they were trained for. This is problematic, as many applications require extensive amounts of labeled training data, which is often difficult or infeasible to collect.
Researchers have attempted to overcome the requirement for labelled data through unsupervised or self-supervised learning [49,57,21,33,6,44,55], primarily using two families of techniques, often in combination: (1) enforcing equivariant representations through synthetic image transformations, such as homographies [6,44] or non-rigid deformations [26,49]; and (2) clustering features in a latent space over a predetermined number of channels, where every channel is meant to represent a single object or a single part of an object. The first objective ensures that keypoints can deal with changes due to camera pose or body articulation, but synthetic transformations do not capture any high-level semantics. The second objective can encode semantic information into the keypoints, but methods assume that every keypoint (a specific object or part, such as the corner of the mouth) appears only once in a given image, such that a 'specialized' heatmap can be learned for each semantic keypoint. This is an effective way to constrain optimization, but is unrealistic in many scenarios, and restricts the applicability of these methods to isolated instances-for example, a learned landmark detector for Ours LatentKeypointGAN [18] Figure 1: Landmark detection on images with multiple faces. (Left) Our model is trained on CelebA [31], which contains single faces, but can be trivially applied here by simply increasing the number of keypoints. (Right) Methods based on the multi-channel-detection paradigm fail to generalize, and would require processing each instance separately. Images are in the public domain.
human faces will break down if we simply introduce a second face to the image: see Fig. 1. In other words, these methods require each object instance to be isolated before keypoint extraction-which is itself a non-trivial problem, and makes them unsuitable as a generic keypoint detector.
We propose a method to solve this. Similarly to previous works, we estimate a heatmap representing the likelihood of each pixel being a keypoint, and a feature map encoding a descriptor for each pixel, should that pixel be a keypoint. To avoid being limited to single instances, unlike most existing works [49,57,21,33,18], we predict a single heatmap, from which we extract multiple keypoints. We then focus on the properties that we want keypoints to have, and show how to accomplish them without supervision or task-specific assumptions. Specifically, we focus on the two key aspects outlined above: (a) keypoints should be equivariant to geometric transformations, i.e., they should follow the same transformation that the image does; and (b) group semantically similar things together. The former can be easily achieved by enforcing equivariance of both keypoint locations and their descriptors under random perturbations, following previous efforts like [49]-with additional considerations, because, unlike existing methods, we do not assume that the images contain a single instance.
The latter is more difficult to solve under the multiple-instance scenario. We address it with a novel solution based on unsupervised clustering of the feature space. We learn a latent representation for each keypoint that we use to reconstruct the image, similarly to [57,21,33,15,2,32], and show how to discover the re-occurring structure or semantic information by clustering these features in the latent space. Specifically, we learn a codebook of cluster centers-which we call prototypes-and enforce that the distribution of descriptors at keypoint locations corresponds to the distribution of these prototypes through a novel sliced Wasserstein objective [5] that utilizes Gaussian Mixture Models (GMM). This naturally leads to descriptors clustering around the prototypes.
Our approach allows us to achieve both objectives without task-specific labels, including implicit assumptions such as that the images contain the same objects, or that each object appears only once [49,57,21,33,18]: see Fig. 1 for examples. Our main contributions are as follows: • We show how to encode both geometry and semantics into keypoints and descriptors, without supervision. For the latter, we propose a novel technique based on prototype clustering. • Unlike existing methods, our method does not require knowing the exact number of keypoints a priori, and can naturally deal with multiple instances and occlusions. • Our approach does not require domain-specific knowledge and can be applied to multiple taskswe showcase landmark detection, multiple-instance classification, and object discovery, with performance on part with the state of the art-while being able to deal with multiple instances. Other than hyperparameters, the only component we tailor to the task is the loss used to reconstruct the image from the keypoints (L recon , Sec. 3.3). Our method is the only one applicable to all tasks.

Related work
Landmark learning. The term landmark is used in the literature to refer to localized, compact representations of object structures such as human body parts. We favor the term keypoint instead, as our method can be applied to a variety of tasks, and our keypoints may represent both whole objects or their parts. There is a vast amount of literature on this topic, especially for human bodies [46,45,39] and faces [52,58], the majority of which use manual annotations for supervision. Among unsupervised methods, Thewlis et al. [49] enforce landmarks to be equivariant by training a siamese network with synthetic transformations of the input image. They introduce a diversity loss that penalizes overlapping landmark distributions to prevent different channels from collapsing into the same landmark. Zhang et al. [57] improve on this design by adding a reconstruction task solved with an autoencoder, which imbues the landmarks with semantics but still relies on separation constraints to prevent landmark collapse. Our approach does not require such an additional regularization, as our formulation prevents collapse by design. Given a pair of source and target images, Jakab et al. [21] detect landmarks in the target image and combine them with a global representation of the source image to reconstruct the target, and learn without supervision through either synthetic deformations or neighboring frames from video sequences. Lorenz et al. [33] use a two-stream autoencoder to disentangle shape and appearance and learn to reconstruct an object from its parts. We use similar strategies to enforce equivariance in the location of the landmarks (keypoints), but additionally learn distinctive features with semantic meaning. In concurrent work, He et al. [18] pushed performance in unsupervised keypoint detection further with a GAN-based framework. All of these methods, however, use dedicated heatmap channels for each landmark to be detected and thus assume that the image contains one and only one object instance-by contrast, our approach learns a single-channel heatmap and can handle a variable number of objects (see Fig. 1 for a comparison with [18]).
Unsupervised object discovery. Our method can be applied to object discovery, by representing objects as keypoints. Decomposing complex visual scenes with a variable number of potentially repeating and partially occluded objects into semantically meaningful components is a fundamental cognitive ability, and remains an open challenge. Finding objects in images has been historically addressed through bounding-box-based detection or segmentation methods trained on manual annotations [17,47,13]-learning object-centric representations of scenes without supervision has only recently seen breakthroughs [23,2,15,32,9,10]. MONet [2] trains a Variational Auto-Encoder and a recurrent attention network to segment objects and reconstruct the image. IODINE [15] similarly relies on VAEs to infer a set of latent object representations which can be used to represent and reconstruct the scene-and like MONet uses multiple encode-decode steps. Locatello et al. [32] replace this procedure with a single encoding step and iterate on the attention mechanism instead, and introduce the concept of 'slot attention' to map feature representations to discrete sets of objects. All of them share the same paradigm, representing the scene as a collection of latent variables where each variable corresponds to one object. However, the decoder may produce different latent codes for objects at different locations in order to reconstruct all present objects, disregarding their possibly similar semantics, which requires higher network capacity to model complex scenes [32]. By contrast, we explicitly supervise our keypoints to form compact clusters in feature space.
Multiple-instance classification. This problem differs from the above in that the objects in the scene are known a priori (e.g. numbers, or letters) and need not be 'discovered'. It is also related to landmark learning, but it may contain repeated objects, which breaks down the one-object-per-channel paradigm used by most methods devised for that problem. MIST [1] solves this with a differentiable top-k operation which is used to extract a predetermined number of patches, which are then fed to a task-specific network, such as a classifier. It needs to be supervised with class labels indicating how many times each instance appears, which we do not require, and unlike our approach it does not generalize to the case of an unknown or variable number of objects. We show that our method obtains similar performance to MIST and another state-of-the-art top-k selection method based on optimal transport [53].
Unsupervised clustering. The performance of traditional clustering methods such as k-means and Gaussian Mixture Models is highly dependent on the quality of the hand-crafted features, and it is non-trivial to incorporate them into deep networks. A popular solution is to train an autoencoder for reconstruction, and regularize the bottleneck feature to form clusters [54,37]. These methods, however, are mostly only applicable to image classification, because they do not have the notion of locality-i.e., they produce one global description per image.
Local features. Local feature is the preferred term for 'classical' keypoint methods, which with the introduction of SIFT [34] became one of the fundamental building blocks of computer vision.
They remain an open area of research [8,50,51,55], but are now relegated to applications such as 3D reconstruction and re-localization [35,48] or SLAM [38,11]. These problems require matching arbitrary textures (corners, blobs) across a wide range of rotations and scale changes. They are thus more open-ended than the other methods presented in this section and are typically supervised from known pixel-to-pixel correspondences. Relevant to our approach are ideas such as using homographic adaptation to learn equivariant features [29,6], or a soft-argmax operator [56] to select keypoints in a differentiable manner. While recent papers have drastically reduced supervision requirements [51,50,55], truly unsupervised learning for general-use local features remains a 'holy grail'. Figure 2: Framework -We generate a heatmap H and a feature map F through an encoder E. We then select K keypoints at locations L and retrieve K descriptors D from the feature map F, by performing non-maxima suppression on the heatmap and picking the top responses. We then decode these keypoints with the decoder D and reconstruct the input image. We use multiple losses for training. L eqv enforces equivariance of heatmap and features w.r.t. a random thin plate spline transformation T . L sw and L km use learned prototypes P to force descriptors to form semantic clusters. L recon minimizes the difference between image and reconstruction, similarly to an autoencoder.

Method
We illustrate our approach in Fig. 2. Given an image, we auto-encode it to obtain a latent representation of sparse keypoints. In more detail, the encoder E receives an input image I ∈ R H×W ×3 and generates a feature map F ∈ R H×W ×C , preserving the spatial dimensions of the input. Unlike previous works that utilize one heatmap per keypoint, our method outputs a single heatmap H ∈ R H×W on which we detect all keypoints, allowing for multiple instances of the same keypoint (or object) to be detected. From the heatmap H, we extract K keypoint locations L ∈ R K×2 and their corresponding descriptors D ∈ R K×C from F. This sparse, latent representation of our autoencoder, i.e., L and D, is then fed into the decoder D, which generatesĨ, a reconstruction of the input image. To force these individually detected D to form semantic clusters in the latent space, we learn M prototypes P ∈ R M ×C , which act as 'cluster centers' that guide the clustering process via losses to be discussed in Sec. 3.3.
This network architecture is similar to previous autoencoder-based landmark discovery papers in that we use an encoder to extract keypoint locations and descriptors and a decoder to reconstruct the image [57, 21, 33]-but our method differs in two key aspects: (a) our encoder generates a single heatmap instead of multiple ones, allowing us to operate beyond the assumption that each feature appears once and only once in every image; and (b) we regulate the latent space by clustering features around prototypes. Note that while we still need to choose a value for K (and M ), our approach can handle a variable number of objects at inference time as long as the number of extracted keypoints is sufficiently large, by filtering out low confidence detections: see CLEVR (Sec. 4.2) for results on images with a variable number of objects while using a fixed number of keypoints.

Encoder E -Keypoint and descriptor extraction
We use a U-Net-like [47] network that takes as input a unit-normalized image I ∈ [0, 1] H×W ×3 to generate a single-channel heatmap H ∈ R H×W and a multi-channel feature map F ∈ R H×W ×C .
(We found C=32 to be sufficient for all datasets and tasks.) To detect multiple instances from H we apply non-maximum suppression (NMS) to H and select the top K maxima. Note that this is a non-differentiable process-inspired by LIFT [56,41], we apply soft-argmax in the neighborhood of each maximum to obtain the final keypoint locations L, restoring differentiability w.r.t. locations. We extract the descriptors D at each keypoint location from the feature map F via simple bilinear interpolation. Please refer to Supplementary Material for further details.

Decoder D -Image reconstruction from sparse features
To decode the sparse, keypoint-based representation into an image, we use another U-Net, identical to the encoder except for the number of input/output channels. We transform the sparse keypoints and descriptors into a dense representation by creating a feature map. Specifically, in order to make keypoints 'local', we generate one feature map per keypoint by copying the descriptors D to the locations L and convolving them with a Gaussian kernel with a pre-defined variance, generating K feature maps. We then aggregate them with a weighted sum, according to the detection score. With this input, we train the decoder D to reconstruct the image I intoĨ.
If the object mask is necessary-e.g. for the object discovery task, in Sec. 4.2-instead of a weighted sum we simply reconstruct one RGB-A image for each of the K feature maps and apply alphablending to obtain the final reconstructionĨ. The alpha channel can be used as the object mask, required by the evaluation metrics. See Supplementary Material for more architectural details.

Learning objectives
We present various learning objectives to encourage our keypoint-based representation to have certain desirable properties. All objectives are critical to the final performance, see ablation in Sec. 4.4.
Learning meaningful features -L recon . Similarly to a standard auto-encoder, our network learns to extract high-level features by minimizing the reconstruction loss between the input image I and the reconstructed imageĨ. We use either the mean squared error (MSE) or the perceptual loss [22]. We find that the perceptual loss is advantageous on real data with complex backgrounds such as CelebA [31], but the MSE loss does better on simpler datasets with no background such as CLEVR [25] or Tetrominoes [25], commonly used by object discovery papers [2,15,32], as well as H36M after removing the background [20].
For the perceptual loss we use a pre-trained, frozen VGG16 network to extract features. V l is the output of layer l and c l its weighting factor, as presented in [22].
Unsupervised clustering -L sw , L km . In order to embed semantic meaning into our learned descriptors, we encourage descriptors to cluster in feature space. Descriptors of the same cluster should move closer together and descriptors of different clusters further apart, which is similar to the objective of contrastive learning papers [16,4,3]. More specifically, we learn M prototypes P ∈ R M ×C , where each prototype represents the 'center' of a class. Our goal is two-fold. First, we wish the prototypes to follow the same distribution as the features. Second, we wish the features to cluster around the prototypes. For the first objective, we minimize the sliced Wasserstein distance [5], and for the second objective, we use a k-means loss.
-Sliced Wasserstein loss L sw . To compute the sliced Wasserstein distance [5], we randomly sample N projections w i from a unit sphere, each projecting the C-dimensional descriptor space into a single dimension. Then, in one-dimensional space, the Wasserstein distance W can be calculated as the square distance between sorted, projected descriptors, and prototypes. The average of these N distances is the sliced Wasserstein loss we seek to minimize.
The difficulty in applying this method, however, is that the number of samples from the two distributions must be equal, which is not the case for our descriptors and prototypes. We thus represent the prototypes as modes in a Gaussian Mixture Model (GMM) [7], so we can draw as many samples from the prototype distribution as the number of keypoints we extract (K) to compute the sliced Wasserstein loss (see Supplementary Material for more details). Denoting the sampled, sorted prototypes asP ∈ R K×C and the sorted descriptors asD, we write -k-means loss L km . With the k-means loss, we encourage descriptors to become close to the closest prototype, relatively to the other prototypes. Inspired by [4], we minimize the log soft-min of the 2 distances between prototypes and descriptors, where the logarithm avoids the vanishing gradients of the soft-min. We write Making detections equivariant -L eqv . Keypoints should be equivariant to image transformations such as affine or non-rigid deformations. This idea appears in several local feature papers [6,44], and unsupervised landmark learning papers [49,57,21]. Following [26,49] we apply a thin plate spline transformation T and use the l 2 distance as a loss to ensure that both heatmap and feature map are equivariant under non-rigid deformations. Note that the encoder E produces a heatmap and a feature map and we apply the loss to both separately-we use a single equation for simplicity: Iterative training. We found that straightforward training of our framework leads to convergence difficulties. We thus alternate training of descriptors and prototypes. Specifically, we train only the encoder and the decoder for one step via λ recon L recon + λ km L km + λ eqv L eqv (see Supplementary Material for details). We then train the prototypes and their GMM for eight steps using L sw . We repeat this iterative training process until convergence.

Results
We apply our method to three tasks: multiple-instance object detection, object discovery, and landmark detection. We use five different datasets.

MNIST-Hard [1]
contains synthetically-generated images composed of multiple MNIST digits. We follow the procedure in [1] and put 9 MNIST digits on a black canvas at random positions, sampling from a Poisson distribution. Variations of the same digit may co-occur in the same image. We generate 50K such images for training and testing, respectively.
CLEVR [25] contains visual scenes with a variable number of objects in each scene. Following [15,2,32], we only use the first 70K images that contain at most 6 objects (also known as CLEVR6). We train our model with the first 60K images and evaluate with the last 10K.
Tetrominoes [25] contains 80K images of Tetrominoes, a geometric shape composed of four squares. Each image contains three randomly sampled, non-overlapping Tetrominoes. There are 19 unique tetrominoes shapes (including rotations). Each Tetromino can be in one of the 6 possible colors. We train our model using the first 50K images and evaluate with the last 10K.
CelebA [31] contains 200k images of human faces. Earlier efforts [33] used aligned images (with centered faces) for both training and evaluation, a dataset that is now heavily saturated. We follow [30] instead, which uses CelebA with unaligned images, filtering out those where the face accounts for less than 30% of the pixels. We follow the standard procedure and train our model without supervision using all images except for the ones in the MAFL (Multi-Attribute Facial Landmark) test set, and train a linear regressor which regresses the 5 manually annotated landmarks, on the MAFL training set. We report pixel accuracy normalized by inter-ocular distance on the MAFL test set.
Human 3.6M (H36M) [20] contains 3.6M captured human images with ground truth joint locations from 11 actors (7 for training, 4 for test) and 17 activities. We follow [57], using 6 actors in the training set for training and the last one for evaluation. We subtract the background with a method provided by the dataset. As for CelebA, we first train our model on the training set and then learn a linear regressor that predicts the location of 36 joints. Since the front and back views are often similar, we follow [57] and swap left and right joint labels on images where the subject is turned around.

Multiple-instance classification and detection: MNIST-Hard [1]
MNIST-Hard contains nine instances of ten different digits-we thus extract K=9 keypoints and M =10 prototypes to represent them. The main challenge for this task is that the images may contain multiple instances of the same digit. Note that all the baseline methods used in the landmark detection task (Sec. 4.3) follow the one-class-per-channel paradigm and will thus fail in this scenario. This limitation is also mentioned in MIST [1], which reports that a channel-wise approach, using the same auto-encoder as MIST while replacing their single-heatmap approach with K separately-predicted channels, results in poor performance: we list this baseline as 'Ch.-wise'. We also consider [53], a differentiable top-k operator based on optimal transport, and a fully supervised version of MIST with both location and class annotation available during training that acts as a performance upper bound for the MIST architecture [1].  To evaluate detection accuracy we compare the ground-truth bounding boxes to our predictions by placing bounding boxes the approximate size of a digit (20 × 20px) around each keypoint. We consider a prediction to be correct with an intersection over union (IoU) of at least 50%. For the classification task, we assign prototypes to each class through bipartite matching, while maximizing accuracy. We also report the intersection of the two metrics. We report results in Table 1. Our method outperforms all baselines on the detection task and is close to MIST on the classification task. Note that ours is the only unsupervised method: all baselines have access to class labels, or class labels and location. Qualitative results for detections, learned prototypes, and t-SNE [36] visualizations of descriptors and prototypes are shown in Fig. 3.

Object Discovery: CLEVR [25] and Tetrominoes [25]
We also consider a multi-object discovery task using two different datasets: CLEVR and Tetrominoes. Performance is evaluated with the Adjusted Rand Index (ARI) [43,19], which computes the similarity between two clusterings, adjusted for the chance grouping of elements, ranging from 0 (random) to 1 (perfect match). This requires semantic masks, so for this task, we modify our network as outlined in Sec. 3.2. Our masks are generated by applying an argmax operator over the alpha channels of the K per-keypoint reconstructions, and are subsequently used to compute ARI. In addition to the standard metric, we report the classification performance of our method, based on shape, color, and size (the latter only for CLEVR)-which none of the baselines are capable of. To calculate classification accuracy, as our method is fully unsupervised, we simply assign, for each prototype, the properties (shape, color, size) that best represent the prototype. We use K=6 keypoints and M =48 prototypes for CLEVR and K=3 and M =114 for Tetrominoes, which are the maximum number of objects and the number of combinations of shape, color, and size, respectively. We report the results in  . We do not report standard deviation for our method, which is deterministic at inference. Our approach is on par with the state of the art in ARI, and can additionally perform unsupervised classification. Note that for CLEVR, shape classification performance is lower because the dataset contains 3D motion (near/far), and some prototype capacity is used to explain this rather than shape, as more pixels change by moving an object from the near to the back plane than e.g. moving between 'sphere' and 'cube'. This lack of 3D spatial information is currently a limitation of our method.

Landmark Detection:
CelebA-in-the-wild [31] and H36M [20] We train our model separately on CelebA and H36M. We extract K=4 keypoints with M =32 prototypes for CelebA and K=16 keypoints with M =32 prototypes for H36M. As noted previously, we follow [49,57,21,33] for evaluation and train a linear layer using ground truth landmark annotations to map the learned landmarks to a set of human-defined landmarks. In existing landmark detection methods, the linear layer has an input size K × 2 and an output size T × 2, where K is the number of heatmap channels and T is the number of human-defined landmarks, which is possible, as they assume a keypoint to be detected precisely once. This is not directly applicable to our method, as we may detect more than one keypoint belonging to the same prototype-e.g., in an extreme case, all K keypoints could belong to a single prototype. To support multiple detections, we form the input to our linear layer as a K × M × 2 tensor, which is populated while preserving prototype information through the indexing: for instance, if keypoints i and j belong to prototype k, we will place their coordinates at [i, k] and [j, k]. Unused entries are simply set to zero-note that only K coordinates are non-zero, the same number that is used for the baseline methods. -4.91% Jakab [21] 19.42% -Lorenz [33] 15.49% 2.79% He [18] 25.81% -He [18]  We report quantitative results in Tab. 3 and show detections and clustering results in Fig. 6. Our method provides comparable performance to the methods targetting single objects but is still able to deal with multiple instances, which none of the baseline methods can do; see Fig. 1. While there is still a gap between our method and the best performing single-instance landmark detector, which is a limitation of our method, consid-  Figure 5: Tetrominoes dataset. Each prototype represents unique shape-color combination.   Fig. 7 shows the corresponding results in the quality of the clustering with t-SNE [36]. Each ingredient has a significant impact on learning semantically meaningful descriptors.

Conclusion
We introduce a novel method able to learn keypoints without any supervision signal. We improve on prior work by removing the key limitation of predicting separate channels for each keypoint, which requires the image to contain a single instance. We achieve this with a novel feature clustering strategy based on prototypes, and show that our method can be applied to a wide range of tasks and datasets. Training and inference code will be made available upon publication. We also provide a discussion on the potential societal impact of our method in the Supplementary Material.
Limitations. While our approach is able to model occlusions (see highlights in Fig. 4), it remains unaware of 3D structure, which limits its applicability to scenarios where scene understanding requires a notion of compositionality (see shape classification performance on Tab. 2). It is also unable to model the background. We intend to address these limitations in future work.

A Supplementary material
The appendix includes a detailed description of the soft-argmax used in encoder E in Sec.

A.1 Detailed Description of the Soft-Argmax in Encoder E
In the encoder E, we first generate a scoremap S the same shape as input image I with a U-Net-style network, followed by a sigmoid activation, whose output is in [0, 1]. Then, we apply NMS on the scoremap S and select top K keypoints P from it. However, this top-k selection is non-differentiablespecifically, the gradients only flow back through the K locations on the H × W scoremap S. To avoid sparsity in the gradient, we construct a kernel of size L × L centered at each keypoint location and use a soft-argmax [56] function to compute the final keypoint location as a weighted sum of the region around the center of the kernel. By using these weighted centers, the gradient flows back through L × L locations for each keypoint, rather than just one. As a by-product, we also obtain sub-pixel accuracy on the keypoints.
where kernel{P i } are the points within a kernel of size L × L centered at point P i , and τ is a hyper parameter to control the hardness of the softmax operation.

A.2 Sampling for Sliced Wasserstein Loss
In order to apply the sliced Wasserstein loss to match the distribution of the prototypes to the distribution of the descriptors, we require an equal number of prototypes and descriptors. Thus, we model the descriptors as a Gaussian Mixture Model (GMM) with a small predefined variance Σ (a hyperparameter), since we would ideally want the descriptors to form very sharp modes around the Gaussian centers. Thus, denoting the sampled descriptors from the prototype GMM asD we sample where π j are the mixture weights of the GMM for the j-th prototype. For the mixture weight π j , note that each prototype may have a different number of descriptors associated with it. We thus define it according to the ratio of descriptors that are associated with each prototype-we associate via finding the nearest prototype with the 2 norm-but with a term that encourages exploration when the compactness of the prototypes is not equal. For example, when a certain prototype dominates but is widely spread, we would ideally want to explore using not just this single prototype. Mathematically, denoting the ratio of descriptors associated with prototype m as r m , and the variance of the descriptors associated with the prototype as σ m , we write and For prototypes without any descriptors assigned, we simply set σ m = 1.

A.3 Generalization to CLEVR10
We demonstrate that our approach generalizes to a larger number of instances without any retraining in Tab. 5, where we follow [32] to evaluate our CLEVR6-trained model on CLEVR10 (which contains up to 10 objects). Our model can deal with a larger number of instances and generalizes very well, with only a very small drop in performance. We do not report numbers for the baselines as they are not available.

A.4 Architecture Detail
We summarize all the main components in our pipeline in Tab. 6.

A.5 Hyperparameter Table
We report the hyperparameters we use for each dataset/task in Tab. 7. The number of prototypes and keypoints are set according to the dataset characteristics. The softmax kernel size was also adjusted to match the dataset image size and the rough size of the object of interest. Batch sizes were determined according to the memory limit of our GPU. Except for CelebA, all settings are similar 1 .

A.6 Computation Resources
We train all our models on NVIDIA V100 GPUs with 32GB of RAM. Training on MNIST converges within 12 hours on a single GPU. For the other four datasets, we use two GPUs, which allows for larger batch size, for about 24 hours.

A.7 Dataset Licenses
A.7.1 MNIST http://yann.lecun.com/exdb/mnist/ The MNIST-Hard dataset is derived from the MNIST dataset. Below is the license for the original MNIST dataset: without warranty or condition of any nature. The authors disclaim all other warranties, expressed or implied, including but not limited to implied warranties of merchantability and fitness for a particular purpose, with respect to the data and any accompanying materials.

RESTRICTION AND LIMITATION OF LIABILITY In no event shall the authors be liable
for any other damages whatsoever arising out of the use of, or inability to use this dataset and its associated software, even if the authors have been advised of the possibility of such damages. 4. RESPONSIBLE USE It is YOUR RESPONSIBILITY to ensure that your use of this product complies with these terms and to seek prior written permission from the authors and pay any additional fees or royalties, as may be required, for any uses not permitted or not specified in this agreement. 5. ACCEPTANCE OF THIS AGREEMENT Any use whatsoever of this dataset and its associated software shall constitute your acceptance of the terms of this agreement. By using the dataset and its associated software, you agree to cite the papers of the authors, in any of your publications by you and your collaborators that make any use of the dataset, in the following format (NOTICE THAT CITING THE DATASET URL INSTEAD OF THE PUBLICA-TIONS, WOULD NOT BE COMPLIANT WITH THIS LICENSE AGREEMENT): A.7.6 multi-face images All multi-face images are in the public domain.

A.8 Societal Impacts
Our method would be a front-end for a typical computer vision pipeline and thus is not immediately linked to any particular application with significant societal impact. Nonetheless, our method would facilitate unsupervised learning from images which could have significant impact down the road. Unsupervised learning would significantly reduce the need for human effort within any automated workflow, which would bring both convenience and a certain amount of change-the latter may require careful consideration when adopting these technologies more widely.


Alternating Phase Langevin Sampling with Implicit Denoiser Priors for Phase Retrieval

Phase retrieval is the nonlinear inverse problem of recovering a true signal from its Fourier magnitude measurements. It arises in many applications such as astronomical imaging, X-Ray crystallography, microscopy, and more. The problem is highly ill-posed due to the phase-induced ambiguities and the large number of possible images that can fit to the given measurements. Thus, there's a rich history of enforcing structural priors to improve solutions including sparsity priors and deep-learning-based generative models. However, such priors are often limited in their representational capacity or generalizability to slightly different distributions. Recent advancements in using denoisers as regularizers for non-convex optimization algorithms have shown promising performance and generalization. We present a way of leveraging the prior implicitly learned by a denoiser to solve phase retrieval problems by incorporating it in a classical alternating minimization framework. Compared to performant denoising-based algorithms for phase retrieval, we showcase competitive performance with Fourier measurements on in-distribution images and notable improvement on out-of-distribution images.

INTRODUCTION
Phase retrieval is the reconstruction problem of obtaining lost phase information of a signal given only its Fourier magnitude. It arises in many physical measurement applications including X-ray crystallography, transmission electron microscopy, astronomical imaging, and more. More formally, the phase retrieval problem asks to recover a signal x * ∈ R n from phaseless measurements of the form where A ∈ C m×n is the measurement matrix and w is noise.
In the Fourier case, A is replaced with F, the Fourier transform, that operates on x * . Due to the phase ambiguity and noise present in the problem, phase retrieval is an ill-posed problem. There are many possible images that can fit to the given measurements, most of them lacking natural structure. Image priors combat this by constraining the space of possible solutions to images that do contain the structure that the object of interest is believed to obey. There is a rich history of image priors across computational imaging, and we highlight previous priors utilized in phase retrieval.
A common structural prior traditionally enforced in imaging applications is sparsity with respect to a wavelet basis. Several works have considered the use of sparsity priors in phase retrieval to improve reconstruction and yield more efficient algorithms [1]. However, sparsity priors are limiting as they inherently enforce a linear signal model, which may potentially lack the richness to model more complicated data of interest. On the learning side, there has been tremendous progress on priors in the form of generative models given by a deep neural network. For example, the works [2,3] have shown that generative networks can potentially better represent natural signals compared to sparsity priors and also demonstrate strong recovery performance. One drawback of this approach, however, is that generative models require an abundance of clean data to be trained. Moreover, such generators are often tailored specifically to a data distribution of interest, and struggle to reconstruct out-of-distribution images. These points ultimately limit their use in practical applications.
An interesting class of (implicit) priors are in the form of neural network-based denoisers, which are models that are trained remove additive noise from images. Denoisers implicitly learn natural image structure through many paired examples of noisy and clean images. Such models have shown impressive performance as regularizers in several linear inverse problems [4]. Additionally, these models have shown strong recovery performance as regularizers in phase retrieval and showcase the ability to recover images outside of their training distribution [5]. Recently, the work of [6] showed that it was possible to sample from the prior implicit in a denoiser, and demonstrated strong performance in using this prior to help solve linear inverse problems.
Our contributions: In this work, we propose an algorithm called Alternating Phase Langevin Sampling (APLS) to solve phase retrieval problems by combining classical alternating projection methods with modern denoiser priors. Our algorithm incorporates a posterior sampling method to access the implicit prior of a denoiser in an alternating minimiza-tion framework. Then, we empirically study our approach in the challenging Fourier measurement regime and showcase competitive results for in-distribution images and significantly outperform other approaches on out-of-distribution images. Our work showcases robustness and generalization performance for such priors, and highlights the potential for incorporating implicit denoiser priors more broadly for nonlinear inverse problems.

RELATED WORK
Phase retrieval methods can be classified as either convex or non-convex approaches. Convex techniques include semidefinite programming methods, such as PhaseLift [7], and methods based on linear programming, such as PhaseMax [8]. Non-convex methods for phase retrieval generally fall under the alternating projections framework, originating from the Gerchberg-Saxton [9] and Fienup [10] algorithms, or direct first-order methods, such as Wirtinger Flow [11]. We now detail some of the alternating minimization approaches that our work is inspired by.

Alternating Minimization
Alternating minimization provides an optimization framework for solving nonlinear inverse problems [12]. In particular, the nonlinear system b = |Ax * | can be reduced to a linear system if we had access to the true phase p of Ax * . However, p is not actually known and so the solution to the linear system is obtained by solving arg min Alternating minimization then alternates between updating the phase p and the solution x to iteratively minimize (2). Denoting the phase of z ∈ C by Ph(z) := z |z| , the phase p t and signal x t are updated via the following iteration for each t 1: The authors of [12] showcase geometric convergence of the algorithm with proper initialization and establish sample complexity guarantees for Gaussian measurements.

Denoisers for Regularization
Denoisers were first used as regularizers in solving inverse problems by [13]. This method is termed plug-and-play regularization because a denoiser can be "plugged in" as a regularizer when solving an inverse problem to produce more accurate recoveries [4].
The current state-of-the-art method for solving phase retrieval using a denoiser as a prior is prDeep [5]. They use the Regularization by Denoising framework (RED) [14] that explicitly defines the regularizer as where D(x) is the denoiser. For the denoiser, a powerful neural network-based denoiser DnCNN [15] is used.

Prior Implicit in a Denoiser
Why are denoisers such strong priors for inverse problems? When training to remove additive noise, denoisers must develop an understanding of what the important properties of structured, natural images are. This implicit prior embedded within denoisers was exposed by [6] through Empirical Bayes estimation and was used to achieve competitive performance in solving linear inverse problems. Specifically, it was shown in [6] that there is an interesting connection [16] between the score of the density p(y) induced by images contaminated with Gaussian noise y = x + z with variance σ 2 and MMSE denoisers: wherex(y) = E[x|y] is the MMSE estimate, which can be approximated by a denoiser trained on a large paired dataset of clean and noisy images. Based on this observation, [6] proposed an iterative algorithm that samples from a denoiser to take a step in the direction specified by the observation density gradient and gradually reduce the effective noise σ to eventually converge to a high-probability region of p(x). To achieve this, [6] employed a Langevin sampling algorithm derived from Langevin Dynamics to inject a decreasing amount of noise into the gradient ascent procedure.
Since the effective noise is decreasing, the denoiser must be able to adjust to the amount of noise as the stochasticity of the gradient changes. Most denoisers such as the one used in prDeep, are trained to only handle a specific noise value. Instead, the algorithm of [6] uses a bias-free CNN denoiser. By removing all the additive bias terms of the convolutional layers, [17] discovered that the resulting denoiser would become robust to the amount of additive noise.
Finally, this Langevin sampling algorithm was used to solve linear inverse problems such as deblurring, superresolution, and inpainting. Updating the gradient step to handle the conditional density p(x|Ax = Ax * ) allows the algorithm to sample the denoiser in a way that also fits to linear measurements Ax = Ax * .

APPROACH
We aim to incorporate an implicit denoiser prior into an efficient algorithm to solve phase retrieval problems. The de-  noiser prior has learned universal properties of natural images, and, thus, could potentially provide strong generalization guarantees to different datasets. One issue, however, is that the approach of [6] only allowed for sampling from the posterior conditioned on linear measurements p(x|Ax = Ax * ), rendering an extension to nonlinear inverse problems challenging. Inspired by alternating projection approaches, we propose an alternating minimization-sampling procedure that 1) fixes a candidate phase p t and then 2) updates the signal by sampling from the posterior x t ∼ p(x|Ax = p t b). We first describe the Langevin sampling algorithm for sampling from the posterior of the image conditioned on the linear part of the measurements.

Langevin Sampling
As noted in [6], the residual of the denoiser D(x) :=x(y)−y is proportional to the gradient of log observation density. This result can be utilized to perform the following update steps to approximately sample from the conditional density p(x|Ax = x c ), as outlined in Algorithm 1. The main update step is where hd x is the gradient step with step size h and γz is injected noise. The quantity γ controls the proportion of injected noise z ∼ N (0, I) and is given by where β ∈ (0, 1] is a hyperparameter. σ 2 = d 2 2 N (N is the size of the image) adjusts the amount of injected noise in accordance to the size of the gradient step. The gradient step, d x , is and uses D, the residual of the denoiser, to approximately sample from the conditional distribution p(x|Ax = x c ). See Algorithm 1 for the full details. Note that the quantity h t = h0t 1+h0(t−1) acts as an accelerator for the algorithm by increasing the step size and decreasing the amount of injected noise over time.
Algorithm 1: Posterior Sampling via the Implicit Denoiser Prior Data: We use the Alternating Minimization framework to decouple the non-linearity of the phase retrieval problem to incorporate a denoiser prior with Langevin Sampling. Specifically, we perform an alternating minimization-sampling algorithm that first fixes the phase and then approximately samples from the posterior conditioned on the phase-induced measurements by utilizing the implicit denoiser prior. That is, for initial points p 0 and x 0 , we compute for t 1, Additionally, we enforce the gradient of the Langevin Sampling to account for the phase p t by providing it the measurements x c = p t b. We outline more details in Algorithm 2.

EXPERIMENTAL RESULTS
The classical phase retrieval problem is to recover true signal x * from measurements b = |F(x * )| + w. Poisson shot noise is a common noise model that appears in many phase retrieval applications [18]. We approximate this as where α is a parameter that increases with the noise level. We use 4× over-sampled Fourier measurements by placing an n × n image in the center of a 2n × 2n square and taking its 2D Fourier transform. The location of the image in this square is the support which we assume is known a priori. As in [5], for the initial input x 0 we use the output of the Hybrid Input-Output (HIO) algorithm [10]. We take 50 random initializations x 1 , x 2 , . . . x 50 and run HIO for 50 iterations. We take the x i with the lowest residual b − |F(x i )| 2 and run HIO for another 1000 iterations with the initial iterate x i . The result is used as the input x 0 for APLS. Our algorithm is initialized with h 0 = 0.1, β = 0.0001, and T 2 = 500. Empirically, we achieved best results when the number of iterations to run the Langevin Sampling algorithm between phase updates was just T 1 = 1.
We compare to prDeep which is the state-of-the-art using the regularization by denoising framework. We use 4 DnCNN denoisers trained on the Berkeley Segmentation Dataset with noise standard deviations of 50, 25, 15, and 5 to implement the pipeline used in the original paper.
We evaluate the reconstructions of HIO, prDeep, and our algorithm on a set of 6 natural grayscale images (Figure 1) and a set 6 unnatural grayscale images (Figure 2). We calculate the Peak Signal to Noise Ratio, PSNR = 10 log 10 ( 255 2 mean((x−x * )) ), to assess the reconstruction quality between reconstructed imagex and original image x * . Table  1 displays the average PSNR across both test sets. We see that our method is comparable on natural images and we showcase considerable improvement for unnatural images over all noise levels.

CONCLUSION
We proposed the Alternating Phase Langevin Sampling (APLS) algorithm for phase retrieval, which incorporates an implicit denoiser prior in an alternating minimization framework. We demonstrated competitive performance with prDeep, a strong baseline incorporating a denoiser as a regularizer for phase retrieval. For natural images, our performance was comparable while for unnatural, out-ofdistribution images, we showcased superior performance. This work showcases the benefits of extracting the prior in a denoiser more directly. There are many interesting avenues of future work, including applying our approach to phase retrieval problems in more scientific domains, such as astronomical imaging, and theoretically analyzing the convergence and sample complexity of our method.


Evaluation of phenotype-driven gene prioritization methods for Mendelian diseases

Abstract Yuan et al. recently described an independent evaluation of several phenotype-driven gene prioritization methods for Mendelian disease on two separate, clinical datasets. Although they attempted to use default settings for each tool, we describe three key differences from those we currently recommend for our Exomiser and PhenIX tools. These influence how variant frequency, quality and predicted pathogenicity are used for filtering and prioritization. We propose that these differences account for much of the discrepancy in performance between that reported by them (15–26% diagnoses ranked top by Exomiser) and previously published reports by us and others (72–77%). On a set of 161 singleton samples, we show using these settings increases performance from 34% to 72% and suggest a reassessment of Exomiser and PhenIX on their datasets using these would show a similar uplift.

We were pleased to see the recent publication in this journal on phenotype-driven gene prioritization methods for Mendelian disease [1]. This evaluation of the methods on two real clinical datasets (N = 305 and 209) will be welcomed by the clinical genetics community, especially as it is performed by an independent group that have not been involved in developing any of the assessed tools. However, we were surprised by the reported figures for our Exomiser tool, e.g. 15% and 26% of diagnoses identified as the top ranking hits in their two datasets compared with 38-58% for AMELIE, xRare and LIRICAL. As reported recently, we have used Exomiser to identify 77% of diagnoses as the top-ranking candidate for the 100 000 Genomes Project and this has led to its adoption in the Genomic Medicine Service of the UK National Health Service [2]. Similarly, for a retinal disease dataset 74% of diagnoses were ranked first by Exomiser [3] and a completely independent group from the Los Angeles Children Hospital reported a performance of 72% [4].
The authors state that default settings were used for each assessed software, but a closer inspection of the settings used by them when assessing Exomiser 12.1.0 revealed three key differences from those we routinely use. Here, we investigated whether these differences are likely to account for the above discrepancies. On a set of 161 diagnosed singleton cases from the 100 000 Genomes project (100KGP), using the same Exomiser version (12.1.0) they tested, we observe that 34% are identified as the top candidate with their settings versus 72% with ours. Incrementally reintroducing our settings to theirs highlights the differences these make: (1) Including mode of inheritance (MOI) specific frequency filtering (0.1% for dominant or homozygous recessive modes, 2% for compound-heterozygous, recessive) increases performance to 44%. This is the recommended default in the version of Exomiser 12.1.0 they assessed and is closer to the minor allele frequency (MAF) filtering settings used for the higher performing tools in their hands such as LIRICAL and AMELIE. Note this filter, as used in the default recommendation, runs every possible MOI using appropriate MAF settings rather than restricting to preselected modes of inheritance that may not be known ahead of analysis.
(2) Including the failedFrequencyFilter further increases performance to 62%. This setting filters out any low-quality variants that are not f lagged as PASS in the FILTER column of the VCF file. The settings Yuan et al. [1] used for LIRICAL and AMELIE did include filtering of non-PASS variants and this likely explains much of the improved performance they reported for these tools. Although this filter was documented in the example settings for Exomiser 12.1.0 it was not explicitly stated as a default recommendation. Since release 13 of Exomiser (September 2021) this is explicitly specified in the preset options. (3) Using REVEL and MVP as more modern sources of predicted pathogenicity data than Polyphen2, Muta-tionTaster and SIFT further increases performance to 72%. Again, although REVEL and MVP were available for the version of Exomiser tested (12.1.0), they were not clearly f lagged as default settings but since release 13 they are now in the recommended preset options. Figure 1 summarizes these results alongside the previously reported performance of Exomiser on clinical datasets.
In conclusion, we recommend all users to use the preset options provided in Exomiser, or if they are altering configuration to make sure that the above three options are retained for optimal performance: e.g. the provided example VCF file (Pfeiffer.vcf.gz) and clinical data (pfeiffer-phenopacket.yml) will run with these recommended settings if simply run as java -jar exomiser-cli-13.0.1.jar-sample examples/ pfeiffer-phenopacket.yml-vcf examples/Pfeiffer.vcf.gzassembly hg19. PhenIX, another tool they assessed, runs as part of the Exomiser framework and will similarly benefit from using these recommended settings. We suspect that most, if not all, of the differences in performance the authors observe between xRare, AMELIE and LIRICAL and Exomiser/PhenIX are due to these variant filtering settings rather than the underlying gene prioritization algorithms they set out to evaluate. We would be very interested to see if this is the case in a future assessment of their datasets.

Key Points
• Variant filtering settings, as well as phenotype similarity approaches, are critical for the performance of phenotype-driven gene prioritization approaches. • Mode of inheritance specific frequency and variant quality filtering alongside use of REVEL and MVP are the latest recommended options for Exomiser and PhenIX. • Use of these latest recommended settings, versus those used in the study of Yuan et al., increased detection of the diagnosis as the top-ranked candidate from 34% to 72% in our hands.