Kendall Shape-VAE : Learning Shapes in a Generative Framework
Editors: Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo, Nina Miolane
Abstract
Learning an interpretable representation of data without supervision is an important precursor for the development of artificial intelligence.
In this work, we introduce Kendall
Shape -VAE, a novel Variational Autoencoder framework for learning shapes as it disentangles the latent space by compressing information to simpler geometric symbols.
In Kendall
Shape -VAE, we modify the Hyperspherical Variational Autoencoder such that it results in an exactly rotationally equivariant network using the notion of landmarks in the Kendall shape space.
We show the exact equivariance of the model through experiments on rotated
MNIST.
1. Introduction
The learning of reliable, compressed, and information-rich representations is essential in machine learning.
Among the many unsupervised likelihood-based frameworks, Variational
Autoencoders (Kingma and Welling, 2013; Rezende et al., 2014) have proven to be an intuitive framework with compelling results and have thus gained popularity in representation learning and generative image modeling.
VAEs allow for learning of low dimensional latent representations, and can, e g., be adapted for learning disentangled representations; (Chen et al., 2018) or modeling of dependencies between random variables in separated latent spaces (Ilse et al., 2020).
There has been considerable interest in making the latent space invariant to a transformation, but it has been tricky to do so with a VAE framework as it involves separate training of invariant latent spaces for specific transformations (Bepler et al., 2019) which cannot be generalized well.
Adding a geometric structure to latent space has improved on the generation capacity of VAEs (Chadebec and Allassonni`ere, 2021), and could improve on the notion of disentanglement if the geometry is biased towards it (Davidson et al., 2018).
Alternatively, a few works provide post-hoc explanations via classifiers that map the latent space to outcome, but they come with their limitations (Liu et al., 2019).
For the tasks of domain generalization (Ilse et al., 2020; Mitko, 2019) and compressed sensing (Kuzina et al., 2021) it is important to have explicit control over how the latent space is encoded.
We see a lot of potential in using geometry to structure VAE latent spaces as to obtain interpretable latents.
Towards this end, we present a new VAE framework, the Kendall
Shape VAE.
It encodes each image as a geometric symbol, or shape , and does so in a rotation equivariant manner.
Our approach is motivated by our daily use of symbols, pictograms, and ideograms in general, to represent concepts in a concise geometry form.
As such, the learnt symbols in our Kendall Shape VAEs should be thought of as neural ideograms 1 .
We illustrate the interpretability of our method by showing how similar the latents of the two different images of the same class look as well as show that the shape remains consistent.
Our contributions in this paper can be summarized as:
• We introduce Kendall Shape VAE, a fully equivariant VAE with a hyperspherical latent space, equivalent to a Kendall shape space.
• Points in latent space correspond to landmarks that define shapes in Kendall shape space.
We show that equivalence classes of images are mapped to invariant shapes.
• We show that this shape is truly invariant to rotation and translation, by the use of equivariant encoder-decoder, and that our method disentangles shape from pose .
• We perform experiments in a simple unsupervised manner without any special equivariant loss or conditional alignment.
2. Related Work
Variational Autoencoders, as introduced by Kingma and Welling (2013) and Rezende et al.
(2014), have a great advantage of fast and tractable sampling compared to other likelihood based frameworks.
Since then a lot of work has been done to improve on VAE models by reducing the gap between approximate and true posterior distributions (Rezende and
Mohamed, 2015; Kingma et al., 2016; Cremer et al., 2018), formulating tighter bounds (Li and Turner, 2016; Burgess et al., 2018b), for extending VAEs to discrete variables (Rolfe, 2017) and addressing posterior collapse (Lucas et al., 2019).
Hierarchical VAEs with advanced neural network architectures with multi-scale dependency between compressed representations, were presented in (Vahdat and Kautz, 2020;
Child, 2020).
In (Razavi et al., 2019) the prior is learnt and hence results in higher quality of generation.
Many works have focused on disentanglement of latent spaces (Burgess et al., 2018a;
Chen et al., 2018) and explaining latent space of VAEs (Liu et al., 2019).
Our work marries two classes of VAEs in a framework of Kendall shape space VAEs; equivariant and hyperspherical .
Lafarge et al.
(2020) proposes an equivariant VAE, where the latents are equivariant without having information of pose or orientation.
Bepler et al.
(2019) disentangles translation and rotation in an image by training a VAE on a spatial grid.
Kuzina et al.
(2021) introduces an equivariant prior and results in approximate rotational equivariant latents.
These equivariant methods have the usual perks of equivariance being an important inductive bias for image data that significantly reduces model complexity (Elesedy, 2022).
They do, however, suffer from the same limitations as plain VAEs; namely, the latent space is Euclidean and the common Gaussian priors limits utilization of the space only locally around the origin.
Hyperspherical VAEs (Davidson et al., 2018, 2019), on the contrary, allow for uniform priors on the sphere, which makes that the entire space (which is compact) is utilized.
Our work inherits best of both equivariant and hyperspherical VAEs.
3. Background
As a preliminary to our method, we first review VAEs and Hyperspherical VAEs (SVAEs).
VAEs : Variational autoencoders consist of a generator p θ ( x | z ), a prior p θ ( z ) and an approximate posterior (encoder?) q ϕ ( z | x ).
When we parameterize the joint distribution by a neural network the marginalization over the latent variables (as to obtain p θ ( x )) is untractable.
So instead, our objective is to maximize the evidence lower bound (ELBO), using an approximate posterior q ϕ ( z | x ).
The ELBO is given by
[Equation]
Hyperspherical VAEs : In the original VAE, the prior and posterior are both defined as a normal distribution.
In SVAE, however, we need to work with distributions on the sphere, for which the von Mises Fisher (vMF) distribution is the natural choice as it is the stationary distribution of a convection-diffusion process on the hypersphere S m − 1 , just like the normal is on R m .
The probability density function of the vMF distribution for a random unit vector z ∈ S m − 1 is defined as
[Equation]
where
[Equation]
|| µ ||
= 1,
and I ν denotes a modified Bessel function of the first kind at order ν .
The
Kullback–Leibler divergence can be analytically computed between a vMF distribution vMF ( z | µ, κ ) and a uniform distribution on a hypersphere S m − 1 U ( x ) via
[Equation]
4. Kendall Shape VAE (KS-VAE) Framework
In (Davidson et al., 2018), the use of von Mises Fisher (vMF) distribution as a prior is motivated to support learning distributions lying on non-euclidean manifolds.
Consider a case where the data lies on a circle; the normal VAE tries to force the embedding to be mapped into an approximated posterior distribution that has support on the entire S 1 .
But this fails, as in the low dimensions, Gaussian distribution presents probability mass around the origin, while in high dimensions, it tends to resemble a uniform distribution on the surface of a hypersphere, with majority of its mass concentrated on the hyperspherical shell.
By approximating the posterior using vMF in SVAE, the resulting latents are separated on the surface of the hypersphere and in addtition the model learns more efficiently.
In this work we leverage the benefits of SVAEs, and additionally give a new interpretation by considering hyperspherical latent vectors as meaningful aligned set of points by relating them to landmarks in Kendall shape spaces.
This new viewpoint on hyperspherical VAEs further gives insight in how to make them equivariant to rotations , as we show in the next section.
4.1. Kendall Shape Spaces
(Kendall et al., 1999) defined pre-shapes by k labelled points in Euclidean space R m .
Two configurations of k labelled points are regarded as the same shape if either of their preshapes can be transformed into the other by a rotation about a shared centroid.
The quotient structure was extended in (Kendall et al., 2009) by defining shape space as the quotient of the space of landmark configurations by the group of translations, scale, and rotations.
We follow the notation of (Guigui et al., 2021) which defines a pre-shape space S km with k points ∈
R m as the space of m × k matrices, M ( m, k ).
We are interested in mapping the equivalence class formed by invariant transformations on this set of points or landmarks.
To do so we remove the effects of translation (subtract the mean) and remove effects of scaling (divide by Frobenius norm).
This defines the pre-shape space
S
[Equation]
= {
[Equation]
k
[Equation]
i
[Equation] x [Equation] = 1 } which is identified with the hypersphere of dimension S m ( k − 1) − 1 .
To remove the effect of rotations, we define the equivalence relation ∼ on S km by x ∼ y ⇐⇒ ∃ R ∈ SO ( m ) such that y
= R x for all x, y ∈ S km .
For x ∈ S km , let [ x ] then denote its equivalence class for ∼ .
This equivalence relation results from the group action of SO ( m ) on R m .
From computational perspective, we need a vector, which we label as x 0 ∈ S km , to represent the entire equivalent class [ x 0 ] =
{ R x 0 | R ∈ SO ( m ) } .
As a flattened vector x 0 ∈ R km , this vector transforms via the SO ( m ) group representation ρ ( R ) :
= ⊕ k R , i.e., a block diagonal matrix with R along the diagonal.
4.2. The Kendall Shape VAEs Requires Equivariant Encoders/Decoders
In our KS-VAE setting, we consider mappings from (square integrable) images f ∈ L 2 ( R m ) to a latent pre-shape representation z ∈ S km via an encoder Enc :
L 2 ( R m )
→ S km .
This mapping should be equivariant if we want the landmarks in S km to be geometrically meaningful with respect to the content in the input image.
That is, the content of the image may exhibit features or objects at a certain pose/orientation R , like the rotation of a digit in the rotatedMNIST dataset, or the arbitrary rotation under which a cell might present itself under a microscope in histopathology (Lafarge et al., 2020).
We would want that if the input image merely rotates that the landmarks rotate accordingly and the actual shape remains invariant (its content stays the same but only the orientation changes).
As often we are interested in the invariant content of the image f , we can talk about an equivalence class of images [ f ], where each image this class is related via a rotation.
It is important to understand that the equivalence classes themselves are invariant to scaling, translation and rotations, but the individual members z ∈ [ z ] and f ∈ [ f ] are covariant vectors under the action of SO ( m ) and only invariant to translation and scaling.
That is, an encoder Enc :
L 2 ( R m )
→ S km should be equivariant via
[Equation]
[Equation] with SIM ( m ) as a group of rotations, translations and scale on SO ( M ) and L g the leftregular representation of SIM ( m ) on m -dimensional images, given by L g f ( x ) =
f ( g − 1 x ).
Similarly, we want the decoder Dec : S km
→ L 2 ( R m ) to be equivariant.
This can be ensured if we let them be parameterized by group equivariant CNNs (Cohen and Welling, 2016), specifically by SIM ( m )-CNNs as proposed in (Knigge et al., 2022).
5. Implementation 5.1. Rotation Equivariant Layers
In this work we focus on the notion of rotation equivariance, and assume the datasets do not exhibit scale variations 2 such that our networks do not have to be scale invariant.
We then require SE ( m ) equivariant neural networks.
We rely on the escnn library (Cesa et al., 2022; Weiler and Cesa, 2019), which provide a general program for building E ( m ) equivariant neural networks.
It is based on the construction of feature fields with fibers (vector-valued feature channels) that transform under specific representations of SO ( m ).
In light of the escnn library, it is worth noting that our input images are feature fields with fibers that transform under the trivial representation of SO ( m ), i.e., each RGB value in a color image stays invariant under rotation.
The output of the encoder Enc (the preshapes) transform via type-1 irreducible representations, which are just the rotation matrices ρ 1 ( R ) =
R .
That is, a vector z ∈ S km of k landmarks in R m transforms via k rotation matrices as specified in Subsec.
The escnn library further requires the specification for the hidden feature fields, for which we use regular representations (cyclic permutations), as these fields are compatible with usual activation functions such as ReLU.
5.2. The Encoder: Equivariant vMF Parameter Extraction
In the VAE setting, the encoder should encode for the parameters of vMF distribution on the (( k − 1) m − 1)-dimensional hypersphere.
The vMF is parametrized by a mean µ ∈ S km , i.e., k landmarks with zero mean and unit norm.
As such the encoder equivariantly obtains µ as k type-1 features which we normalize to unit norm.
During training we add a loss on the squared mean as to ensure that the encoder predicts centered landmarks.
The vMF also requires a concentration parameter κ ∈ R > 0 , which the equivariant encoder predicts as a type-0 scalar ˜ k that defines κ via κ = softmax(˜ κ ) + 1 as to ensure a minimum variance of 1.
Finally, the encoder also needs to predict a pose R ∈ SO ( m ).
In 2D, this pose is encoded as a type-1 vector ˜ r ∈ R 2 , which we normalize to r = (cid:18) r x r y (cid:19) =
˜
r ∥ ˜ r ∥ from which we can extract a rotation angle via the arctan, or directly parametrize a rotation matrix
R = (cid:18) r x − r y r y r x
(cid:19) .
5.3. The Posterior Shape-vMF is Invariant to Rotations
The three components (pose R , shape µ , concentration parameter κ ) allow us to obtain a posterior distribution on the equivalence class of shapes
[ z ] by mapping the inferred µ to a µ 0 which represents the canonical shape, or class representative, via µ 0 =
ρ
( R − 1 ) µ .
It is noteworthy that the shape representative µ 0 is invariant to input rotations, since due to equivariance of Enc we have that the parameters as a function of input image f satisfy ∀
[Equation]
and R [ L [Equation]
= ˜ [Equation] [Equation] [Equation]
Thus, vector representation of the equivalence class [ µ 0 ] represented by µ 0 is truly invariant: ∀ ˜
[Equation]
SO (
[Equation]
[Equation]
Give the above, we are thus able to learn a posterior distribution for the vector representative z 0 , which represents the entire equivalence class [ f ] via q ( z 0 | f )
=
vMF ( µ 0 , κ ), since ∀ f,f 0 ∈ [ f ] : q ( z 0 | f )
=
q ( z 0 | f 0 ).
In the generative mode, we sample from this distribution a representative latent ˆ z 0 ∼ q ( z 0 | f ), map it to the corresponding pose ˆ z = ρ ( R [ f ])
z 0 using the estimated pose R , after which we pass ˆ z through the equivariant decoder to obtain the reconstructed image ˆ f = Dec(ˆ z ).
5.4. Summary: The KS-VAE recipe
• Enc( f ) → ( R , µ, κ ) • Extract the representative latent z 0 , using z 0 =
R − 1 z , where z is latent with a canonical orientation.
Sampling :
ˆ z 0 ∼ q ( z 0 | [ f ]) • Dec(ˆ z 0 ) → transforms to ˆ z and then reconstructs the image ˆ f .
6. Experiments
In this section, we perform experiments with KS-VAE and compare it with standard VAE and SVAE.
Here all the models are trained in an unsupervised way, without any prior knowledge on the equivalence classes or the poses of every instance.
6.1. Qualitative results
For qualitative results, we demonstrate that KS-VAE is rotationally equivariant.
For each equivalence class (or each orbit), KS-VAE can learn a representative f 0 .
The KS-VAE is able to encode every input image to its orbit, as well as extract the corresponding pose
R ∈ SO ( m ) that it has relative to the representative.
We observe that the KS-VAE seems to map the digits 6 and 9 to the same equivalence class.
This is efficient whilst still allowing to discern the two via the estimated pose within the class.
For each image, it generates a shape which remains consistent when the image is transformed.
In Figure 2, the first column corresponds to a rotated digit from MNIST dataset, the second column is its reconstruction, the third column is the reconstruction using a canonical orientation and the last column contains a shape made using four landmark points for the different digits.
6.2. Quantitative Results
Here in Table 1, we quantify and characterise the differences in various VAE models in the task of equivariant learning and show that Shape-VAE outperforms standard VAE and
SVAE.
All VAEs use the same architectures for the encoder and decoder; in KS-VAE, however, the conv-layers are replaced by their group convolutional counter parts for the C 8 cyclic permutation group (8 rotations covering SO (2)).
The input images are padded with ones at the bottom and right in order to obtain images of size 29 × 29.
The encoder performsconvolutions (followed by ReLU) with 5x5 kernels, without padding as to shrink the image to a single pixel.
The decoder also consists of 7 layers with 5x5 convolutions, but with padding size 4 as to iteratively expand t he image size to 29 × 29.
The feature channels increase (decrease) linearly in the encoder (decoder) from 8 in the first layer to 96 in the last layer. [Table]
7. Discussion
Like any other hypserspherical model, this one too faces with the vanishing surface area problem.
It shows unstable behavior of hyperspherical models in high dimensions.
The surface area of a hypersphere for dimension m is defined as
[Equation]
As S ( m − 1) → 0 as m → ∞ , this leads to unstable behavior.
Additionally different number of landmarks leads to different shape resulting to different shapes for the same class/image.
However, this could be remedied using a conditional VAE and figuring out the optimal number of landmarks for the most intuitive shape .
It is noteworthy that the shapes , i.e., the collection of landmarks, do not exclusively contain information about shape in the common sense, but also includes variations in writing style and line thickness.
In that sense, it is not any different to hypershperical VAEs except that our method is equivariant to rotations.
In future work, we intend to disentangle the representation in a pure shape component, and a separate appearance component which could encode for style, texture etc.
As such, our method could provide the fundamentals for learning disentangled representations in terms of shape-pose akin to classical active appearance models.
8. Conclusion
In this paper, we proposed a novel VAE framework that successfully learns consistent symbols for equivalence classes of images, which disentangles pose from shape.
In addition to the interpretable nature of the latent space, KS-VAEs outperform both standard VAEs as well as hyperspherical VAEs in terms of representational efficiency.
This underpins the efficiency in representing abstract concepts in the form of geometric symbols.
Towards this end, we presented a Kendall Shape-VAE framework for learning neural ideograms.
Acknowledgments
Sharvaree Vadgama is funded by the Hybrid Intelligence Center, a 10-year programme funded by the Dutch Ministry of Education, Culture and Science through the Netherlands
Organisation for Scientific Research.
Erik J. Bekkers is financed by the research programme
VENI (grant number 17290) funded by the Dutch Research Council.
All authors sincerely thank everyone involved in funding this work