SAG-VAE: End-to-end Joint Inference of Data Representations and Feature Relations

The ability to capture relations within data can provide the much needed inductive bias for robust and interpretable Machine Learning algorithms. Variational Autoencoder (VAE) is a promising candidate for such purpose thanks to their power in data representation inference, but its vanilla form and common variations cannot process feature relations. In this paper, inspired by recent advances in relational learning with graph neural networks, we propose the Self-Attention Graph Variational AutoEncoder (SAG-VAE) model which can simultaneously learn feature relations and data representations in an end-to-end manner. The SAG-VAE is trained by jointly inferring the posterior distribution of two types of latent variables, which respectively represent the data and the feature relations. The feature relations are represented as a graph structure, and the presence of each edge is determined by a Gumbel-Softmax distribution. The generative model is accordingly parameterized by a graph neural network with a special attention mechanism we introduced in the paper. Therefore, the SAG-VAE model can generate via graph convolution and be trained via backpropagation. Experiments based on graphs show that SAG-VAE is capable of approximately retrieving edges and links between vertices based entirely on feature observations. Furthermore, experiments on image data illustrate that the learned feature relations can provide the SAG-VAE robustness against perturbations in image reconstruction and sampling. The learned feature relations as graph adjacency matrices are observed to be structured, which provides intuitive interpretability of the models.