# Neural Congealing: Aligning Images to a Joint Semantic Atlas

Dolev Ofri-Amar1 Michal Geyer1 Yoni Kasten2 Tali Dekel1
1Weizmann Institute of Science 2NVIDIA Research
Project webpage: https://neural-congealing.github.io/

###### Abstract

We present Neural Congealing – a zero-shot self-supervised framework for
detecting and jointly aligning semantically-common content across a given set
of images. Our approach harnesses the power of pre-trained DINO-ViT features
to learn: (i) a joint semantic atlas – a 2D grid that captures the mode of
DINO-ViT features in the input set, and (ii) dense mappings from the unified
atlas to each of the input images. We derive a new robust self-supervised
framework that optimizes the atlas representation and mappings per image set,
requiring only a few real-world images as input without any additional input
information (e.g., segmentation masks). Notably, we design our losses and
training paradigm to account only for the shared content under severe
variations in appearance, pose, background clutter or other distracting
objects. We demonstrate results on a plethora of challenging 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 large-scale training data is
scarce (e.g., coffee mugs). We thoroughly evaluate our method and show that
our test-time optimization approach performs favorably compared to a state-of-
the-art method that requires extensive training on large-scale datasets.

Figure 1: Given a set of input images, our method automatically detects and
jointly aligns semantically-common content across the images. This is achieved
through a test-time training approach that estimates a unified 2D atlas that
represents the common semantic content, and dense mappings from the joint
atlas to each of the input images. Our atlas and mappings are optimized per
input set in a self-supervised manner by leveraging a pre-trained DINO-ViT
model. Our method can be applied to diverse image sets, without requiring any
additional training data, and allows us to automatically propagate an edit
applied to a single image across the entire set.

##  1 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 [58, 38], browsing images through canonical
primitives, and 3D reconstruction (e.g., [49, 8]).

The task of joint image alignment dates back to the seminal congealing [22,
14, 30, 13], which aligns a set of images into a common 2D space. Recently,
GANgealing [38] 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 [36]. 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 [4]. DINO-ViT features have been shown to serve as an
effective visual descriptor, capturing localized and semantic information
(e.g., [51, 2]). 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 [38] 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.

##  2 Related Work

##### Joint Image Set Alignment.

Congealing [30, 14, 22] 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 [13], which showed
improvement in face verification accuracy. Further, various methods have been
proposed to generalize congealing to several modes, e.g., through clustering
[11, 10, 28, 37, 23, 53], and make it more robust to occlusions [39]. Other
methods are based on pairwise optical flow, requiring either consistent
matching across image pairs in a collection [57, 44], or factorizing the
collection into simpler subspaces, allowing to simplify the matching task [20,
35]. AverageExplorer [58] presented a user interactive framework for browsing
Internet photo collections through average images representing modes in the
collections.

Recently, GANgealing [38] used a Spatial Transformer Network (STN) [16] 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 [17] 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 [56], 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.

##### Semantic Correspondences.

Prior to the deep-learning era, various methods have been proposed to tackle
the task of establishing sparse point correspondences between an image pair
[25, 50]. 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., [9, 48, 1]), or by training a
model end-to-end for establishing semantic correspondences (e.g., [42, 43]).
Many of them use groundtruth supervision for training [5], while others aim to
tackle the task in a weakly-supervised [42, 43] or entirely unsupervised
fashion [2]. 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 [4].

##### DINO-ViT Features as Local Semantic Descriptors.

Recent works showed the power of ViT (Vision Transformer) features as local
and global semantic descriptors, specifically features of a pre-trained DINO-
ViT [4]. Several works [4, 47, 55] showed the use of these features for
various applications such as instance segmentation, object discovery or
transfer learning on downstream tasks. GCD [52] categorize all unlabeled
images given only a partially labeled dataset. Other works used the pre-
trained features for both object localization and segmentation [29], and even
object part discovery and segmentation [7]. STEGO [12] 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 [2] 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 [51] 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 [2, 51] and use DINO-ViT’s spatial features as
dense descriptors for aligning images from semantically-related categories.

##  3 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 [4, 2,
51].

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 [2]. This approach allows us to
robustly align only the common regions in highly cluttered scenes.

Figure 2: _Neural Congealing._ Each of the input images $I_{i}$ in the set (a)
is fed into a pre-trained and fixed DINO-ViT model. (b) Spatial features,
$K_{i}$, are extracted and processed to estimate initial rough saliency masks
$S_{i}$ (see [2]), both of which are used to train: (i) a mapping network
(STN), which (c) transforms each of the input images into a joint 2D space,
and (ii) (d) a joint learnable 2D feature atlas, $K_{\mathcal{A}}$, and an
atlas saliency map $S_{\mathcal{A}}$. Note that while the input saliency masks
are noisy and contain objects from the background, our method is able to
refine them and focus on the most common object in the set. See SM for the
full set of images.

###  3.1 Semantic Joint Image Alignment

Given an input image set $\\{I_{i}\\}_{i=1}^{N}$ and a pre-trained (and fixed)
DINO-ViT model, we extract for each image its keys features from the last
layer, denoted by $K_{i}\in\mathbb{R}^{H\times W\times D}$; an initial per-
image saliency mask $S_{i}$ is estimated by applying a simple clustering-based
method directly to the extracted features (see [2]). These maps roughly
capture salient foreground regions, but are often noisy and contain uncommon
regions across the images, as seen in Fig. 4.

We define a learnable atlas $\mathcal{A}$ as a 2D grid of latent features
$K_{\mathcal{A}}\in\mathbb{R}^{H_{\mathcal{A}}\times W_{\mathcal{A}}\times
D}$, and a saliency mask $S_{\mathcal{A}}\in\mathbb{R}^{H_{\mathcal{A}}\times
W_{\mathcal{A}}}$. We define the 2D mapping of each atlas point
$\mathbf{x}_{\mathcal{A}}$ to each of the input images $I_{i}$ as follows:

$\mathbf{x}_{i}=\mathcal{M}(I_{i},\mathbf{x}_{\mathcal{A}})$ (1)

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 [38], $\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:

$\mathcal{M}_{r}(I_{i},\mathbf{x})=sR\mathbf{x}+\mathbf{t}$ (2)

where $R\in SO(2)$, $\mathbf{t}\in\mathbb{R}^{2}$ and $s\in\mathbb{R}^{+}$ are
2D rotation, translation and global scale, respectively. Our method also
supports horizontal flips (see Appendix A.4.1).

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

$\mathcal{M}_{f}(I,\mathbf{x})=\mathbf{x}+\mathbf{w}$ (3)

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.

###  3.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:

$\mathcal{L}=\mathcal{L}_{keys}+\lambda_{s}\mathcal{L}_{saliency}+\lambda_{r}\mathcal{L}_{reg_{\mathcal{M}}}+\lambda_{a}\mathcal{L}_{reg_{\mathcal{A}}}$
(4)

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

##### Semantic loss ${\bf\mathcal{L}_{keys}}$.

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:

$\displaystyle\mathcal{L}_{keys}=\sum_{i=1}^{N}\frac{1}{N\cdot\Sigma_{S_{\mathcal{A}}}}\sum_{\mathbf{x}_{\mathcal{A}}}S_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\cdot\mathcal{D}\left(K_{i}(\mathbf{x}_{i}),K_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\right)$
(5)

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}}$.

##### Saliency loss ${\bf\mathcal{L}_{saliency}}$.

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:

$\displaystyle\mathcal{L}_{saliency}=\frac{1}{N\cdot
N_{\mathcal{A}}}\sum_{i=1}^{N}\sum_{\mathbf{x}_{\mathcal{A}}}\rho_{0.7}\left(S_{i}(\mathbf{x}_{i}),S_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\right)$
(6)

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

$\rho_{\delta}(a,b)=\begin{cases}\begin{array}[]{lr}\frac{1}{2}(a-b)^{2},&\text{if
}|a-b|<\delta\\\
\delta\cdot\left(|a-b|-\frac{1}{2}\delta\right)&\text{otherwise
}\end{array}\end{cases}$ (7)

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

##### Transformation regularization ${\bf\mathcal{L}_{reg_{\mathcal{M}}}}$.

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:

$\mathcal{L}_{reg_{\mathcal{M}}}=\lambda_{s_{1}}\mathcal{L}_{scale}+\lambda_{s_{2}}\mathcal{L}_{mag}+\mathcal{L}_{smooth}$
(8)

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:

$\mathcal{L}_{scale}=\frac{1}{N}\sum_{i=1}^{N}|1-s_{i}|^{2}$ (9)

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:

$\mathcal{L}_{mag}=\frac{1}{N\cdot
N_{\mathcal{A}}}\sum_{i=1}^{N}\sum_{x_{\mathcal{A}}}\left\lVert\mathbf{w}_{i}\right\rVert_{2}^{2}$
(10)

where $\mathbf{w}$ is the per-pixel flow vector defined in Eq. (3). The term
$\mathcal{L}_{smooth}$, defined as in [19], 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:

$\mathcal{L}_{smooth}=\frac{1}{N\cdot
N_{\mathcal{A}}}\sum_{i=1}^{N}\sum_{x_{\mathcal{A}}}\left(\left\lVert
J^{T}J\right\rVert_{F}+\left\lVert\left(J^{T}J\right)^{-1}\right\rVert_{F}\right)$
(11)

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

Figure 3:  _Neural Congealing results._ We show results on a variety of image
sets, each containing images with large variations in pose, appearance, scale,
different domains, and cluttered backgrounds. We show (a) the original images
together with the congealed images, and (b) the average image in atlas space
together with the joint atlas saliency mask. See SM for results on all full
sets.

##### Atlas regularization ${\bf\mathcal{L}_{reg_{\mathcal{A}}}}$.

This term controls the localization and sparsity of the atlas:

$\mathcal{L}_{reg_{\mathcal{A}}}=\mathcal{L}_{center}+\lambda_{p}\mathcal{L}_{sparsity}$
(12)

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:

$\mathcal{L}_{center}=\left\lVert\frac{1}{\Sigma_{S_{\mathcal{A}}}}\cdot\sum_{\mathbf{x}_{\mathcal{A}}}S_{\mathcal{A}}(\mathbf{x}_{\mathcal{A}})\cdot\mathbf{x}_{\mathcal{A}}\right\rVert_{2}^{2}$
(13)

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:

$\mathcal{L}_{sparsity}=\mathcal{L}_{sparsity}^{S_{\mathcal{A}}}+\lambda_{k}\mathcal{L}_{sparsity}^{K_{\mathcal{A}}}$
(14)

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

$\mathcal{L}_{sparsity}^{S_{\mathcal{A}}}=\gamma\left\lVert
S_{\mathcal{A}}\right\rVert_{1}+\Psi_{0}(S_{\mathcal{A}})$ (15)

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:

$\mathcal{L}_{sparsity}^{K_{\mathcal{A}}}=(1-S_{\mathcal{A}})\cdot\left\lVert
K_{\mathcal{A}}\right\rVert_{1}$ (16)

###  3.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 [38], 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}$.

##  4 Results

We tested our method on a variety of image sets, containing 5-25 images, from
LSUN dataset [56], AFHQ dataset [6], Pixabay [40], Shutterstock [46], 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.

###  4.1 Qualitative Results.

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 [38] that was
trained on naturally-looking images, and thus struggles to generalize to other
domains.

Figure 4: _Refined saliency masks (overlay)._ Initial saliency masks extracted
using [2], together with our final learned atlas saliency mask. Our method is
able to refine the masks capturing the common object, while ignoring cluttered
background content or irrelevant objects. Figure 5: _Editing results._ Sample
results for editing for two different sets. Edits are propagated to the same
semantic parts to all images. See SM for more results. Figure 6: _Across-
domain set._ We compare results to GANgealing [38]. While they struggle with
such images, our method is able to semantically align them.

##### Edit results.

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.

##### Refined masks.

As discussed in Sec. 3.1, the initial saliency masks extracted from DINO-ViT
features using [2] 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.

Figure 7: _Comparison to GANgealing[38]._ We compare results from sets used in
evaluation (Sec. 4.2). Since our method is based on alignment according to
semantic DINO-ViT features, even in hard examples, our method focuses on
aligning the most common salient object part, e.g., the birds in (c). Note
that as in [38], our method supports horizontal flips (see Appendix A.4.1 for
technical details).

###  4.2 Quantitative Evaluation and Comparison.

We evaluate our framework on the task of semantic point correspondences on
SPair-71K [33] and CUB-200-2011 [54]. 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 [38] and set
$h,w$ to be the image size.

##### Results on SPair-71K.

We use the same pre-processing as in [38], 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 [38]; 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.

Method | Supervision | Cat | Dog | Bicycle
---|---|---|---|---
HPF [32] | supervised | 52.9 | 32.8 | 18.9
DHPF [34] | supervised | 61.6 | 46.1 | 23.8
SCOT [24] | supervised | 63.1 | 42.5 | 20.7
CHM [31] | supervised | 64.9 | 56.1 | 29.3
CATs [5] | supervised | 66.5 | 56.5 | 34.7
WeakAlign [42] | weakly-supervised | 31.8 | 22.6 | 17.6
NC-Net [43] | weakly-supervised | 39.2 | 18.8 | 12.2
CNNgeo [41] | self-supervised | 32.7 | 22.8 | 16.7
A2Net [45] | self-supervised | 35.6 | 24.3 | 18.5
GANgealing [38] | GAN-supervised | 67.0 | 23.1 | 37.5
Ours | self-supervised |  54.5/70.7* | 35.8 | 29.1
Table 1: _PCK-Transfer on categories from SPair-71K with $\alpha=0.1$._
*Results for relaxed setting, aligning only the cats’ heads.

##### Results on CUB-200-2011.

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 [38]. 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).

###  4.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.

Method | Supervision |<EMAIL_ADDRESS>
---|---|---
Ours w/o $\mathcal{L}_{keys}$ | self-supervised | 30.7
Ours w/o saliency masks | self-supervised | 53.4
Ours w/o $\mathcal{L}_{reg_{\mathcal{M}}}$ | self-supervised | 36.9
Ours w/o $\mathcal{L}_{reg_{\mathcal{A}}}$ | self-supervised | 64.9
GANgealing [38] | GAN-supervised | 56.8
Ours | self-supervised | 63.6
Table 2: _PCK-Transfer on subsets of CUB-200-2011._ Comparison to [38],
together with ablation for our different loss terms.  Figure 8: Ablation of
each of the loss terms in our objective function. See Sec. 4.3.

##  5 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).

##  6 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.

Figure 9: _Limitations._ Sets containing symmetric objects under large
rotations may converge partially due to relative position between semantic
parts. In addition, our method is not designed to align images depicting more
than one instance of the shared mode.

##### Acknowledgements:

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

## References

  * [1] Kfir Aberman, Jing Liao, Mingyi Shi, Dani Lischinski, Baoquan Chen, and Daniel Cohen-Or.  Neural best-buddies: Sparse cross-domain correspondence.  ACM Transactions on Graphics (TOG), 2018.
  * [2] Shir Amir, Yossi Gandelsman, Shai Bagon, and Tali Dekel.  Deep vit features as dense visual descriptors.  ECCVW What is Motion For?, 2022.
  * [3] Omer Bar-Tal, Dolev Ofri-Amar, Rafail Fridman, Yoni Kasten, and Tali Dekel.  Text2live: Text-driven layered image and video editing.  In Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XV, pages 707–723. Springer, 2022.
  * [4] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.  Emerging properties in self-supervised vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650–9660, 2021.
  * [5] Seokju Cho, Sunghwan Hong, Sangryul Jeon, Yunsung Lee, Kwanghoon Sohn, and Seungryong Kim.  Cats: Cost aggregation transformers for visual correspondence.  Advances in Neural Information Processing Systems, 2021.
  * [6] Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha.  Stargan v2: Diverse image synthesis for multiple domains.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2020.
  * [7] Subhabrata Choudhury, Iro Laina, Christian Rupprecht, and Andrea Vedaldi.  Unsupervised part discovery from contrastive reconstruction.  Advances in Neural Information Processing Systems, 2021.
  * [8] Gabriele Facciolo, Carlo De Franchis, and Enric Meinhardt-Llopis.  Automatic 3d reconstruction from multi-date satellite images.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pages 57–66, 2017.
  * [9] Philipp Fischer, Alexey Dosovitskiy, and Thomas Brox.  Descriptor matching with convolutional neural networks: a comparison to sift.  arXiv preprint arXiv:1405.5769, 2014.
  * [10] Brendan J Frey and Nebojsa Jojic.  Estimating mixture models of images and inferring spatial transformations using the em algorithm.  In Proceedings. 1999 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (Cat. No PR00149). IEEE, 1999.
  * [11] Brendan J. Frey and Nebojsa Jojic.  Transformation-invariant clustering using the em algorithm.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2003\.
  * [12] Mark Hamilton, Zhoutong Zhang, Bharath Hariharan, Noah Snavely, and William T Freeman.  Unsupervised semantic segmentation by distilling feature correspondences.  arXiv preprint arXiv:2203.08414, 2022.
  * [13] Gary Huang, Marwan Mattar, Honglak Lee, and Erik Learned-Miller.  Learning to align from scratch.  Advances in neural information processing systems, 25, 2012.
  * [14] Gary B Huang, Vidit Jain, and Erik Learned-Miller.  Unsupervised joint alignment of complex images.  In 2007 IEEE 11th international conference on computer vision, pages 1–8. IEEE, 2007.
  * [15] Peter J. Huber.  Robust Estimation of a Location Parameter.  The Annals of Mathematical Statistics, 35(1):73 – 101, 1964.
  * [16] Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al.  Spatial transformer networks.  Advances in neural information processing systems, 28, 2015.
  * [17] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila.  Analyzing and improving the image quality of stylegan.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8110–8119, 2020.
  * [18] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila.  Analyzing and improving the image quality of stylegan.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
  * [19] Yoni Kasten, Dolev Ofri, Oliver Wang, and Tali Dekel.  Layered neural atlases for consistent video editing.  ACM Transactions on Graphics (TOG), 40(6):1–12, 2021.
  * [20] Ira Kemelmacher-Shlizerman and Steven M Seitz.  Collection flow.  In 2012 IEEE Conference on Computer Vision and Pattern Recognition, pages 1792–1799. IEEE, 2012.
  * [21] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [22] Erik G Learned-Miller.  Data driven image models through continuous joint alignment.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 28(2):236–250, 2005.
  * [23] Xiaoming Liu, Yan Tong, and Frederick W Wheeler.  Simultaneous alignment and clustering for an image ensemble.  In 2009 IEEE 12th International Conference on Computer Vision. IEEE, 2009.
  * [24] Yanbin Liu, Linchao Zhu, Makoto Yamada, and Yi Yang.  Semantic correspondence as an optimal transport problem.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020.
  * [25] David G Lowe.  Distinctive image features from scale-invariant keypoints.  International journal of computer vision, 60(2):91–110, 2004.
  * [26] Erika Lu, Forrester Cole, Tali Dekel, Weidi Xie, Andrew Zisserman, David Salesin, William T. Freeman, and Michael Rubinstein.  Layered neural rendering for retiming people in video.  ACM Trans. Graph., 2020.
  * [27] Erika Lu, Forrester Cole, Tali Dekel, Andrew Zisserman, William T. Freeman, and Michael Rubinstein.  Omnimatte: Associating objects and their effects in video.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
  * [28] Marwan A Mattar, Allen R Hanson, and Erik G Learned-Miller.  Unsupervised joint alignment and clustering using bayesian nonparametrics.  arXiv preprint arXiv:1210.4892, 2012.
  * [29] Luke Melas-Kyriazi, Christian Rupprecht, Iro Laina, and Andrea Vedaldi.  Deep spectral methods: A surprisingly strong baseline for unsupervised semantic segmentation and localization.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
  * [30] Erik G Miller, Nicholas E Matsakis, and Paul A Viola.  Learning from one example through shared densities on transforms.  In Proceedings IEEE Conference on Computer Vision and Pattern Recognition. CVPR 2000 (Cat. No. PR00662). IEEE, 2000.
  * [31] Juhong Min and Minsu Cho.  Convolutional hough matching networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
  * [32] Juhong Min, Jongmin Lee, Jean Ponce, and Minsu Cho.  Hyperpixel flow: Semantic correspondence with multi-layer neural features.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019.
  * [33] Juhong Min, Jongmin Lee, Jean Ponce, and Minsu Cho.  Spair-71k: A large-scale benchmark for semantic correspondence.  arXiv prepreint arXiv:1908.10543, 2019.
  * [34] Juhong Min, Jongmin Lee, Jean Ponce, and Minsu Cho.  Learning to compose hypercolumns for visual correspondence.  In European Conference on Computer Vision. Springer, 2020.
  * [35] Hossein Mobahi, Ce Liu, and William T Freeman.  A compositional model for low-dimensional image set representation.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1322–1329, 2014.
  * [36] Ron Mokady, Omer Tov, Michal Yarom, Oran Lang, Inbar Mosseri, Tali Dekel, Daniel Cohen-Or, and Michal Irani.  Self-distilled stylegan: Towards generation from internet photos.  In ACM SIGGRAPH 2022 Conference Proceedings, pages 1–9, 2022.
  * [37] Tom Monnier, Thibault Groueix, and Mathieu Aubry.  Deep transformation-invariant clustering.  Advances in Neural Information Processing Systems, 2020.
  * [38] William Peebles, Jun-Yan Zhu, Richard Zhang, Antonio Torralba, Alexei A Efros, and Eli Shechtman.  Gan-supervised dense visual alignment.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 13470–13481, 2022.
  * [39] Yigang Peng, Arvind Ganesh, John Wright, Wenli Xu, and Yi Ma.  Rasl: Robust alignment by sparse and low-rank decomposition for linearly correlated images.  IEEE transactions on pattern analysis and machine intelligence, 34(11):2233–2246, 2012.
  * [40] Pixabay.  https://pixabay.com/.
  * [41] Ignacio Rocco, Relja Arandjelovic, and Josef Sivic.  Convolutional neural network architecture for geometric matching.  In Proceedings of the IEEE conference on computer vision and pattern recognition, 2017.
  * [42] Ignacio Rocco, Relja Arandjelović, and Josef Sivic.  End-to-end weakly-supervised semantic alignment.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6917–6925, 2018.
  * [43] Ignacio Rocco, Mircea Cimpoi, Relja Arandjelović, Akihiko Torii, Tomas Pajdla, and Josef Sivic.  Neighbourhood consensus networks.  Advances in neural information processing systems, 2018.
  * [44] Michael Rubinstein, Armand Joulin, Johannes Kopf, and Ce Liu.  Unsupervised joint object discovery and segmentation in internet images.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1939–1946, 2013.
  * [45] Paul Hongsuck Seo, Jongmin Lee, Deunsol Jung, Bohyung Han, and Minsu Cho.  Attentive semantic alignment with offset-aware correlation kernels.  In Proceedings of the European Conference on Computer Vision (ECCV), 2018.
  * [46] Shutterstock.  https://www.shutterstock.com/.
  * [47] Oriane Siméoni, Gilles Puy, Huy V Vo, Simon Roburin, Spyros Gidaris, Andrei Bursuc, Patrick Pérez, Renaud Marlet, and Jean Ponce.  Localizing objects with self-supervised transformers and no labels.  arXiv preprint arXiv:2109.14279, 2021.
  * [48] Edgar Simo-Serra, Eduard Trulls, Luis Ferraz, Iasonas Kokkinos, Pascal Fua, and Francesc Moreno-Noguer.  Discriminative learning of deep convolutional feature point descriptors.  In Proceedings of the IEEE international conference on computer vision, pages 118–126, 2015.
  * [49] Xiao Tang, Tianyu Wang, and Chi-Wing Fu.  Towards accurate alignment in real-time 3d hand-mesh reconstruction.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 11698–11707, October 2021.
  * [50] Engin Tola, Vincent Lepetit, and Pascal Fua.  Daisy: An efficient dense descriptor applied to wide-baseline stereo.  IEEE transactions on pattern analysis and machine intelligence, 32(5):815–830, 2009.
  * [51] Narek Tumanyan, Omer Bar-Tal, Shai Bagon, and Tali Dekel.  Splicing vit features for semantic appearance transfer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10748–10757, 2022.
  * [52] Sagar Vaze, Kai Han, Andrea Vedaldi, and Andrew Zisserman.  Generalized category discovery.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
  * [53] Andrea Vedaldi, Gregorio Guidi, and Stefano Soatto.  Joint data alignment up to (lossy) transformations.  In 2008 IEEE Conference on Computer Vision and Pattern Recognition. IEEE, 2008.
  * [54] C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie.  .  Technical Report CNS-TR-2011-001, California Institute of Technology, 2011\.
  * [55] Yangtao Wang, Xi Shen, Shell Xu Hu, Yuan Yuan, James L Crowley, and Dominique Vaufreydaz.  Self-supervised transformers for unsupervised object discovery using normalized cut.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
  * [56] Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop.  arXiv preprint arXiv:1506.03365, 2015.
  * [57] Tinghui Zhou, Yong Jae Lee, Stella X Yu, and Alyosha A Efros.  Flowweb: Joint image set alignment by weaving consistent, pixel-wise correspondences.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1191–1200, 2015.
  * [58] Jun-Yan Zhu, Yong Jae Lee, and Alexei A Efros.  Averageexplorer: Interactive exploration and alignment of visual data collections.  ACM Transactions on Graphics (TOG), 33(4):1–11, 2014.

##  Appendix A Implementation Details

###  A.1 General

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 [2]. 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 [38]. Both architectures are
based on the design of the ResNet-based discriminator from StyleGAN2 [18].
Tables 3 and 4 detail the layers of the rigid and non-rigid STN respectively,
and 5 and 6 detail the building blocks.

##### Rigid STN.

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:

$\displaystyle\theta=\pi\cdot\tanh(o_{1})\quad,\quad
R=\begin{bmatrix}\cos\theta&-\sin\theta\\\ \sin\theta&\cos\theta\end{bmatrix}$
(17) $\displaystyle s=\exp(o_{2})$ (18)
$\displaystyle\mathbf{t}=\begin{bmatrix}o_{3}\\\ o_{4}\end{bmatrix}$ (19)

##### Non-Rigid STN.

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.

block | layer | output size
---|---|---
0 | bilinear downsample (using conv2d) | $3\times 128\times 128$
1 | convL(64, 1, 1, 0, fusedLeakyReLU) | $64\times 128\times 128$
2 | ResBlock((64, 3, 1, 1), (128, 3, 2, 0), (128, 1, 2, 0)) | $128\times 64\times 64$
3 | ResBlock((128, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 32\times 32$
4 | ResBlock((512, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 16\times 16$
5 | ResBlock((512, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 8\times 8$
6 | ResBlock((512, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 4\times 4$
7 | convL(512, 1, 2, 0, fusedLeakyReLU) |  $512\times 4\times 4$ (flattened)
8 | linear + fusedLeakyReLU | $1\times 512$
9 | linear | $1\times 4$
Table 3: _Architecture of rigid STN._ block | layer | output size
---|---|---
0 | convL(64, 1, 1, 0, fusedLeakyReLU) | $64\times 128\times 128$
1 | ResBlock((64, 3, 1, 1), (128, 3, 2, 0), (128, 1, 2, 0)) | $128\times 64\times 64$
2 | ResBlock((128, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 32\times 32$
3 | ResBlock((512, 3, 1, 1), (512, 3, 2, 0), (512, 1, 2, 0)) | $512\times 16\times 16$
4 | ResBlock((512, 3, 1, 1), (512, 3, 1, 1), (512, 1, 1, 0)) | $512\times 16\times 16$
5 | convL(512, 3, 1, 1, fusedLeakyReLU) | $512\times 16\times 16$
6 | conv2d(512, 3, 1, 1) + ReLU + conv2d(2, 3, 1, 1) |  $2\times 16\times 16$ (coarse flow)
7 | conv2d(512, 3, 1, 1) + ReLU + conv2d(576, 3, 1, 1) |  $576\times 16\times 16$ (upsampling weights)
Table 4: _Architecture of non-rigid STN._ ResBlock((channels1, kernel1,
stride1, padding1), (channels2, kernel2, stride2, padding2), (channels3,
kernel3, stride3, padding3))
---
order | type | layer
0 | conv1 | convL(channels1, kernel1, stride1, padding1, fusedLeakyReLU)
1 | conv2 | convL(channels2, kernel2, stride2, padding2, blur, fusedLeakyReLU)
2 | skip | convL(channels3, kernel3, stride3, padding3, blur)
Table 5: _Architecture of a ResBlock._ convL(channels, kernel, stride,
padding, blur, fusedLeakyReLU)
---
order | layer
0 | blur (upfirdn2d)
1 | conv2d(channels, kernel, stride, padding)
2 | fusedLeakyReLU
Table 6: _Architecture of a convL layer._

###  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).

##### Rigidity Loss $\mathcal{L}_{smooth}$.

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

$\mathcal{L}_{smooth}=\frac{1}{N\cdot
N_{\mathcal{A}}}\sum_{i=1}^{N}\sum_{x_{\mathcal{A}}}\left(\left\lVert
J^{T}J\right\rVert_{F}+\left\lVert\left(J^{T}J\right)^{-1}\right\rVert_{F}\right)$
(20)

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:

$J=\frac{1}{\Delta}\begin{bmatrix}\mathbf{j}_{1}&\mathbf{j}_{2}\end{bmatrix}\in\mathbb{R}^{2\times
2}$ (21)

where

$\displaystyle\mathbf{j}_{1}=\mathcal{M}\left(I_{i},\mathbf{x}_{\mathcal{A}}+\Delta\cdot\begin{bmatrix}1\\\
0\end{bmatrix}\right)-\mathcal{M}(I_{i},\mathbf{x}_{\mathcal{A}}),$ (22)
$\displaystyle\mathbf{j}_{2}=\mathcal{M}\left(I_{i},\mathbf{x}_{\mathcal{A}}+\Delta\cdot\begin{bmatrix}0\\\
1\end{bmatrix}\right)-\mathcal{M}(I_{i},\mathbf{x}_{\mathcal{A}})$ (23)

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
[21] 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.

##### Loss coefficients.

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

$\mathcal{L}=\mathcal{L}_{keys}+\lambda_{s}\mathcal{L}_{saliency}+\lambda_{r}\mathcal{L}_{reg_{\mathcal{M}}}+\lambda_{a}\mathcal{L}_{reg_{\mathcal{A}}}$
(24)

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

$\mathcal{L}_{keys}=\lambda_{l}L_{2}+\mathcal{D}_{cos}$ (25)

with $\lambda_{l}=0.875$.

$\mathcal{L}_{reg_{\mathcal{M}}}=\lambda_{s_{1}}\mathcal{L}_{scale}+\lambda_{s_{2}}\mathcal{L}_{mag}+\mathcal{L}_{smooth}$
(26)

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:

$\mathcal{L}_{smooth}=\mathcal{L}_{smooth}^{\Delta=1}+\lambda_{s_{3}}\mathcal{L}_{smooth}^{\Delta=20}$
(27)

with $\lambda_{s_{3}}=3.5$.

$\mathcal{L}_{reg_{\mathcal{A}}}=\mathcal{L}_{center}+\lambda_{p}\mathcal{L}_{sparsity}$
(28)

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:

$\mathcal{L}_{sparsity}=\mathcal{L}_{sparsity}^{S_{\mathcal{A}}}+\lambda_{k}\mathcal{L}_{sparsity}^{K_{\mathcal{A}}}$
(29)

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 [38], 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.

Figure 10: _Results without ${\bf\mathcal{L}_{reg_{\mathcal{A}}}}$._ For each
set, we show the original images, and the congealed images with and without
${\bf\mathcal{L}_{reg_{\mathcal{A}}}}$; on the right are the average image in
atlas space and the atlas saliency. This regularization encourages cleaner
atlases, e.g., ignoring background clutter in _Guitars_ , and allows us to
better capture the dominant shared content, e.g., focusing only on the face in
_Art Cats_ allows us to better align the third cat from the right.

##  Appendix B Point Correspondence Between A Pair of Images

As in [38], 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 [38], 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).