Attribute Graphs Underlying Molecular Generative Models: Path to Learning with Limited Data
Samuel C. Hoffman, Payel Das, Karthikeyan Shanmugam, Kahini Wadhawan, Prasanna Sattigeri
TL;DR
This work tackles robust molecular property prediction under distribution shift by learning a latent-to-attribute graphical model from a pre-trained generative autoencoder. It introduces PerturbLearn, which perturbs latent dimensions to infer a sparse, DAG-structured attribute graph and identifies a Markov blanket of attributes that best predict a target property with limited data. The approach demonstrates improved out-of-distribution generalization and interpretability by leveraging latent-mediated relations rather than retraining the generator, using pharmacokinetic datasets and RDKit descriptors. Overall, the method offers a data-efficient pathway to domain adaptation in molecular design with partial explainability via the inferred latent–attribute structure.
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.