# Inference for Generative Capsule Models

Alfredo Nazabal The Alan Turing Institute, London, UK  Christopher K.I.
Williams  The Alan Turing Institute, London, UK  School of Informatics,
University of Edinburgh, UK

###### Abstract

Capsule networks (see e.g. Hinton et al., 2018) aim to encode knowledge and
reason about the relationship between an object and its parts. In this paper
we focus on a clean version of this problem, where data is generated from
multiple geometric objects (e.g. triangles, squares) at arbitrary
translations, rotations and scales, and the observed datapoints (parts) come
from the corners of all objects, without any labelling of the objects. We
specify a generative model for this data, and derive a variational algorithm
for inferring the transformation of each object and the assignments of points
to parts of the objects. Recent work by Kosiorek et al. [2019] has used
amortized inference via stacked capsule autoencoders (SCA) to tackle this
problem—our results show that we significantly outperform them. We also
investigate inference for this problem using a RANSAC-type algorithm.

##  1 Introduction

An attractive way to set up the problem of object recognition is
_hierarchically_ , where an object is described in terms of its parts, and
these parts are in turn composed of sub-parts, and so on. For example a face
can be described in terms of the eyes, nose, mouth, hair, etc.; and a teapot
can be described in terms of a body, handle, spout and lid parts. This
approach has a long history in computer vision, see e.g. Recognition-by-
Components by Biederman [1987]. Advantages of recognizing objects by first
recognizing its constituent parts include tolerance to the occlusion of some
parts, and that parts may vary less under a change of pose than the appearance
of the whole object.

One big advantage of a parts-based setup is that if the pose111i.e. the
location and rotation of the object in 2D or 3D. of the object changes, this
can have very complicated effects on the pixel intensities in an image, but
the geometric transformation between the object and the parts can be described
by a simple linear transformation (as used in computer graphics).
Convolutional neural networks can allow a recognition system handle 2D shifts
of an object in the image plane, but attempts to generalize such equivariances
across more general transformations lead to cumbersome constructions, and have
not been very successful.

Recent work by Hinton and collaborators [Sabour et al., 2017, Hinton et al.,
2018] has developed _capsule networks_. The key idea here is that a part in a
lower level can vote for the pose of an object in the higher level, and an
object’s presence is established by the agreement between votes for its pose.
Hinton et al. [2018, p. 1] use an iterative process called “routing-by-
agreement” which “updates the probability with which a part is assigned to a
whole based on the proximity of the vote coming from that part to the votes
coming from other parts that are assigned to that whole”. Subsequently
Kosiorek et al. [2019] framed inference for a capsule network in terms of an
autoencoder, the Stacked Capsule Autoencoder (SCA). Here, instead of the
iterative routing-by-agreement algorithm, a neural network $h^{\mathrm{caps}}$
takes as input the set of input parts and outputs predictions for the object
capsules’ instantiation parameters $\\{{\mathbf{y}}_{k}\\}_{k=1}^{K}$. Further
networks $h_{k}^{\mathrm{part}}$ are then used to predict part candidates from
each ${\mathbf{y}}_{k}$.

The objective function used in Hinton et al. [2018] (their eq. 4) is quite
complex (involving four separate terms), and is not derived from first
principles. In this paper we argue that the description in the paragraph above
is backwards—it is more natural to describe the generative process _by which
an object gives rise to its parts_ , and that the appropriate routing-by-
agreement inference algorithm then falls out naturally from this principled
formulation. Below we focus on a single-layer of part-object relations, as we
need to show first that this is working properly; but the generative model
should be readily extensible to deeper hierarchies.

The contributions of this paper are to:

  * •

Derive a variational inference algorithm for a generative model of object-part
relationships, including a relaxation of the permutation-matrix formulation
for matching object parts to observations;

  * •

Obtain a principled routing-by-agreement algorithm from this formulation;

  * •

Evaluate the performance of our model vs. competitors (including a RANSAC-type
algorithm, Fischler and Bolles 1981) on the inference task in scenes
consisting of 2D points.

##  2 Overview

Our images consist of a set of objects in different poses. Examples include
images of handwritten digits, faces, or geometric shapes in 2D or 3D. An
object can be defined as an instantiation of a specific model object (or
template) with a particular pose (or geometric transformation). Furthermore,
objects, and thus templates, are decomposed in parts, which are the basic
elements that comprise the objects. For example, a digit can be decomposed in
different strokes, or a geometric shape can be decomposed into vertices.

More formally, let $T=\\{T_{k}\\}_{k=1}^{K}$ be the set of $K$ templates that
are used to generate a scene. Each template
$T_{k}=\\{{\mathbf{p}}_{n}\\}_{n=1}^{N_{k}}$ is composed of $N_{k}$ parts
${\mathbf{p}}_{n}$. We assume that scenes can only be generated using the
available templates. Furthermore, every scene can present a different
configuration of objects, with some objects missing from some scenes. For
example, in scenes that could potentially contain all digits from 0 to 9 once,
if only the digits 2 and 5 are in the image, we consider that the other digits
are missing.

Each observed template in a scene is then transformed by an independent
transformation ${\mathbf{y}}_{k}$, different for each template, generating a
transformed object $X_{k}=\\{{\mathbf{x}}_{n}\\}_{n=1}^{N_{k}}$

$T_{k}\xrightarrow{{\mathbf{y}}_{k}}X_{k}.$ (1)

Assume that we are given a scene $X=\\{{\mathbf{x}}_{m}\\}_{m=1}^{M}$ composed
of $M$ observed parts. The _inference problem_ for $X$ involves different
tasks. We need to determine which objects from the set of templates were
employed to generate the scene222We will only consider fully missing objects
in this paper, but the methods described below are also applicable to
partially occluded objects with some parts missing.. Also, we need to infer
what transformation ${\mathbf{y}}_{k}$ was applied to each template to
generate the objects. This then allows us to infer the _correspondences_
between the template points and the scene points. If all the templates were
employed in the scene, then the number of observed parts $M$ is equal to the
sum of all the parts of all the templates $N=\sum_{k=1}^{K}N_{k}$.

Figure 1: Scenes composed of ${\mathbb{R}}^{2}$ points (upper figures) and
their corresponding objects matching (lower figures).

In this paper we particularly focus on a clean version of this problem
involving "constellations", where the data consists of 2D points derived from
the vertices of geometric shapes (e.g. squares and triangles), as used in
[Kosiorek et al., 2019, §2.1]. The transformations ${\mathbf{y}}_{k}$ are
restricted to linear similarity transformations that translate, rotate and
scale the templates to generate the objects of the scene. In general parts can
have appearance information that helps to distinguish them (e.g. for face
parts), but the constellation data is particularly hard as the datapoints are
just points in $\mathbb{R}^{2}$ with no appearance information. The extension
of the model to cover appearance features is described in supp. mat. A. Figure
1 shows several examples of shapes located at multiple positions with
different rotations and scaling factors. Given a set of unlabelled points
comprising a scene (upper figures), we seek to find the objects that were used
to generate such scene (lower figures).

##  3 Generative model for capsules

We propose a generative model to describe our problem. We will consider a
Gaussian mixture model approach. Given a set of observations (points) $X$ that
were generated by $K$ different components, we want to identify the points in
$X$ that belong to each component. Each component is a different object with
transformation ${\mathbf{y}}_{k}$, and the template model is used to generate
the data. Every observation in $X$ is a part that belongs uniquely to one of
the objects.

In this work we only use translation, rotation and scaling in 2D as possible
transformations for each object. Consider a template with parts
${\mathbf{p}}_{n}=(p_{nx},p_{ny})$ for $n=1,\ldots,N_{k}$ that we want to
scale by a factor $s$, rotate through with a clockwise rotation angle $\theta$
and translate by $(t_{x},t_{y})$, to obtain a transformed object with parts
${\mathbf{x}}_{n}=(x_{nx},x_{ny})$. For each part in the template, the
transformation works as follows:

$\begin{pmatrix}x_{nx}\\\ x_{ny}\end{pmatrix}=\begin{pmatrix}t_{x}\\\
t_{y}\end{pmatrix}+\begin{pmatrix}s\cos{\theta}&s\sin{\theta}\\\
-s\sin{\theta}&s\cos{\theta}\end{pmatrix}\begin{pmatrix}p_{nx}\\\
p_{ny}\end{pmatrix},$ (2)

which can also be expressed as

$\begin{pmatrix}x_{nx}\\\
x_{ny}\end{pmatrix}=\begin{pmatrix}1&0&p_{nx}&p_{ny}\\\
0&1&p_{ny}&-p_{nx}\end{pmatrix}\begin{pmatrix}t_{x}\\\ t_{y}\\\
s\cos{\theta}\\\ s\sin{\theta}\end{pmatrix}\stackrel{{\scriptstyle
def}}{{=}}F_{kn}{\mathbf{y}}_{k},$ (3)

where ${\mathbf{y}}_{k}$ is the ${\mathbb{R}}^{4}$ column vector and
$F_{kn}\in{\mathbb{R}}^{2\times 4}$ is the matrix to its left. From (3) we can
completely characterize a template $T_{k}$ by a set of matrices
$\\{F_{kn}\\}_{n=1}^{N_{k}}$ containing all the information of its parts.
Also, the information about the transformation applied to template $T_{k}$ is
fully contained in the ${\mathbb{R}}^{4}$ vector ${\mathbf{y}}_{k}$. Notice
that we can recover each of the terms of the transformation from the elements
of vector ${\mathbf{y}}_{k}$

$\displaystyle t_{x}=y_{k1},$ $\displaystyle t_{y}=y_{k2},$ (4) $\displaystyle
s=\sqrt{y_{k3}^{2}+y_{k4}^{2}},$
$\displaystyle\theta=\arctan{\frac{y_{k4}}{y_{k3}}}.$ (5)

Now, let $z_{mnk}\in\\{0,1\\}$ indicate whether the point ${\mathbf{x}}_{m}$
belongs to the part $n$ of object $k$. Every point $m$ belongs uniquely to a
tuple $(k,n)$, or in other words, a point ${\mathbf{x}}_{m}$ belongs uniquely
to the part defined by ${\mathbf{y}}_{k}$ acting on the template matrix
$F_{kn}$. Interestingly, the opposite is also partially true. Every tuple
$(k,n)$ belongs uniquely to a point $m$ or it is unassigned if object $k$ is
missing in the scene.

The joint distribution of the variables in the model is given by

$p(X,Y,Z)=p(X|Y,Z)p(Y)p(Z),$ (6)

where $p(X|Y,Z)$ is a Gaussian mixture model explaining how the points in a
scene were generated from the templates

$p(X|Y,Z)=\prod_{m=1}^{M}\prod_{k=1}^{K}\prod_{n=1}^{N_{k}}N\left({\mathbf{x}}_{m}|F_{kn}{\mathbf{y}}_{k},\lambda^{-1}I_{2}\right)^{z_{mnk}},$
(7)

with $\lambda$ being the precision of the Gaussian distribution for each point
and $I_{2}$ the identity matrix of dimension $2$. Note that such a model for a
_single object_ was developed by Revow et al. [1996].

The distribution over similarity transformations ${\mathbf{y}}_{k}$ is
modelled with a ${\mathbb{R}}^{4}$ Gaussian distribution with mean
${\boldsymbol{\mu}}_{0}$ and precision matrix $\Lambda_{0}$:

$p(Y)=\prod_{k=1}^{K}N({\mathbf{y}}_{k}|{\boldsymbol{\mu}}_{0},\Lambda_{0}^{-1}).$
(8)

Notice that modelling ${\mathbf{y}}_{k}$ with a Gaussian distribution implies
that we are modelling the translation $(t_{x},t_{y})$ in ${\mathbb{R}}^{2}$
with a Gaussian distribution.

If $\mu_{0}=\mathbf{0}$ and $\Lambda_{0}=I_{4}$ then
$s^{2}=y^{2}_{k3}+y^{2}_{k4}$ has a $\chi^{2}_{2}$ distribution, and $\theta$
is uniformly distributed in its range $[-\pi,\pi]$ by symmetry. For more
complex linear transformations (e.g. an affine transformation), we need only
to increase the dimension of ${\mathbf{y}}_{k}$ and change the form of
$F_{kn}$, but the generative model in (8) would remain the same.

In a standard Gaussian mixture model, the assignment matrix $Z$ is
characterized by a Categorical distribution, where each point
${\mathbf{x}}_{m}$ is assigned to one part

$p(Z)=\prod_{m=1}^{M}\text{Cat}({\mathbf{z}}_{m}|{\boldsymbol{\pi}}),$ (9)

with ${\mathbf{z}}_{m}$ being a 0/1 vector with only one 1, and
${\boldsymbol{\pi}}$ being the probability vector for each tuple $(k,n)$.
However, the optimal solution to our problem occurs when each part of a
template belongs uniquely to one point in a scene. This means that $Z$ should
be a permutation matrix, where each point $m$ is assigned to a tuple $(k,n)$
and vice versa. Notice that a permutation matrix is a square matrix, so if
$M\leq N$, we add dummy rows to $Z$, which are assigned to missing points in
the scene. The set of permutation matrices of dimension $N$ is a discrete set
containing $N!$ permutation matrices. They can be modelled with a discrete
prior over permutation matrices, assigning each matrix $Z_{i}$ a probability
$\pi_{i}$, such as:

$p_{perm}(Z)=\sum_{i=1}^{N!}\pi_{i}\;I[Z=Z_{i}]$ (10)

with $\sum_{i=1}^{N!}\pi_{i}=1$ and $I[Z=Z_{i}]$ the indicator function, being
equal to $1$ if $Z=Z_{i}$ and $0$ otherwise. The number of possible
permutation matrices increases as $N!$, which makes exact inference over
permutations intractable. An interesting property of $p_{perm}(Z)$ is that its
first moment ${\mathbb{E}}_{p_{perm}}[Z]$ is a doubly-stochastic (DS) matrix,
a matrix of elements in $[0,1]$ whose rows and columns sum to 1. We propose to
relax $p_{perm}(Z)$ to a distribution $p_{DS}(Z)$ that is characterized by the
doubly-stochastic matrix $A$ with elements $a_{mnk}$, such that
${\mathbb{E}}_{p_{DS}}[Z]=A$:

$p_{DS}(Z)=\prod_{m=1}^{N}\prod_{k=1}^{K}\prod_{n=1}^{N_{k}}a_{mnk}^{z_{mnk}}.$
(11)

$A$ is fully characterized by $(N-1)^{2}$ elements. We will use a uniform
prior over $Z$ with $a_{mnk}=\frac{1}{N}$. However, note that $p_{DS}$ can
also represent a particular permutation matrix $Z_{i}$ by setting the
appropriate entries of $A$ to 0 or 1, and indeed we would expect this to occur
during variational inference (see sec. 3.1) when the model converges to a
correct solution.

###  3.1 Variational Inference

Variational inference for the above model can be derived similarly to the
Gaussian Mixture model case [Bishop, 2006, Chapter 10]. The variational
distribution under the mean field assumption is given by

$q(Z,Y)=q(Z)q(Y),$ (12)

where the optimal solutions for $\log{q(Z)}$ and $\log{q(Y)}$ under the
generative model can be expressed as

$\displaystyle\log{q(Z)}$ $\displaystyle\propto$
$\displaystyle{\mathbb{E}}_{q(Y)}[\log{p(X,Y,Z)}]$ (13)
$\displaystyle\log{q(Y)}$ $\displaystyle\propto$
$\displaystyle{\mathbb{E}}_{q(Z)}[\log{p(X,Y,Z)}]$ (14)

For $q(Z)$ we obtain an expression with the same distribution model as the
prior in (11)

$\displaystyle
q(Z)\propto\prod_{m=1}^{N}\prod_{k=1}^{K}\prod_{n=1}^{N_{k}}\rho_{mnk}^{z_{mnk}},$
(15)

where $\rho_{mnk}$ represents the unnormalized probability of point $m$ being
assigned to tuple $(k,n)$ and vice versa. These unnormalized probabilities
have a different expression depending on whether we are considering a point
that appears in the scene ($m\leq M$)

$\displaystyle\log{\rho_{mnk}}$
$\displaystyle=\log{a_{mnk}}-\log{2\pi}+\log{\lambda}-$
$\displaystyle\frac{\lambda}{2}{\mathbb{E}}_{{\mathbf{y}}_{k}}[({\mathbf{x}}_{m}-F_{kn}{\mathbf{y}}_{k})^{T}({\mathbf{x}}_{mt}-F_{kn}{\mathbf{y}}_{k})],$
(16)

or whether we are considering a dummy row of the prior ($m>M$),

$\displaystyle\log{\rho_{mnk}}$ $\displaystyle=\log{a_{mnk}}.$ (17)

When a point is part of the scene (16), and thus $m\leq M$ the update of
$\rho_{mnk}$ is similar to the Gaussian mixture model case. However, if a
point is not part of the scene (17), and thus $m>M$ then the matrix is not
updated and the returned value is the prior $a_{mnk}$. This specific choice of
distribution for both $p_{DS}(Z)$ and $q(Z)$ leads naturally to the addition
of dummy rows that are employed to make the matrix $\rho$ square. Now, the
normalized distribution $q(Z)$ becomes:

$\displaystyle
q(Z)=\prod_{m=1}^{N}\prod_{k=1}^{K}\prod_{n=1}^{N_{k}}r_{mnk}^{z_{mnk}},$ (18)

where ${\mathbb{E}}_{q(Z)}[z_{mnk}]=r_{mnk}$. The elements $r_{mnk}$ of matrix
$R$ represent the posterior probability of each point $m$ being uniquely
assigned to the part-object tuple $(n,k)$ and vice-versa. This means that $R$
needs to be a DS matrix. This can be achieved by employing the Sinkhorn-Knopp
algorithm [Sinkhorn and Knopp, 1967], which updates a square non-negative
matrix by normalizing the rows and columns alternately, until the resulting
matrix becomes doubly stochastic (see supp. mat. B for more details). The use
of the Sinkhorn-Knopp algorithm for approximating matching problems has also
been described by Powell and Smith [2019] and Mena et al. [2020], but note
that in our case we also need to alternate with inference for $q(Y)$.
Furthermore, the optimal solution to the assignment problem occurs when $R$ is
a permutation matrix itself. When this happens we exactly recover a discrete
posterior (with the same form as (10)) over permutation matrices where one of
them has probability one, with the others being zero.

The distribution for $q(Y)$ is a Gaussian with

$\displaystyle q(Y)$
$\displaystyle=\prod_{k=1}^{K}N({\mathbf{y}}_{k}|{\boldsymbol{\mu}}_{k},\Lambda_{k}^{-1}),$
(19) $\displaystyle\Lambda_{k}$
$\displaystyle=\Lambda_{0}+\lambda\sum_{m=1}^{M}\sum_{n=1}^{N_{k}}r_{mnk}F_{kn}^{T}F_{kn},$
(20) $\displaystyle\mu_{k}$
$\displaystyle=\Lambda_{k}^{-1}(\Lambda_{0}{\boldsymbol{\mu}}_{0}+\lambda\sum_{m=1}^{M}\sum_{n=1}^{N_{k}}r_{mnk}F_{kn}^{T}{\mathbf{x}}_{m}),$
(21)

where the updates for both $\Lambda_{k}$ and ${\boldsymbol{\mu}}_{k}$ depend
explicitly on the templates employed in the model through the transformation
matrices $F_{kn}$ and the precision term $\lambda$. Note that the prediction
from datapoint $m$ to the mean of ${\mathbf{y}}_{k}$ is given by
$\sum_{n}r_{mnk}F^{T}_{kn}{\mathbf{x}}_{m}$, i.e. a weighted sum of the
predictions of each part $n$ with weights $r_{mnk}$. These expressions remain
the same when considering a Gaussian mixture prior such as (9).

Finally, the expectation term in (16) under this model becomes

$\displaystyle{\mathbb{E}}_{{\mathbf{y}}_{k}}[({\mathbf{x}}_{m}-F_{kn}{\mathbf{y}}_{k})^{T}({\mathbf{x}}_{m}-F_{kn}{\mathbf{y}}_{k})]=$
(22)
$\displaystyle({\mathbf{x}}_{m}-F_{kn}{\boldsymbol{\mu}}_{k})^{T}({\mathbf{x}}_{m}-F_{kn}{\boldsymbol{\mu}}_{k})+\text{trace}(F_{kn}^{T}F_{kn}\Lambda_{k}^{-1}).$

The evidence lower error bound (ELBO) $L(q)$ for this model is decomposed in
three terms:

$\displaystyle L(q)$ $\displaystyle={\mathbb{E}}_{q}[\log{p(X|Y,Z)}]$ (23)
$\displaystyle-KL(q(Y)||p(Y))-KL(q(Z)||p(Z)),$

with $KL(q||p)$ being the Kullback-Leibler divergence between distributions
$q$ and $p$. Expressions for these two terms are given in supp. mat. B. The
first term indicates how well the generative model $p(X|Y,Z)$ fits the
observations under our variational model $q(Y,Z)$:

$\displaystyle{\mathbb{E}}_{q}[\log{p(X|Y,Z)}]=M\left(-\log{2\pi}+\log{\lambda}\right)$
(24)
$\displaystyle-\frac{\lambda}{2}\sum_{m=1}^{M}\sum_{k=1}^{K}\sum_{n=1}^{N_{k}}r_{mnk}({\mathbf{x}}_{m}-F_{kn}{\boldsymbol{\mu}}_{k})^{T}({\mathbf{x}}_{m}-F_{kn}{\boldsymbol{\mu}}_{k})$
$\displaystyle-\frac{\lambda}{2}\sum_{m=1}^{M}\sum_{k=1}^{K}\sum_{n=1}^{N_{k}}r_{mnk}\text{trace}(F_{kn}^{T}F_{kn}\Lambda_{k}^{-1}).$

Finally, in Algorithm 1 we present the inference procedure for this model. The
alternating updates of $q(Y)$ and $q(Z)$ carry out routing-by-agreement. We
describe how the annealing of $\lambda$ (line 10) is implemented in Section 6.

Algorithm 1 Variational Inference

1:Initialize $\lambda$, $\lambda_{max}$ and $R\sim U[0,1]^{N\times N}$,
$\forall m,n,k$

2:$R=\text{SinkhornKnopp}(R)$

3:while not converged do

4: Update $\Lambda_{k}$ (20), $\forall k$

5: Update ${\boldsymbol{\mu}}_{k}$ (21), $\forall k$

6: Update $\log{\rho_{mnk}}$ (16)(17), $\forall m,n,k$

7: Update $R=\text{SinkhornKnopp}(\rho)$

8: if ELBO has converged then

9: if $\lambda<\lambda_{max}$ then

10: Anneal $\lambda$

11: else

12: converged = True return $R,\mu_{k},\Lambda_{k}$

###  3.2 Comparison with other objective functions

In Hinton et al. [2018] an objective function $cost^{h}_{k}$ is defined (their
eq. 1) which considers inference for the pose of a higher-level capsule $k$ on
pose dimension $h$. Translating to our notation, $cost^{h}_{k}$ combines the
votes from each datapoint $m$ as $cost^{h}_{k}=\sum_{m}r_{mk}\ln P^{h}_{m|k}$,
where $P^{h}_{m|k}$ is a Gaussian, and $r_{mk}$ is the “routing softmax
assignment” between $m$ and $k$. It is interesting to compare this with our
equation (21). Firstly, note that the vote of ${\mathbf{x}}_{m}$ to part $n$
in object $k$ is given explicitly by $F^{T}_{kn}{\mathbf{x}}_{m}$, i.e. we do
not require introduction of an independent voting mechanism, this falls out
directly from the inference. Secondly, note that our $R$ must keep track not
only of assignments to objects, but also to parts of the objects. In contrast
to Hinton et al. [2018], our inference scheme is derived from variational
inference on the generative model, rather than introducing an _ad hoc_
objective function.

The specialization of the SCA method of Kosiorek et al. [2019] to
constellation data is called the “constellation capsule autoencoder (CCAE)”
and discussed in their sec. 2.1. Under their equation 5, we have that

$\displaystyle
p({\mathbf{x}}_{1:M})=\prod_{m=1}^{M}\sum_{k=1}^{K}\sum_{n=1}^{N}\frac{a_{k}a_{k,n}}{\sum_{i}a_{i}\sum_{j}a_{ij}}N({\mathbf{x}}_{m}|\mu_{k,n},\lambda_{k,n}),$
(25)

where $a_{k}\in[0,1]$ is the presence probability of capsule $k$,
$a_{k,n}\in[0,1]$ is the conditional probability that a given candidate part
$n$ exists, and $\mu_{k,n}$ is the predicted location of part $k,n$. The
$a_{k}$s are predicted by the network $h^{\mathrm{caps}}$, while the
$a_{k,n}$s and $\mu_{k,n}$s are produced by separate networks
$h_{k}^{\mathrm{part}}$ for each part $k$.

We note that (25) provides an autoencoder style reconstructive likelihood for
${\mathbf{x}}_{1:M}$, as the $a$’s and $\mu$’s depend on the data. To handle
the arbitrary number of datapoints $M$, the network $h^{\mathrm{caps}}$
employs a Set Transformer architecture [Lee et al., 2019]. In comparison to
our iterative variational inference, the CCAE is a “one shot” inference
mechanism. This may be seen as an advantage, but in scenes such as Figure 1
with overlapping objects, humans may perform reasoning like “if that point is
one of the vertices of a square, then this other point needs to be explained
by a different object” etc, and it may be rather optimistic to believe this
can be done in a simple forward pass. Also, CCAE cannot exploit prior
knowledge of the geometry of the objects, as it relies on an opaque network
$h^{\mathrm{caps}}$ which requires extensive training.

###  3.3 Related work

Above we have discussed the work of Sabour et al. [2017], Hinton et al. [2018]
and Kosiorek et al. [2019]. There is some more recent work, for example Li et
al. [2020] present an undirected model, based on the capsule RBM of Li and Zhu
[2019]. Compared to our work, this is an undirected as opposed to directed
generative model. Also Smith et al. [2021] propose a layered directed
generative model, where the children of higher level capsules inherit its pose
(but transformed). Both Li et al. [2020] and Smith et al. [2021] consider
tree-structured relationships between parents and children, similar to Hinton
et al. [2000] and Storkey and Williams [2003]. However, most importantly,
these recent capsule models do not properly handle the fact that the input to
a capsule should be a _set_ of parts; instead in their work the first layer
capsules that interact with the input image model specific location-dependent
features/templates in the image, and their second layer capsules have
interactions with the specific first layer capsules (e.g. the fixed affinity
$\rho^{k}_{ij}$ of Smith et al. [2021] parent $i$ in layer $k$ for child $j$).
But if we consider a second-layer capsule that is. for example, detecting the
conjunction of the 3 strokes comprising a letter “A”, then at different
translations, scales and rotations of the A it will need to connect to
different image level stroke features, so these connection affinities must
depend on transformation variables of the parent.

##  4 A RANSAC approach to inference

A radical alternative to “routing by agreement” inference is to make use of a
“random sample consensus” approach (RANSAC, Fischler and Bolles, 1981). We can
derive a RANSAC approach to the constellations problem that follows the same
principles described in our generative model.

For any pair of points on one of the templates we can try matching these to
every possible pair of $M(M-1)$ points in the scene. The key insight is that a
pair of known points is sufficient to compute the 4-dimensional
${\mathbf{y}}_{k}$ vector in the case of similarity transformations. In
general the minimal number of points needed is known as the size of the basis;
for example a three-point basis would be required for a general affine
transform in 2D. Using the transformation $\hat{{\mathbf{y}}}_{k}$, we can
then _predict_ the location of the remaining parts of the template, and
_check_ if these are actually present. If so, this provides evidence for the
existence of $T_{k}$ and $\hat{{\mathbf{y}}}_{k}$. After considering the
$M(M-1)$ subsets, it then combines the identified instantiations to give an
overall explanation of the scene. A summary of the algorithm is given in
Algorithm 2.

In more detail, assume we have chosen parts $n_{1}$ and $n_{2}$ as the basis
for object $k$, and that we have selected datapoints ${\mathbf{x}}_{i}$ and
${\mathbf{x}}_{j}$ as their hypothesized counterparts in the image. Let
${\mathbf{x}}_{ij}$ be the vector obtained by stacking ${\mathbf{x}}_{i}$ and
${\mathbf{x}}_{j}$, and $B_{k}$ be the $4\times 4$ square matrix obtained by
stacking $F_{kn_{1}}$ and $F_{kn_{2}}$. Then
$\hat{{\mathbf{y}}}_{k}=B_{k}^{-1}{\mathbf{x}}_{ij}$. Finally,
$\text{SubsetMatch}(\hat{X}_{k},X)$ selects those points in $X_{k}$ that are
close to $X$ with a given tolerance and add them to the output. Among them,
the solution is given by the one that minimizes
$\sum_{n=1}^{N_{k}}(\hat{x}_{nk}-x_{nk})^{2}$.

The above algorithm chooses a specific basis for each object, but one can
consider all possible bases for each object. It is then efficient to use a
hash table to store the predictions for each part, as used in Geometric
Hashing [Wolfson and Rigoutsos, 1997]. Geometric Hashing is traditionally
employed in computer vision to match geometric features against previously
defined models of such features. This technique works well with partially
occluded objects, and is computationally efficient if the basis dimension is
low.

Algorithm 2 RANSAC approach

1:$T$: $K$ templates of the scene

2:$B_{k}$: base matrix for template $T_{k}$

3:$X$: $M$ points of the scene

4:$out=[]$

5:for ${\mathbf{x}}_{i}\in X$ do

6: for ${\mathbf{x}}_{j}\in X\setminus x_{i}$ do

7: ${\mathbf{x}}_{ij}=\text{Vectorize}({\mathbf{x}}_{i},{\mathbf{x}}_{j})$

8: for $k=1:K$ do

9: $\hat{{\mathbf{y}}}_{k}=B_{k}^{-1}{\mathbf{x}}_{ij}$

10: $T_{k}\xrightarrow{\hat{{\mathbf{y}}}_{k}}\hat{X}_{k}$

11: if SubsetMatch($\hat{X}_{k},X)$ then

12: Add $(T_{k},\hat{{\mathbf{y}}}_{k},\hat{X}_{k})$ to $out$ return $out$

##  5 Evaluation

In order to allow fair comparisons, we use the same dataset generator for
geometric shapes employed by Kosiorek et al. [2019]. We create a dataset of
scenes, where each scene consists of a set of 2D points, generated from
different geometric shapes. The possible geometric shapes (templates) are a
square and an isosceles triangle, with parts being represented by the 2D
coordinates of the vertices. We use the same dimensions for the templates as
used by Kosiorek et al. [2019], side 2 for the square, and base and height 2
for the triangle. All templates are centered at $(0,0)$. In every scene we can
have at most two squares and one triangle. Each shape is transformed with a
random transformation to create a scene of 2D points given by the object
parts. To match the evaluation of Kosiorek et al. [2019], all scenes are
normalized so as the points lie in $[-1,1]$ on both dimensions. While creating
the scene, we select randomly (with probability 0.5) whether an object is
going to be present or not, but delete empty scenes. A test set used for
evaluation is comprised of 450-460 non-empty scenes, based on 512 draws.

Additionally, we study how the methods compare when objects are created from
noisy templates. We consider that the original templates used for the creation
of the images are corrupted with Gaussian noise with standard deviation
$\sigma$. Once the templates are corrupted with noise, a random transformation
${\mathbf{y}}_{k}$ is applied to obtain the object $X_{k}$ of the scene. As
with the noise-free data, the objects are normalized to lie in [-1,1] on both
dimensions.

###  5.1 Evaluation Metrics

Three metrics are used to evaluate the performance of the methods: variation
of information, adjusted Rand index and segmentation accuracy. They are based
on partitions of the datapoints into those associated with each object, and
those that are missing. Compared to standard clustering evaluation metrics,
some modifications need to be made to handle the missing objects.

In a given scene $X$ there are $M$ points, but we know that there are $N\geq
M$ possible points that can be produced from all of the templates. Assume that
$K^{\prime}\leq K$ templates are active in this scene. Then the points in the
scene are labelled with indices $1,\ldots,K^{\prime}$, and we assign the
missing points index $0$. Denote the ground truth partition as
$V=\\{V_{0},V_{1},\ldots V_{K^{\prime}}\\}$. An alternative partition output
by one of the algorithms is denoted by
$\hat{V}=\\{\hat{V}_{0},\hat{V}_{1},\ldots\hat{V}_{\hat{K}^{\prime}}\\}$. The
predicted partition $\hat{V}$ may instantiate objects or points that were in
fact missing, thus it is important to handle the missing data properly.

In Information Theory, the variation of information (VI) [Meilă, 2003] is a
measure of the distance between two partitions of elements (or clusterings).
For a given set of elements, the variation of information between two
partitions $V$ and $\hat{V}$, where $N=\sum_{i}|V_{i}|=\sum_{j}|\hat{V}_{j}|$
is defined as:

$VI(V,\hat{V})=-\sum_{i,j}r_{ij}\left[\log{\frac{r_{ij}}{p_{i}}}+\log{\frac{r_{ij}}{q_{j}}}\right]$
(26)

where $r_{ij}=\frac{|V_{i}\cap\hat{V}_{j}|}{N}$, $p_{i}=\frac{|V_{i}|}{N}$ and
$q_{j}=\frac{|\hat{V}_{j}|}{N}$. In our experiments we report the average
variation of information of the scenes in the dataset.

The Rand index [Rand, 1971] is another measure of similarity between two data
clusterings. This metric takes pairs of elements and evaluates whether they do
or do not belong to the same subsets in the partitions $V$ and $\hat{V}$

$RI=\frac{TP+TN}{TP+TN+FP+FN},$ (27)

where TP are the true positives, TN the true negatives, FP the false positives
and FN the false negatives. The Rand index takes on values between 0 and 1. We
use instead the adjusted Rand index (ARI) [Hubert and Arabie, 1985], the
corrected-for-chance version of the Rand index. In our experiments, we compute
the average adjusted Rand index of the scenes in our dataset.

The segmentation accuracy (SA) is based on obtaining the maximum bipartite
matching between $V$ and $\hat{V}$, and was used by Kosiorek et al. [2019] to
evaluate the performance of CCAE. For each set $V_{i}$ in $V$ and set
$\hat{V}_{j}$ in $\hat{V}$, there is an edge $w_{ij}$ with the weight being
the number of common elements in both sets. Let $W(V,\hat{V})$ be the overall
weight of the maximum matching between $V$ and $\hat{V}$. Then we define the
average segmentation accuracy as:

$SA=\sum_{i=1}^{I}\frac{W(V_{i},\hat{V}_{i})}{W(V_{i},V_{i})}=\frac{1}{N}\sum_{i=1}^{I}W(V_{i},\hat{V}_{i}),$
(28)

where $I$ is the number of scenes. Notice that $W(V_{i},V_{i})$ represents a
perfect assignment of the ground truth, both the observed and missing subsets,
and thus $W(V_{i},V_{i})=N$.

There are some differences on how we compute the SA metric compared to
Kosiorek et al. [2019]. First, they do not consider the missing points as part
of their ground truth, but as we argued above this is necessary. They evaluate
the segmentation accuracy in terms of the observed points in the ground truth,
disregarding possible points that were missing in the ground truth but
predicted as observed in $\hat{V}$. Second, they average the segmentation
accuracy across scenes as

$SA=\frac{\sum_{i=1}^{I}W(V_{i},\hat{V}_{i})}{\sum_{i=1}^{I}W(V_{i},V_{i})}.$
(29)

For them, $W(V_{i},V_{i})=M_{i}$, where $M_{i}$ is the number of points
present in a scene. In our case, both averaging formulae are equivalent since
our $W(V_{i},V_{i})$ is the same across scenes.

##  6 Experiments

CCAE is trained by creating random batches of 128 scenes as described in the
previous section and optimizing the objective function in (25). The authors
run CCAE for $300K$ epochs, and when the parameters of the neural networks are
trained, they use their model on the test dataset to generate an estimation of
which points belong to which capsule, and where the estimated points are
located in each scene.

For the variational inference approach, $p(X|Y,Z)$ in (7) allows us to model
scenes where the points are corrupted with some noise. The precision $\lambda$
controls the level of noise allowed in the model. We use an annealing strategy
to fit $\lambda$, increasing it every time the ELBO has converged, up to a
maximum value of $\lambda_{max}=10000$. We set the hyperparameters of the
model to $\mu_{0}=\mathbf{0}$, $\Lambda_{0}=I_{4}$ and $a_{mnk}=\frac{1}{N}$.
We run Algorithm 1 with 5 different random initializations of $R$ and select
the solution with the best ELBO. Similarly to Kosiorek et al. [2019], we
incorporate a sparsity constraint in our model, that forces every object to
explain at least two parts. Once our algorithm has converged, for a given $k$
if any $r_{mnk}>0.9$ and $\sum_{m}\sum_{n}r_{mnk}<2$ it means that the model
has converged to a solution where object $k$ is assigned to less than 2 parts.
In these cases, we re-run Algorithm 1 with a new initialization of $R$. Notice
that this is also related to the minimum basis size necessary in the RANSAC
approach for the types of transformations that we are considering.

The implementation of the RANSAC algorithm (Algorithm 2) considers all matches
between the predicted and the scene points where the distance between them is
less than $0.1$. Among them, it selects the matching with minimum distance
between scene and predicted points. For both the variational inference
algorithm and the RANSAC algorithm, a training dataset is not necessary. These
algorithms are applied directly to each test scene and return a set of objects
describing the points in the scene. Unfortunately we do not have access to the
code employed by Hinton et al. [2018], so we have been unable to make
comparisons with it.

###  6.1 Results

Figure 2: Reconstruction examples from CCAE and GCM-DS. The upper figures show
the ground truth of the test images. The middle figures show the
reconstruction and the capsule assignments (by colours) of CCAE. The lower
figures show the reconstruction and the parts assignment of GCM-DS. Datapoints
shown with a given colour are predicted to belong to the reconstructed object
with the corresponding colour.

In Table 1 we show a comparison between CCAE, the variational inference method
with a Gaussian mixture prior (9) (GCM-GMM), with a prior over permutation
matrices (10) (GCM-DS), and the RANSAC approach 333Code at:
https://github.com/anazabal/GenerativeCapsules. For GCM-GMM and GCM-DS we show
the results where the initial $\lambda=500$, relegating a comparison across
different initializations of $\lambda$ to the supp. mat. C. We provide results
for the metrics described above, and an average scene accuracy metric, where a
scene is correct if the method returns the full original scene, and is
incorrect otherwise. Scenes without noise and with noise levels of
$\sigma=0.1,\;0.25$ are considered.

Table 1: Comparison between the different methods. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Different levels of Gaussian noise with standard deviation $\sigma$ are considered. Metric | Model |  $\sigma$=0 |  $\sigma$=0.1 |  $\sigma$=0.25
---|---|---|---|---
SA $\uparrow$ | CCAE | 0.828 | 0.754 | 0.623
GCM-GMM | 0.753 | 0.757 | 0.744
GCM-DS | 0.899 | 0.882 | 0.785
RANSAC | 1 | 0.992 | 0.965
ARI $\uparrow$ | CCAE | 0.599 | 0.484 | 0.248
GCM-GMM | 0.586 | 0.572 | 0.447
GCM-DS | 0.740 | 0.699 | 0.498
RANSAC | 1 | 0.979 | 0.914
VI $\downarrow$ | CCAE | 0.481 | 0.689 | 0.988
GCM-GMM | 0.478 | 0.502 | 0.677
GCM-DS | 0.299 | 0.359 | 0.659
RANSAC | 0 | 0.034 | 0.135
| CCAE | 0.365 | 0.138 | 0.033
Scene | GCM-GMM | 0.179 | 0.173 | 0.132
Acc $\uparrow$ | GCM-DS | 0.664 | 0.603 | 0.377
| RANSAC | 1 | 0.961 | 0.843

We can see that GCM-DS improves over CCAE and GCM-GMM in all of the metrics,
with GCM-GMM being comparable to CCAE. Interestingly, for the noise-free
scenarios, the RANSAC method achieves a perfect score for all of the metrics.
Since there is no noise on the observations and the method searches over all
possible solutions of ${\mathbf{y}}_{k}$, it can find the correct solution for
any configuration of geometric shapes in a scene. For the noisy scenarios, all
the methods degrade as $\sigma$ increases. However, the relative performance
between them remains the same, with RANSAC performing the best, followed by
GCM-DS and then GCM-GMM. Note that the GCM models here assume that the
template for each model is unperturbed; however, as per supp. mat. A one could
allow a prior on the part locations which would take into account that they
are now noisy.

We have conducted paired t-tests between CCAE and GCM-GMM, GCM-DS and RANSAC
on the three metrics. The differences between CCAE and GCM-DS are
statistically significant with p-values less than $10^{-7}$, and between CCAE
and RANSAC with p-values less than $10^{-28}$. For CCAE and GCM-GMM the
differences are not statistically significant.

Figure 2 shows some reconstruction examples from CCAE and GCM-DS for the
noise-free scenario. We show additional examples for the noisy cases in the
supp. mat. E. In columns (a) and (b) we can see that CCAE recovers the correct
parts assignments but the object reconstruction is amiss. In (a) one of the
squares is reconstructed as a triangle, while in (b) the assignment between
the reconstruction and the ground truth is not exact. For GCM-DS, if the parts
are assigned to the ground truth properly, and there is no noise, then the
reconstruction of the object is perfect. In column (c) all methods work well.
In column (d), CCAE fits the square correctly (green), but adds an additional
red point. In this case GCM-DS actually overlays two squares on top of each
other. Both methods fail badly on column (e). Note that CCAE is not guaranteed
to reconstruct an existing object correctly (square or triangle in this case).
In column (f) we can see that CCAE fits an irregular quadrilateral (blue) to
the assigned points, while GCM-DS obtains the correct fit.

##  7 Discussion

Above we have described a principled generative capsules model. It leads
directly to a variational inference algorithm which can handle either a
mixture formulation or incorporation of the DS matrix constraint. We have also
shown that this formulation outperforms the CCAE of Kosiorek et al. [2019] on
the constellation data generator that they used.

In this paper we have not discussed the learning of the model (i.e. the
$F_{kn}$ matrices), but in all such work, inference for $q(Z,Y)$ forms an
inner loop, with model learning as an outer loop. Also, the interpretable
nature of the model means that we can exploit the geometric knowledge where it
is available; in contrast the CCAE’s opaque inference networks do not allow
this.

One challenge for the future is to extend this work beyond purely geometric
situations, e.g. to images with multiple digits as in Kosiorek et al. [2019].
The theory for this is presented in supp. material A. The other notable
challenge is around generalizing RANSAC-style inference. For the geometric
testbed, this is remarkably effective, and it leads to the question if similar
techniques could work in a more general setting.

###### Acknowledgements.

This work was supported in part by The Alan Turing Institute under EPSRC grant
EP/N510129/1.

## References

  * Biederman [1987] I. Biederman.  Recognition-by-components: A theory of human image understanding.  _Psychological Review_ , 94:115–147, 1987.
  * Bishop [2006] C. M. Bishop.  _Pattern Recognition and Machine Learning_.  Springer, 2006.
  * Cootes et al. [1995] T. F. Cootes, C. J. Taylor, D. H. Cooper, and J. Graham.  Active Shape Models—Their Training and Application.  _Computer Visoin and Image Understanding_ , 61(1):38–59, 1995.
  * Fergus et al. [2003] R. Fergus, P. Perona, and A. Zisserman.  Object Class Recognition by Unsupervised Scale-Invariant Learning.  In _IEEE Computer Society Conference on Computer Vision and Pattern Recognition_ , 2003.
  * Fischler and Bolles [1981] M. A. Fischler and R. C. Bolles.  Random sample consensus: a paradigm for model fitting with applications to image analysis and automated cartography.  _Communications of the ACM_ , 24(6):381–395, 1981\.
  * Hinton et al. [2000] G. E. Hinton, Z. Ghahramani, and Y. W. Teh.  Learning to Parse Images.  In S. A. Solla, T. K. Leen, and K.-R. Müller, editors, _Advances in Neural Information Processing Systems 12_ , pages 463–469. MIT Press, Cambridge, MA, 2000.
  * Hinton et al. [2018] G. E. Hinton, S. Sabour, and N. Frosst.  Matrix capsules with EM routing.  In _International Conference on Learning Representations_ , 2018\.
  * Hubert and Arabie [1985] L. Hubert and P. Arabie.  Comparing partitions.  _Journal of Classification_ , 2(1):193–218, 1985\.
  * Kosiorek et al. [2019] A. Kosiorek, S. Sabour, Y. W. Teh, and G. E. Hinton.  Stacked capsule autoencoders.  In _Advances in Neural Information Processing Systems_ , pages 15512–15522, 2019.
  * Lee et al. [2019] J. Lee, Y. Lee, J. Kim, A. R. Kosiorek, S. Choi, and Y.-W. Teh.  Set Transformer: A Framework for Attention-based Permutation-invariant Neural Networks.  In _Proc. of the 36th International Conference on Machine Learning_ , pages 3744–3753, 2019.
  * Li and Zhu [2019] Y. Li and X. Zhu.  Capsule Generative Models.  In I. V. Tetko et al., editors, _ICANN_ , pages 281–295. LNCS 11727, 2019.
  * Li et al. [2020] Y. Li, X. Zhu, R. Naud, and P. Xi.  Capsule Deep Generative Model That Forms Parse Trees.  In _2020 International Joint Conference on Neural Networks (IJCNN)_. 2020.
  * Meilă [2003] M. Meilă.  Comparing clusterings by the variation of information.  In _Learning Theory and Kernel Machines_ , pages 173–187. Springer, 2003.
  * Mena et al. [2020] Gonzalo Mena, Erdem Varol, Amin Nejatbakhsh, Eviatar Yemini, and Liam Paninski.  Sinkhorn permutation variational marginal inference.  In _Symposium on Advances in Approximate Bayesian Inference_ , pages 1–9. PMLR, 2020.
  * Powell and Smith [2019] B. Powell and P. A. Smith.  Computing expectations and marginal likelihoods for permutations.  _Computational Statistics_ , pages 1–21, 2019.
  * Rand [1971] W. M. Rand.  Objective criteria for the evaluation of clustering methods.  _Journal of the American Statistical Association_ , 66(336):846–850, 1971.
  * Revow et al. [1996] M. Revow, C. K. I. Williams, and G. E. Hinton.  Using Generative Models for Handwritten Digit Recognition.  _IEEE Trans. on Pattern Analysis and Machine Intelligence_ , 18(6):592–606, 1996.
  * Sabour et al. [2017] S. Sabour, N. Frosst, and G.E. Hinton.  Dynamic routing between capsules.  In _Advances in Neural Information Processing Systems_ , pages 3856–3866, 2017.
  * Sinkhorn and Knopp [1967] R. Sinkhorn and P. Knopp.  Concerning nonnegative matrices and doubly stochastic matrices.  _Pacific Journal of Mathematics_ , 21(2):343–348, 1967.
  * Smith et al. [2021] L. Smith, L. Schut, Y. Gal, and M. van der Wilk.  Capsule Networks—A Generative Probabilistic Perspective, 2021.  https://arxiv.org/pdf/2004.03553.pdf.
  * Storkey and Williams [2003] A. J. Storkey and C. K. I. Williams.  Image Modelling with Position-Encoding Dynamic Trees.  _IEEE Trans Pattern Analysis and Machine Intelligence_ , 25(7):859–871, 2003.
  * Wolfson and Rigoutsos [1997] H. J. Wolfson and I. Rigoutsos.  Geometric hashing: An overview.  _IEEE Computational Science and Engineering_ , 4(4):10–21, 1997.

Supplementary Material for _Inference for Generative Capsule Models_

##  Appendix A Adding appearance features to the generative model

The generative model in sec. 3 only considers the geometry, but not the
appearance each part. In this appendix we show how to extend the model to
include appearance features, and also include a prior distribution $p(P)$ over
the geometry and appearance of the parts.

Consider a template (or model) $T_{k}$ for the $k$th object. $T_{k}$ is
composed of $N_{k}$ parts $\\{{\mathbf{p}}_{n}\\}_{n=1}^{N_{k}}$.444For
simplicity of notation we suppress the dependence of ${\mathbf{p}}_{n}$ on $k$
for now. Each part ${\mathbf{p}}_{n}$ is described in its reference frame by
its _geometry_ ${\mathbf{p}}^{g}_{n}$ and its _appearance_
${\mathbf{p}}^{a}_{n}$. Each object also has associated latent variables
${\mathbf{y}}_{k}$ which transform from the reference frame to the image
frame; ${\mathbf{y}}_{k}$ is split into geometric variables
${\mathbf{y}}^{g}_{k}$ and additional variables ${\mathbf{y}}^{a}_{k}$. The
geometric latent variables ${\mathbf{y}}^{g}_{k}$ transform the locations of
the parts ${\mathbf{p}}_{n}^{g}$ into locations in the the image frame
${\mathbf{x}}^{g}_{n}$. Below we consider 2D templates and a similarity
transformation (translation, rotation and scaling) for each object, but this
can be readily extended to allow 3D templates and a scene-to-viewer camera
transformation.

The appearance ${\mathbf{x}}_{n}^{a}$ of part $n$ in the image may depend on
both ${\mathbf{y}}^{g}_{k}$ and ${\mathbf{y}}^{a}_{k}$. Consider a handwritten
digit model, where each part is a stroke. The orientation of stroke $n$ will
depend not only on its orientation in the reference frame (described by
${\mathbf{p}}^{a}_{n}$) but also the rotation of the geometric transformation
${\mathbf{y}}^{g}_{k}$. As another example, the width of the stroke of part
$n$ in the template will depend on ${\mathbf{y}}^{a}_{k}$, which will describe
the thickness of the pen used to draw the digit; note that this will be shared
across parts. The stroke width observed in the image will also depend on the
scale $s$ in the transformation ${\mathbf{y}}^{g}_{k}$, so in general
${\mathbf{x}}^{a}_{n}$ will depend on both ${\mathbf{y}}^{g}_{k}$ and
${\mathbf{y}}^{a}_{k}$. Parts can be extracted from real images using e.g. the
Part Capsule Autoencoder (PCAE) of Kosiorek et al. [2019].

In more detail, the 2D similarity transformation is defined as follows:
Consider a template with the location of part $n$ in the reference frame given
by ${\mathbf{p}}^{g}_{n}=(p^{g}_{nx},p^{g}_{ny})^{T}$, for $n=1,\ldots,N_{k}$.
As in sec. 3 we wish to scale by a factor $s$, rotate through with a clockwise
rotation angle $\theta$ and translate by $(t_{x},t_{y})^{T}$, to obtain a
transformed object with the location of part $n$ being
${\mathbf{x}}^{g}_{n}=(x^{g}_{nx},x_{ny})^{T}$. For each part in the template,
the transformation works as follows:

$\begin{pmatrix}x^{g}_{nx}\\\
x^{g}_{ny}\end{pmatrix}=\begin{pmatrix}1&0&p^{g}_{nx}&p^{g}_{ny}\\\
0&1&p^{g}_{ny}&-p^{g}_{nx}\end{pmatrix}\begin{pmatrix}t_{x}\\\ t_{y}\\\
s\cos{\theta}\\\ s\sin{\theta}\end{pmatrix}\stackrel{{\scriptstyle
def}}{{=}}F^{g}_{kn}{\mathbf{y}}^{g}_{k},$ (30)

where ${\mathbf{y}}^{g}_{k}$ is the ${\mathbb{R}}^{4}$ column vector, and
$F^{g}_{kn}\in{\mathbb{R}}^{2\times 4}$ is the matrix to its left.

Let $X$ denote the matrix of observed parts
${\mathbf{x}}_{1},\ldots,{\mathbf{x}}_{M}$, including both the geometric and
appearance features for each part. Let $Y$ describe the matrix of $K$ vectors
${\mathbf{y}}_{1},\ldots,{\mathbf{y}}_{K}$, including both the geometric and
additional features. Let $P$ denote the matrix of $K$ templates, comprised of
$P_{1},\ldots,P_{K}$. Each $P_{k}$ holds the geometric and appearance
information for template $T_{k}$. As in sec. 3, $p(Z)$ describes the
distribution over assignments of observations to objects and parts, so that
$z_{mnk}=1$ indicates that ${\mathbf{x}}_{m}$ is assigned to part $n$ of
object $k$.

Then the joint distribution of the variables in this model is given by

$p(P,X,Y,Z)=p(P)p(Y)p(Z)p(X|Y,Z).$ (31)

We describe each of these terms in turn: $p(P)$ describes the prior
distribution on the templates, which can either be a delta function on a
specific geometry and appearance, or modelled e.g. with a Gaussian
distribution, as in an active shape model, see e.g. Cootes et al. [1995].

$p(Y^{g})$ describes the distribution over similarity transformations. As in
sec. 3, each ${\mathbf{y}}^{g}_{k}$ is modelled with a ${\mathbb{R}}^{4}$
Gaussian distribution with mean ${\boldsymbol{\mu}}_{0}$ and precision matrix
$\Lambda_{0}$.

$Y^{a}$ describes the additional latent variables, which describe for each
template common appearance information like the width of the pen stroke used
to draw a handwritten character, or the shared colour of two eyes in a face.
$p(Y^{a})$ is assumed to be Gaussian, and $p(Y)=p(Y^{a})p(Y^{g})$.

$p(X|P,Y,Z)$ is assumed to have the form

$p(X|P,Y,Z)=\prod_{m=1}^{M}\prod_{k=1}^{K}\prod_{n=1}^{N_{k}}p({\mathbf{x}}_{m}|P_{k},{\mathbf{y}}_{k})^{z_{mnk}},$
(32)

where $z_{mnk}$ identifies the match between datapoint ${\mathbf{x}}_{m}$ and
the corresponding part $n$ of template $k$. For the geometric part we have

$p({\mathbf{x}}^{g}_{m}|P_{k},{\mathbf{y}}_{k})\sim
N\left({\mathbf{x}}^{g}_{m}|F^{g}_{kn}{\mathbf{y}}^{g}_{k},\lambda^{-1}I_{2}\right)$
(33)

with $\lambda$ being the precision of the Gaussian noise distribution for each
point, and $I_{2}$ the 2D identity matrix. For the appearance features we
assume a similar factor analysis model as for the geometry, so that

$p({\mathbf{x}}^{a}_{m}|P_{k},{\mathbf{y}}_{k})\sim
N\left({\mathbf{x}}^{a}_{m}|F^{a}_{kn}{\mathbf{y}}_{k},D^{a}\right),$ (34)

where $F^{a}_{kn}$ maps from ${\mathbf{y}}_{k}$ to the predicted appearance
features in the image, and $D^{a}$ is a diagonal matrix of variances. Note
that ${\mathbf{x}}^{a}_{m}$ can depend on both ${\mathbf{y}}^{g}_{k}$ and
${\mathbf{y}}^{a}_{k}$, while ${\mathbf{x}}^{g}_{m}$ only depends
${\mathbf{y}}^{g}_{k}$. Stacking ${\mathbf{x}}^{g}_{m}$ and
${\mathbf{x}}^{a}_{m}$ into ${\mathbf{x}}_{m}$, we can combine equations 33
and 34 to obtain

$p({\mathbf{x}}_{m}|P_{k},{\mathbf{y}}_{k})\sim
N\left({\mathbf{x}}_{m}|F_{kn}{\mathbf{y}}_{k},D\right)$ (35)

where $D$ consists of the diagonal matrices of $\lambda^{-1}I_{2}$ and
$D^{a}$. Note that $F_{kn}$ has a block of zeros so that
${\mathbf{x}}^{g}_{m}$ does not depend on ${\mathbf{y}}^{a}_{kn}$.

The model described above is similar to some in the literature. For example,
Revow et al. [1996] described a geometric model similar to eq. 33 for a single
object, but without the DS matrix constraint on $Z$. Fergus et al. [2003]
model the shape and appearance of the parts of a single object, but again
without a DS matrix constraint on $Z$.

##  Appendix B Details for Variational Inference

The Sinkhorn-Knopp algorithm is given in Algorithm 3.

Algorithm 3 Sinkhorn-Knopp algorithm

1:procedure SinkhornKnopp($X$)

2: while $X$ not doubly stochastic do

3: Normalize rows of $X$: $x_{ij}=\frac{x_{ij}}{\sum_{j}x_{ij}}$, $\forall{i}$

4: Normalize columns of $X$: $x_{ij}=\frac{x_{ij}}{\sum_{i}x_{ij}}$,
$\forall{j}$ return $X$

The Kullback-Leibler divergence between the two Gaussian distributions $q(Y)$
and $p(Y)$ in our model has the following expression:

$\displaystyle-
KL(q(Y)||p(Y))=\sum_{k=1}^{K}\left(2+\frac{1}{2}\log{|\Lambda_{0}|}-\frac{1}{2}\log{|\Lambda_{k}|}\right.$
$\displaystyle\left.-\frac{1}{2}({\boldsymbol{\mu}}_{k}-{\boldsymbol{\mu}}_{0})^{T}\Lambda_{0}({\boldsymbol{\mu}}_{k}-{\boldsymbol{\mu}}_{0})-\frac{1}{2}\text{trace}(\Lambda_{0}\Lambda_{k}^{-1})\right).$
(36)

The expression for $KL(q(Z)||p(Z))$ is given by

$\displaystyle KL(q(Z)||p(Z))$
$\displaystyle=\sum_{m=1}^{N}\sum_{k=1}^{K}\sum_{n=1}^{N_{k}}r_{mnk}\log{\frac{r_{mnk}}{a_{mnk}}}.$
(37)

##  Appendix C Different initializations of $\lambda$

In this section, we show the effect of varying the initial value of $\lambda$
for both the GCM-GMM and the GCM-DS methods. We have tested 6 different values
of $\lambda$ (50, 100, 500, 1000, 2000, 5000). In Tables 2, 3 and 4 we show
the results for all metrics, methods and initial values of $\lambda$ for
$\sigma=0$, $\sigma=0.1$, and $\sigma=0.25$ resp.

Table 2: Different initializations of $\lambda$. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Noise-free case. Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.828 | 0.599 | 0.481 | 0.365
GCM-DS $\lambda=50$ | 0.869 | 0.667 | 0.350 | 0.581
GCM-DS $\lambda=100$ | 0.876 | 0.678 | 0.328 | 0.624
GCM-DS $\lambda=500$ | 0.899 | 0.740 | 0.299 | 0.664
GCM-DS $\lambda=1000$ | 0.904 | 0.752 | 0.302 | 0.653
GCM-DS $\lambda=2000$ | 0.899 | 0.728 | 0.332 | 0.622
GCM-DS $\lambda=5000$ | 0.886 | 0.679 | 0.402 | 0.576
GCM-GMM $\lambda=50$ | 0.823 | 0.449 | 0.592 | 0.402
GCM-GMM $\lambda=100$ | 0.823 | 0.487 | 0.555 | 0.362
GCM-GMM $\lambda=500$ | 0.753 | 0.586 | 0.478 | 0.179
GCM-GMM $\lambda=1000$ | 0.771 | 0.550 | 0.543 | 0.214
GCM-GMM $\lambda=2000$ | 0.764 | 0.513 | 0.613 | 0.199
GCM-GMM $\lambda=5000$ | 0.762 | 0.488 | 0.650 | 0.194
Table 3: Different initializations of $\lambda$. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Gaussian noise with $\sigma=0.1$. Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.754 | 0.484 | 0.689 | 0.138
GCM-DS $\lambda=50$ | 0.840 | 0.608 | 0.451 | 0.518
GCM-DS $\lambda=100$ | 0.859 | 0.644 | 0.389 | 0.559
GCM-DS $\lambda=500$ | 0.883 | 0.699 | 0.359 | 0.603
GCM-DS $\lambda=1000$ | 0.888 | 0.715 | 0.352 | 0.601
GCM-DS $\lambda=2000$ | 0.874 | 0.682 | 0.406 | 0.544
GCM-DS $\lambda=5000$ | 0.872 | 0.659 | 0.443 | 0.515
GCM-GMM $\lambda=50$ | 0.793 | 0.468 | 0.566 | 0.322
GCM-GMM $\lambda=100$ | 0.800 | 0.497 | 0.547 | 0.318
GCM-GMM $\lambda=500$ | 0.757 | 0.573 | 0.502 | 0.173
GCM-GMM $\lambda=1000$ | 0.768 | 0.521 | 0.595 | 0.204
GCM-GMM $\lambda=2000$ | 0.765 | 0.460 | 0.684 | 0.178
GCM-GMM $\lambda=5000$ | 0.753 | 0.403 | 0.780 | 0.140
Table 4: Different initializations of $\lambda$. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Gaussian noise with $\sigma=0.25$. Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.623 | 0.248 | 0.988 | 0.033
GCM-DS $\lambda=50$ | 0.771 | 0.475 | 0.669 | 0.338
GCM-DS $\lambda=100$ | 0.781 | 0.486 | 0.652 | 0.362
GCM-DS $\lambda=500$ | 0.785 | 0.498 | 0.659 | 0.377
GCM-DS $\lambda=1000$ | 0.779 | 0.476 | 0.699 | 0.351
GCM-DS $\lambda=2000$ | 0.768 | 0.470 | 0.708 | 0.291
GCM-DS $\lambda=5000$ | 0.766 | 0.435 | 0.752 | 0.272
GCM-GMM $\lambda=50$ | 0.736 | 0.469 | 0.564 | 0.265
GCM-GMM $\lambda=100$ | 0.730 | 0.467 | 0.581 | 0.221
GCM-GMM $\lambda=500$ | 0.744 | 0.447 | 0.677 | 0.132
GCM-GMM $\lambda=1000$ | 0.735 | 0.387 | 0.801 | 0.102
GCM-GMM $\lambda=2000$ | 0.727 | 0.310 | 0.921 | 0.088
GCM-GMM $\lambda=5000$ | 0.717 | 0.242 | 1.036 | 0.064

We can observe that no matter the initialization of $\lambda$, GCM-DS is
always better than CCAE and GCM-GMM. As we increase the initial lambda, GCM-DS
performs better across all the metrics. We have found that the performance of
GCM-DS and GCM-GMM degrades with $\lambda>1000$.

##  Appendix D Comparison across methods using the ground truth mask

The results in the previous section include the missing objects in the
creation of $\hat{V}$, as described in Section 5.2 in the main paper. Failing
to recognize exactly which points in a scene are missing leads to a decrease
in performance for all the metrics. However, Kosiorek et al. [2019] instead
used the ground truth missing indicators for the objects in order to compute
the segmentation accuracy of their method (see equation (27) in the main
paper). If we use this strategy to all the models and metrics, we obtain
Tables 5, 6 and 7. When we compare these results with the tables in the
previous section we can see that this alternative leads to a clear improvement
in the scores of all the methods, but with GCM-DS still performing better than
CCAE or GCM-GMM.

Table 5: Different initializations of $\lambda$ using the ground truth missing mask. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Noise-free case. Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.950 | 0.897 | 0.157 | 0.736
GCM-DS $\lambda=50$ | 0.985 | 0.967 | 0.052 | 0.915
GCM-DS $\lambda=100$ | 0.989 | 0.973 | 0.042 | 0.941
GCM-DS $\lambda=500$ | 0.971 | 0.938 | 0.095 | 0.886
GCM-DS $\lambda=1000$ | 0.966 | 0.929 | 0.108 | 0.860
GCM-DS $\lambda=2000$ | 0.966 | 0.930 | 0.104 | 0.854
GCM-DS $\lambda=5000$ | 0.958 | 0.912 | 0.132 | 0.825
GCM-GMM $\lambda=50$ | 0.886 | 0.805 | 0.209 | 0.646
GCM-GMM $\lambda=100$ | 0.905 | 0.830 | 0.191 | 0.683
GCM-GMM $\lambda=500$ | 0.942 | 0.884 | 0.160 | 0.718
GCM-GMM $\lambda=1000$ | 0.937 | 0.879 | 0.179 | 0.657
GCM-GMM $\lambda=2000$ | 0.927 | 0.861 | 0.213 | 0.603
GCM-GMM $\lambda=5000$ | 0.925 | 0.858 | 0.216 | 0.585
Table 6: Different initializations of $\lambda$ using the ground truth missing mask. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Gaussian noise with $\sigma=0.1$ Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.867 | 0.729 | 0.403 | 0.386
GCM-DS $\lambda=50$ | 0.957 | 0.917 | 0.135 | 0.794
GCM-DS $\lambda=100$ | 0.972 | 0.940 | 0.094 | 0.857
GCM-DS $\lambda=500$ | 0.957 | 0.908 | 0.142 | 0.833
GCM-DS $\lambda=1000$ | 0.956 | 0.906 | 0.144 | 0.831
GCM-DS $\lambda=2000$ | 0.946 | 0.887 | 0.172 | 0.794
GCM-DS $\lambda=5000$ | 0.942 | 0.875 | 0.194 | 0.759
GCM-GMM $\lambda=50$ | 0.856 | 0.752 | 0.272 | 0.544
GCM-GMM $\lambda=100$ | 0.875 | 0.768 | 0.272 | 0.561
GCM-GMM $\lambda=500$ | 0.922 | 0.837 | 0.232 | 0.621
GCM-GMM $\lambda=1000$ | 0.912 | 0.824 | 0.268 | 0.548
GCM-GMM $\lambda=2000$ | 0.905 | 0.822 | 0.278 | 0.474
GCM-GMM $\lambda=5000$ | 0.881 | 0.792 | 0.334 | 0.379
Table 7: Different initializations of $\lambda$ using the ground truth missing mask. For SA, ARI and Scene Accuracy the higher the better. For VI the lower the better. Gaussian noise with $\sigma=0.25$ Model |  SA $\uparrow$ |  ARI $\uparrow$ |  VI $\downarrow$ |  Sc. Acc. $\uparrow$
---|---|---|---|---
CCAE | 0.774 | 0.621 | 0.540 | 0.033
GCM-DS $\lambda=50$ | 0.897 | 0.811 | 0.302 | 0.556
GCM-DS $\lambda=100$ | 0.908 | 0.822 | 0.281 | 0.611
GCM-DS $\lambda=500$ | 0.900 | 0.805 | 0.302 | 0.623
GCM-DS $\lambda=1000$ | 0.895 | 0.796 | 0.317 | 0.609
GCM-DS $\lambda=2000$ | 0.891 | 0.784 | 0.333 | 0.578
GCM-DS $\lambda=5000$ | 0.881 | 0.770 | 0.359 | 0.528
GCM-GMM $\lambda=50$ | 0.852 | 0.743 | 0.287 | 0.545
GCM-GMM $\lambda=100$ | 0.865 | 0.756 | 0.282 | 0.558
GCM-GMM $\lambda=500$ | 0.887 | 0.773 | 0.337 | 0.468
GCM-GMM $\lambda=1000$ | 0.869 | 0.748 | 0.392 | 0.355
GCM-GMM $\lambda=2000$ | 0.855 | 0.734 | 0.419 | 0.278
GCM-GMM $\lambda=5000$ | 0.834 | 0.704 | 0.470 | 0.196

##  Appendix E Noisy cases examples

In Figures 3 and 4 we show several examples of objects generated from noisy
templates with corruption levels of $\sigma=0.1$ and $\sigma=0.25$ resp. GCM-
DS and RANSAC does not allow for deformable objects to try to fit the points
exactly, contrary to CCAE. Both methods try to find the closest reconstruction
of the noisy points in the image by selecting the geometrical shapes that are
a best fit to those points. Nonetheless, both methods can determine that which
parts belong together to form a given object, even when the matching is not
perfect.

Figure 3: Reconstruction examples from CCAE, GCM-DS and RANSAC with Gaussian
noise $\sigma=0.1$. Figure 4: Reconstruction examples from CCAE, GCM-DS and
RANSAC with Gaussian noise $\sigma=0.25$.