TLDR We propose a new formalization for disentanglement on graphs using tools from Lie Algebras. We then provide a mechanism to encode graph properties into subgroups within a Variational Autoencoder framework.

Introduction to disentanglement

  • Disentangled representations [1,2] encode information about the salient (or explanatory) factors of variation in the data, and isolate each specific factor in only a few dimensions, unraveling the underlying interactions of complex data.
  • Graph disentanglement [3,4] aims to factorize the semantic factors using graph convolutional networks, e.g. distinguishing the neighborhood type (work, family, hobby) in a social network.

Formalization of disentanglement

In this work we provide a notion of symmetry-induced disentanglement for graphs, following the formalization from Higgins et al. 2018 [5], leveraging tools from Lie algebras.

Lie group and Lie Algebra
A Lie group \(G\) is a group of symmetries where the symmetries are continuous. A Lie group is associated with a Lie algebra \(\mathfrak{g}\), which is the tangent space to the identity element of \(G\). A Lie algebra is parameterized with a basis \(\{\mathbb{A}_i\}_{i=1}^k\), where every element in \(\mathfrak{g}\) can be written as \(\mathbb{A} = \mathbb{A}_1 t_1 + \ldots + \mathbb{A}_k t_k\), where \(t_i\) are the coordinates. The elements of the Lie algebra can be mapped into the Lie group with an exponential function: \(\exp: \mathfrak{g} \rightarrow G\).

Intuition
The latent representation \(\hat{Z}\) is disentangled with respect to the Lie group \(G\), if a change in the coordinate \(t_i\) is associated with a change in \(\hat{z}_i\).

Unconditional disentanglement
The graph embedding \(\hat{Z}\) obtained by \(\textstyle f(\hat{Z}|T)\) is unconditionally disentangled with respect to the Lie group coordinates \(\textstyle T=\{t_j\}_{j=1}^k\) if: (1) there is a group action on \(\textstyle \hat{Z}\): \(\cdot : G \times \hat{Z} \rightarrow \hat{Z}\), (2) the map \(\textstyle f = \exp \mathbb{A}(T): T \rightarrow \hat{Z}\) is equivariant between actions on \(T\) and \(\hat{Z}\), (3) there is a decomposition \(\textstyle \hat{Z} = \hat{z}_1 \times \hat{z}_2 \ldots \times \hat{z}_k\), where each coordinate \(t_i\) affects only the corresponding component \(\hat{z}_i\).

Conditional disentanglement
The graph embedding \(\textstyle \hat{Z}\) obtained by \(\textstyle f(\hat{Z}|T, Z)\) is conditionally disentangled with respect to the Lie group coordinates \(\textstyle T=\{t_j\}_{j=1}^k\) conditioned on \(Z\) if (1) there is a group action on \(\textstyle \hat{Z}\): \(\cdot : G \times \hat{Z} \rightarrow \hat{Z}\), (2) the map \(\textstyle f:\exp \mathbb{A} (T) \times Z \rightarrow \hat{Z}\) is equivariant between actions on \(T\) and \(\textstyle \hat{Z}\) for any fixed \(Z\), and (3) \(\textstyle \exp \mathbb{A} (T)\) factorizes into the product \(\textstyle \prod_{i=1}^k \exp (t_i \mathbb{A}_i) \times Z\) such that each component \(\textstyle \exp (t_i \mathbb{A}_i) \times Z\) is affected only by the corresponding coordinate \(\textstyle t_i\) for \(i \in \{1, 2, \ldots, k\}\).

Proposed Lie group VAE model

  • We design a generative model based on the Variational Autoencoder framework.
  • We provide conditional or unconditional versions of the ELBO.
  • The ELBO is implemented with 4 modules: 2 encoders, 2 decoders.

Probabilistic model

We consider the graph data \(\mathcal{G}=(X,\mathcal{A})\) and two latent variables \(Z\) (graph encoding) and \(T\) (group encoding).

Unconditional ELBO
\(\mathcal{L_{\text{unc}}} = \mathbb{E}_{q(Z|\mathcal{G}) q(T|Z)} \log p(\mathcal{G}|\hat{Z}) p(\hat{Z}|T) - \mathbb{E}_{q(Z|\mathcal{G})} \text{KL} (q(T|Z) || p(T)) - \mathbb{E}_{q(Z|\mathcal{G})} \log q(Z|\mathcal{G})\)

Conditional ELBO
\(\mathcal{L_{\text{cond}}} = \mathbb{E}_{q(Z|\mathcal{G}) q(T|Z)} \log p(\mathcal{G}|\hat{Z}, Z) p(\hat{Z}|T) - \mathbb{E}_{q(Z|\mathcal{G})} \text{KL} (q(T|Z) || p(T)) - \mathbb{E}_{q(Z|\mathcal{G})} \log q(Z|\mathcal{G})\)

Architecture

We provide a graph VAE with an additional Lie group autoencoder.

Graph encoder (\(E_{\text{graph}}\))
VAE reparameterization that maps to latent variable \(Z\)
\(q(Z | X, \mathcal{A}) = \mathcal{N}(Z | \mu, \text{diag}(\sigma^2))\) where \(\mu = \text{GCN}_{\mu}(\overline{X}, \mathcal{A})\), \(\log \sigma = \text{GCN}_{\sigma}(\overline{X}, \mathcal{A})\) and \(\overline{X}=\text{Pool}(\text{GCN}(X, \mathcal{A}))\)

Group encoder (\(E_{\text{group}}\))
VAE reparameterization that maps to coordinates \(\{T_i\}_{i=1}^k\):
\(q(T|Z, \mathcal{A}) = \prod\nolimits_{i=1}^k q(t_i | Z, \mathcal{A})\) with \(q(t_i | Z, \mathcal{A}) = \mathcal{N}(t_i | \mu_i, \text{diag}(\sigma_i^2))\)

Group decoder (\(D_{\text{group}}\))
For each coordinate \(t_i \in T\), we learn a Lie algebra basis element \(\{\mathbb{A}_i\}_{i=1}^k \in \mathfrak{g}\). The coordinates and basis are then aggregated and fed into an exponential mapping layer, to achieve a group representation.
\(p(\hat{Z} | T) = g(T) = \exp \mathbb{A}(T)\) where \(\mathbb{A}(T) = \sum\nolimits_{i=1}^k t_i \mathbb{A}_i \quad \text{for} \ g \in G, \ \mathbb{A} \in \mathfrak{g}\) The conditional case involves conditioning on \(Z\):
\(\mathbb{A}(T) = p(\hat{Z} | Z, T) = \sigma(\text{agg}[\exp \mathbb{A}(T), Z])\), where \(Z = \{z_j\}_{j=1}^k\)

Graph decoder (\(D_{\text{graph}}\))
We map the latent \(\hat{Z}\) into \(N\) latent variables using sigmoid functions \(\sigma_j\), and then aggregates the \(N\) latent variables to obtain the full graph representation. \(p(\mathbb{A}, X | \hat{Z}) = \prod\nolimits_{j=1}^N \big[ p(\mathbb{A}, X| \sigma_j (\hat{Z}) ) \big]\)

Discussion

Limitations

  • Currently quantitative evaluation is limited to synthetic graphs (same as in images)
  • Further discussions on identifiability and causality of disentanglement can be explored.

Conclusion

  • We formalize the notion of conditional disentanglement on graphs and propose a novel framework for graph disentanglement by leveraging tools from Lie algebras.
  • We design a graph VAE based on a Lie group parameterization, and provide a novel ELBO criteria for optimizing conditional disentanglement.
  • Experiments on quantitative disentanglement, few-shots classification and molecular generation provide evidence of the capabilities of our model.

References

  1. Van Steenkiste et al. Are disentangled representations helpful for abstract visual reasoning? NeurIPS 2019
  2. Bengio et al . Representation learning: A review and new perspectives. TPAMI 2013
  3. Ma et al, Disentangled graph convolutional networks, ICML 2019
  4. Yang et al. Factorizable graph convolutional networks NeurIPS 2020
  5. Higgins et al. Towards a definition of disentangled representations. Arxiv 2018