Attribute Graphs Underlying Molecular Generative Models: Path to Learning with Limited Data

TL;DR

This paper introduces a simple perturbation-based algorithm to uncover attribute graphs underlying molecular generative models, enabling robust property prediction under limited data and distribution shifts.

Contribution

It presents a novel method to infer attribute graphs from pre-trained autoencoders, revealing how latent variables influence multiple attributes and improving transferability in molecular property prediction.

Findings

The attribute graph effectively models relationships between latent codes and molecular attributes.

Predictors based on the derived Markov blanket are robust to distribution shifts.

The method outperforms existing causal discovery and feature selection techniques in limited data scenarios.

Abstract

Training generative models that capture rich semantics of the data and interpreting the latent representations encoded by such models are very important problems in un-/self-supervised learning. In this work, we provide a simple algorithm that relies on perturbation experiments on latent codes of a pre-trained generative autoencoder to uncover an attribute graph that is implied by the generative model. We perform perturbation experiments to check for influence of a given latent variable on a subset of attributes. Given this, we show that one can fit an effective graphical model that models a structural equation model between latent codes taken as exogenous variables and attributes taken as observed variables. One interesting aspect is that a single latent variable controls multiple overlapping subsets of attributes unlike conventional approaches that try to impose full independence. Using a pre-trained generative autoencoder trained on a large dataset of small molecules, we demonstrate that the graphical model between various molecular attributes and latent codes learned by our algorithm can be used to predict a specific property for molecules which are drawn from a different distribution. We compare prediction models trained on various feature subsets chosen by simple baselines, as well as existing causal discovery and sparse learning/feature selection methods, with the ones in the derived Markov blanket from our method. Results show empirically that the predictor that relies on our Markov blanket attributes is robust to distribution shifts when transferred or fine-tuned with a few samples from the new distribution, especially when training data is limited.