Symmetry-induced Disentanglement on Graphs
Anonymous Author(s)
AffiliationAddress email
Abstract
Learning disentangled representations is important for unraveling the underlyingcomplex interactions between latent generative factors.
Disentanglement has beenformalized using a symmetry-centric notion for unstructured spaces, however,graphs have eluded a similarly rigorous treatment.
We fill this gap with a newnotion of conditional symmetry based disentanglement, and leverage tools from
Lie algebras to encode graph properties into subgroups using suitable adaptationsof generative models such as Variational Autoencoders.
Unlike existing workson disentanglement, proposed models can learn to segregate the latent space intouncoupled and entangled parts.
Experiments on synthetic and real datasets revealthe ability of these models to learn effective disengaged representations, andimprove performance on downstream tasks such as few-shot classification andmolecular generation.
Introduction
Disentanglement represents a fundamental desideratum in Artificial Intelligence, because it capturesinformation about the salient (or explanatory) factors of variation in the data, and it isolates information about each specific factor in only a few dimensions, unraveling the underlying interactions ofcomplex data [2, 56].
Disentangled representations have often been deemed responsible for providingneural models with the ability to improve real-word downstream tasks [2, 55, 17, 56, 32], and havedemonstrated empirically the ability to enable critical properties such as sample efficiency [18] andsystematic generalization [65, 28, 39].
Higgins et al.
[17] have recently presented a symmetry-based definition for disentanglement, whichuses tools from group theory to formalize the idea that disentanglement is achieved when the abstractand the vector spaces are equivariant.
Such symmetry-based definition presents various advantages:first it aligns with the idea that auxiliary information can be exploited for imposing a structure on thelatent space [21, 36], secondly, it has been successfully used as a theoretical framework for recentdisentangled Variational Autoencoder (VAE) works [23, 47] [60, 66], finally, differently from otherdefinitions of disentanglement [10, 12, 2] it may enable a fundamental connection with several graphneural network (GNN) approaches that leverage group symmetries [5, 35, 9, 3].
Disentanglement approaches have been initially developed on vision-based datasets [16, 6, 22, 8, 14],however, numerous recent methods have improved GNNs performance using disentanglement [11, 37,34, 27, 61].
For example, [34] propose to use a convolutional layer to identify the factor responsiblefor linking a node to its neighbors via a neighborhood routing mechanism, while [61] enable multirelational disentanglement, by encoding all the possible links between nodes with attention.
While previous work in the GNN domain has demonstrated the advantages of exploiting groupsymmetries [35, 3], investigating disentanglement from a symmetry group perspective is still anunexplored area in the GNN literature.
Furthermore, a rigorous formalization of disentanglement for graphs, which aligns the formalism provided by [17] within a GNN architecture is still missing in thecontemporary disentangled GNN literature [34, 61, 27].
In this paper, we make several contributions to fill the mentioned gaps in disentangled GNNs.
• We introduce two formal definitions for unconditional
and conditional symmetry
-based
disentanglement on graphs, which provide the tools for encoding graph properties intosubgroups by leveraging a Lie algebra parameterization.
Our intuition is that each generativefactor can correspond to a Lie algebra coordinate element, and control a single graphproperty (Section 2).
• Based on the proposed Lie group disentanglement framework, we design two disentangledGNN modules (unconditional and conditional) that auto-encode both a graph representationand a Lie group, within a VAE framework (Section 3).
• We provide the first systematic evaluation on several disentanglement metrics [22, 16, 8,12, 42]
and demonstrate that the proposed models can outperform all contemporary GNN
baselines [34, 61]
and successfully learn to segregate the latent space into uncoupledand entangled parts (Section 5.1).
Finally, we show that our models are able to improve
performance on downstream tasks such as few-shot classification, compression (Section 5.2)and molecular graph generation (Section 5.3).
Proposed framework
2.1 Symmetry-based disentanglement
We consider the definition for disentangled representations from Higgins et al.
Definition 1 (Symmetry-based disentanglement) .
Let us consider a generative process b = W
→ Omapping a set of world states W into a set of observations (i.e. data) O , and an inference process h =
O → Z mapping the data into the vector space.
By composing b and h , we obtain f = h ◦ b , whichis equivalent to f : W → Z .
Given group actions on W · : G × W → W and on Z · : G × Z → Z ,the function f is defined equivariant between the actions on W and Z if the action commutes with f ,which can be expressed as an equivariant map: g · f ( w ) =
f ( g · w ) , ∀ g ∈ G, ∀
w ∈
W. Assuming thatthe symmetry group G can be decomposed into subgroups G 1 × . . .
× G n , a vector representation
Z is disentangled with respect to the decomposition of G if: (1) there is an action defined on therepresentation set · : G ×
Z → Z , (2) the map f : W → Z is equivariant between actions on W and
Z , (3) there is a decomposition Z
= Z 1 × Z 2 . . .
×
Z n such that each Z i is affected only by G i andfixed for G j , j ̸
= i .
Following the notation of Definition 1, we propose the notion of conditional disentanglement.
Definition 2 (Conditional symmetry-based disentanglement) .
We consider: a generative processb =
W × E → O mapping a set of disentangle-able world states W and non-disentangle states Einto a set of observations (i.e. data) O , and an inference process h
= O → Z mapping the data intothe vector space.
By composing b and h , we obtain f = b ◦ h , which maps W × E to Z .
Givengroup actions on W · : G × W → W and on Z · : G × Z → Z , the function f is defined equivariantbetween the actions on W and Z if the action commutes with f according to the following equivariantmap: g · f ( w, e ) =
f ( g · w, e ) , ∀ g ∈ G, ∀ w ∈ W, ∀ e ∈ E .
2.2 Lie algebras for graph networks
While previous works explored disentanglement on GNNs [34, 61], they lack formal frameworksfor designing disentangled GNNs, which is necessary for aligning with the notions of symmetrygroups from Definition 1.
In the next two paragraphs, we introduce Lie algebras and leverage them toprovide a disentangled GNN architecture.
Lie groups and Lie algebras A Lie group G is a group of symmetries where the symmetries arecontinuous [49].
A Lie group is associated with a Lie algebra g , which is the tangent space to theidentity element of G .
A Lie algebra is parameterized with a basis { A i } ki =1 , where every element ing can be written as A =
A 1 t 1 + . . .
+ A k t k , where t i are the coordinates.
Elements of the Lie algebracan be mapped into the Lie group with an exponential function: exp : g → G .
In the next paragraph, we show how we can connect the notion of a Lie group G with the definition of disentanglement in a
GNN model, by associating each Lie algebra coordinates t i element with a generative factor.
GNN architecture Given two latent variables Z and T , we instantiate the unconditional disentanglement from Definition 3 with a probabilistic model that maximize the log-likelihood of the graphdata G = ( X, A ) , as follows: We consider an input graph G = ( X, A ) where X denotes the matrix obtained by stacking thefeatures of the nodes, and A denotes the adjacency matrix.
Let N be the number of nodes in G andd the number of features, so X ∈ R N × d . We consider a Lie group of dimension L associated withk Lie algebra basis elements.
First, we compute the graph features f = Pool ( GCN ( G )) ∈ R N 1 × L ,where the GCN maps the features to a group matrix of size N × L and Pool reduces the number ofnodes to N 1 using a certain pooling rate.
Next, the second layer of the GCN maps the features f to alatent code Z of dimension N 1 × k .
To enhance disentanglement, Z is then fed into k GCNs, whereeach the output of each GCN is another N 1 × k matrix.
We then perform an aggregation operationon each GCN as: T = agg ( { GCN j ( Z ) =
t j } kj =1 ) .
The vector T represents the k coordinates of a Lie algebra g , and models the generative factors of the dataset.
Finally, T is fed into the exponentialmap, which reconstructs the variable ˆ Z , while imposing a group structure on it, by learning the basis
A ∈ g , i.e., ˆ Z = exp
A ( T ) where A ∈ g and A ( T ) =
(cid:80) k
[-=
-]i =
1 A i t i .
This model is further explained in Section 3.
2.3 Formalizing disentanglement on graphs
Intuitive definition
The latent representation ˆ Z is disentangled with respect to the Lie group G ,if a change in the coordinate t i is associated with a change in ˆ z i .
In other words, disentanglementmeans that the semantics ( t 1 , t 2 , . . .
, t k ) are equivariant with respect to the properties that are beingpredicted (ˆ z 1 , ˆ z 2 , . . .
, ˆ z k ) .
Next, we provide the formal definitions for unconditional and conditionalgraph disentanglement, where the former expresses ˆ Z in function of T without accounting for Z ,while the latter expresses ˆ Z based on T , after fixing an encoding Z .
Definition 3 (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: (1) there isa group action on ˆ Z : · : G × ˆ Z → ˆ Z , (2) the map f = exp A
( T ) :
T →
ˆ Z is equivariant betweenactions on T and ˆ Z , (3) there is a decomposition ˆ Z =
ˆ z 1 × ˆ z 2 . . .
× ˆ z k , where each coordinate t iaffects only the corresponding component ˆ z i .
Definition 4 (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 conditionedon Z if (1) there is a group action on ˆ Z : · : G × ˆ Z → ˆ Z , (2) the map f : exp A ( T ) ×
Z → ˆ Z isequivariant between actions on T and ˆ Z for any fixed Z , and (3) exp A ( T ) factorizes into the product 119 (cid:81)
ki =1 exp( t i A i ) ×
Z such that each component exp( t i A i )
× Z is affected only by the correspondingcoordinate t i for i ∈ { 1 , 2 , . .
. , k } .
2.4 Disentanglement constraints
Since a Lie group parameterization is not sufficient for a group to be decomposed into subgroups thatare parameterized independently by a single coordinate [66, 60], we provide the necessary constraintsfor enforcing disentanglement.
Commutative penalty First, disentanglement can be achieved by enforcing a coordinate t i to beidentified by single group exp( t i A i ) , representing a single property variation in the data ˆ z i .
Suchproperty is expressed by the equivalence between the exponential map of the sum, and the productover exponential maps, which [66] proved to be achieved when there is commutativity between the
Lie algebra basis.
Namely, if A i
A j =
A j A i then exp (cid:80) ki =1 t i A i
=
(cid:81) ki
=1 exp t
i A i .
Hessian penalty Furthermore, disentanglement can be strengthened using the hessian penalty [42],which states that the Hessian matrix with respect to a disentangled representation is always zero,since there should be independence between different dimensions.
Zhu et al.
[66] adapt this penalty in a Lie algebra parameterization setup, as H ij =
δ 2 g ( T ) δt i δt j where g ( T ) =
exp (cid:80) ki
=1 t i A i , and alsoshow 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
Lie Group GVAE
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
Method
We design a Lie group graph VAE framework [24], as depicted in Figure 1, which considers the latentgraph encoding variable Z , the graph features X , the symmetry Lie group structure T , and the graphadjacency matrix A .
We propose two decoders, the unconditional one reconstructs the properties ˆ Zgiven T , while the conditional one reconstructs the properties ˆ Z given Z and T .
3.1 Probabilistic model
Proposition 1 (Unconditional ELBO) .
Given two latent variables Z and T , we instantiate theunconditional disentanglement from Definition 3 with a probabilistic model that maximize thelog-likelihood of the graph data G =
( X, A ) , as follows:
[Equation]
Proposition 2 (Conditional ELBO) .
We extend Proposition 1 with a conditional version, corresponding to the conditional disentanglement from Definition 4.
[Equation]
Propositions 1 and (2) are proved respectively in Appendix B and C.
Equations (1) and (2) are implemented with four modules, namely, two stochastic networks q ( Z |G )( E graph ), encoding graph G into vector Z and q ( T | Z ) ( E group ), encoding the group structure T givengraph encoding Z , and two deterministic networks p ( ˆ Z | T ) ( D group ), decoding ˆ Z from the groupstructure T , and p ( G| ˆ Z )
( D graph ), decoding the graph G from ˆ Z .
Note that the decoder D groupis responsible for the group structure injection, because it imposes p ( ˆ Z | T ) to be a Lie algebraparameterization.
Furthermore, in the conditional ELBO model the D group decoder is given byp ( G| ˆ Z, Z ) , where the reconstructed latent code ˆ Z , is conditioned on the fixed latent variable Z .
3.2 Encoders
We first encode the graph representation in the latent space Z , by decomposing the features Xand adjacency matrix A using an inference network E graph q
( Z | X, A ) .
We use a single 2-layer
GCN, where the first layer reduces the graph features to the symmetry group dimension, andperforms a pooling operation to reduce the number of nodes: X = Pool ( GCN ( X, A )) , while the second layer reduces the number of features to k while computing the matrices µ = GCN µ ( X, A ) ,log σ = GCN σ ( X, A ) , which are then combined into a Z with the reparameterization trick, asq ( Z | X, A ) = N ( Z | µ, diag ( σ 2 )) .
Secondly, we use E group to map the latent vector Z into a vector T , which represents the k Lie algebracoordinates.
Each { t i } ki =1 is the output of a graph convolutional encoder (GCN) over Z .
Each GCNconsists of two layers, where first one computes Z = GCN
( Z, A ) , while the second one computesthe matrices µ = GCN µ ( Z, A ) and log σ = GCN σ ( Z, A ) , which are then combined to compute t iwith the reparameterization trick.
[Equation]
3.3 Decoders
We first leverage D group for decoding the latent variable ˆ Z using a Lie group G and a Lie algebrag , with a Lie algebra parameterization.
Specifically, for each coordinate t i ∈ T , we learn a Liealgebra basis element { A i } ki =1 ∈ g .
The coordinates and basis are then aggregated and fed into anexponential mapping layer, to achieve a group representation, as follows.
[Equation]
For the conditional version of the ELBO, D group ˆ Z is computed by combining Z and A ( T ) , asp ( ˆ Z | Z, T )
=
σ ( agg [exp A ( T ) , Z ]) , where Z = { z j } kj =1 , agg is the mean, and σ is the ReLUnon-linear activation function.
Secondly, we apply D graph , which first maps the latent ˆ Z into N latent variables using sigmoidfunctions σ j , and then aggregates the N latent variables to obtain the full graph representation.
[Equation]
3.4 Training
The proposed unconditional and conditional models are trained by optimizing the loss given byβ L + λh + γc , where L is the ELBO, from Equation (1) (unconditional) or (2) (conditional), whileh and c are respectively the hessian and commutative penalties, obtained as regularization terms.
Theparameters are found empirically, and we provide further insights on their effect on the model andits performance by providing an ablation study in the experiments in Section 5.
In terms of neuralarchitecture we set each encoder and decoder to be made of 3 of the described layers, where eachlayer takes as input the graph features of the previous layer.
Related work
Symmetry-based disentanglement Initial works on symmetry disentanglement [7, 45] have proposed use reinforcement learning to achieve irreducible representations of group elements throughobservation of action transitions, while [41] does not require labelled action-transition pairs.
In arelated line of work, [66] propose to decompose a Lie group in multiple subgroups, and enforcecommutativity and independence between subgroups, while [60] impose a cyclic group structure,while encouraging such group to be isomorphic to the ground truth.
Other works focus on manifoldsproperties, for example, [43], propose a nonparametric algorithm that aims to discover a decomposition of a data manifold by investigating its holonomy group, while [13] model disentanglementas a product of the low-dimensional submanifolds underlying the data space.
Finally, [58] designa method which iteratively updates a partition, modeled as a symmetry group, until the consideredfactor is invariant, while [4] uses group representation theory for defining distributed latent operators.
Differently, from the reviewed methods, we propose to analyze the graph data modality.
Disentangled GNNs [34] propose a routing mechanism aimed at separating the underlying factorsin a graph dataset, while [61] use an attention mechanism for disentangling the representationof GNNs, and provide two evaluation metrics for disentanglement on graphs.
In a related line of work, [31] achieve disentanglement by encouraging independence in the latent space, while[1] propose a method for disentangling multi-relational graphs with an application on pedestriantrajectory prediction.
Furthermore, [11] propose a generative model for spatio-temporal graphs viadisentangled representations, while [27] leverage a contrastive framework for disentangling the latentfactors in GNNs.
Differently from the mentioned GNNs methods, we provide a formal definition ofdisentanglement for graphs, and do leverage group structures for modeling.
Graph generative models We also relate to VAEs models on graph data modality, for example [59,64] propose VAE-based Information Bottleneck methods for improving GNNs performance onclassification tasks, by learning minimal sufficient representations.
Another line of work aimsto generate molecular graphs, by proposing graph-VAE models that are provided with encoderand decoders able to handle graph structures, and to condition molecule generation on specificproperties [30, 53].
Similarly, [63] propose a deep autoregressive model that learns to generate graphsby training on a representative set of graphs and decomposes the graph generation process into asequence of node and edge formations.
We differ from these methods as we explicitly investigate thedisentanglement capabilities of generative models.
Experiments
Section 5.1 evaluates the disentanglement capabilities of the models, Section 5.2 provides a compression and few-shots classification experiment, Section 5.3 assesses the generation capabilities onmolecular datasets.
5.1 Disentanglement evaluation
Experimental setup Vision-based disentanglement [16, 22] requires the generative factors to beknown, is based on synthetic datasets that have such property.
We follow the protocol, and constructtwo graph datasets based respectively on random graphs, where we can sample a combination ofgenerative factors that are required to synthesize node features and adjacency matrices in each graph.
More details on the datasets are in Appendix A.
We measure 5 disentanglement metrics from theliterature, including β VAE metric [16] (Beta), FactorVAE metric (FVM) [22], Mutual Information
Gap (MIG) [8], DCI Disentanglement [12], and Factor Leakage (FL) [41].
Ablation study We start by evaluating the impact of our model components on the disentanglementscores, by maintaining all parameters fixed, except from one that varies.
We probe group size, hessianpenalty, commutative penalty, subspace size, beta penalty, and model depth.
In Figure 2 we report themean over 10 runs for our unconditional model on the synthetic Watts-Strogatz dataset.
We observethat 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 bythe Lie algebra coordinate t i can control different data variations.
The hessian penalty results to beeffectively enhancing disentangled representations, while the commutativity penalty is less effective,and we can explain this result, because the former requires the Lie algebra basis elements to havemutual products of zeros while subgroup decomposition only requires their commutators to be zeros,which also confirms the results obtained for image datasets [66].
Quantitative disentanglement benchmark We compare our models with several baselines, over10 random seed runs, where we measure the mean and variance, by training for 200 epochs on theproposed synthetic datasets and measuring the quantitative metrics on 1000 data points.
Baselinesinclude GAT [57] and GCN [25] based autoencoders, as well as disentangled graph methods such as
FAC [61] and DIS [34].
Our unconditional and conditional Lie group models are named UNC and
CON and set with group size of 81, a hessian penalty of 40, and a commutativity penalty of 5.
In Table 1 we observe that our Lie group models are able to improve on previous baselines on allmetrics on the evaluated datasets, and the conditional version achieves the best performance in mostcases.
The FL and FVM metrics show the highest performance increase.
FL is designed to capturethe disentanglement on each generative factor, which, in our models is controlled by the group size,while in previous approaches is not considered.
FVM is the most used metrics for vision datasets, [Table] and provides a more robust assessment than the Beta metric.
We note that while previously proposeddisentangled approaches DIS and FAC demonstrate higher disentanglement when compared to GCNand GAT , they are outperformed by our methods.
5.2 Compression and few-shots classification
Experimental setup Following the setup from [15], we consider the task of learning a compressedmodel via multiple pooling operations, and test on few-shots classification capabilities on datasetswith multiple classes, including FRANKENSTEIN [40] (2 classes), COLORS-3 [26] (11 classes),
Mutagenicity [48] (2 classes), NCI1 [51] (2 c).
The baselines include MIA [15], which propose anattention-based pooling for compression, as well as GCN [25], and GAT [57] based autoencoders,equipped with the same pooling layer.
All models are set to have a pooling rate of 0.8, and a depthof 3.
Our models are set with a group size of 81, hessian penalty of 20 and commutativity penaltyof 5.
The task consists in training first without labels to learn the representation, and then usingthe compressed model as input for training a GCN classifier a labelled subset of the dataset, whichinvolves 100 samples for training and 100 for testing.
Results In Table 2 we report the mean and variance across 10 runs for MSE, the classificationaccuracy, and the size of the trained model.
The results show that CON and UNC outperform theconsidered baselines in terms of few-shots classification accuracy on all datasets, while achieving asmaller compressed size.
We hypothesize that the generalization performance is due to more efficientlearning due to disentanglement, in fact disentangled representations have been studied in the contextof generalization in previous work [38, 65, 28].
In terms of MSE, our models show a decrease inperformance compared to MIA , but improve on GCN and GAT based autoencoders.
[Table] [Table]
5.3 Molecular graph generation
The molecular generation experiment is two-fold, the first part involves learning the distribution of themolecular data while the second one involves fine-tuning the pre-trained model with a Reinforcement
Learning (RL) procedure for a downstream generation task.
We employ our disentangled modelsduring the training phase, and then follow the RL method from GCPN [62] for fine-tuning tasks,which include random generation, property optimization and constrained optimization.
For all tasks,we follow the setups from previous GNN approaches [20, 52, 33] designing our experiments based on TorchDrug 2 and DIG [29].
We set our disentangled models with a group size of 81, a hessian penaltyof 40 and commutativity penalty of 5.
In terms of baselines, we include JT-VAE [20], GraphAF [52],
GraphDF [33], GCPN [62].
Random generation This task evaluates the quality of randomly generated samples on 4 standardmetrics in percentage, including: valid molecules with resampling (Val), valid molecules withoutresampling (Res), unique molecules (Uni), novel molecules that not appear the training data (Nov).
Following the baselines’ setup, we evaluate on ZINC [19], QM9 [46], and MOSES [44].
We train ourmodels for 10 epochs, with a batch size of 32 and the learning rate of 0.001, and compute the metricsover 10K generated molecules.
In Table 3 we observe that our CON model achieves the highest scoreof validity withour resampling on the ZINC dataset, and we note that both UNC and CON are ableto improve the performance of GraphAF and GCPN on most metrics, while achieving comparableperformance with GraphDF.
Since our model is not optimized for molecule generation, the positiveresults of UNC and CON indicate that disentangled representations may represent way to achieve moreexpressive generation, and that encouraging the separation of generative factors into different latentcodes, 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 scoresfor the penalized logP property, and the quantitative estimation of drug-likeness property (QED).
Thesetup involves pretraining our models for 300 epochs on ZINC random generation, applying the RLfine-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 penalizedlogP and QED, and achieve comparable performance to the more advanced GraphDF.
In Figure 3 weshow the molecules generated using CON .
While our model does not outperform GraphDF, this resultis important because it demonstrates that the disentangled representations learned from our modelscan be leveraged to improve the performance of the downstream task. [Table]
Constrained Optimization This task aims to modify the input molecular graph for improving itspenalized logP score while keeping the similarity between the input and modified molecules higherthan the threshold δ .
We follow the setup of GraphAF and GraphDF, which select 800 moleculesfrom ZINC with low penalized logP scores as the input molecules to be optimized.
We pretrain ourmodel 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), andsimilarities (Sim) between them and their corresponding input molecules, as well as the success rate(Suc).
In Table 5 we compare the results for our CON model with GraphAF and GraphDF.
We observethat CON is able to outperform GraphAF
3 out of 4 times for improvement values, and achievescomparable results with GraphDF, which we motivate by the fact that GraphDF provides furtherfine-tuning in order to improve the baseline.
In terms of similarity and success rate, our model hascomparable results with the baselines for all δ setups.
[Table]
Conclusion
We formalize the notion of conditional disentanglement on graphs and propose a novel framework forgraph 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 ELBOcriteria for optimizing conditional disentanglement.
Our method achieves superior performance onquantitative disentanglement benchmarks when compared to contemporary disentangled GNNs andother convolutional layers.
Finally, we demonstrate strong capabilities on few-shots classification andmolecular generation experiments.
Limitations We rely on synthetic dataset for evaluating disentanglement, because current quantitative metrics are based on the ability to sample each factor and combine them into an observation.
Since this is an open problem in representation learning [22], future work will investigate how toprovide metrics that can evaluate disentanglement on real-worl d datasets.
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 tothem?
[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] Provided in Appendix3.
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 as zip
file in the supplemental material
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] We specify the experimental setups in Section 5(c)
Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?
[Yes] See beginning of Section
(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] See beginning of Section 54.
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 in the supplemental material as zip file.
(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]
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 , basedrespectively on Watts-Strogatz, and Erdos-Renyi random graphs.
The datasets consist in a list ofgraphs, where we control the generative factors necessary for synthesizing node feature and adjacencymatrices in each graph.
This allows us to create labels and latent codes for each graph, similarly tothe canonical dataset for disentanglement such as Dsprites 5 .
Node feature factors The node feature matrix is computed for both datasets as a random tensorof dimension N nodes × N features .
The number of node is fixed to 20, while the features are 64 .
Theelements 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, wherewe 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 ringtopology, 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-Renyidatasets 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]
B Unconditional Disentanglement
The unconditional ELBO (Proposition 1), states that given two latent variables Z and T modeling thelog-likelihood of the graph data G is bounded by the following L unc ELBO:
[Equation]
Proof.
Using Jensen inequality, we have:
[Equation]
=
[Equation]
|
[Equation]
By applying Jensen inequality again, we obtain:
[Equation]
Assuming Markov settings, which imply q ( T | Z )
=
q ( T |G , Z ) and p ( G| Z, T ) =
p ( G| Z ) , we have:
[Equation]
[Equation]
p
[Equation]
[Equation]
q
[Equation]
Z |G )
[Equation]
Z
[Equation]
G )
C Conditional Disentanglement
Proof.
Using Jensen’s inequality, we note that
[Equation]
Using our conditional independence assumption q
[Equation]
, Z
[Equation]
=
q
[Equation]
and applying Jensen’s inequality again
[Equation]
we get
[Equation]
Finally, assuming
E q
[Equation]
, we getlog P ( G ) ≥ E q ( Z |G )
E q ( T | Z )
[Equation]
=
[Equation]
Z
[Equation]
G
[Equation]
| Z )
[Equation]
| ˆ
[Equation]
|
[Equation]
[Equation]
|
[Equation]
|G )
[Equation]
Z |
[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 CON model on social network datasets including IMDB binary
IMDB-multi, COLLAB, as well as macro molecules such as PROTEINS and MUTAG.
The baselinesfor this experiment include FactorGCN [61] and DGCL [27].
For our model we set group-size to 81,hessian penalty to 40 and commutative penalty to 5.
We follow the classification setup from Yanget al.
[Table]
with a ten-fold cross-validation procedure
[Table]
and report accuracy and standard deviation.
The results show that both UNC and CON are able to improve the performance of the baselines.
Inparticular, CON achieves the highest accuracy on all the datasets
[Table]
[Table]
[Table]
E Correlation analysis
Following previous work [61, 34], we report the correlation matrix of the hidden graph features.
The correlation matrix aims to show whether the features capture mutually exclusive information,which is achieved when the there is a block-wise correlation pattern.
Specifically, we follow the setupfrom Yang et al.
[61], which analyse the correlation matrix on a dataset composed of graphs, wherethe generative factors are synthetic factor graphs.
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 4 we note that both our UNC and
CON models achieve a block-wise correlation pattern, the conditional version has also some otherhighlighted regions, capturing the entangled factors.
F Learned graph variations
In this experiment we train the CON model on Watts-Strogatz (WS) random graphs, which have 3generative 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 then follow Stoehr et al.
[54]and project latent variables in pairs to see how the factors vary.
In Figure 5 we report the traversal of the CON model trained on 3 latent variables, and we observe thefollowing.
• The variable z 0 controls the node attribute value which is indicated by the color in the rightcolumn.
For example, in the top row on the right, we see the color changing from blue towhite along the horizontal axis.
• The variable z 1 controls connectivity, for example on top row, following the y axis, we seethat the edges decrease, and in second row, along the x axis, edges increase.
• The variable z 2 controls the number of nodes, and we see that in the middle row and the lastrow the number of nodes decreases along the vertical axis.
G Molecular experiment
In order to have a fair comparison with Flow models, such as GraphAF [52] and GraphDF [33],we incorporate our parameterization into the training process of a Flow model.
Both GraphAFand GraphDF are designed using relational GCN [50] (RGCN) as building block to compute anembedding of a graph using a layer R-GCN: H Li = R-GCN ( G i ) , ˜ h
i = sum ( H Li ) where sumdenotes the sum-pooling operation, and H Li,j ∈ R k denotes the embedding of the j -th node in theembeddings 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 networkfrom Section 3.
We then apply the lie algebra exponential mapping to reconstruct the embedding Has p ( H | T ) , following the D group network from Section 3.
Furthermore, we incorporate the hessianpenalty 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 the random graph generation task, which involves quantifying 4 metrics inpercentage, 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 thatwith the inclusion of a lie algebra reparameterization and the hessian penalty, the performance onrandom generation is enhanced, and the Flow + CON model achieves the best performance on allmetrics a part from uniqueness, where it achieves comparable results.
[Table]
Secondly
[Table]
we report the results for the property optimization task in Table
[Table]
which show the scoresfrom the top 3 generated molecules for two selected properties (penalized logp and QED).
Thebaselines are the same ones as in the random generation task.
We observe that our Flow + CON modelis the top performing method on penalized logP for the top 3 scores
[Table]
and it achieves the top results,together with the other methods, for QED.
[Table]
Finally, we show in Table 11 the results for the constrained optimization task.
We report only GraphDFas baseline, which is the top performing model among the baselines.
The evaluation involves reportingthe mean and standard deviation of metrics including: the largest property improvement (Imp), andsimilarities (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-basedmodel, we are able to further improve the results, and achieve the top performance.
[Table]