Unsupervised Learning of Group Invariant and Equivariant Representations
Anonymous Author(s)
AffiliationAddress email
Abstract
Equivariant neural networks, whose hidden features transform according to representations of a group G acting on the data, exhibit training efficiency and animproved generalisation performance.
In this work, we extend group invariant andequivariant representation learning to the field of unsupervised deep learning.
Wepropose a general learning strategy based on an encoder-decoder framework inwhich the latent representation is separated in an invariant term and an equivariantgroup action component.
The key idea is that the network learns to encode anddecode data to and from a group-invariant representation by additionally learningto predict the appropriate group action to align input and output pose to solve thereconstruction task.
We derive the necessary conditions on the equivariant encoder,and we present a construction valid for any G , both discrete and continuous.
Wedescribe explicitly our construction for rotations, translations and permutations.
We test the validity and the robustness of our approach in a variety of experimentswith diverse data types employing different network architectures.
Introduction
An increasing body of work has shown that incorporating knowledge about underlying symmetries inneural networks as inductive bias can drastically improve the performance and reduce the amount ofdata needed for training Cohen & Welling (2016a); Bronstein et al. (2021).
For example, the equivariant design with respect to the translation symmetry of objects in images proper of convolutional neuralnetworks (CNNs) has revolutionized the field of image analysis LeCun et al.
Message Passingneural networks, respecting permutation symmetries in graphs, have enabled powerful predictivemodels on graph-structured data Gilmer et al.
(2017); Defferrard et al.
Recently, much workhas been done utilizing 3D rotation and translation equivariant neural networks for point cloudsand volumetric data, showing great success in predicting molecular ground state energy levels withhigh fidelity Miller et al.
; Anderson et al.
; Klicpera et al.
Schütt et al.
Invariant models take advantage of the fact that often properties of interest, such as the class label ofan object in an image or the ground state energy of a molecule, are invariant to certain group actions(e.g., translations or rotations), while the data itself is not (e.g., pixel values, atom coordinates).
There are several approaches to incorporate invariance into the learned representation of a neuralnetwork.
The most common approach consists of teaching invariance to the model by data augmentation: during training, the model must learn that a group transformation on its input does not affectits label.
While this approach can lead to improved generalization performance, it reduces trainingefficiency and quickly becomes impractical for higher dimensional data Thomas et al.
Asecond technique, known as feature averaging, consists of averaging model predictions over grouptransformations of the input Puny et al.
While feasible with finite groups, this method requires,for instance, sampling for infinite groups Lyle et al.
A third approach is to impose invariance
[Figure]
[Equation]
[Figure]
[Equation]
𝑍 as a model architectural design.
The simplest option is to restrict the function to be learned to be acomposition of symmetric functions only Schütt et al.
Such choice, however, can significantlyrestrict the functional form of the network.
A more expressive variation of this approach consistsof an equivariant neural network, followed by a symmetric function.
This allows the network toleverage the benefits of invariance while having a larger capacity due to the less restrictive natureof equivariance.
In fact, in many real-world application, equivariance is beneficial if not necessary
Smidt (2020); Miller et al.
For example, the interaction of a molecule (per se rotationalinvariant) with an external magnetic field is an intrinsically equivariant problem.
All aforementioned considerations require some sort of supervision to extract invariant representationsfrom data.
Unsupervised learning of group invariant representations, despite its potential in the fieldof representation learning, has been impaired by the fact that the representation of the data in generaldoes not manifestly exhibit the group as a symmetry.
For instance, in the case of an encoder-decoderframework in which the bottleneck layer is invariant, the reconstruction is only possible up to a grouptransformation.
Nevertheless, the input data is typically parametrized in terms of coordinates in somevector space X , and the reconstruction task can only succeed by employing knowledge about thegroup action on X .
Following this line of thought, this work is concerned with the question: Can we learn to extract boththe invariant and equivariant representations of data in an unsupervised way?
To this end, we introduce a group-invariant representation learning method that encodes data in agroup-invariant latent code and a group action.
By separating the embedding in a group-invariant anda group-equivariant part, we can learn expressive lower-dimensional group-invariant representationsutilizing the power of autoencoders (AEs).
In particular, all our results trivially extend to variational
AEs (VAEs), thus advancing generative approaches of group-invariant data.
We can summarize themain contributions of this work as follows:
• We introduce a novel framework for learning group equivariant representations.
Our representations are by construction separated in an invariant and group action function (equivariant) component.
• We characterize the mathematical conditions of the group action function component andwe propose an explicit construction suitable for any group G .
To the best of our knowledge, this is the first method for unsupervised learning of separated invariant-equivariantrepresentations valid for any group.
• We show in various experiments the validity and flexibility of our framework by learningrepresentations of diverse data types with different network architectures.
Method2.1 Background
We begin this section by introducing the basic concepts which will be central in our work.
A group G is a set equipped with an operation (here denoted · ) which is associative as well as havingan identity element e and inverse elements.
In the context of data, we are mainly interested in howgroups represent geometric transformations by acting on spaces and, in particular, how they describethe symmetries of an object or of a set.
In either case, we are interested in how groups act on spaces.
This is represented by a group action : given a set X and a group G , a (left) action of G on X is amap ρ : G × X → X such that it respects the group property of associativity and identity element.
If
X is a vector space, which we will assume for the remainder of the text, we refer to group actions ofthe form
ρ X : G → GL ( X ) as representations of G , where the general linear group of degree n
GL ( X ) is represented by the set of n × n invertible matrices.
Given a group action, a concept whichwill play an important role in our discussion is given by the fixed points of such an action.
Formally,given a point x ∈ X and an action (representation) ρ X of G on X , the stabilizer of G with respect tox is the subgroup G x = { g ∈ G | ρ X ( g ) x
= x } ⊂ G .
In the context of representation learning, we assume our data to be defined as the space ofrepresentation-valued functions on some set V ,
i.e., X = { f | f : V → W } .
For instance, apoint cloud in three dimensions can be represented as the set of functions f : R 3 → Z 2 , assigning toevery point r ∈ R 3 the value f ( r ) =
(the point is not included in the cloud) or f ( r ) = 1 (the pointis included in the cloud).
Representations ρ V of a group G on V can be extended to representationson f , and therefore on X , ρ X : G → GL ( X ) , as follows
[Equation]
In what follows, we will then only refer to representations for the space X , implicitly referring toequation (1) for mapping back to how the various components transform.
A map φ : V → W is saidto be G -equivariant with respect to the actions (representations) ρ V , ρ W if φ ( ρ V ( g ) v ) =
ρ W ( g ) φ ( v )for every g ∈ G and v ∈ V .
Note that G -invariance is a particular case of the above, where wetake ρ V , ρ W to be the trivial representations.
An element x ∈ X can be described in terms of a
G -invariant component and a group element g ∈ G , as follows: let φ inv : X → X/G , be an invariantmap mapping each element x ∈ X to a corresponding canonical element ˆ x in the orbit in the quotientspace X/G .
Then for each x ∈ X there exist a g ∈ G such that x =
ρ X ( g ) φ inv ( x ) .
2.2 Problem Definition
We consider a classical autoencoder framework with encoding function η : X → Z and decodingfunction δ : Z → X , mapping between the data domain X , and latent domain Z , minimizing thereconstruction objective d ( δ ( η ( x )) , x ) , with a difference measure d (e.g., L p norm).
As discussedabove, we wish to learn the invariant map η ( φ inv in the previous paragraph), thus
Property 2.1.
The encoding function η : X → Z is G -invariant, i.e., η ( ρ X ( g ) x ) =
η ( x ) ∀ x ∈
X, ∀ g ∈ G .
The decoding function δ maps the G -invariant representation
z ∈ Z back to the data domain X .
However, as z is G -invariant, δ can at best map η ( x ) ∈
Z back to an element ˆ x ∈ X such thatˆ x ∈ { ρ X ( g ) x |∀ g ∈ G } , i.e., an element in the orbit of x through G .
This is depicted in Figure 1a.
Thus, the task of the decoding function δ : Z → X is to map encoded elements z
= η ( x ) ∈
Z to anelement ˆ x ∈ X such that ∃ ˆ g x ∈ G such that
[Equation]
We call ˆ x the canonical element of the decoder δ .
We can rewrite the reconstruction objective witha G -invariant encoding function η as d ( ρ X (ˆ g − 1 ) δ ( η ( x )) , x ) .
One of the main results of this workconsists in showing that ˆ x and ˆ g x can be simultaneously learned by a suitable neural network.
That is,we have the following property of our learning scheme:
Property 2.2.
There exists a learnable function ψ : X → G such that, given suitable η, δ asdescribed above the relation ρ X ( ψ ( x ))
δ ( η ( x )) =
x , holds for all x ∈ X .
We call any function ψ satisfying (2.2) a suitable group function.
Figure 1b describes schematicallyour proposed framework.
In what follows, we will first characterize the defining properties of suitablegroup functions.
Subsequently, we will describe our construction, valid for any group G .
2.3 Predicting Group Actions
In the following we further characterize the properties of ψ .
We begin by stating two key results,while we refer to the Appendix A for the proofs.
Proposition 2.3.
Any suitable group function ψ : X → G is G -equivariant at a point x ∈ X up thestabilizer G x , i.e., ψ
( ρ X ( g ) x ) ⊆ g · ψ ( x ) G x .
Proposition 2.4.
The image of any suitable group function ψ : X → G is surjective into GG X , where
G X is the stabilizers of all the points of X .
Let us briefly discuss an example.
Suppose X = { x = ( x 0 , x 1 , x 2 , x 3 ) ∈ R 4 × 2 | x i =
ρ R 2 ( g θ = π/ 2 ) i x 0 , x 0 ∈ R 2 } and G = SO (2) .
X describes all collections of vertices of squarescentered at the origin of R 2 , and it is easy to check that g X
=
Z 4 , generated by a π/ 2 rotation aroundthe origin.
In this case, any such square can be brought to any other square (of the same radius) by arotation of an angle θ < π/ 2 , thus Im ψ ⊇ { g θ ∈ SO (2) | 0 ≤ θ ≤ π/ 2 } =
SO (2) /
Z 4 .
Combining the two propositions above we have the following
Lemma 2.5.
Any suitable group function ψ is an isomorphism
O x ≃ G/G x for any x ∈ X , where
O x ⊂ X is the orbit of x with respect to G in X .
2.4 Proposed Construction
Next, we turn to our proposed construction of a class of suitable group functions that satisfy Property2.2 for any data space X and group G .
As we described above, these functions must be learnable.
Property 2.6 ( Proposed construction ) .
Without loss of generality, we write our target functionψ =
ξ ◦ µ , where µ : X → Y is a learnable map between the data space X and the embeddingspace Y , while ξ : Y → G is a deterministic map.
Our construction is further determined by thefollowing properties:
• We impose µ : X → Y to be G -equivariant, that is, µ ( ρ X ( g ) x ) =
ρ Y ( g ) µ ( x ) for allx ∈ X and g ∈ G .
• We ask that Y is an homogeneous space, that is, given any element y 0 ∈ Y , every elementy ∈ Y can be written as y =
ρ Y ( g ) y
for some g ∈ G .•
The map ξ : Y → G is defined as follows: ξ ( y ) =
g such that
y
=
ρ Y ( g ) y
for any chosenpoint y 0 ∈ Y .
In what follows we will show that our construction satisfies the properties of the previous section.
For proofs see Appendix.
We begin with the following
Proposition 2.7.
Let ψ = ξ ◦ µ be a suitable group function and let µ : X → Y be G -equivariant.
Then, G x = G µ ( x ) for all x ∈ X .
The result of the above proposition is crucial for our desired decomposition of the learned embedding,as it ensures that no information about the group action on X is lost through the map µ : if a groupelement acts non-trivially in X , it will also act non-trivially in Y .
Proposition 2.8.
Given y, y 0 , the element g such that y ≡ ρ Y ( g ) y 0 is unique up to the stabilizer
G y 0 .
This proposition establishes the equivariant properties of the map ξ .
Finally, we have
Proposition 2.9.
Let ψ = ξ ◦ µ where µ and ξ are as described above.
Then, ψ is a suitable groupfunction.
2.5 Intuition Behind the Proposed Framework
We conclude this rather technical section with a comment on the intuition behind our construction.
Assuming for simplicity that the domain set V admits the structure of vector space, Y representsthe space spanned by all basis vectors of V .
The point y 0 represent a canonical orientation of suchbasis, and the element ξ ( y )
=
g is the group element corresponding to a basis transformation.
Asall elements can be expressed in terms of coordinates with respect to a given basis, it is natural toconsider a canonical basis for all orbits, justifying the assumption of homonogeneity of the space Y .
Further, let us assume that the invariant autoencoder correctly solves its task, x ∼ δ ( η ( x )) .
Nowlet ˆ x ∈ O x such that ˆ x = δ ( η ( x )) , and by definition, ˆ x = ρ V ( g ) x for some g ∈ G .
Now, the correct orbit element is identified when ψ (ˆ x ) =
e , since ψ ( x ) = g − 1 · ψ (ˆ x ) =
g − 1 and thusρ X
( g − 1 )
δ ( η ( x )) =
ρ X ( g − 1 ) x ′ = x .
Hence, during training ψ needs to learn which orbit elementsare decoded as “canonical”, i.e., without the need of an additional group transformation.
To clarify,here “canonical” does not reflect any specific property of the element, but it simply refers to theorientation learned from the decoder during training.
In fact, different decoder architectures orinitializations will lead to different canonical elements.
Finally, note how the different parts of our proposed framework, as visualized in Figure 1b, can bejointly trained by minimizing the objective d ( ρ X ( ψ ( x ))
δ ( η ( x )) , x ) , which is by construction groupinvariant, i.e., not susceptible to potential group-related bias in the data (e.g. data that only occurs incertain orientations).
Application to Common Groups
In this section we describe how our framework applies to a variety of common groups which we willthen implement in our experiments.
As discussed in Section 2.2 and visualized in Figure 1b, the maincomponents of our proposed framework are the encoding function η , the decoding function
δ and thegroup function ψ .
As stated in Property 2.1, the only constraint for the encoding function
η is that ithas to be group invariant.
This is in general straightforward to achieve for different groups as wewill demonstrate in Section 5.
Our proposed framework does not constrain the decoding function δother than that it has to map elements from the latent space Z to the data domain X .
Hence, δ can bedesigned independently of the group of interest.
The main challenge is in defining the group functionψ =
ξ ◦ µ such that it satisfies Property 2.2.
Following Property 2.6 we now turn to describing ourconstruction of ξ , µ and Y for a variety of common groups.
Orthogonal group SO (2) .
The Lie group SO (2) is defined as the set of all rotations about theorigin in R 2 .
We take Y to be the circle S 1 ⊂ R 2 , that is, the space spanned by unit vectors in R 2 .
Now, S 1 is a homogeneous space: any two points s 0 , s 1 ∈ S 1 are related by a rotation.
Withoutloss of generality, we take the reference vector y
to be the vector (1 , 0) ∈ S 1 .
Then given a vectory ∈ S 1 , we can write
[Equation]
thus, the function ξ : S 1 → SO (2) is determined by ξ ( y ) =
g θ such that θ = arccos( y x ) =
arcsin( y y ) .
Orthogonal group SO (3) .
We assume that X has no fixed points, as this is usually the case forgeneric shapes (point clouds) in R 3 .
It would be tempting to take Y to be the sphere S 2 ⊂ R 3 ,that is, the space spanned by unit vectors in R 3 .
While this space is homogeneous, it does notsatisfy the condition that the stabilizers of G are trivial.
In fact, given any vector y 1 ∈ S 2 , we have
G y 1 = { g ∈ SO (3) | g is a rotation about y 1 } .
In order to construct a space with the desired property, consider a second vector w 2 ∈ S 2 orthogonalto y 1 , y 2 ⊂ y ⊥ 1 .
Taking Y to be the space spanned by y 1 , y 2 ∈ S 2 , it is easy to see that now all thestabilizers are trivial.
Finally, let y
=
y 1 × y 2 ∈ S 2 , then we construct the rotation matrix
R
[Equation]
y
[Equation]
, which satisfies
[Equation]
(cid:18) 1 00 1(cid:19) ⊺ =
Ry 0 .
Symmetric group S n .
A suitable space Y is the set of ordered collections of unique elements of the set M = { 1 , 2 , . . . , n } .
For instance, for n = 3 , we have Y =
It is trivial to see that the action of thepermutation group on the set Y is free, that is, all the stabilizers are trivial.
Explicitely, given anypermutation-equivariant vector w ∈ R n , we obtain an element y = argsort ( w ) ∈ Y .
Moreover, it isalso obvious that any element in Y can be written as P σ (1 , 2 , . . . , n ) =
( σ (1) , σ (2) , . . .
, σ ( n )) , thatis, a group element acting on the canonical y
= (1 , 2 , . . . , n ) .
Translation group T n .
Here we take Y = R n , which is homogeneous with respect to the translationgroup.
In fact, any vector y ∈ Y can be trivially written as y =
y
+ 0 =
y + y 0 , where 0 is the originof R n .
Our group function takes therefore the form ξ ( y ) =
y .
Euclidean group SE ( n ) .
A generic transformation of the Euclidean group on a n -dimensionalrepresentation v ∈ V isv (cid:55)→
[Equation]
Let µ = ( µ 1 , µ 2 , . . .
, µ n +1 ) be a collection of n + 1 n -dimensional SE ( n ) -equivariant vectors, thatis, µ i ( ρ X ( g ) x ) =
ρ Y ( g ) µ ( x ) , i = 1 , . . . , n .
We construct (cid:98) y
a = ( µ a − µ n +1 ) /
|| µ a − µ n +1 || ∈ S n ,a = 1 , . . . , n , where S n is the unit n -dimensional sphere.
These n ortho-normal vectors aretranslation invariant but rotation equivariant, and are suitable to construct the rotation matrix
[Equation]
(5) while the extra vector (cid:98) y n +1 = µ n +1 can be used to predict the translation action.
Putting all together,the space Y is described by n vectors y a = (cid:98)
y a + (cid:98) y n +1 , and y
= I n is the n × n unit matrix, as
[Equation]
Related Work
Group equivariant neural networks.
Group equivariant neural networks have shown great successfor various groups and data types.
There are two main approaches to implement equivariance in alayer and, hence, in a neural network.
The first, and perhaps the most common, imposes equivarianceon the space of functions and features learned by the network.
Thus, the parameters of the modelare constrained to satisfy equivariance Thomas et al.
(2018); Weiler & Cesa (2019a); Weiler et al.(2018a); Esteves et al.
The disadvantage of this approach consists in the difficulty of designingsuitable architectures for all components of the model, transforming correctly under the group action
Xu et al.
The second approach to equivariance consists in lifting the map from the spaceof features to the group G , and equivariance is defined on functions on the group itself Romero &
Hoogendoorn (2020); Romero et al.
Hoogeboom et al.
Although this strategy avoidsthe architectural constraints, applicability is limited to homogeneous spaces Hutchinson et al.
(2021)and involves an increased dimensionality of the feature space, due to the lifting to G .
Equivariancehas been explored in a variety of architecture and data structures: Convolutional Neural Networks
Cohen & Welling (2016a); Worrall et al.
; Weiler et al.
(2018c)
; Bekkers et al.
; Thomaset al.
; Dieleman et al.
Kondor & Trivedi (2018); Cohen & Welling (2016b); Cohenet al.
; Finzi et al.
Transformers Vaswani et al.
Fuchs et al.
Hutchinsonet al. (2021); Romero & Cordonnier (2020), Graph Neural Networks Defferrard et al.
; Brunaet al. (2013); Kipf & Welling (2016); Gilmer et al.
Satorras et al.
(2021) and Normalizing
Flows Rezende & Mohamed (2015); Köhler et al.
; Boyda et al.
These methodsare usually trained in a supervised manner and combined with a symmetric function (e.g. pooling) toextract group-invariant representations.
Group equivariant autoencoders.
Another line of related work is concerned with group equivariantautoencoders.
Such models utilize specific network architectures to encode and decode data in anequivariant way, resulting into equivariant representations only Hinton et al.
(2011); Sabour et al.(2017); Kosiorek et al.
; Guo et al.
Feige (2019) use weak supervision in an AE toextract invariant and equivariant representations.
Winter et al.
(2021) implement a permutationinvariant AE to learn graph embeddings, in which the permutation matrix for graph matching islearned during training.
In that sense, the present work can be seen as a generalization of theirapproach for a generic data type and any group.
Experiments
In this section we present a diverse set of experiments for the various symmetry groups discussed in
Section 3.
5.1 Rotated MNIST
[Figure]
In the first experiment, we train an SO(2)-invariant autoencoder on the original (nonrotated) MNIST dataset and validate the trainedmodel on the rotated MNIST dataset mni whichconsists of randomly rotated versions of theoriginal MNIST dataset.
For the functions ηand ψ
we utilize SO(2)-Steerable Convolutional
Neural Networks Weiler & Cesa (2019b).
Formore details about the network architecture andtraining, we refer to Appendix B.
In Figure 2we show images in different rotations and therespective reconstructed images by the trainedmodel.
The model decodes the different rotated versions of the same image (i.e., elements from thesame orbit) to the same canonical output orientation (second row in Figure 2).
The trained modelmanages to predict the right rotation matrix (group action) to align the decoded image with the inputimage, resulting in an overall low reconstruction error.
Note that the model never saw rotated imagesduring training but still manages to encode and reconstruct them due to its inherent equivariant design.
We find that the encoded latent representation is indeed rotation invariant (up to machine precision),but only for rotations of an angle θ = n · π 2 , n ∈ N .
For all other rotations, we see slight variationsin the latent code, which, however, is to be expected due to interpolation artifacts for rotations ona discretized grid.
Still, inspecting the 2d-projection of the latent code of our proposed model in
Figure 3, we see distinct clusters for each digit class for the different images from the test dataset,independent of the orientation of the digits in the images.
In contrast, the latent code of a classicalautoencoder exhibits multiple clusters for different orientations of the same digit class.
5.2 Set of Digits
Next, we train a permutation-invariant autoencoder on sets of digits.
A set with N digits is representedby concatenating one-hot vectors of each digit in a N × D -dimensional matrix, where we take D = 10 .
Notice that this matrix-representation of a set is not permutation invariant.
We randomly sampled1.000.000 different sets for training and 100.000 for the final evaluation with N = 20 , 30 , 40 , 100 ,respectively, removing all permutation equivariant sets (i.e., there are no two sets that are the sameup to a permutation).
For comparison, we additionally trained a classical non-permutation-invariantautoencoder with the same number of parameters and layers as our permutation-invariant version.
Formore details on the network architecture and training we refer to Appendix C.
Here, we demonstratehow the separation of the permutation-invariant information of the set (i.e., the composition of the set)from the (irrelevant) order-information results in a significant reduction of the space needed to encodethe set.
In Figure 4a, we plot the element-wise reconstruction accuracy of different sized sets for bothmodels for varying embedding (bottleneck) sizes.
As the classical autoencoder has to store both thecomposition of digits in the set (i.e., number of elements for each of the 10 digits classes) as well astheir order in the permutation-dependent matrix representation, the reconstruction accuracy drops forincreasing size of the set N for a fixed embedding size.
For the same reason, perfect reconstructionaccuracy is only achieved if the embedding dimension is at least as large as the number of digits in theset.
On the contrary, our proposed permutation invariant autoencoder achieves perfect reconstructionaccuracy with a significant lower embedding size.
Crucially, as no order information has to be storedin the embedding, this embedding size for perfect reconstruction accuracy also stays the same forincreasing size N of the set.
In Figure 4b we show one example for a set x with N = 100 digits, withthe predicted canonical orbit element ˆ x
and the predicted permutation matrix.
As perhaps expected,the canonical element clusters together digits with same value, while not using the commonly used
[0,0,100] [100,0,0] [0,100,0] [50,50,0] [50,0,50] [0,50,50] [8, 6, 2, 7, 9, 2, 8, 5, 8, 0, 9, 7, 9, 6, 7, 3, 9, 4, 6, 9, 1, 7, 5, 9, 3, 7, 3, 7, 6, 8, 8, 1, 1, 4, 5, 9, 2, 0, 1, 8, 8, 3, 2, 5, 5, 7, 6, 7, 4, 1, 1, 0, 2, 8, 7, 8, 4, 9, 6, 4, 0, 4, 2, 9, 0, 5, 8, 8, 3, 4, 3, 9, 9, 7, 4, 4, 3, 3, 7, 3, 9, 6, 1, 0, 1, 7, 2, 3, 9, 9, 8, 8, 7, 5, 1, 9, 7, 6, 8, 0] [1, 1, 1, 1, 1, 1, 1, 1, 1, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 4, 4, 4, 4, 4, 4, 4, 4, 4, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 2, 2, 2, 2, 2, 2, 2, 5, 5, 5, 5, 5, 5, 5] 𝑥= 𝑥′= a) b) c)
P= order of Arabic numerals.
This learned order (here [1,9,4,0,3,6,8,7,2,5]) stays fixed for the trainednetwork for different inputs but changes upon re-initialization of the network.
In Figure 4c we show the two-dimensional embedding of a permutation invariant autoencoder trainedon set of N = 100 elements chosen from D = 3 different classes (e.g. digits 0,1,2).
As the setsonly consists of 3 different elements (but in different compositions and order) we can visualize the 308 (cid:0) D + N −
N (cid:1) =
(cid:0) 102100 (cid:1)
= 5151 elements in the two-dimensional embedding and colour them accordingto their composition.
As our proposed auteoncoder only needs to store the information about the setcomposition and not the order, the embedding is perfectly structured with respect to the compositionas can be seen by the colour gradients in the visualization of the embedding.
5.3 Point Cloud
Point clouds are a common way to describe objects in 3D space, such as the atom positions of amolecule or the surface of an object.
As such, they usually adhere to 3D translation and rotation symmetries and are unordered, i.e., permutation invariant.
Hence, we investigate in the next experiment acombined SE( 3 )- and S N -invariant autoencoder for point cloud data.
We use the Tetris Shape toydataset Thomas et al.
(2018) which consists of 8 shapes, where each shape includes N = 4 points in3 D space, representing the center of each Tetris block.
To generate various shapes, we augment the 8shapes by adding Gaussian noise with σ = 0 .
01 standard deviation on each node’s position.
Differentorientations are obtained by rotating the point cloud with a random rotation matrix R ∈ SO (3) andfurther translating all node positions with the same random translation vector t ∈ R 3 ≃ T 3 .
Foradditional details on the network architecture and training we refer to Appendix D. In Figure 5 wevisualize the input points and output points before and after applying the predicted rotation.
Themodel successfully reconstructs the input points with high fidelity (mean squared error of ∼ 4 × 10 − 5 )for all shapes and arbitrary translations and rotations.
Figure 5b shows the two-dimensional embedding of the trained SE( 3 )- and S N -invariant autoencoder.
Augmenting the points with random noiseresults into slight variations in the embedding, while samples of the same Tetris shape class stillcluster together.
The embedding is invariant with respect to rotations, translation and permutationsof the points.
Notably, the SE(3)-invariant representations can distinguish the two chiral shapes(compare green and violet coloured shapes in the bottom right of Figure 5b).
These two shapes aremirrored versions of themselves and should be distinguished in an SE( 3 ) equivariant model.
Modelsthat achieve SE( 3 ) invariant representations by restricting themselves to composition of symmetricfunctions only, such as working solely on distances (e.g. SchNet Schütt et al.
(2018)) or angles (e.g.
ANI-1 Smith et al.
between points fail to distinguish these two shapes Thomas et al.
Molecular Conformations.
We showcase our learning framework on real-world data by autoencoding the atom types and geometries of small molecules from the QM9 database Ramakrishnan et al.(2014).
We achieved a reconstruction RMSE of 0 .
15 ± 0 .
07 Å for atom coordinates and perfect atomtype accuracy on 5000 unseen test conformations (see Figure 5c for two examples and Appendix E for
[Figure]
We only showcase the effect of aligning by applying the predicted rotation matrix, while the models are still trained in predicting the permutation matrix and translation vector to perform the alignment.
more reconstruction predictions).
Given a point cloud of N nodes, notice that the G -inv.
embeddingz has to store information about the Cartesian coordinates P ∈ R 3 N as well as the 5 distinct atomtypes A ∈ { 0 , 1 } 5 N represented as one-hot encodings.
Given the largest molecule in the QM9database with N max = 29 atoms, the degrees of freedom of the data space 2 X can assumed to be3 · 29 · 5 · 29 = 12615 , making the input data high-dimensional, while we compress the molecularconformations into z ∈ Z ⊂ R 256 dimensions.
Conclusion
In this work we proposed a novel unsupervised learning strategy to extract representations from datathat are separated in a group invariant and equivariant part for any group G .
We defined the sufficientconditions for the different parts of our proposed framework, namely the encoder, decoder and groupfunction without further constraining the choice of a ( G -) specific network architecture.
In fact, wedemonstrate the validity and flexibility of our proposed framework for diverse data types, groups andnetwork architectures.
To the best of our knowledge, we propose the first general framework for unsupervised learningof separated invariant-equivariant representations valid for any group.
Our learning strategy canbe applied to any AE framework, including variational AEs.
It would be compelling to extend ourapproach to a fully probabilistic approach, where the group action function samples from a probabilitydistribution.
Such formalism would be relevant in scenarios where some elements of a group orbitoccur with different frequencies, enabling this to be reflected in the generation process.
For instance,predicting protein-ligand binding sites depends on the molecule’s orientation with respect to theprotein pocket or cavity.
Thus, in a generative approach, it would be highly compelling to generate agroup action reflecting a candidate molecule’s orientation in addition to a candidate ligand.
We planto return to these generalization and apply our learning strategy to non-trivial real-world applicationsin future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
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?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or
(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 Proofs
Proposition A.1.
Any suitable group function ψ : X → G is G -equivariant at a point x ∈ X up thestabilizer G x , i.e., ψ
( ρ X ( g ) x ) ⊆ g · ψ ( x ) G x .
Proof : As the relation (see Property 2.2)
[Equation]
must hold for any x ∈ X , it must hold for any point x ′ =
ρ X ( g ) x in the orbit of x , which then reads
[Equation]
where we used the invariance of η .
On the other hand, applying ρ X ( g ) to both sides of (7) we have
[Equation]
since η ( ρ X ( g ′ ) x ) =
η ( x ) and ρ X ( g 1 ) ρ X ( g 2 ) =
ρ X ( g 1 g 2 ) .
Combining (8) and (9) it follows thatρ X ( ψ (
[Equation]
1 · g
[Equation] [Equation]
g
[Equation]
η ( x
[Equation] [Equation]
x
[Equation]
, (10) that is, ψ ( x ) − 1 · g − 1 · ψ ( ρ X ( g ) x )) ∈
G δ ( η ( x )) .
Now, since x and δ ( η ( x )) by assumption belong tothe same orbit of G , it follows that they have isomorphic stabilizers, G δ ( η ( x )) ≃ G x .
Thus, we haveshown that ψ ( ρ X ( g ) x ))
=
g · ψ ( x ) · g ′ , where g ′ ∈ G x , which proves our claim.
Proposition A.2.
The image of any suitable group function ψ : X → G is surjective into GG X , where
G X is the stabilizers of all the points of X .
Proof : Let x ∈
X be such that x = δ ( η ( x )) , that is, ψ ( x ) =
G x , the stabilizer of x .
Notethat each orbit cointains at least one such element.
For any element g ∈ G we have that, using
Proposition A.1, ψ ( ρ X ( g ) x )
= g · ψ ( x ) · ˜ g , where ˜ g ∈ G x .
Since ψ ( x ) · ˜ g ∈ G x as well, it thenfollows that the image of ψ is G up to an action by an element of the stabilizer G x .
Applyingthe above reasoning to every points x ∈ X , we have that Im ( ψ ) = ∪ x ∈ X GG x =
G ∩ x ∈ X G x , where∩ x ∈ X G x = G X
= { g ∈ G | ρ X ( g ) x
= x, ∀ x ∈ X } , proving our claim.
Lemma A.3.
Any suitable group function ψ is an isomorphism
O x ≃ G/G x for any x ∈ X , where
O x ⊂ X is the orbit of x with respect to G in X .
Proof : Surjectivity follows directly from Proposition A.2.
To show injectivity, consider x, x ′ ∈
O xsuch that ψ ( x ′ )
=
ψ ( x ) · ˜ g , where ˜ g ∈ G x .
From Proposition 2.3 it follows that x ′ = x , which provesthe claim.
Proposition A.4.
Let ψ = ξ ◦ µ be a suitable group function and let µ : X → Y be G -equivariant.
Then, G x = G µ ( x ) for all x ∈ X .
Proof : Let g ∈ G x , that is,
ρ X ( g ) x = x .
Applying µ to both sides of this equation we obtainµ ( x ) = µ ( ρ X ( g ) x )
=
ρ Y ( g ) µ ( x ) , where we used the G -equivariance of µ .
Hence, G x ⊆ G µ ( x ) .
To prove the opposite inclusion, let g ∈ G µ ( x ) but g / ∈ G x , and let x ′ =
ρ X ( g ) x .
Now, µ ( x ′ )
=
ρ Y ( g ) µ ( x ) = µ ( x ) , thus µ , and therefore ψ = ξ ◦ µ , maps the distinct element x, x ′ to the samegroup element ψ ( x ) = ψ ( x ′ ) , in contradiction with Proposition 2.3.
Proposition A.5.
Given y, y 0 , the element g such that y ≡ ρ Y ( g ) y 0 is unique up to the stabilizer
G y 0 .
Proof : Suppose that there exist g 1 , g 2 ∈ G
such that ρ Y ( g 1 ) y
=
ρ Y ( g 2 )
y 0 , then ρ Y ( g − 1 2 g 1 )
y
= y 0 ,which implies g − 1 2 g 1 ∈ G y 0 .
Proposition A.6.
Let ψ = ξ ◦ µ where µ and ξ are as described above.
Then, ψ is a suitable groupfunction.
Proof : We show that our construction describes an isomorphism O x ≃ G/G x for all x ∈ X .
Givenx ∈ X and g ∈ G , Propositions
A.4 and A.5 imply
[Equation]
that is, ψ possesses the G -equivariant property as required in Proposition 2.3, which in turns implyinjectivity, as in Lemma A.3.
Surjectivity follows from the same argument as in Proposition A.2,since the proof only relies on the equivariant properties of ψ , which we showed in (11).
B Model architecture Rotated MNIST
We follow Weiler & Cesa (2019b) and use steerable CNNs to parameterize functions η and µ .
Incontrast to classical CNNs, CNNs with O(2)-steerable kernels transform feature fields respecting thetransformation law under actions of O (2) .
We can define scalar fields s : R 2 → R and vector fieldsv : R 2 → R 2 that transform under group actions (rotations) the following: [Equation]
[Equation]
Thus, scalar values are moved from one point on the plane R 2 to another but are not changed, whilevectors are moved and changed (rotated) equivalently.
Hence, we can utilize steerable CNNs toencode samples in R 2 in O (2) -invariant scalar features and O (2)
-equivariant vector features.
We canuse the scalar features s as G -invariant representation z ∈ Z and following Section 3 (Orthogonalgroup SO (2) ) utilizing a single vector features v to construct the rotation matrix R as:
[Equation]
In our experiments we used seven layers of steerable CNNs as implemented by Weiler & Cesa(2019b).
We did not use pooling layers, as we found them to break rotation equivariance and onlyaveraged over the two spatial dimensions after the final layer to extract the final invariant embeddingand equivariant vector.
In each layer we used 32 hidden scalar and 32 hidden vector fields.
In thefinal layer we used 32 scalar fields (32 dimensional invariant embedding) and one vector feature field.
The Decoding function δ : Z → R 2 can be parameterized by a regular CNN.
In our experiments weused six layers of regular CNNs with 32 hidden channels, interleaved with bilinear upsampling layersstarting from the embedding expanded to a 2 × 2 × 32 tensor.
Training was done on one NVIDIA Tesla V100 GPU in approximately 6 hours.
We also implemented and trained the quotient autoencoder (QAE) approach proposed by Mehr et al.(2018) on the MNIST dataset for the group
SO (2) , discretized in 36 rotations with the loss
[Equation]
(14) where x is a MNIST sample and y is the reconstructed sample.
We evaluated the resulting embeddingson the rotated MNIST test set (in such a way that the evaluation is the same as for our model).
In
Figure 6 we plot TSNE embeddings for this approach, and we can observe that the embedding spaceshows a clearer structure, in comparison with the classical model.
However, in comparison, ourapproach results in a better clustering of the different digits classes.
That shows that the discretizationstep, while it helps in structuring the embedding space in “signal clusters”, still does not capture thefull continuous nature of the group.
To further quantitatively compare the three methods (ours, QAE(a) and classical AE), we evaluated the reconstruction loss as well as the (digit class) classificationaccuracy of a KNN classifier trained on 1000 embeddings of each method.
We present in the tablebelow the results for the reconstruction loss and for the classification accuracy of a KNN classifiertrained on the AE embeddings.
To obtain a fair comparison, we kept the architecture and the traininghyperparameters exactly identical for all the strategies.
We note that our strategy outperforms boththe classical AE as well as the strategy of (a) in both tasks.
We also performed an additional experiment, we trained a fully equivariant AE (that is, the embeddingitself is fully equivariant, i.e. multiple 2-dimensional vectors) on MNIST with G = SO (2) , and weperform an invariant pooling afterwards to extract the invariant part.
Specifically, we have trained KNN classifiers on (a) the invariant embedding corresponding to the norm of the 2-dimensionalvectors forming the bottleneck representation, (b) the angles between the first and all other vectorsand on (c) the full invariant embedding we obtained by combining the the norms and angles. Wechoose the number of vectors in the bottleneck in such a way that the dimensionality of the fullinvariant representation coincides with the one of our model. We visualized the resulting TSNEembeddings in Figure 6 and show the downstream performance of the KNN classifiers in Table 1.
From the results we can see that, in comparison to the approximate invariant (QAE) and our invarianttrained model, the invariant projected equivariant representations perform inferior.
Although weextract a complete invariant representation (which performs better than a subset of this representationlike the norm or angle part), the resulting representation is apparently not as expressive and e.g.useful in a downstream classification task.
This aligns well with our hypothesis, that our proposed
[Table]
framework poses a sensible supervisory signal to extract expressive invariant representations that aresuperior to invariant projections of equivariant features.
We belief that this experiment gives strongempirical evidence for the motivation of this work and thank the reviewer very much for suggestingsuch a comparison, strengthening our manuscript.
C Model architecture Set of Digits
We can rewrite the equation P σ (1 , 2 , . . . , n ) =
( σ (1) , σ (2) , . . .
, σ ( n )) , in vector form by representing set elements by standard n × 1 column vectors e i (one-hot encoding) and σ by a permutationmatrix P σ whose (i,j) entry is 1 if i = σ ( j ) and 0 otherwise, then:
[Equation]
Hence, encoding function η should encode a set of elements in a permutation invariant way and ψshould map a set M to a permutation matrix P σ :
[Equation]
We follow Zaheer et al.
and parameterize η by a neural network γ that is applied element-wiseon the set followed by an invariant aggregation function Σ (e.g. sum or average) and a second neural network β : [Equation]
[Equation]
In our experiments we parameterized γ and β with regular feed-forward neural networks with threelayers respectively, also using ReLU activations and Batchnorm.
The output of function γ is equivariant and can also be used to construct ψ .
We follow Winter et al.(2021) and define a function s : R d
→ R mapping the output of γ for every set element to a scalarvalue.
By sorting the resulting scalars, we construct the permutation matrix P σ with entries p ij thatwould sort the set of elements with respect to the output of s :
[Equation]
As the argsort operation is not differentiable, we utilizes a continuous relaxation of the argsortoperator proposed in (Prillo & Eisenschlos, 2020; Grover et al., 2019):
[Equation]
[Equation] where the softmax operator is applied row-wise, d ( x, y ) is the L 1 -norm and τ ∈ R + a temperatureparameter.
Decoding function δ can be parameterized by a neural network that maps the permutation-invariantset representation back to either the whole set or single set elements.
In the letter case, where the samefunction is used to map the same set representation to the different elements, additional fixed positionembeddings can be fed into the function to decode individual elements for each position/index.
Forthe reported results we choose this approach, using one-hot vectors as position embeddings and a4-layer feed-forward neural network.
Training was done on one NVIDIA Tesla V100 GPU in approximately 1 hours.
D Model architecture Point Cloud - Tetris 3D
We implement a graph neural network (GNN) that transform equivariantly under rotations andtranslations in 3D space, respecting the invariance and equivariance constraints mentioned in Eq.
(5)and (6) for n = 3 .
Assume we have a point cloud of N particles each located at a certain position x i ∈ R 3 in Cartesianspace.
Now given some arbitrary ordering σ ( · ) for the points, we can store the positional coordinatesin the matrix P =
[ x 1 , ..., x N ] ∈ R N × 3 .
Standard Graph Neural Networks (GNNs) perform messagepassing Gilmer et al.
(2017) on a local neighbourhood for each node.
Since we deal with a pointcloud, common choice is to construct neighbourhoods through a distance cutoff c > 0 .
The edges ofour graph are specified by relative positions
[Equation]
and the neighbourhood of node i is defined as N ( i ) =
{ j : d ij :=
|| x ij || ≤ c } .
Now, our data (i.e., the point cloud) lives on a vector space X , where we want to learn an SE(3)invariant and equivariant embedding wrt.
arbitrary rotations and translations in 3 D space.
Let thefeature for node i consist of an invariant (type-0) embedding h i ∈ R F s , an equivariant (type-1)embedding w i ∈ R 3 × F v that transforms equivariantly wrt.
arbitrary rotation but is invariant totranslation.
Such a property can be easily obtained, when operating with relative positions.
Optionally, we can model another equivariant (type-1) embedding t i ∈ R 3 which transformsequivariantly wrt.
translation and rotation.
As our model needs to learn to predict group actionsin the SE(3) symmetry, we require to predict an equivariant translation vector ( b ∈ T 3 ),
as well asa rotation matrix ( A ∈ SO (3) ), where we will dedicate the t vector to the translation and the wvector(s) to the rotation matrix.
As point clouds might not have initial features, we initialize the SE(3)-invariant embeddings asone-hot encoding h i
=
e i for each node i = 1 , . . .
, N .
The (vector) embedding dedicated forpredicting the rotation matrix is initialized as zero-tensor for each particle, i.e., w i = 0 and thetranslation vector is initialized as the absolute positional coordinate, i.e. to,
t i = x i .
We implement following edge function ϕ
e : R 2 F s +
(cid:55)→ R F s +2 F v + k with
[Equation]
20) and set k = 1 if the GNN should model the translation and k = 0 else.
Notice that the message m ijin Eq.
(20) only depends on SE(3) invariant embeddings.
Now, (assuming k = 1 ) we further split themessage tensor into 4 tensors,
[Equation]
which we require to compute the aggregated messages for the SE(3) invariant and equivariant nodeembeddings.
We include a row-wise transform ϕ s : R F s (cid:55)→ R F s for the invariant embeddings using a linear layer:
[Equation]
The aggregated messages for invariant (type-0) embedding h i are calculated using:
[Equation]
22) where ⊙ is the (componentwise) scalar-product.
The aggregated equivariant features are computed using the tensor-product ⊗ and scalar-product ⊙from (invariant) type-0 representations with (equivariant) type-1 representations:
[Equation]
[Equation]
where 1 ∈ R 3 is the vector with 1 ’s as components and ( a × b ) denotes the cross product betweentwo vectors a, b ∈ R 3 .
The tensor in Eq.
(23) is equivariant to arbitary rotations and invariant to translations.
It is easy toprove the translation invariance, as any translation t ∗ ∈ T 3 acting on points x i , x j does not changethe relative position x ij =
( x j + t ∗ ) − ( x i + t ∗ ) =
x j − x i .
To prove the rotation equivariance, we first observe that given any rotation matrix A ∈ SO (3) actingon the provided data, as a consequence relative positions rotate accordingly, since
[Equation]
The tensor product ⊗ between two vectors u ∈ R 3 and v ∈ R F s , commonly also referred to as outerproduct is defined as
[Equation]
, and returns a matrix given two vectors.
For the case that a group representation of SO(3), i.e. arotation matrix R , acts on u , it is obvious to see with the associativity property
[Equation]
692 The cross product ( w i × w j ) ∈ R 3 × F v used in equation (23) between type-1 features w i and w j isapplied separately on the last axis.
The cross product has the algebraic property of rotation invariance,i.e. given a rotation matrix
A acting on two 3-dimensional vectors a, b ∈ R 3 the following holds:
[Equation]
696 Now, notice that the quantities that "transform as a vector" which we call type-1 embeddings are in
S = { x ij , w i , t i } Ni,j =1 .
Given a rotation matrix A acting on elements of S , we can see that the result in (23)(cid:88)
[Equation]
is rotationally equivariant.
We update the hidden embedding with a residual connection
[Equation]
[Equation]
and use a Gated-Equivariant layer with equivariant non-linearities as proposed in the PaiNN architecture Schütt et al.
(2021) to enable an information flow between type-0 and type-1 embeddings.
The type-1 embedding for the translation vector is updated in a residual fashion
[Equation]
where we can replace the tensor-product with a scalar-product, as m t,ij ∈ R .
The result in Eq.
(26)is translation and rotation equivariant as the first summand t i is rotation and translation equivariant,while the second summand is only rotation equivariant since we utilize relative positions.
For the SE(3) Tetris experiment, the encoding function η : X (cid:55)→ Z is a 5-layer GNN encoder with
F = F s = F v = 32 scalar- and vector channels and implements the translation vector, i.e. k = 1 .
The encoding network η outputs four quantities: two SE(3) invariant node embedding matrices(cid:101) H, M ∈ R N × F , one SO(3) equivariant order-3 tensor (cid:102) W ∈ R N × 3 × F as well as another SE(3)equivariant matrix T ∈ R N × 3 .
We use two linear layers 3 to obtain the SE(3) invariant embedding matrix H ∈ R N × 2 as well asthe SO(3) equivariant embedding tensor W ∈ R N × 3 × 2 .
Notice that the linear layer returning the
W tensor can be regarded as the function ψ rot that aims to predict the group action in the SO(3)symmetry, while we use the identity map for the translation vector, i.e. ψ transl = T .
As point clouds can be regarded as sets, we obtain an permutation invariant embedding by averagingover the first dimension of the { H, W, T } tensors,
[Equation]
[Equation]
while we use the M matrix to predict the permutation matrix P σ with the ψ perm function, in similarfashion as described in Eq.
To construct the rotation matrix R out of 2 vectors in R 3 as describedin Section 3, we utilize the SO(3) equivariant embedding w .
The decoding network δ : Z (cid:55)→ X is similar to the encoder a 5 -layer SE( 3 )
-equivariant GNN but doesnot model the translation vector, i.e. k = 0 .
The decoder δ maps the SE( 3 ) as well as S( N )-invariantembedding
h back to a reconstructed point cloud
ˆ P ∈ R N × 3 .
At the start of decoding, we utilize alinear layer to map the G − invariant embedding h ∈ R 2 to a higher-dimension, i.e.
[Equation] [Equation]
Next, to “break” the symmetry and provide the nodes with initial type-0 features, we utilize fixed(deterministic) positional encodings as suggested by Winter et al.
(2021) for each node i = 1 , . . .
, Nto be summed with ˜ h .
Notice that this addition enables us to obtain distinct initial type-0 embeddings{ ˆ h i } Ni =
For the start positions, we implement a trainable parameter matrix P θ of shape ( N × 3) for thedecoder.
Now, given an initial node embedding ˆ H ∈ R N × F s , we apply the S( N ) group action, by multiplyingthe predicted permutation matrix P σ with ˆ H from the left to obtain the canonical ordering as
[Equation]
[Table]
To retrieve the correct orientation required for the pairwise-reconstruction loss, we multiply theconstructed rotation matrix R with the initial start position matrix P θ
[Equation]
With such construction, we can feed the two tensors to the decoder network δ to obtain the reconstructed point cloud as
[Equation]
where t ∈ R 3 is the predicted translation vector from the encoder network, added row-wise for eachnode position.
Architecture QM9.
For the QM9 dataset, we use the same model components as described in the
Tetris experiment, with the difference of including atom species as SE (3) − invariant features andsetting F s = 256 , F v = 32 and increasing the dimensionality of the latent space to 256 .
Finetuning.
We performed additional experiments on the pretrained group-invariant AE on theextended GEOM-QM9 dataset Axelrod & Gomez-Bombarelli (2021) which, as opposed to thestandard QM9 dataset ( ≈ 130 k samples), contains multiple conformations of small molecules.
Wetrained the autoencoder on a reduced set of GEOM-QM9 ( ≈ 641 k ), containing up to 10 conformationsper molecule and utilized this pretrained encoder network to regress (invariant) energy targets, suchas internal energy U or enthalpy H on the original QM9 dataset.
We observed that the pretrained encoder network learns faster and achieves better generalizationperformance than the architectural identical network trained from scratch.
In Figure 7 we illustratesthe learning curves for the two networks on different fraction on 5% and 25% labelled samplesfrom original QM9 dataset to analzye the benefit of finetuning a pre-trained encoder network on alow-data regime, when regressing on the enthalpy H .
On a held-out test dataset of 1000 samples, thepretrained encoder network achieves superior generalization performance in terms of R 2 with 0 .
7529vs.
0970 in the 5% data regime, and 0 .
9908 vs. 0 .
9093 in the 25% data regime compared to theencoder that was trained from scratch.
In Table2 we show additional comparisons of the pretrainednetwork against a network that was trained from scratch for 50 epochs on the restricted dataset.
As shown in Table 2, the pretrained encoder achieves improved generalization performance on thetest dataset compared to its architectural identical model that was trained from scratch.
We believethat training the group-invariant autoencoder on a larger diverse dataset of (high-quality) molecularconformations facilitates new opportunities in robust finetuning on different data-scarse datasets formolecular property prediction.
E Molecular Conformations: Further Examples
We show additional reconstructions of 12 randomly selected small molecules from the QM9 testdataset.
Noticeably, our trained autoencoder is able to reconstruct molecular conformations with [Figure] F N
O
C
H 0. [Figure] complex geometries as depicted in the third column (from the left). We notice that the AE is notable to perfectly reconstruct the conformation shown in the 4th column of the 2nd row.
Althoughthis molecule does not exhibit a complicated geometrical structure, its atomistic composition (ofonly containing nitrogen and carbon as heavy atoms) could be the reason why the encoding of theconformation is pointing into a non-densely populated region in the latent space, as nitrogen does nothave a large count in the total QM9 database, see Figure 9.
Training was done on one NVIDIA Tesla V100 GPU in approximately 1 day.
F ShapeNet
We also run experiments on the ShapeNet dataset.
As the dataset comes in an aligned form (e.g. carsare always aligned in the same orientation), we additionally applied random 90 degree rotations toremove this bias.
We utilized 3D Steerable CNNs proposed by Weiler et al.
(2018b)
as equivariantencoder for the 3d voxel input space.
We utilized the scalar outputs as rotation-invariant embedding( z ) and predict (analogously to our experiments on 3d point clouds) 2 rotation-equivariant vectorsto construct a rotation matrix ( g ).
Similar to our MNIST experiment, we compared the resultingembedding space to the embeddings produced by a non-invariant autoencoder model.
In Figure 10we visualize a TSNE projection of the embeddings of both models.
We can see a well structuredembedding space for our model with distinct clusters for the different shape classes.
On the otherhand, the embeddings produced by the non-invariant autoencoder is less structured and one can makeout different clusters for the same shape label but in different orientations.
Moreover, we comparedthe downstream performance and generalizability of a KNN classifier on shape classification, trainedon 1000 embeddings and tested on the rest.
The classifier based on our rotation-invariant embeddingsachieved an accuracy of 0.81 while the classifier based on the non-invariant embeddings achieved anaccuracy of only 0.63.