Unsupervised Object Representation Learning using Translation and Rotation Group Equivariant VAE
Abstract
In many imaging modalities, objects of interest can occur in a variety of locationsand poses (i.e. are subject to translations and rotations in 2d or 3d), but the locationand pose of an object does not change its semantics (i.e. the object’s essence).
That is, the specific location and rotation of an airplane in satellite imagery, orthe 3d rotation of a chair in a natural image, or the rotation of a particle in acryo-electron micrograph, do not change the intrinsic nature of those objects.
Here, we consider the problem of learning semantic representations of objectsthat are invariant to pose and location in a fully unsupervised manner.
We addressshortcomings in previous approaches to this problem by introducing TARGET-VAE,a translation and rotation group-equivariant variational autoencoder framework.
TARGET-VAE combines three core innovations: 1) a rotation and translationgroup-equivariant encoder architecture, 2) a structurally disentangled distributionover latent rotation, translation, and a rotation-translation-invariant semantic objectrepresentation, which are jointly inferred by the approximate inference network,and 3) a spatially equivariant generator network.
In comprehensive experiments, weshow that TARGET-VAE learns disentangled representations without supervisionthat significantly improves upon and avoids the pathologies of previous methods.
Semantic representations learned by TARGET-VAE on images highly corruptedby rotation and translation, approach those learned on consistently posed objects,dramatically improving clustering and pose inference on multiple datasets.
Introduction
In many imaging modalities, objects of interest are arbitrarily located and oriented within the imageframe.
Examples include airplanes in satellite images, galaxies in astronomy images [1] and evenparticles in single-particle cryo-electron microscopy (cryo-EM) micrographs [2].
However, neitherthe location nor the rotation of an object within an image frame changes the nature (i e. the semantics)of
the object itself.
An airplane is an airplane regardless of where it is in the image, and differentrotations of a particle in a cryo-EM micrograph are still projections of the same protein.
Hence, thereis great interest in learning semantic representations of objects that are invariant to their locations andposes.
In general, unsupervised representation learning methods do not recover representations that disentangle the semantics of an object from its location or its pose.
Popular unsupervised deep learningmethods for images, such as variational autoencoders (VAE) [3], usually use encoder-decoder frameworks in which an encoder network is learned to map an input image to an unstructured latent variableor distribution over latent variables which is then decoded back to the input image by the decodernetwork.
However, the unstructured nature of the latent variables means that they do not separate intoany specific and interpretable sources of variation.
To achieve disentanglement, unstructured methods have been proposed that encourage independence between latent variables [4, 5].
However, thesemethods make no prior or structural assumptions about what these latent variables should encode,even when some sources of variation e.g. an object’s location or pose are prevalent and well-known.
Recently, several methods have proposed structured models that explicitly model rotation or translation within their generative networks [6, 7, 8] by formulating the image generative model as afunction mapping coordinates in space to pixel values.
Although promising, only the generativeportion of these methods is equivariant to rotation and translation, and the inference networks havelackluster performance due to poor inductive bias for these structured latents.
To address this, we propose TARGET-VAE, a Translation and Rotation Group Equivariant Variational
Auto-Encoder.
TARGET-VAE is able to learn semantic object representations that are invariantto pose and location from images corrupted by these transformations by structurally decomposingthe image generative factors into semantic, rotation, and translation components.
We performapproximate inference with a group equivariant convolutional neural network [9] and speciallyformulated approximate posterior distribution that allows us to disentangle the latent variables into arotationally equivariant latent rotation, translationally equivariant latent translation, and rotation andtranslation invariant semantic latent variables.
By combining this with a spatial generator network, ourframework is completely invariant to rotation and translation, unsupervised, and fully differentiable;the model is trained end-to-end using only observed images.
In comprehensive experiments, weshow that TARGET-VAE accurately infers the rotation and translation of objects without supervisionand learns high quality object representations even when training images are heavily confounded byrotation and translation.
We then show that this framework can be used to map continuous variationin 2D images of proteins collected with cryo-EM.
Related Work
In the recent years, there has been a great progress in developing models for unsupervised semanticanalysis of images using VAEs [6, 10, 11, 12], flow-based methods [13], Generative Adversarial
Networks (GAN) [14, 15, 16], and capsules [17].
In these methods, usually the goal is to build alow-dimensional representation of the images, which can be sampled to reconstruct similar imageswith some variations, or to use the representations to identify clusters in the data.
The unsupervisedlearning methods used for clustering objects, need to address the variations caused by the commontransformations in the objects.
One simple approach to building invariant representations, is dataaugmentation.
In this approach, training dataset plus samples generated by applying multipletransformations on them, are used to learn the transformation-invariant representations in a lowerdimensional latent space [18, 19].
This training approach is not efficient and more importantly, it doesnot guarantee invariance among the representations of the transformed versions of the same sample.
Another group of methods that use data augmentation for learning transformation-invariant features,are based on contrastive learning [20, 21, 22].
These models learn to output similar encoding forthe different transformed versions of the same data sample.
Even though this training scheme hasmade great progress in the recent years, it is not an efficient training scheme for identifying invariantrepresentations, since same as data augmentation, it needs to consider a large number of transformedsamples to achieve invariance.
In some methods, the model is trained to learn an invariant representations by constraining thelatent space [4, 5, 15].
In these methods, the latent space is not explicitly structured to separatethe representation variable from the transformation ones.
The latent space that can capture eachtargeted transformation in a structured space, can theoretically disentangle them from the contentof the image. In [17], 3D point cloud representation of objects are decomposed into capsules, in aself-supervised learning scheme, and the model learns transformation-invariant descriptors along withthe transformation-equivariant poses for the capsules, which leads to a more accurate reconstructionand clustering. A structured latent space is proposed in Spatial-VAE [6], where the latent space isdivided to variables encoding the representation and transformations.
However, only the generativemodel of Spatial-VAE is equivariant to rotation and translation, and the inference network does notperform any meaningful inference in these transformations.
To be able to disentangle the representation, from the transformations, the inference model shouldbe equivariant to the transformations.
Convolutional layers are one of the most popular featureextraction methods in image processing, due to their translation-equivariant features.
However, theyare not rotation-equivariant.
There has been many studies in designing transformation-equivariant
[Figure]
Maps
Encoder
Decoder q ( 𝑡,𝑟|𝑦) Inference module for attention q 𝜃 𝑡,𝑟,𝑦) 𝒒 𝜽 𝒕,𝒓,𝒚)~(𝓝 𝝁 𝒊𝜽 ,𝝈 𝒊𝜽 +𝜽 𝒐𝒇𝒇𝒔𝒆𝒕 ) 𝒒 𝒛 𝒕,𝒓,𝒚)~𝓝 𝝁 𝒛𝒊 ,𝝈 𝒛𝒊
Sampled z f(x 1 , x 2 ) pixel-value for (x 1 , x 2 ) x 1 , x
Fully connected layers Inference module for z
Inference module for 𝜃 𝒒 𝒕,𝒓 𝒚 ~𝒑 𝒕 𝒑(𝒓) Rotated and translated q 𝑧 𝑡,𝑟,𝑦) … x 1 x Random Fourier feature expansion z … … …
Fully-connected Layer on coordinates Fully-connected Layer on representation c o n c a t … … … and specially rotation-equivariant neural networks [9, 23, 24]. The Group-equivariant Convolutional
Neural Networks (G-CNNs) [9], are one of these proposed methods where the 2D rotation space isdiscretized, and filters are applied over this rotation dimension, in addition to the spatial dimensionsof the image.
In this method, without adding anymore parameters and by adding small extracomputational cost, the regular convolutional layers can reflect the rotation of the input images.
Using
G-CNNs in the inference model can add rotation-equivariance to the extracted features, which havenot been accounted for in the previous unsupervised object detection methods.
Regarding the equivariance in the image generation process, a number of studies have proposed totrain the generative model as a function that maps 2D or 3D spatial coordinates of pixels to theirvalues [7, 8, 25, 26]. Due this mapping, any transformation in the pixel coordinates can transformthe generated image.
However, in these studies, no direct inference on the transformations isperformed.
To the best of our knowledge, we are the first method to establish rotation and translationequivariance in both the inference and the generative models of our proposed VAE framework, alongwith disentangling the mentioned transformations as structured variables in the latent space from thecontinuous translation-rotation-invariant representation variable of the objects.
Method
In this section, we present in details our proposed VAE framework that enables an unsuperviseddisentanglement of the rotation and translation of the objects from their representations.
Figure 1shows the overall framework.
3.1 Image generation process
Images can be considered as a combination of values for discreetly identifiable pixels.
Transformationof an image, can also be considered as transforming the spatial coordinates of the pixels in the image.
To have a generation process adherent to the transformations identified in the latent space, we defineour generator as a function, which is mapping the spatial coordinates of the pixels to their values.
The generator outputs a probability distribution as p (ˆ y i | x i , z ) , where z , x i , and ˆ y i are the latentrepresentation variable, the spatial coordinates of a point, and the value for that point, respectively.
Similar to [3], for an image with n pixels, we can define the probability of the image generated in this manner as sum over probability of its pixel values:
[Equation]
To define spatial coordinates over the pixels of the image, we use a Cartesian coordinates with theorigin at the center of the image.
The translation and rotation of the coordinates can be done byshifting the origin, and rotating coordinates around it, respectively.
Depending on the value of therepresentation variable z , the generator acts as a function over the pixel coordinates to reconstruct theinput image.
In this setup the generation process is equivariant with regard to the transformations onthe coordinate system.
Assuming R ( θ ) as the rotation matrix for angle θ , and translation value of t ,the probability of the generated image from Formula 1, can be rewritten as
[Equation]
In this study, we focus on disentangling the rotation and translation from the representation of data,but the image generation probability from Formula 2, can be extended the other transformations too.
3.2 Approximate inference on latent variables
We implement a VAE framework to perform approximate inference on representation, rotation,and translation latent variables.
Since rotation optimization is non-convex, we have discretized 2Drotation space to r values, each related to a specific k 2 πr angle, where k ∈ { 0 , 1 , ..., r − 1 } .
Werepresent the approximate posterior as q ( z, θ, t, r | y ) , where z is the representation value, θ is therotation angle, and t and r refer to the translation and discretized rotation values, and y is the inputimage.
This posterior distribution is intractable, and making the simplifying assumption that q(z|t,r,y)and q( θ |t,r,y) are independent, we can factorize the posterior distribution as
[Equation]
, where q ( t, r | y ) identifies a 3D distribution over t and r , and q ( θ | t, r, y ) and q ( z | t, r, y ) , are mixturedistributions conditioned on r , t , and the input y .
We use N (0 , I ) for the prior distribution over z .
The prior over r can be uniform or Gaussian, whichis calculated based on the θ offset of each rotation dimension.
For example, when r = 4 , we haverotation values of θ offset ∈ { 0 , π 2 , π, 3 π 2 } , and assuming a uniform prior over r , means calculatingprobability of θ offset for each rotation dimension.
The prior over translation values t is set to be Gaussian with mean 0, which denotes the center of the 2D coordinate system.
The standard deviationof this prior depends on the number of pixels in our fixed-size coordinate system and whether wewant the model to be flexible with large translation values or not.
Lower standard deviation for theprior over translations, will cause the model to pose higher loss towards larger predicted translationvalues.
Using these priors, we calculate the Kullback-Leibler (KL) divergence for the latent variables:
[Equation]
To simplify this equation, we use the factorization from Equation 3, and based on calculations detailedin Appendix A1, we get:
[Equation]
, where KL t,r = (cid:80) t,r q ( t, r | y ) log q ( t [Equation]
r | y ) p ( t,r ) , KL θ = KL ( q ( θ | t, r, y ) || p ( θ | t, r, y )) , and KL z =
KL ( q ( z | t, r, y ) || p ( z | t, r, y )) .
Multiplication of Kl θ and KL z with q ( t, r ) , weights the KL-divergence on θ and z based on thejoint posterior distributions over t and r. We can calculate p ( t, r ) by multiplying p ( t ) and p ( r ) , sincet and r are independent.
The prior p ( z | t, r, y ) can be considered equal to p ( z | y ) , which is N (0 , I ) ,and p ( θ | t, r, y ) can be simplified as p ( θ | r ) since θ is independent from t .
We define this prior as
N ( θ offset , πr ) , which sets a Gaussian prior over rotation angles θ in each r dimension, with mean equal to the rotation value of the filter in that r dimension and the standard deviation that is dependenton the number of r .
We are tightening this prior by lowering their standard deviation as r grows, sothat the distributions over the different r dimension would not overlap.
The full variational lower-bound for our model in this VAE framework is calculated based on thereconstruction loss and the Kl-divergence as:
[Equation]
, where we calculate the KL-divergence value in closed form using Equation 5 during the training.
3.3 Neural network architecture
Convolutional layers are translation equivariant meaning that the spatial shift in the image causesthe same shift in their output feature maps, but convolutional layers are not rotation-equivariant.
To incorporate rotation-equivariance in the inference model, we choose to use group convolutionallayers from [9].
A symmetry group G is defined as a set X with an operation . , which is associativeon the set, and has inverse and identity elements.
Furthermore, combination of two symmetrytransformations on a set, also creates a symmetry group.
We can consider the regular convolutionallayers as a function f : Z 2 → R k , that maps every pixel to a vector of k values where k is the numberof kernels in the convolutional layer.
The convolutional layer does this by sliding the kernels on theinput image.
Replacing this shifting operation with transformations from group G , creates a functionover group G [27].
We define p r group where r ∈ N , and the operations consist of all translations, and r rotations, aboutany center of rotation in a 2D space.
In the group convolutional layer, each kernel is rotated byk 2 πr angles, where k ∈ { 0 , 1 , ..., r − 1 } and then we calculate its dot product with the input imageor feature map.
Rotating kernels happens in addition to the translation of them across the input.
As a result, the output feature map of the p r group convolutions has one extra dimension whichcorresponds to the rotations of the kernel.
As shown in Figure 1, we define three parallel 1x1 3D convolutional layers on the translation androtation equivariant feature maps of the group convolutional layers.
These layers output three logposterior distributions of log q ( t, r | y ) , log q ( z | t, r, y ) , and log q ( θ | t, r, y ) .
To build the associationbetween the θ values of each r dimension and the rotation of the kernel in that r dimension, we addthe mean values from the q ( θ | t, r, y ) with the respective θ offset values from their r dimension.
During the training, we sample from the joint posterior distribution of log q ( t, r | y ) using GumbelSoftmax [28], and use the soft one-hot output to calculate the weighted sum for the parametersthat describe the approximate posterior distributions over content (z), and rotation ( θ ).
We definea coordinate system where x is in range, and is centered on the middle of the image.
Theoutput of the Gumbel-Softmax directly describes a sample from the approximate posterior over thetranslation, which we get by calculating the weighted sum over the translation grid points.
Then wetranslate and rotate this coordinate system using the sampled t and θ values.
The generator receivesthe sampled z value along with the coordinates a of point at every pass, and outputs the value fromthat point in the image.
We use random Fourier feature expansion [29] on the coordinates of thepixel before feeding it to the fully-connected layers.
As shown in Figure 1, there are two sets ofparallel fully-connected layers each with 512 hidden units, where one receives the Fourier featureexpanded values of the coordinates, and the other one takes the z as input.
The output of these layersare concatenated before being input to a final set of fully-connected layers to output the value for thepixel.
Results4.1 Experiment setup
We run our experiments with p 4 , p 8 , and p 16 group convolutional layers in the inference model.
Theinput images are normalized to the range of [0,1].
We use 128 rectangular filters where size of thekernels in the first layer is set to a value larger than the size of the targeted objects.
We do this becausewe use only one group convolutional layer in our models and we want to make sure that our receptivefield is at least as large as the objects of interest.
We use leaky-ReLU activation functions and the batch size for all the experiments is set to 100.
We use ADAM optimizer [30], with learning rate of2e-4, and the learning rate is reduced by factor of 0.5 after no improvements in the loss for 10 epochs.
We run the training for a maximum of 500 epochs with early-stopping in case of no improvementsin the loss for 20 epochs.
We run all our experiments in a single NVIDIA A100 gpu, with 80 GBmemory.
The fully-connected layers in the generator, each has 512 hidden units, and the output can be oneprobability value in the case of binary valued pixels, or mean and standard deviation if the pixels arereal values.
For the RGB images, the dimension of output for each pixel can be easily extended tothree.
4.2 TARGET-VAE predicts the translation and rotation values of objects with highcorrelation
We evaluate how accurately, our proposed model can identify the translation and rotation values ofthe objects in the images.
For this purpose, we created two datasets using hand-written digits in
MNIST.
In the first one, the digits are randomly rotated by sampling from N (0 , π 2 16 ) and randomlytranslated by sampling from N (0 , 5 2 ) pixels.
We call this dataset MNIST(N), because of the normaldistribution over the rotation angles.
In the second dataset, which we call MNIST(U), the digitsare translated similar to the MNIST(N), but the rotations are randomly sampled from the uniformdistribution of U (0 , 2 π ) .
In both datasets, the train and test sets have 60,000, and 10,000 images,where each image is 50x50 pixels.
Appendix A2 shows some examples of the data and the rotationdistribution in these datasets.
We train TARGET-VAE with p 4 , p 8 , and p 16 on these generated datasets.
For the prior over θ , weuse N (0 , π 4 ) for the MNIST(N) dataset, and U (0 , 2 π ) for the MNIST(U).
For p ( θ | r ) we calculatep ( θ offset ( r ) ) ∼ p ( θ ) , where θ offset ( r ) is the rotation angle of the kernels in each dimension of r. Inthe training for both of the datasets, we set the prior over t to N (0 , 5 2 ) .
After the training, to evaluate our model’s performance on the prediction of the translation and rotationof the objects, we use samples from the test set as the input to the inference model to obtain q ( t, r | y )and q ( θ | t, r, y ) .
To identify the most-likely ( t, r ) , we apply the arg-max function on q ( t, r | y ) , and weuse them to get θ from q ( θ | t, r, y ) .
For the translation values, we calculate Pearson correlation, butbecause of the circular nature of the rotations values, we calculate the circular correlation [31] forthem.
Table 1, shows the correlation values on the test set for MNIST(N) and MNIST(U).
The valuesin these tables are based on the model trained with dimension of the representation space ( z _ dim )set to 2.
To the best of our knowledge, the only other method that performs direct inference on therotation values to identify the invariant representations for the objects, is spatial-VAE [6].
To compareour results, we run spatial-VAE with the same number of parameters and z_dim=2.
As shown in Table1, in both MNIST(N) and MNIST(U) datasets, TARGET-VAE significantly improves the rotationcorrelation value. Our experiments show that increasing dimension of the representation space, doesnot have a significant effect on the translation and rotation correlation values, but increasing thelevel of the discretization over the rotation dimension usually improves the rotation correlation. Thisis confirmed by TARGET-VAE with p 16 achieving the highest rotation correlation value. In thetranslation column, the two values are showing the translation correlation for x 1 and x 2 coordinates.
We believe the reason for the slightly higher translation correlation of spatial-VAE than our proposedmodels might be in the fact that TARGET-VAE can predict the translations only in a discrete pixellevel manner, and this might affect its ability in identifying the translations that have sub-pixel values.
[Table] [Figure] Some of the digits in MNIST datasets, have natural ambiguity with regard to their rotations.
Investigating the correlation between the predicted rotation values and the ground-truth ones, shows adifference about π between the some of the predicted angles and their ground-truth ones.
Figure 2shows this difference.
By drawing digit-wise rotation correlation figures in appendix A.3, we showthat this π difference mostly belongs to the three digits of 0, 1, and 8.
This shows how our model canacknowledge the rotation symmetry in these digits.
4.3 TARGET-VAE improves clustering accuracy by learning translation and rotationinvariant features
By disentangling rotation and translation from the representation values, TARGET-VAE can identifythe different clusters of the objects.
Similar to 4.2, we use arg-max ( q ( t, r )) to identify the mostprobable representation from q ( z | t, r, y ) .
We train TARGET-VAE with p 4 , p 8 , and p 16 and evaluatetheir clustering performance using agglomerative clustering [32] on MNIST(N) and MNIST(U).
Table 2 shows the clustering accuracies on these two datasets, where increasing the dimension ofthe representation space, helps the model with the separation of different digits and achieving higherclustering accuracy.
Based on the accuracies in this table and rotation correlations in Table 1, higherdiscretization of the rotation space, allows the TARGET-VAE to have a better estimate of the rotationof the objects, and this is turn results in more rotation-invariant representation encoding and higherclustering accuracy.
To compare the clustering performance of TARGET-VAE with other similarmethods, we train a standard VAE [3], beta-VAE [4], and spatial-VAE [6] with about same number ofparameters as TARGET-VAE, on MNIST(N) and MNIST(U) datasets.
We set the z_dim=2 for bothspatial-VAE and TARGET-VAE.
Since the standard VAE and the beta-VAE do not have dedicatedlatent variables for identifying translation and rotation, we set the dimension of latent space to 4 forthem, so that all of the models would have equal latent space dimensions.
Also, for the beta-VAE,we set β value to 4 based on the proposed settings in [4].
As shown in the Table 2, all variationsof the TARGET-VAE outperform the other methods.
We believe that the direct inference on therotation and translation, and dedicating structured latent variables to encode these transformations in TARGET-VAE, plus the transformation-equivariant generator, are the main reason for the dramaticincrease in the clustering accuracy with our proposed method.
In MNIST(N), since there are limitedrange for the rotation values, increasing the size of the group convolution, does not necessarilyincrease the rotation correlation or the clustering accuracy, where in MNIST(U) due to the largervariations in the rotation values, increasing the size of the group convolution leads to a better estimateof rotations and improved clustering accuracy.
Figure 3 shows the confusion matrix and t-sne plot for TARGET-VAE with p 8 group-convolution, on MNIST(N).
In appendix A.4, we perform experiments using images with multiple digits to show that TARGETVAE can detect multiple objects in the evaluation stage.
Furthermore, in appendix A.5 and A.6,we train and evaluate TARGET-VAE on the dSprites [33], and galaxy zoo [1].
We also perform anablation study on the important components of our proposed model in Appendix A.7. [Table]
4.4 Identifying class-averages and protein conformations in the electron-microscopy imageswith TARGET-VAE
In the single particle electron microscopy images, biological particles can appear with differentrotations and in different locations of the images.
To be able to identify different types of particlesor different conformations of the protein, the representation vector should be invariant to thesetransformations.
Here, we demonstrate the capacity of TARGET-VAE in capturing different classesand variable conformations of particles in the electron microscopy images.
Identifying different classes of particles in EMPIAR-10025.
The dataset EMPIAR-10025 [34]contains T20S Proteasome micrographs with 2.8 Å resolution.
We use Topaz [35] package to create161292 images of 400x400 pixels, where each image contains a particle.
After downsampling byfactor of 4 and normalizing the images, we train TARGET-VAE with p 8 and z_dim=2 on them.
Dueto the high-complexity of the structures of the proteins, we use 6 fully-connected layers with 512units in the generator.
In the rotation inference our prior over r is uniform.
We randomly select 10%of images as validation set to monitor the training process.
As shown in Figure 4, our model identifiesthree clusters in this dataset.
The image on the left side of the Figure 4 contains instances of inputimages which are extremely noisy.
Identifying different views of particles in EMPIAR-10029.
The EMPIAR-10029 dataset containssimulated EM dataset of GroEL particles.
This dataset has 10000 images of shape 200x200, whichwe downsample to 100x100 and normalize.
Similar to EMPIAR-10025, we train TARGET-VAE withthe same settings as the EMPIAR-10025.
The latent space does not identify separate clusters in thisdataset, but performing a random sampling in the latent space helps us with identifying differentprojections of the biological particle (Fig. 5).
We have also tested our model in identifying theheterogeneity in some other cryo-EM datasets, which we discuss in appendix A.8.
Conclusion
In this paper, we presented TARGET-VAE, a translation and rotation group equivariant variationalautoencoder framework for learning translation and rotation invariant object representation in imageswithout supervision.
By structuring the encoder to be translation and rotation equivariant, the modelcan identify the latent rotation and translation variables and disentangle them from the latent variablesencoding the content of the image.
The spatially equivariant generator samples from the posteriordistribution to reconstruct the image.
We showed that this approach not only accurately estimates therotation and translation of objects, but also learns content representations that reflect known semantics(i.e. , clusters match known semantic labels) across multiple datasets.
Although we only consider 2dimages containing single objects with rotation and translation, this framework can be extended toother object transformations by adopting encoder networks with the proper equivariance and adaptingthe latent distribution.
In the future, we expect this framework can also be extended to multi-objectdetection, object tracking over time, and to 3d environments.
TARGET-VAE lays the groundwork fora new generation of fully unsupervised object detection and semantic analysis methods for cryo- EMand other imaging modalities.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] Our proposed framework can perform inference on the
rotation and translation values of the object and use this inference to learn a rotation
and translation invariant representation value for them.
(b) Did you describe the limitations of your work?
[Yes] We mentioned the limitations inthe conclusion in the form the future directions.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
To the
best of our knowledge, our work does not have any negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes] We have read the ethics review guidelines and we believe that our paper conforms to that.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] In the partsthat we made any assumptions, we clearly explained them.(b)
Did you include complete proofs of all theoretical results?
[Yes] We included proof for one of our equations in section 3.2 which needs proof.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] The code isavailable at —————————–(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] All the details are either explained in the paper or can be found inthe code.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] The variation in the results is not large enough to justify using error bars.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] We included all the processing information.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A] All the data that we are using is the public domain.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
The data that we are using does not contain any personally identifiable information or offensive content.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Appendix
A.1 Calculating Kullback-Leibler divergence
Based on the standard definition for the KL-divergence, we have:
[Equation]
To simplify this equation
[Equation]
we define the first part of the result from Equation 7 as:
[Equation]
The second part of the result from Equation 7 can be further expanded as:
[Equation]
Assuming these definitions for KL θ and KL z :
[Equation]
[Equation]
, we can rewrite the the Equation 7 using Equations 8, 9, 10, and 11 as:
[Equation]
A.2 MNIST(N) and MNIST(U) datasets
We generated two datasets, MNIST(N) and MNIST(U), by rotating and translating digits in MNIST.
The rotations of digits in MNIST(N) are randomly sampled from N (0 , π 2 16 ) , and in MNIST(U) arerandomly sampled from U (0, 2 π ) (Figure 6).
A.3 Digit-wise rotation correlation, and RMSE of the predicted rotations
Following our discussion about the difference between the predicted angle and the ground-truthvalues, we measure the digit-wise rotation correlation for each dataset (Figure 7).
We find that theerrors in the predicted rotations are concentrated around π for digits 0, 1, and 8.
We suspect that thisis because these digits have rough 180 degrees symmetry.
To study the accuracy of the predicted rotation angles by TARGET-VAE, we calculate the meanstandard deviation of the predicted rotations, introduced in [17].
This metric basically measures themean square error between the rotation of the object in the input image and the predicted rotationfor that object.
We rotate each image in the test set of MNIST(U), 160 times using angles uniformlypicked from [0, 2 π ], and then pass them through the encoder to get the predicted rotation angle forthem.
Since assigning predicted rotation of zero to an object is arbitrary in our framework (modelmight assign predicted rotation angle of zero to the object that is actually rotated 90 degrees), we subtract the predicted rotation for each one of 160 rotated objects from the predicted rotation valuewhen the input object is not rotated.
Table 3, shows the root mean squared error (RMSE) calculatedin average for each digit’s predicted rotation.
As expected, the RMSE of rotations is highest whenthere is no inference done on rotations as in Spatial-VAE [6], and using finer discretization in rotationinference results in more accurate prediction of the rotation values.
Some digits such as 0, 1, and 8,due to their symmetry, have less accurate rotation predictions compared to the other ones.
[Table]
A.4 TARGET-VAE identifies multiple objects without supervision
We create a new dataset using multiple rotated and translated digits from MNIST(U).
We selectimages at random from MNIST(U) and insert them in random locations of a 150x150 pixel frame.
We call this new dataset MNIST(multi).
We use the model trained on MNIST(U) to identify thelocation, rotation, and the content latent for images in MNIST(multi).
We find that the modelcorrectly identifies and reconstructs the objects (Figure 8).
Even though the model, can identifynon-overlapping objects in this approach, it struggles to identify the overlapping objects.
We believethese preliminary results show the potential of our model to be trained on images with multipleobjects and to identify the objects independent of their location and rotation.
[Figure]
for visualization purposes.
The high probability values in this attention map show the predicted locations of the objects.
We use the peaks in q ( t, r | y ) to sample from q ( θ | t, r, y ) and q ( z | t, r, y ) , the predicted rotation and content values for each object.
The sampled θ , t , and z values are used to reconstruct each individual object.
A.5 TARGET-VAE predicts the rotation of dSprites shapes and clusters them with highaccuracy
The dSprites dataset was introduced to benchmark unsupervised disentangled representation learningmethods.
It contains 64x64 images of three shapes: squares, ellipses, and hearts.
Each shape isrotated by one of 40 values in [0,2 π ], translated across both x and y dimensions, and scaled usingone of six linearly spaced values in [0.5,1].
We train TARGET-VAE with p 8 group convolution andz_dim=2 on dSprites.
We use a uniform prior over r and N (0 , π ) prior over θ .
TARGET-VAE learnscontent representations that capture the type and scale of objects, but not their location and rotation(9).
Furthermore, TARGET-VAE accurately predicts the rotation of the shapes, and, interestingly,the symmetry groups of the shapes become immediately apparent when examining the correlationbetween the predicted and ground truth rotations.
A.6 TARGET-VAE learns transformation-invariant representations in the galaxy dataset
Galaxy zoo contains images of galaxies gathered by Sloan Digital Sky Survey.
This dataset containsmore than 61000 RGB images, which we cropped and downsampled to 64x64 and split into trainand test sets.
The galaxies appear with different rotations and translations in the images.
We train
TARGET-VAE with p 8 group convolution on the train set.
TARGET-VAE learns to accurately predict the rotation and translation of the galaxies and reconstructs centered and aligned galaxy images whendecoding with t and θ set to zero (Figure 10).
A.7 Ablation studies
We conducted extensive ablation study to validate the effectiveness of the main components of ourproposed method.
Variant 1 - Inference only on translations We evaluated the importance of the rotation equivariantfeatures, by modifying the inference model to use regular convolutional layers instead of the groupconvolutional ones.
As a result, the posterior distributions will be q(t|y), q(z|t,y), and q( θ |t,y),where they only depend on the input y and the translation value t. We observed that, as expected,
[Table]
[Table]
eliminating inference on the discretized rotation dimension has a negative effect on identifyingtransformation-invariant representations (Table 4).
Variant 2 - Inference only on translations + using group convolutional layers In this variant, weuse the group convolutional layers in the inference model, but we do not perform any inference on therotation dimension.
We apply a fully-connected layer on the output of the group convolutional layersto map the rotation values at every location, to a single value.
As a result the posterior distribution arethe same as variant 1.
We experiment with this variant to identify how effective is the use of the groupconvolution layers.
It turns out that if we do not perform any inference on the rotation dimension, justusing group convolutional layers for feature extraction only slightly improves the clustering accuracy.
Variant 3 - Inference on both translations and rotations without adding θ offset : In this variant,we use group convolutional layers and we perform inference on both translations and rotations.
Thedifference between this variant and our proposed framework is that we are not adding θ offset ofkernels in each rotation dimension to the q( θ |t,r,y).
Adding θ offset to the posterior on θ , allows us tobreak down the rotation space among the r rotation dimension, and without it, the model proves notcapable of identifying the rotations of the digits in the MNIST(U) dataset.
Variant 4 to 6 - Increasing the level of discretization of the rotation space
In these variants,we perform inference on rotation and translations and we add θ offset of each rotation dimension toits corresponding q( θ |t,r,y) distribution.
In these variants, we show that by increasing the level ofdiscretization of the rotation space, model can have a better estimate of the actual rotation values, andthis in turn helps the model with improving the clustering accuracy.
Table 4, shows the translation, rotation correlation along with the clustering accuracy of the mentionvariants above, on the MNIST(U) dataset.
Performing inference on rotations plus adding θ offset tothe posterior on θ , significantly increases the rotation correlation and the clustering accuracy.
A.8 Learning transformation-invariant representations of proteins with TARGET-VAE
In cryo-EM images, rotation and translation are the major transformations that cause variations inparticles in the image.
Here, we show the result of our experiments with TARGET-VAE to learntransformation-invariant representations, for some more cryo-EM datasets.
Identifying hinge motion of 5HDB We train our model on a dataset of 20000 simulated projectionsof integrin α -IIb in complex with integrin β
(5HDB).
We aim to identify the rotation invariantrepresentations of the protein to be able to identify the variations in its structure.
We train TARGETVAE with p 8 and z_dim=1 and uniform prior over θ .
Figure 11 shows some examples of the inputimages and their reconstructed versions with translation and rotation set to 0.
After training, wesample from the representation latent space and reconstruct the images with no rotation and translation.
We observe that the reconstructed images identifies different conformations of the particle.
Learning arm motion of particle in CODH/
ACS We have about 14000 40x40 pixels images ofthe CODH/ACS protein complex.
We train TARGET-VAE with p 8 and z_dim=2 for 100 epochs.
Theprior over θ is set to π
and we set the generator to have 6 fully-connected layers with 512 hiddenunits in each.
After training, we sample from the transformation-invariant representation space andreconstruct the particle to identify different conformations of the protein.
Figure 12 shows some ofthese different conformations.