Neural Congealing: Aligning Images to a Joint Semantic Atlas

Introduction

Humans can easily associate and match semantically-related objects across images, even under severe variations in appearance, pose and background content. For example, by observing the images in Fig. 1, we can immediately focus and visually compare the different butterflies, while ignoring the rest of the irrelevant content. While computational methods for establishing semantic correspondences have seen a significant progress in recent years, research efforts are largely focused on either estimating sparse matching across multiple images (e.g., keypoint detection), or establishing dense correspondences between a pair of images. In this paper, we consider the task of joint dense semantic alignment of multiple images. Solving this long-standing task is useful for a variety of applications, ranging from editing image collections , browsing images through canonical primitives, and 3D reconstruction (e.g., ).

The task of joint image alignment dates back to the seminal congealing , which aligns a set of images into a common 2D space. Recently, GANgealing has modernized this approach for congealing an entire domain of images. This is achieved by leveraging a pre-trained GAN to generate images that serve as self-supervisory signal. Specifically, their method jointly learns both the mode of the generated images in the latent space of the GAN, and a network that predicts the mappings of the images into the joint mode. GANgealing demonstrated impressive results on in-the-wild image sets. Nevertheless, their method requires a StyleGAN model pre-trained on the domain of the test images, e.g., aligning cat images requires training StyleGAN on a large-scale cat dataset. This is a challenging task by itself, especially for unstructured image domains or uncurated datasets . Furthermore, they require additional extensive training for learning the mode of the generated images and their mapping network (e.g., training on millions of generated images). In this work, we take a different route and tackle the joint alignment task in the challenging setting where only a test image set is available, without any additional training data. More specifically, given only a few images as input (e.g., $<$25 images), our method estimates the mode of the test set and their joint dense alignment, in a self-supervised manner.

We assume that the input images share a common semantic content, yet may depict various factors of variations, such as pose, appearance, background content or other distracting objects (e.g., Mugs in Fig. 3). We take inspiration from the tremendous progress in representation learning, and leverage a pre-trained DINO-ViT – a Vision Transformer model trained in a self-supervised manner . DINO-ViT features have been shown to serve as an effective visual descriptor, capturing localized and semantic information (e.g., ). Here, we propose a new self-supervised framework that jointly and densely aligns the images in DINO-ViT features space. To the best of our knowledge, we are the first to harness the power of DINO-ViT for dense correspondences between in-the-wild images. More specifically, given an input image set, our framework estimates, at test-time: (i) a joint latent 2D atlas that represents the mode of DINO-ViT features across the images, and (ii) dense mappings from the atlas to each of the images. Our training objective is driven by a matching loss encouraging each image features to match the canonical learned features in the joint atlas. We further incorporate additional loss terms that allow our framework to robustly represent and align only the shared content in the presence of background clutter or other distracting objects.

Since our atlas and mappings are optimized per set, our method works in a zero-shot manner and can be applied to a plethora of image sets, including sets of mixed domains (e.g., aligning images depicting sculpture and artwork of cats), sets depicting related yet different object categories (e.g., dogs and tigers), or domains for which a dedicated generator is not available (e.g., coffee mugs). We thoroughly evaluate our method, and demonstrate that our test-time optimization framework performs favorably compared to and on-par with state-of-the-art self-supervised methods. We further demonstrate how our atlas and mappings can be used for editing the image set with minimal effort by automatically propagating edits that are applied to a single image to the entire image set.

Related Work

Congealing introduced the task of jointly aligning images into a common 2D image, representing the geometric mode of the set. This has been done by minimizing the entropy of intensity values in a pixel stack after the alignment. They have demonstrated the use of their method for several applications, including image classification given a few labels per class. These seminal works were extended by incorporating deep learning, combining unsupervised alignment with unsupervised feature learning , which showed improvement in face verification accuracy. Further, various methods have been proposed to generalize congealing to several modes, e.g., through clustering , and make it more robust to occlusions . Other methods are based on pairwise optical flow, requiring either consistent matching across image pairs in a collection , or factorizing the collection into simpler subspaces, allowing to simplify the matching task . AverageExplorer presented a user interactive framework for browsing Internet photo collections through average images representing modes in the collections.

Recently, GANgealing used a Spatial Transformer Network (STN) to predict a transformation from any image sampled from a predefined domain (e.g. cat images) into a shared aligned space. They leverage the style-pose disentanglement in a pre-trained StyleGAN2 model to provide training supervision. Specifically, during training, their method simultaneously learns the mode of object pose across a large collection of images, and trains the STN to map each image to this mode. Their method demonstrates impressive results on complex data such as LSUN , yet requires extensive compute and large-scale training data. In contrast, we take the congealing task to the realm of test-time optimization, where only a small test image set is available (e.g., $<$25 images). Thus, our method can be applied on diverse domains, or image sets that comprise of images from mixed yet related domains, as shown in Fig. 1, 3, 6 and 7.

Prior to the deep-learning era, various methods have been proposed to tackle the task of establishing sparse point correspondences between an image pair . However, due to their local nature, and the lack of global context, they cannot handle significant color and shape variations. Later, data-driven based descriptors opened the door to establishing correspondences based on higher level information by learning from data representations that encode semantic global information; such methods either work by extracting features from a pre-trained classification model (e.g., ), or by training a model end-to-end for establishing semantic correspondences (e.g., ). Many of them use groundtruth supervision for training , while others aim to tackle the task in a weakly-supervised or entirely unsupervised fashion . Our method also aims at learning semantic correspondences, however, our focus is on aligning multiple images jointly by leveraging descriptors extracted from a pre-trained DINO-ViT model .

Recent works showed the power of ViT (Vision Transformer) features as local and global semantic descriptors, specifically features of a pre-trained DINO-ViT . Several works showed the use of these features for various applications such as instance segmentation, object discovery or transfer learning on downstream tasks. GCD categorize all unlabeled images given only a partially labeled dataset. Other works used the pre-trained features for both object localization and segmentation , and even object part discovery and segmentation . STEGO showed a method for unsupervised semantic segmentation, which unlike most works that learn a feature per pixel, they use contrastive loss that produces low rank representation for the DINO features.

It was shown by that the deep features of a pre-trained DINO-ViT encode semantic information at fine spatial granularity and capture semantic object parts. Furthermore, they showed that this information is shared across domains of semantically-related categories, e.g. cats and dogs, allowing them to design dense descriptors which are used for various applications of co-segmentation, part co-segmentation and point correspondences between images. Splice established a method for semantic appearance transfer using DINO-ViT features, and presented additional powerful properties such as keys inversion, which shows the amount of details they hold of the original image.

We build upon the findings from and use DINO-ViT’s spatial features as dense descriptors for aligning images from semantically-related categories.

Neural Atlas Congealing

The input to our method is a collection of images that share a common semantic content, e.g., a set of images depicting different types of guitars in natural scenes. The shared content may significantly differ across the images in appearance, structure, pose, and appear in cluttered scenes containing complex backgrounds or other distracting objects (e.g., Fig. 3). Our goal is to automatically detect the common content across the images and estimate a geometric transformation that maps each image into a joint 2D space. Our key idea is to harness the power of deep features extracted from a pre-trained (and fixed) DINO-ViT model, which have been shown to capture localized semantic information under significant appearance and pose variations .

Specifically, our framework, illustrated in Fig. 2, jointly aligns the images in DINO-ViT feature space using two learnable components: (i) a unified latent 2D atlas that represents the common semantic mode in DINO-ViT space across the images, and (ii) a Spatial Transformer Network (STN) that aligns each of the input images to the joint latent atlas. Our method is fully automatic and self-supervised: each of the input images is first fed into DINO-ViT and features (keys) are extracted from the last layer, and serve as our spatial semantic descriptors. The atlas representations, in which each pixel stores a latent feature, and the STN parameters are then optimized such that the transformed features of each image are aligned with the joint atlas features.

We further define a saliency value in each atlas pixel that takes a continuous value between zero and one. We optimize the atlas saliency using a voting-based loss w.r.t. rough initial image saliency masks that are estimated from DINO-ViT features in a pre-processing step . This approach allows us to robustly align only the common regions in highly cluttered scenes.

where $\mathbf{x}_{i}$ is the estimated corresponding point of $\mathbf{x}_{\mathcal{A}}$ in image $I_{i}$. Note that applying $\mathcal{M}$ on each of the atlas coordinates allows us to backward warp $I_{i},S_{i}$ or $K_{i}$ into the atlas space.

Similarly to , $\mathcal{M}$ is modeled as a composition of rigid and non-rigid transformations $\mathcal{M}_{r}\circ\mathcal{M}_{f}$, each is estimated by a separate STN. That is, $\mathcal{M}_{r}$ is a global 2D similarity transformation defined by:

Given an image $I$, the non-rigid transformation $\mathcal{M}_{f}$ is defined by a dense flow field:

where $\mathbf{w}$ is a per-pixel 2D offset.

In practice, each of the transformations is obtained by an STN model: one takes $I_{i}$ as input and predicts the parameters of the similarity transformation $\mathcal{M}_{r}$; the image $I_{i}$ is then backward warped using $\mathcal{M}_{r}$ and fed to the second STN that predicts $\mathcal{M}_{f}$. See network architecture details in Appendix A.2.

2 Training

Given $\{I_{i},S_{i},K_{i}\}_{i=1}^{N}$, we now turn to the task of learning the joint atlas $\mathcal{A}=(K_{\mathcal{A}},S_{\mathcal{A}})$ and 2D mappings $\mathcal{M}=(\mathcal{M}_{r},\mathcal{M}_{f})$. Our objective function incorporates four main loss terms and takes the following form:

where $\lambda_{s},\lambda_{a}$ and $\lambda_{r}$ control the relative weights between the terms, and are fixed throughout training.

The semantic loss is our driving loss and it encourages an alignment in DINO-ViT feature space between the atlas and each of the congealed images:

where $\Sigma_{S_{\mathcal{A}}}=\sum_{\mathbf{x}_{\mathcal{A}}}S_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})$, $K_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})$ is the atlas feature at location $\mathbf{x}_{\mathcal{A}}$, and $K_{i}(\mathbf{x}_{i})$ is the corresponding feature in $I_{i}$ (see Eq. 1). The distance metric $\mathcal{D}$ is defined by: $\lambda_{l}\left\lVert K_{i}(\mathbf{x}_{i})-K_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\right\rVert_{2}^{2}+\mathcal{D}_{cos}\left(K_{i}(\mathbf{x}_{i}),K_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\right)$, where $\mathcal{D}_{cos}$ is the cosine distance. Note that since we are interested in aligning only the semantically common content across the images, our loss is weighted according to the atlas saliency $S_{\mathcal{A}}$.

This term serves as a saliency-voting loss that allows us to capture the common content across the images by the atlas saliency $S_{\mathcal{A}}$. Formally, the atlas saliency is optimized to match the initial congealed images’ saliency masks. Since the image saliency masks are often rough and contain clutter or irrelevant salient objects of the scene (Fig. 2), we use a robust loss:

where $N_{\mathcal{A}}$ is the number of pixels in the atlas, and $\rho_{\delta}(a,b)$ denotes the Huber loss with parameter $\delta$:

Intuitively, each image “votes” for the regions that should be salient in the atlas, and the aggregated common salient content is estimated.

For obtaining a shared representation that is as undistorted as possible, while containing some distortions for aligning objects with different proportions of semantic parts, we apply regularization on both mapping networks:

where $\lambda_{s_{1}}$ and $\lambda_{s_{2}}$ are the relative weights.

$\mathcal{L}_{scale}$ regularizes $\mathcal{M}_{r}$ from changing the scale of the original images in the atlas space:

where $s_{i}$ is the scale parameter of the learned rigid transformation for image $I_{i}$.

The non-rigid transformation is encouraged to be as small as possible:

where $\mathbf{w}$ is the per-pixel flow vector defined in Eq. (3). The term $\mathcal{L}_{smooth}$, defined as in , is used to prevent the non-rigid mapping from distorting the shared content by encouraging as rigid as possible mapping. Formally, this term is defined by:

where $J$ is the Jacobian matrix of $\mathcal{M}$ at $\mathbf{x}_{\mathcal{A}}$. See Appendix A.3 for more details.

This term controls the localization and sparsity of the atlas:

where $\lambda_{p}$ is the relative weight.

Since the position of the shared content in the atlas space is arbitrary, we define $\mathcal{L}_{center}$ to encourage the common object to be mapped to the center of the atlas space:

where $\mathbf{x}_{\mathcal{A}}$ coordinates are normalized to be in the range $(-1,1)\times(-1,1)$. By minimizing the norm of the saliency’s center of mass, we encourage it to be as close as possible to the atlas center that is located at the origin.

We further observe that without any sparsity regularization, the atlas often contains non-common information. $\mathcal{L}_{sparsity}$ encourages both $S_{\mathcal{A}}$ and $K_{\mathcal{A}}$ to be sparse:

where $\lambda_{k}$ is the relative weight. We follow and define the sparsity loss term for the atlas saliency as a combination of L1- and L0-approximation regularization terms

where $\Psi_{0}(x)\equiv 2Sigmoid(5x)-1$ is a smooth L0 approximation that penalizes non zero elements, and $\gamma$ is the relative weight between the terms.

For the atlas features, we apply L1 sparsity loss on non-salient parts only:

3 Editing

Once we have the atlas representation, we can use the average image of all congealed images as a template for editing. Then, the edit in the atlas space is automatically propagated back to all original images. One can also apply an edit on one of the images and propagate it to all the rest by passing through the atlas space.

As in , given an image $I_{i}$ and an RGBA edit image in atlas space we apply forward warping using $\mathcal{M}(I_{i},\mathbf{x}_{\mathcal{A}})$ to the image space and then apply alpha blending of the warped edit image with $I_{i}$.

Results

We tested our method on a variety of image sets, containing 5-25 images, from LSUN dataset , AFHQ dataset , Pixabay , Shutterstock , and the general Internet. Each set contains objects that share the same semantic parts from different categories including painted/animated/real animals, mugs, guitars, etc. The domain, pose/orientation and appearance, as well as the amount of irrelevant salient objects in the background, may significantly change between the images in each set. We pre-process the images to $256\times 256$ px, with border padding in case of an unsuited aspect-ratio. See Appendix A for full implementation details.

Sample examples for input sets along with our joint alignment result can be seen in Fig. 1, 2, 3, 6 and 7. Fig. 3 also shows visualizations of the average image in the atlas space and the atlas saliency. The full set of results is included in the Supplementary Materials (SM). As can be seen, our method successfully aligns diverse in-the-wild sets under significant differences in object scale (e.g. Mugs in Fig. 3), proportions between semantic parts (e.g., Garfield’s ears in Mix Cats, or the Corgi’s ears in Mix Animals in Fig. 3), slight differences in out-of-plane rotation (e.g., white tiger in Mix Animals in Fig. 3), and under non-rigid deformations (e.g., the butterflies in Fig. 1).

In addition, Fig. 6 shows an example where our method aligns images across different domains (paintings, food, etc.). This demonstrates the flexibility of our test-time training approach compared to GANgealing that was trained on naturally-looking images, and thus struggles to generalize to other domains.

Fig. 1 and 5 show sample editing results automatically applied to the input set (Sec. 3.3). As can be seen, the edits are mapped correctly and accurately to the same semantic regions in all images, under significant variation in scale, pose and appearance. More editing results are included in the SM.

As discussed in Sec. 3.1, the initial saliency masks extracted from DINO-ViT features using are typically very coarse and may contain irrelevant content such as other objects. Fig. 2 and 4 show examples of how our method manages to congeal these rough estimates into an accurate and refined mask that captures the shared content, while robustly filtering out cluttered background content or non-shared objects. The full set of saliency masks is included in the SM.

2 Quantitative Evaluation and Comparison.

We evaluate our framework on the task of semantic point correspondences on SPair-71K and CUB-200-2011 . Specifically, given a source image $I_{A}$ and a target image $I_{B}$ together with their ground truth point correspondences, we transfer the points from image $I_{A}$ to the atlas space and map them back to image $I_{B}$ to obtain its predicted points (see Appendix B for technical details). We then measure for each set the PCK-Transfer, i.e., the percentage of keypoints that are mapped within the threshold of $\alpha\cdot\max(h,w)$ from the ground truth. We follow previous works and set $\alpha=0.1$ for both benchmarks, and $h,w$ to be the dimensions of the object’s bounding box for SPair-71K. For CUB, we follow and set $h,w$ to be the image size.

We use the same pre-processing as in , applying border padding for non-square images and resizing to $256\times 256$. We apply our method on each test set separately, each includes 25-26 images. See Appendix A.4.1 for further technical details.

Table 1 reports the results for our method, GANgealing, and a number of leading methods for semantic correspondences. As seen, our method outperforms GANgealing on most sets, and outperforms other self-supervised methods on all sets. Our method performs very well on the Cat and Dog sets, yet in the Bicycle set, due to the large deformations and the symmetric shape of the object, our performance decreases. Nevertheless, even in this challenging case our performance is on-par with most supervised methods on this set. Note that all supervised methods have been directly trained or fine-tuned using ground-truth supervision on the SPair-71K training set.

Figure 7 shows qualitative comparisons to ; GANgealing converges to a global mode across a large dataset, which allows them to achieve good alignment for Bicycle and Birds (Fig. 7(c),(d)). However, in the Dog set this restricts their alignment to capturing only the head, while our method detects the common mode in a given set and can align the full body of the dogs (Fig. 7(a)).

There is an inherent tradeoff between aligning highly articulated content (e.g., the bodies of the non-rigid animals) and maintaining undistorted atlas representation. We demonstrate this tradeoff for the Cat set, by controlling the effective relative weight of our rigidity loss. As seen, with our default parameters, the method aligns the bodies of the cats yet their faces are not accurately aligned. By increasing the relative weight of the rigidity, we can encourage the model to focus on the most rigid part across the set, while disregarding the cats’ body (even though it is a shared salient part across the set). In this setting, fine facial details are accurately aligned and our method outperforms all previous methods, including supervised methods.

Since our method works with small image sets, we randomly sample 14 sets of 25 images each, and train separately on each of them. For fair comparison, we apply the same pre-processing as in . As seen in Table 2, our method achieves better results compared to GANgealing. As seen in Fig. 7(c), GANgealing tends to hallucinate object parts, and struggles with aligning the heads, especially when the object pose differs significantly from the canonical pose learned from the entire domain. Our method, by optimizing the representation and mappings per set, manages to align the heads of the birds even under unusual poses, e.g., distant bird spreading wings (second column from the left).

3 Ablation Study

We ablate the different loss terms of our objective function (Eq. 4), both quantitatively in Table 2 and qualitatively in Fig. 8. Without our driving loss $\mathcal{L}_{keys}$, we notice a significant drop in performance. Fig. 8 shows that even though the saliency masks help a great deal in bringing the birds one on top of the other, there is no semantic alignment between them. Without saliency masks (no atlas saliency), our framework attempts to align all observed content, and thus struggles to converge, or converges only to small parts of the object, in cases of significant background clutter.

Removing $\mathcal{L}_{reg_{\mathcal{M}}}$ provides too much freedom to the non-rigid mapping, which converges either to a single point in the atlas, or spreads in disorder. Thus, the performance drops dramatically and the visual results are not appealing. Our method performs on-par without $\mathcal{L}_{reg_{\mathcal{A}}}$, however, we note that the initial saliencies for the CUB subsets are quite accurate, thus there is no much need for regularizing the atlas in this case. More generally, this loss allows us to obtain cleaner atlases (see Appendix C). Finally, we see in Fig. 8 that $\mathcal{L}_{center}$ of $\mathcal{L}_{reg_{\mathcal{A}}}$ encourages the shared content to be centered, allowing us to keep the birds within atlas borders.

Limitations

Our method relies on semantic similarity in the space of DINO-ViT feature space. Hence, in cases where these features do not capture the semantic association across the images, our method would not work well (e.g., image domains that are not well-represented in DINO’s training data). In addition, in cases of extreme topological changes in the common object across the set, our method struggles to converge to a good alignment due to the strong rigidity constraints, e.g., Fig. 7(b) and Fig. 7(c). Furthermore, we notice that in sets containing symmetric objects with large rotation differences, the relative position between parts may affect the convergence, and may lead to partial alignment, e.g., Fig. 9(a), where the left eye of the leftmost cat is matched to the right eye of the rest of the cats. In general, our framework is not designed to align images depicting more than one instance of the shared mode. In this case, our method may align arbitrarily one of the objects, and in other cases may fail to converge, e.g., Fig. 9(b).

Conclusions

We tackled the congealing task in a particularity challenging setting – jointly aligning a small set of in-the-wild images, without any additional training data other than the test set itself. We showed how to leverage the power of pre-trained DINO-ViT features for this task in a new test-time training framework. We demonstrated the key advantages of our approach w.r.t. existing state-of-the-art methods in its applicability to diverse image domains, lightweight training and overall performance. We further showed that our method can be used for automatically propagating edits to the entire set by simply editing a single image. We believe that our approach – combining test-time optimization with semantic information learned by external large-scale models – holds great promise for dense alignment tasks, and can motivate future research in this direction.

We thank Shai Bagon and Shir Amir for their insightful comments. This project received funding from the Israeli Science Foundation (grant 2303/20).

References

Appendix A Implementation Details

We use images of size $256\times 256$ (using border padding for non-square image), which is the resolution the STN takes as input. The rigid STN resizes the images to the resolution of the atlas, which is set to be $W_{\mathcal{A}}=H_{\mathcal{A}}=128$. For feature extraction, we use dino_vits8 ($D=384$) with stride 4, as in . We extract the keys from the original images and bilinearly upsample them to the atlas resolution.

A.2 Spatial Transformer Architecture Details

We use the same architectures for the STNs as in . Both architectures are based on the design of the ResNet-based discriminator from StyleGAN2 . Tables 3 and 4 detail the layers of the rigid and non-rigid STN respectively, and 5 and 6 detail the building blocks.

The rigid mapping network consists of a ResNet backbone with a fully-connected layer at the end, which outputs four logits $(o_{1},o_{2},o_{3},o_{4})$, to which the following activations are applied to obtain the transformation parametrization:

The non rigid mapping network consists of a ResNet backbone that outputs a $16\times 16$ feature grid which is then fed to two small convolutional networks: the first outputs a $16\times 16$ coarse flow field, and the second outputs weights which are used to perform $\times 8$ upsampling of the coarse flow field to the size of $128\times 128=H_{\mathcal{A}}\times W_{\mathcal{A}}$. The final flow is bilinearly upsampled in case of applying backward warp on inputs of resolution higher than $128$.

When composing both networks, the affine matrix given by the rigid STN is applied to the non-rigid flow which results in the final sampling grid used to congeal the original input image, DINO-ViT features and saliency mask.

A.3 Loss Terms

All losses except for $\mathcal{L}_{sparsity}$ and $\mathcal{L}_{scale}$ are applied in the atlas space within the boundaries of the backward warped images. Formally, for $\mathcal{L}_{keys}$, $\mathcal{L}_{saliency}$, $\mathcal{L}_{mag}$, $\mathcal{L}_{smooth}$ and $\mathcal{L}_{center}$ the sum in the atlas space is taken over the indices $\{\mathbf{x}_{\mathcal{A}}\,\,|\,\,\mathcal{M}(I_{i},\mathbf{x}_{\mathcal{A}})\in I_{i}\}$.

We detach the atlas saliency $S_{\mathcal{A}}$ for losses which should not have an impact on the joint saliency, which are $\mathcal{L}_{keys}$ and local $\mathcal{L}_{smooth}$ (see details next).

Recall the term $\mathcal{L}_{smooth}$, defined as in , is formally defined by:

where $N$ is the number of images, $N_{\mathcal{A}}$ is the number of pixels in the atlas, and $J$ is the Jacobian matrix of $\mathcal{M}$ at $\mathbf{x}_{\mathcal{A}}$. The term is used to prevent the non-rigid mapping from distorting the shared content by encouraging as rigid as possible mapping. The Jacobian matrix is defined by:

and $\Delta$ corresponds to the offset in pixels. This encourages both singular values of $J$ to be 1, which is what is required for a rigid mapping. We apply both local and global constraints, with $\Delta=1$ and $\Delta=20$ respectively.

In practice, similarly to the other losses, we apply the local rigidity loss only within the salient parts of the atlas (similarly to Eq. (5)).

A.4 Training

We train the atlas and the STNs jointly. We first bootstrap the rigid STN for 1000 epochs, and then train both the rigid and non-rigid components on their own (separate objective function), when only the non-rigid network affects the atlas training. We train for a total of 8000 epochs, and use Adam optimizer with a learning rate of $1\cdot 10^{-4}$ for the STNs and $8\cdot 10^{-4}$ for the atlas. Training on a set of 10 images on a Tesla V100-SXM2-32GB takes 1.2 hrs and uses 3.4GB of GPU memory, or 1.8 hrs and 5.3GB of GPU memory in case of including horizontal flips.

The loss coefficients we used for all experiments are as follows:

with $\lambda_{s}=1.25,\lambda_{r}=0.025,\lambda_{a}=0.75$.

with $\lambda_{s_{1}}=8,\lambda_{s_{2}}=80$. As mentioned in Appendix A.3, we apply $\mathcal{L}_{smooth}$ both locally and globally, namely:

with $\lambda_{p}=0.075$. As mentioned in Sec. 3.2, we apply $\lambda_{p}\mathcal{L}_{sparsity}$ both to the atlas saliency and the atlas keys, namely:

with $\lambda_{k}=0.044$. For $\mathcal{L}_{sparsity}^{S_{\mathcal{A}}}$ we set the relative weight between the L1- and L0-approximation terms to be $\gamma=2$.

The entire objective function is multiplied by a scalar $c\cdot\mathcal{L}$, where $c=4000$.

A.4.1 Congealing under extreme deformations.

Similarly to , we include an option of allowing horizontal flips in a given set, which is also used for the training of subsets of SPair 71K and CUB-200-2011 (Sec. 4.2). The flipping is done during training: we train the STNs with both the original images and the flipped images, and update the atlas only according to the orientation that currently has a lower semantic loss (keys loss).

Due to the extreme deformations present in the subsets of SPair 71K and CUB-200-2011, to increase robustness, we further reduce the local and global rigidity coefficients to be $\times 0.25$ its original value and the global rigidity to be $\lambda_{s_{3}}=0.9$. In addition, the atlas representation is gradually updated during training, i.e., the images used to update the atlas are added one-by-one where every 100 epochs the image with the lowest key loss is added. We observed that this training scheme is more stable and allows faster convergence for these sets. For the Bicycle set of SPair-71K, since many images contain only one wheel, we fix the atlas with the image that is most semantically similar to the average keys of the set and train the set with a fixed atlas.

Appendix B Point Correspondence Between A Pair of Images

As in , our method can find dense correspondences between a pair of images. For each image pair $\{I_{A},I_{B}\}$, we transfer the ground truth keypoints $\mathbf{k}_{A}\in I_{A}$ to $I_{B}$. This is done by mapping $\bf{k}_{A}$ to the atlas, obtaining $\mathbf{k}_{\mathcal{A}}$ then mapping it to $I_{B}$. Recall that our mapping $\mathcal{M}=\mathcal{M}_{r}\circ\mathcal{M}_{f}$ is defined from the atlas to each image. For mapping $\mathbf{k}_{A}$ to the atlas, we first compute the inverse of the rigid transformation, which has a closed-form solution (inverse of an affine matrix). Then, since there is no closed form for obtaining the inverse of $\mathcal{M}_{f}$, we follow , and approximate the inverse using nearest neighbors. Finally, we map $\mathbf{k}_{\mathcal{A}}$ to $I_{B}$ by bilinearly sampling the mapping grid of $I_{B}$.

Appendix C Ablation Study: No Atlas Regularization

As discussed in the main paper (Sec. 4.3), the sparsity regularization on the atlas assists our framework in capturing the dominant shared content, while ignoring noise and background clutter.

Sample cases can be seen in Fig. 10: for Guitars, some background content is captured by the atlas w/o this regularization. In Art Cats, the sparsity regularization allows us to only focus on the face, while ignoring unshared regions even if they are initially considered to be salient (cat’s body, third column from the right).