Symmetry-induced Disentanglement on Graphs
Abstract
Learning disentangled representations is important for unraveling the underlying complex interactions between latent generative factors.
Disentanglement has been formalized using a symmetry-centric notion for unstructured spaces, however, graphs have eluded a similarly rigorous treatment.
We fill this gap with a new notion of conditional symmetry for disentanglement, and leverage tools from Lie algebras to encode graph properties into subgroups using suitable adaptations of generative models such as Variational Autoencoders.
Unlike existing works on disentanglement, the proposed models segregate the latent space into uncoupled and entangled parts.
Experiments on synthetic and real datasets suggest that these models can learn effective disengaged representations, and improve performance on downstream tasks such as few-shot classification and molecular generation.
Introduction
Disentanglement represents a fundamental desideratum in learning with limited supervision because it captures information about the salient (or explanatory) factors of variation in the data and isolates information about each specific factor in only a few dimensions, thus unraveling the interactions underlying complex data [2, 60].
Disentangled representations have often been deemed responsible for providing neural models with the ability to improve performance on real-world downstream tasks [2, 20, 36, 59, 60].
Empirically, they have been shown to enable, or facilitate, critical properties such as sample efficiency [21] and generalization [32, 42, 70].
An interesting perspective on disentanglement has been recently proposed by Higgins et al.
Specifically, they advocate viewing disentanglement as decomposition of the latent space of embeddings into subspaces, each of which can be transformed independently under the action of a single subgroup specific to the subspace, without affecting others.
These subgroups arise, in turn, from the decomposition of a symmetry group.
This formalism bestows several benefits, e g., it (a) aligns with the idea that auxiliary information can be exploited for imposing a structure on the latent space [25, 40], and (b) leads to a rigorous theoretical framework for analyzing disentangled models [65, 71], e.g., those based on Variational Autoencoders (VAEs) [27, 50].
While symmetry groups provide a suitable tool both for formalizing [20] and learning [58, 65, 71] disentangled representations in unstructured domains, these techniques have not been investigated in the context of more general and complex data such as graphs.
Furthermore, most works on disentanglement for graphs [1, 16, 31, 35, 38, 66] are not designed for generative settings, and the existing deep generative models (DGMs) for graph disentanglement such as [10, 17, 57] do not consider group symmetries.
Note that differently from other definitions of disentanglement [2, 9, 11], leveraging group symmetries for disentanglement may unravel fundamental connections across several approaches based on graph neural networks (GNNs), i.e., the state-of-the-art models for embedding graphs [3, 5, 8, 13, 15, 39, 64].
We, therefore, pursue two main goals in this work: (a) providing a rigorous symmetry-based formalism for disentanglement on graphs, and 2) designing novel neural architectures that are able to leverage symmetry groups to learn disengaged representations for graphs.
However, accomplishing these objectives requires overcoming some challenges.
In particular, graphs often abstract complex interactions, so the underlying latent space may not factorize completely into only disentangled subspaces.
Thus, the symmetry-based formalism introduced in [20] does not suffice for graphs.
Therefore, we introduce a new, more flexible notion of conditional disentanglement , advocating segregating the latent space into uncoupled and entangled parts.
Translating this idea into an efficient algorithm requires further work, and we appeal to a Lie algebra based parameterization to encode the graph properties into subgroups.
Contributions.
Our contributions can be summarized as follows:
• a novel notion of symmetry-induced conditional disentanglement (Section 3.1) that generalizes the previous definition from [20] (which we call unconditional disentanglement); • a parameterization centered on Lie algebra based on the intuition that each separable generative factor can correspond to a Lie algebra coordinate element, and control a single graph property (Section 3.2); • two algorithms for S ymmetry-
I nduced D isentanglement under unconditional (SID U ) and conditional (SID C ) settings, both using a two-level variational auto-encoding mechanism, namely, one level for the data and other for symmetry group representation (Section 4); and • a systematic evaluation on several disentanglement metrics [7, 11, 19, 26, 44].
The results demonstrate that the proposed models can outperform contemporary GNN baselines [38, 66]; successfully learn to decouple the entangled part of latent space from the disentangled parts (Section 5.1); and improve performance on downstream tasks such as few-shot classification, compression (Section 5.2) and molecular graph generation (Section 5.3).
We begin with a review of the relevant literature, and then proceed to the proposed framework.
Related work
Symmetry-based disentanglement (SBD).
Initial works on SBD, such as Caselles-Dupré et al.
[6], Quessard et al.
[48], propose to use reinforcement learning to achieve irreducible representations of group elements through observation of action transitions.
Painter et al.
[44] introduce a VAE method that does not require labelled action-transition pairs, and demonstrate its validity within the environment from Caselles-Dupré et al.
More recently, some algebraic approaches have been introduced.
Zhu et al.
[71] propose to decompose a Lie group into multiple subgroups, and enforce commutativity and independence between subgroups.
Yang et al.
[65] impose a cyclic group structure, encouraging the group to be isomorphic to the ground truth.
Bouchacourt et al.
[4] invoke group representation theory for defining distributed latent operators that are able to learn disentangled representations.
Tonnaer et al.
[58] provide a new metric for evaluation, and a VAE-based model for learning symmetry-based disentangled representations.
Some works leverage data manifolds for SBD.
For example, Fumero et al.
[12] view disentanglement as a product of low-dimensional sub-manifolds underlying the data space, such that each sub-manifold encodes an explainable data factor.
Pfau et al.
[46] propose a non-parametric algorithm that aims to discover a decomposition of the data manifold by investigating its holonomy group.
Other works explore SBD as well.
For instance, Higgins et al.
[22] demonstrate the importance of symmetry transformations and disentangled representation learning in the context of neuroscience.
Wang et al.
[62] design a contrastive self-supervised learning method that iteratively updates a partition, modeled as a symmetry group, until the considered factor is invariant.
Disentangled GNNs.
To our knowledge, Ma et al.
[38] provide the first graph disentanglement approach, which involves a routing mechanism aimed at separating the underlying factors.
More recently, Yang et al.
[66] exploit an attention mechanism for GNNs, and provide two evaluation metrics for disentanglement on graphs.
Liu et al.
[35] achieve disentanglement by encouraging independence in the latent space, while Bae and Jeon [1] propose a method for disentangling multirelational graphs motivated by an application on pedestrian trajectory prediction.
Li et al.
[31] provides a contrastive framework for disentangling the latent factors in GNNs.
Differently from these methods, we provide a formal definition of disentanglement for graphs, and leverage symmetry groups in a generative model setting.
Graph DGMs.
Wu et al.
[63], Yu et al.
[69] propose VAE-based Information Bottleneck methods for improving the performance of GNNs on classification tasks, by learning minimal sufficient representations.
The work by Simonovsky and Komodakis [56] aims to generate molecular graphs using graph VAEs.
Liu et al.
[34] adapt VAEs for conditioning molecule generation on specific properties.
You et al.
[68] propose a deep autoregressive model that learns to generate graphs by training on a representative set of graphs and decomposes the graph generation process into a sequence of node and edge formations.
Du et al.
[10] propose a DGM for spatio-temporal graphs, and Guo et al.
[18] investigate DGMs with spatial networks.
Finally, Guo et al.
[17] and Stoehr et al.
[57] investigate disentangled representations with graph VAEs.
Proposed framework
In Section 3.1 we review the unconditional symmetry-induced disentanglement, and propose a new definition of conditional disentanglement.
In Section 3.2 we extend these notions to graphs.
Finally, in Section 3.3, we describe additional loss terms to encourage disentanglement.
3.1 Symmetry-based disentanglement
Definition 1 (Unconditional symmetry-based disentanglement [20]) .
Consider a generative process b = W
→ O mapping the world states W into observations (i.e. data) O , and an inference process h
= O → Z mapping the data into a latent vector space.
By composing b and h , we obtain f = h ◦ b such that f : W → Z .
Given group actions on W , i.e., · W : G × W → W and on Z , i.e., · Z : G × Z
→ Z , the function f is said to be equivariant between the actions on W and Z if the actions commute with f .
We can express this equivariance mathematically as g · Z f ( w ) =
f ( g · W w ) , ∀ g ∈ G, ∀ w ∈
W.
Assuming that the symmetry group G can be decomposed into subgroups G 1 × . . .
× G n , a vector representation Z is disentangled with respect to this decomposition if all of the following conditions hold:
• there is an action defined on the representation set · Z : G × Z
→ Z , • the map f : W → Z is equivariant between the actions on W and Z , and • there is a decomposition Z
= Z 1 × Z 2 . . .
×
Z n such that each Z i is affected only by G i and is invariant to (i.e., does not change due to) G j for all j ̸ = i .
Definition 1 does not allow any part of the latent vector space to be entangled, which is likely to be restrictive for complex applications.
Therefore, we propose conditional disentanglement.
Definition 2 (Conditional symmetry-based equivariance) .
Consider a generative process b = W ×
E → O mapping the disentangle-able world states W and non-disentangled states E into observations (i.e. data) O , and an inference process h
= O → Z mapping the data into a latent vector space.
By composing b and h , we obtain f = b ◦ h , which maps W × E to Z .
Given group actions on
W , namely, · W : G × W → W and on Z , namely, · Z : G × Z
→ Z , the function f is said to be conditionally equivariant between the actions on W and Z given E if the actions commute with f .
That is, g · Z f ( w, e ) = f
( g · W w, e ) , ∀ g ∈ G, ∀ w ∈ W, ∀ e ∈ E .
This definition of conditional equivariance naturally leads to a new notion of conditional disentanglement using a characterization similar to the three conditions for the unconditional case.
Specifically, we need to impose the condition that f is conditionally equivariant between the actions of W and Z given E (instead of simply being equivariant between the actions of W and Z ).
Note that while this definition results in a generalized framework for disentanglement, it does not directly translate into a learning algorithm.
Toward that goal, we invoke tools from Lie algebra as described in the next section.
3.2 Disentanglement on graphs
Here, we aim to instantiate Definition 1 and 2 in the context of models that encode graph-structured data.
Drawing inspiration from the previous work on Lie groups and algebras for symmetry-induced disentangled VAEs by Zhu et al.
[71], we develop a formalism for disentanglement on graphs induced by Lie groups.
Lie groups and Lie algebras.
A Lie group G is a group of continuous symmetries [52], and associated with a Lie algebra g , which is the tangent space to the identity element of G .
We can thus parameterize a Lie algebra with a basis { A i } ki =1 , where every element in g can be written as A =
A 1 t 1 + . . .
+ A k t k using coordinates t i .
Elements of the Lie algebra can be mapped back into the Lie group with a matrix exponential map exp : g → G .
We now provide some intuition into how we can connect the notion of a Lie group G with the definition of disentanglement in a graph network, by associating each Lie algebra coordinate t i with a generative factor.
A latent representation ˆ Z , obtained, e.g., from a graph encoder, is disentangled with respect to a Lie group G , if a change in the coordinate t i is associated with a change in only the i th component of ˆ Z , i.e., only ˆ z i .
In other words, disentanglement entails that the semantics ( t 1 , t 2 , . . .
, t k ) are equivariant with respect to the properties (ˆ z 1 , ˆ z 2 , . . .
, ˆ z k ) that are being predicted.
Next, we provide the formal definitions for unconditional and conditional graph disentanglement.
The former expresses ˆ Z as a function of T without accounting for Z , whereas the latter expresses ˆ Z based on T after fixing an encoding Z .
Definition 3 (Lie-algebra based unconditional graph disentanglement) .
The graph embedding ˆ Z obtained by f ( ˆ Z | T ) is unconditionally disentangled with respect to the Lie group coordinates
T = { t j } kj =1 if the following hold: (a) there is a group action · ˆ Z : G × ˆ Z → ˆ Z on ˆ Z , (b) the map f = exp
( A ( T )) :
T → ˆ Z is equivariant between actions on T and ˆ Z , and (c) there is a decomposition ˆ
Z = ˆ z 1 ×
ˆ z 2 . . .
× ˆ z k , where each coordinate t i affects only the corresponding component ˆ z i .
Definition 4 (Lie-algebra based conditional graph disentanglement) .
The graph embedding ˆ Z obtained by f ( ˆ Z | T, Z ) is conditionally disentangled with respect to the Lie group coordinates
T = { t j } kj =1 given Z if (a) there is a group action · ˆ Z : G × ˆ Z → ˆ Z on ˆ Z , (b) the map f : exp A ( T )
×
Z → ˆ Z is equivariant between actions on T and ˆ Z for any fixed Z , and (c)
exp A ( T )
× Z factorizes into the product (cid:81) ki
=1 exp( t i A i ) ×
Z such that each component exp( t i A i ) is affected only by the corresponding coordinate t i for i ∈ { 1 , 2 , . .
. , k } for any fixed Z .
3.3 Disentanglement constraints
It is known [65, 71] that a symmetry group parameterization is usually not sufficient for a group to be decomposed into subgroups that are parameterized independently by a single coordinate.
Therefore, following Zhu et al.
[71], we add two additional constraints for enforcing disentanglement.
These constraints can be relaxed, and included in the loss function as regularization terms.
Commutative penalty.
We would like to enforce a coordinate t i to be identified by single group exp( t i A i ) , and thus represent a single property variation in the data ˆ z i .
Mathematically, one can state this desideratum as an equivalence between the exponential map of the sum, and the product over exponential maps.
[71] proved this to be achieved under commutativity over the Lie algebra basis.
Namely, if A i
A j =
A j A i then exp (cid:16)(cid:80)
ki =1 t i A i (cid:17) =
(cid:81) ki
=1 exp( t i A i ).
Hessian penalty Furthermore, disentanglement can be encouraged using a Hessian penalty [45] based on the fact the Hessian matrix with respect to a disentangled representation is always zero (owing to the independence between different dimensions).
Zhu et al.
[71] adapt this penalty in a Lie algebra parameterization setup, as H ij =
δ 2 g ( T ) δt i δt j where g ( T ) =
exp (cid:16)(cid:80) ki
=1 t i A i (cid:17) , and show that if
A i A j = 0 ∀ i, j then H ij = 0 .
A ,X
E G r a ph Z
E G r o up t. . .
t k
Symmetry-Induced Disentanglement
D G r o up ˆ Z D G r a ph ˆ A , ˆ X
Conditional disentanglement
Z
GCN 1 µσt. . .
E Group
GCN k µ k σ k t k t. . .
t k
A
D Group
A k J
J
L
A ˆ Z exp µ σ
Z ˆ Z conditional case (
) we further condition
Z , by taking mean ( µ ) of Z and ˆ Z , and passing µ through an function σ (
Method
We propose Symmetry-Induced Disentanglement (SID), a VAE framework for graphs [28, 56] which encourages disentangled representations using a Lie group parameterization and two regularization losses (commutative and hessian penalties).
The architecture is depicted in Figure 1.
4.1 Probabilistic formulation
We use the following notation for model variables: latent graph embeddings Z (and ˆ Z ), graph features
X , symmetry Lie group structure T , and the graph adjacency matrix A .
We next derive lower bounds on log-likelihood for both the unconditional and conditional settings (see Appendix for details).
Proposition 1 (Unconditional lower bound) .
Given two latent variables Z and T , we instantiate the unconditional disentanglement from Definition 3 with a probabilistic model that maximizes the log-likelihood of the graph data G =
( X, A ) by optimizing the following lower bound:
[Equation]
Proposition 2 (Conditional lower bound) .
We extend Proposition 1 to account for conditional disentanglement as defined in Definition 4.
We obtain the following lower bound
[Equation]
Architecture.
Propositions 1 and (2) are derived respectively in Appendix B and C.
The respective losses (1 and 2) are implemented with four neural network modules as follows.
Two stochastic encoders ( E graph and E group ) embed respectively the graph data G into a latent variable Z via q ( Z |G ) , and Z into the Lie group coordinates T via q ( T | Z ) .
Two deterministic decoders ( D group and D graph ) reconstruct, respectively, the graph embedding ˆ Z from T via p ( ˆ Z | T ) for the unconditional case, and the graph G from ˆ Z via p ( G| ˆ Z ) .
The decoder network D group also plays a role in implementing the conditional lower bound, via p ( ˆ Z | Z, T ) , which computes ˆ Z conditioned on Z and T .
4.2 Encoders
Graph encoder.
We first encode the graph features X and the adjacency matrix A in the latent space Z , using an inference network E graph q
( Z | X, A ) .
Similarly to Kipf and Welling [28], we use a single 2-layer network, where the first layer leverages a graph convolutional network (GCN) to reduce the dimension of the graph features to the symmetry group size: X = GCN ( X, A ) .
The second layer computes the mean and variance vectors ( µ = GCN µ ( X, A ) , log σ = GCN σ ( X, A ) ), which are then used to sample the latent variable Z from a Gaussian distribution with the reparameterization trick : q ( Z | X, A ) =
N ( Z | µ, diag ( σ 2 )) .
Group encoder.
We employ E group to map the latent vector Z into a vector T of k Lie algebra coordinates, where k is the subspace size (e.g., the number of disentangled elements).
In our settings, each lie algebra coordinate { t i } ki =1 is the output of a GCN with two layers.
The first one computes a graph embedding Z = GCN ( Z, A ) , and the second one produces the vectors µ = GCN µ ( Z, A ) , log σ = GCN σ ( Z, A ) , which in turn yield produce { t i } ki =1 with the reparameterization trick.
Finally, { t i } ki =1 are combined into a single tensor T using a multilayer perceptron (MLP): q ( T | Z, A ) =
MLP ( q ( t i | Z, A )) with q ( t i | Z, A ) =
N ( t i | µ i , diag ( σ 2 i )) for i = 1 . . .
, k (3)
4.3 Decoders
Group decoder.
The module D group reconstructs the latent variable ˆ
Z using a Lie group G and a
Lie algebra g .
Specifically, we first learn a Lie algebra basis element { A i } ki =1 ∈ g for each coordinate t i .
We then aggregate the coordinates t i and basis elements
A i with an exponential map, to obtain a group representation as p
[Equation]
[Equation]
[Equation]
exp( A ( T )) where A ( T ) =
(cid:88) k
i =
1 t i A i for g ∈ G, A ∈ g .
The conditional version of D group is obtained with a slight modification of Eq.
4: ˆ Z is computed by feeding the mean between Z (e.g. the graph embeddings obtained from E graph ) , and exp( A ( T )) into a non-linear activation function σ such as ReLU:
[Equation]
|
[Equation] [Equation]
[Equation]
[Equation] [Equation]
, where Z = { z j } kj =1 .
Graph decoder.
The network D graph reconstructs the graph data ˆ G from the latents ˆ Z .
This is achieved by first passing the latent feature to a GCN layer that maps Z into a vector of the same dimension as the original feature vectors, and subsequently into a sigmoid activation function:
[Equation]
4.4 Training
The proposed unconditional and conditional models are trained by optimizing the loss given by β L + λh + γc , where L is the lower bound from Eq.
1 (unconditional) or Eq.
2 (conditional), and h and c are respectively the hessian and commutative penalties, obtained as regularization terms adapting the implementation from Zhu et al.
We provide insights into the effect of parameters on the models and their performance with an ablation study in Section 5.1.
For neural architecture, we set each encoder and decoder to consist of 3 layers, where each layer takes as input the graph features from the previous layer.
Experiments
We conducted extensive experiments that we decribe now.
Section 5.1 evaluates the disentanglement capabilities of the models, Section 5.2 provides a compression and few-shots classification experiment, and Section 5.3 assesses the generation capabilities on molecular datasets.
5.1 Disentanglement evaluation
We evaluate disentanglement with a) an ablation study for assessing our models components, b) qualitative investigations, and c) quantitative metrics.
For metrics, we follow the evaluation protocols from
Locatello et al.
[36], which involve being able to randomly combine a number of generative factors to sample a new data.
In our case, we provide the necessary generative factors to generate two types of synthetic random graphs (e.g. Erdos-Renyi and Watts-Strogatz).
More details about our dataset construction procedure are provided in Appendix A.
We compute 5 metrics, including β -VAE
Ablation study.
We start by evaluating the impact of our model components on the disentanglement metrics, by maintaining all parameters fixed, except from one that varies.
We probe the effect from different group sizes, hessian penalties, commutative penalties, and subspace sizes.
In Figure 2 we report the mean over 10 runs for our unconditional model on the synthetic Watts-Strogatz dataset.
We observe that when the group size is increased, there is an improvement in disentanglement across all metrics, and this can be explained by the idea that with a larger group size, the subgroups, represented by the Lie algebra coordinate t i can control different data variations.
The hessian penalty results to be effectively enhancing disentangled representations, while the commutative penalty is less effective, and we can explain this result because the former requires the Lie algebra basis elements to have mutual products of zeros while subgroup decomposition only requires their commutators to be zeros, which also confirms the results obtained for image datasets [71].
Correlation analysis.
Following previous work [66, 38], we report the correlation matrix of the hidden graph features, the goal of which is to show whether the features capture mutually exclusive information, which is achieved when the there is a block-wise correlation pattern.
In our models we use a group-size of 32 (to match with hidden features size of the other methods), a hessian penalty of 40 and commutative penalty of 5.
In Figure 3 we note that both our SID U and SID C models achieve a block-wise correlation pattern, however the conditional version has also some other highlighted regions, because in this model involves fixing a latent vector Z to condition the reconstruction.
As a result, the conditioning may be responsible for the visualized correlation on the hidden features.
Learned graph variations.
In this experiment, we train the SID C model on Watts-Strogatz (WS) random graphs, which have 3 generative factors, n for number of nodes, p for connectivity (the probability of rewiring each edge), and k for neighborhood ( k nearest neighbors in a ring topology).
We adapt our model to the visualization implementations provided by Stoehr et al.
[57], which report adjacency matrices (black) and graph (white) information, where variations in node attribute values are indicated in blue.
In Figure 4 we report the results for the SID C model by projecting pairs of latent variable traversals to see how the factors vary, and we observe the following.
The variable zcontrols the node attribute value.
For example, in adj 01 we see the color changing from blue to white along the horizontal axis.
The variable z 1 controls connectivity, for example in graph 01 , in the zaxis (top to bottom) we see that the number of edges decrease, while in graph 12 , along the z 1 axis (left to right), the number of edges increase.
The variable z 2 controls the number of nodes, and we see that in graph 12 and graph 02 the number of nodes decreases along the z 2 axis (bottom to top).
Disentanglement metrics.
We perform 10 random seed runs, by training for 200 epochs on the proposed synthetic datasets and measure the quantitative metrics on 1000 data points.
In terms of
[Table]
baselines, we consider graph VAEs with GAT [61] and GCN [29] encoders, as well as disentangled encoders such as FAC [66] and DIS [38].
We set our models to have a group size of 81, a hessian penalty of 40, and a commutative penalty of 5.
In Table 1 we observe that our models are able to improve on previous baselines on all metrics on the evaluated datasets, and the conditional version achieves the best performance in most cases.
We notice that the FL metric shows a significant performance increase, justified by its ability to capture the disentanglement on each generative factor, which, in our models is controlled by the group size, while in previous approaches is not considered.
In terms of the other previously proposed disentangled approaches, DIS and FAC demonstrate higher disentanglement in most cases, when compared to GCN and GAT , however they are outperformed by our methods.
5.2 Compression and few-shots classification
Following the setup from Ge et al.
[14], we consider the task of learning a compressed model via multiple pooling operations, and test on few-shots classification capabilities on datasets with multiple classes, including FRANKENSTEIN [43] (2 classes), COLORS-3 [30] (11 classes), Mutagenicity [51] (2 classes), NCI1 [54] (2 c).
The baselines include MIA [14], which propose an attention-based pooling for compression, as well as GCN [29], and GAT [61] based autoencoders, equipped with the same pooling layer.
All models are set to have a pooling rate of 0.8, and a depth of 3 layers.
Our models are set with a group size of 81, hessian penalty of 20 and commutative penalty of 5.
The task consists in training first without labels to learn the representation, and then using the compressed model as input for training a GCN classifier a labelled subset of the dataset, which involves 100 samples for training and 100 for testing.
In Table 2 we report the mean and variance across 10 runs for MSE, the classification accuracy, and the size of the trained model.
The results show that SID C and SID U outperform the considered baselines in terms of few-shots classification accuracy on all datasets, while achieving a smaller compressed size.
We hypothesize that the superior accuracy performance of our models is due to a more efficient learning provided by disentangled representations, which have previously demonstrated the ability to enhance the accuracy of predictions in the context of generalization tasks [41, 70, 32].
In terms of MSE, our models show an improvement compared to GCN - and GAT -based autoencoders, but they are outperformed by MIA in most cases.
[Table] [Table]
5.3 Molecular graph generation
Molecular generation involves three tasks, random generation, property optimization and constrained optimization, and a suitable setup can be obtained following the approaches implemented via TorchDrug 2 and DIG [33].
The generation process involves 1) learning the distribution of the molecular data and 2) fine-tuning the pre-trained model on one of the tasks.
We employ our models during the training phase, using a group size of 81, a hessian penalty of 40 and commutative penalty of 5, and then we follow the reinforcement learning fine-tuning method from GCPN [67].
Similar setups are used from previous molecule generation methods, which we include as baselines: JT-VAE [24],
GraphAF
[55], GraphDF [37], GCPN [67].
Random generation.
This task evaluates the quality of randomly generated samples on 4 standard metrics in percentage, including: valid molecules with resampling (Val), valid molecules without resampling (WR), unique molecules (Uni), novel molecules that not appear in the training data (Nov).
We train our models for 10 epochs, with a batch size of 32 and the learning rate of 0.001, and compute the metrics over 10K generated molecules.
In Table 3 we report results on ZINC [23], QM9 [49], and
MOSES [47].
We observe that both SID U and SID C are able to improve the performance of GraphAF and GCPN on most metrics, while achieving comparable performance with GraphDF, furthermore the SID C model achieves the highest score of validity without resampling on the QM9 dataset.
Since our model is not optimized for molecule generation, the positive results of our models indicate that disentangled representations may represent a flexible tool for generating more realistic graphs, and that encouraging the separation of semantic factors into different latent codes, may be beneficial for improving the quality of the samples.
Property optimization The goal of the task is to generate novel molecules with high property scores for the penalized logP property, and the quantitative estimation of drug-likeness property (QED).
The setup involves pretraining our models for 300 epochs on ZINC random generation, applying the RL fine-tuning procedure from GCPN, and measuring the scores from the top 3 generated molecules.
In
Table 4 we observe that our models are able to outperform GCPN and GraphAF for both penalized logP and QED, and achieve comparable performance to the more advanced GraphDF model.
In
Figure 5 we show the molecules generated using SID C .
While our model does not outperform
GraphDF, our results are relevant because they demonstrate that the disentangled representations learned from our models can be leveraged to improve the performance of the downstream task of property optimization.
[Table] [Table]
Constrained Optimization.
This task aims to modify the input molecular graph for improving its penalized logP score while keeping the similarity between the input and modified molecules higher than the threshold δ .
We follow the setup of GraphAF and GraphDF, which select 800 molecules from ZINC with low penalized logP scores as the input molecules to be optimized.
We pretrain our model on the random generation task for 300 epochs, and then fine-tune with the RL procedure from
GraphAF.
We report the mean and standard deviation of the largest property improvement (Imp), and similarities (Sim) between them and their corresponding input molecules, as well as the success rate (Suc).
In Table 5 we compare the results for our SID C model with GraphAF and GraphDF.
We observe that SID C is able to outperform GraphAF 3 out of 4 times for improvement values, and achieves comparable results with GraphDF, which we motivate by the fact that GraphDF provides further fine-tuning in order to improve the baseline.
In terms of similarity and success rate, our model has comparable results with the baselines for all δ setups.
Discussion
Limitations.
For quantitative evaluation, we rely on synthetic data, since current metrics are based on the ability to sample the factors and combine them into observations Kim and Mnih [26].
Future work should investigate how to provide metrics that can evaluate disentanglement on real-world datasets.
A recent line of work from Khemakhem et al.
[25] shows that disentanglement is closely related to model identifiability, i.e., a representation is disentangled when a set of learned parameters is uniquely identified with the parameters in the true distribution.
Our approach has not considered or provided proofs for, identifiability, and we leave this as future work.
Conclusion.
We formalize the notion of conditional disentanglement on graphs and propose a novel framework for graph disentanglement by leveraging tools from Lie algebras.
Based on the new definition of disentanglement, we design a graph VAE based on a Lie group parameterization, and provide a novel ELBO criteria for optimizing conditional disentanglement.
Our method achieves superior performance on quantitative disentanglement benchmarks when compared to contemporary disentangled GNNs and other convolutional layers.
Finally, we demonstrate strong capabilities on few-shots classification and molecular generation experiments.
Acknowledgements
We thank the anonymous reviewers for their helpful comments.
GM’s research was partially funded by EPSRC and the BBC under iCASE.
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] See Section(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] 2.
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] Provided in Appendix 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] Provided at https://mercatali.gitlab.io/sid/ (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] We specify the experimental setups in Section(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See beginning of Section(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] See beginning of Section4.
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] See Section 5.
(b) Did you mention the license of the assets?
[Yes] See Section 5.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We include our assets at https://mercatali.gitlab.io/sid/ (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 offensive content?
[N/A] 5.
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 Synthetic datasets
We construct two datasets with the PyTorch Geometric library 3 and the NetworkX library 4 , based respectively on Watts-Strogatz, and Erdos-Renyi random graphs.
The datasets consist in a list of graphs, where we control the generative factors necessary for synthesizing node feature and adjacency matrices in each graph.
This allows us to create labels and latent codes for each graph, similarly to the canonical dataset for disentanglement such as Dsprites 5 .
Node feature factors The node feature matrix is computed for both datasets as a random tensor of dimension N nodes × N features .
The number of node is fixed to 20, while the features are 64 .
The elements of the feature matrix are sampled from a normal distribution with a given mean and variance, which are two of the generative factors that we control.
Adjacency matrix factors The adjacency matrices are computed via the NetworkX library, where we are able to control the factors responsible for generating the graph.
For Watts-Strogatz graphs we have (1) a factor k , which represents the k nearest neighbors in a ring topology, and (2) factor p , which is the probability of rewiring each edge.
For Erdos-Renyi graphs, we control only one factor p , that represents the probability for edge creation.
Combination Watts-Strogatz data has 2 adjacency factors, while Erdos-Renyi has only one, thus, in order to achieve the same number of graphs in both datasets, we fix the variance to 2 on WattsStrogatz, while we provide 20 values in Erdos-Renyi.
The final Watts-Strogatz and Erdos-Renyi datasets have 40 × 40 × 20 = 32 .
000 graphs.
The combinations for the factors: mean, variance, k , and p are provided in Table 6.
Note that number of nodes is not considered a factor.
[Table]
n/
[Equation]
32K
B Unconditional Disentanglement
The unconditional ELBO (Proposition 1), states that given two latent variables Z and T modeling the log-likelihood of the graph data G is bounded by the following L u ELBO:
[Equation]
Proof.
Using Jensen inequality
[Equation]
we have:
[Equation]
By applying Jensen inequality again, we obtain:
[Equation]
The results follows assuming p ( G| Z, T ) =
p ( G| Z ) , and noting that for computing T , we only use the adjacency information A from G .
C Conditional Disentanglement
Using Jensen’s inequality, we note that for any fixed ˆ Z
[Equation]
Since we use only A from G for computing T
[Equation]
we can write q ( T |G , Z ) =
q ( T | Z, A )
[Equation]
and applying
Jensen’s inequality again
[Equation]
we get
[Equation]
Z |G ) log q ( Z |G ) .
Assuming G is independent of Z and T given ˆ Z , we immediately get log p ( G , ˆ Z )
[Equation]
−
[Equation]
D Scalability experiment
We perform a graph classification task on large scale datasets, and report their dimensions in Table 7.
In Table 8 we report the results for our SID C model on social network datasets including IMDB binary
IMDB-multi, COLLAB, as well as macro molecules such as PROTEINS and MUTAG.
The baselines for this experiment include FactorGCN [66] and DGCL [31].
For our model we set group-size to 81, hessian penalty to 40 and commutative penalty to 5.
We follow the classification setup from Yang et al.
[66], with a ten-fold cross-validation procedure
[Table]
and report accuracy and standard deviation.
The results show that both SID U and SID C are able to improve the performance of the baselines.
In particular, SID C achieves the highest accuracy on all the datasets.
[Table] [Table]
E Molecular experiment
In order to have a fair comparison with Flow models, such as GraphAF [55] and GraphDF [37]
[Table]
we incorporate our parameterization into the training process of a Flow model.
Both GraphAF and GraphDF are designed using relational GCN [53] (RGCN) as building block to compute an embedding of a graph using a layer R-GCN: H Li = R-GCN ( G i ) ,
˜ h
i = sum ( H Li ) where sum denotes the sum-pooling operation, and H Li,j ∈ R k denotes the embedding of the j -th node in the embeddings H Li .
We incorporate some parts of our models after the RGCN encoding, as follows.
From the embedding
H , we compute the lie algebra coordinates, with an encoder q ( T | H ) , following our E group network from Section 4.
We then apply the lie algebra exponential mapping to reconstruct the embedding H as p ( H | T ) , following the D group network from Section 4.
Furthermore, we incorporate the hessian penalty as a regularization.
In all the three tasks, we follow the experimental setup from GraphAF, we set the model with group-size of 81, and a hessian penalty of 40.
We first perform the random graph generation task, which involves quantifying 4 metrics in percentage, including: valid molecules with resampling (Val), valid molecules without resampling (Res), unique molecules (Uni), and novel molecules not appearing in the training data (Nov).
In Table 9 we report results for GCPN, GraphAF, GraphDF and our models.
The results indicate that with the inclusion of a lie algebra reparameterization and the hessian penalty, the performance on random generation is enhanced, and the Flow + SID C model achieves the best performance on all metrics a part from uniqueness, where it achieves comparable results.
Secondly, we report the results for the property optimization task in Table 10, which show the scores from the top 3 generated molecules for two selected properties (penalized logp and QED).
The baselines are the same ones as in the random generation task.
We observe that our Flow + SID C model is the top performing method on penalized logP for the top 3 scores, and it achieves the top results, together with the other methods, for QED.
[Table] [Table]
Finally
[Table]
we show in Table 11 the results for the constrained optimization task.
We report only GraphDF as baseline, which is the top performing model among the baselines.
The evaluation involves reporting the mean and standard deviation of metrics including: the largest property improvement (Imp), and similarities (Sim) between them and their corresponding input molecules, as well as the success rate (Suc).
We observe that by combining the disentangled principles developed in our models into a Flow-based model, we are able to further improve the results
[Table]
and achieve the top performance.
[Table]