Attribute Graphs Underlying Molecular Generative Models: Path to Learning with Limited Data
Authors: Samuel C Hoffman
Payel Das
Karthikeyan Shanmugam
Kahini Wadhawan
Prasanna Sattigeri
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (08/2024)
Attribute Graphs Underlying Molecular Generative Models: Path to Learning with Limited Data
Samuel C. Hoffman IBM Research
sho shoffman@ibm.ccomm
Payel Das IBM Research
daspa
Karthikeyan Shanmugam
korthikkeyamsss@google.com
IBM Research
Kahini Wadhawan IBM Research
kahini, wadhawan @ibm. com
Prasanna Sattigeri
psattig us. ibm. com
IBM Research
Reviewed on OpenReview: https: /ttps://openreee openrviiew......... net/ forum? id= APON4bs
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 simple algorithm that relies on perturbation experiments on latent codes of pre-trained generative autoencoder to uncover 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 pre-trained generative autocncoder trained demonstrate that the graphical model between various molecular attributes and latent codes learned by our algorithm can be used to predict specific property for molecules which are drawn from a different distribution. We compare prediction models trained on varions feature subsets chosen by simple baselines, as well as existing causal discovery and sparse learning /fetuur 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 fine-tuned with few samples from the new distribution, especially when training data is limited.
large datasct of small molecules,
1 Introduction
While deep learning models demonstrate impressive performance in different prediction tasks, they often leverage correlations to learn model of the data. As result, accounting for the spurious among those correlations present in the data, which correspond to biases, weaken the robustness of predictive model on unseen domains with distributional shifts. Such failure can lead to scrious consequences when deploying machine learning models in the real world.
"currently Google Research India

---

Published in Transactions on Machine Learning Research (08/2024)
One way obtain better generalization across domains is by learning causal relations in data, the core idea being to decipher the causal mechanism that is invariant across domains of interest. Structural causal models have provided a principled framework for inferring causality from the data alone which yields a better way to
understand the data as well as the mechanisms underlying the generation of the data (Pearl 2009 Shimizu
et al.
20060. In this framework, causal relationships among variables are represented with a Directed Acyclic (DAG).
Graph
While learning structured causal models from observational data alone is possible for tabular data (Chick- ering
2020 Kalisch & Bühlman 2007 Tsamardinos et al 2006), in many applications involving complex modalities, such as images, text, and biological sequencing data, it may not be reasonable to directly learn a causal graph on the observed covariates. For example, in the case of image data, the observed pixels may hide the structure of the causal generative factors (Chalupka et al., 2014). In those cases, it may suffice to learn a structural model of the data to serve the purpose of domain generalization, which may not capture the true causal factors. Further, in datasets with complex modalitics, limited knowledge about metadata (additional attributes) of the data samples may be readily available that can guide construction of such structural model.
For real-world applications with complex modalities it is critical to learn a model of the high-dimensionss dataset by using the limited labeled metadata available. This also helps in deriving a model of the data that provides better interpretability. We focus on extracting a graphical model between the latent codes of generative model and attributes present in the metadata (even if limited) by leveraging a generative model trained on vast amounts of unsupervised data. Notably, since any accompanying metadata never exhaustively lists all underlying generating factors of data point (e.g., image), our procedure allows for unobserved latents between attributes in the metadata.
In this work, we propose simple alternative method to extract graph structure relating different data attributes to latent representations learned by a given generative model trained on observational data. The generative model is pre-trained on large amount of unlabeled data. The domain-specific graph is then extracted using smaller set of data with easily obtained labels. We then estimate the sensitivity of each of the data attributes by perturbing one latent dimension at a time and use that sensitivity information to construct graphical model of attributes. The model's generalizability is finally tested by predicting altribule on samples from a larget distribution different from the source training data. Additionally, we are interested in scenarios where there is only limited data available from the target/tes distribution. Therefore, we intend to learn the most data-efficient and accurate predictor, given a pre-trained generative model, to solve the downstream task.
To showcase the usefulness of the proposod framework, we focus on predicting a set of pharmacokinetic properties of organic small molecules. Specifically, the out-of-distribution (OOD) small molecules differ in term of chemical scaffolds (core components).
Below, we list our contributions:
1. We propose a novel method to learn the underlying graph structure between latent representations of a pre-trained generative model and attributes accompanying the dataset allowing for confounding between altributes.
2. We use the structure of the learnt model to do feature selection for predicting propertics of out-of- distribution samples in limited data setting.
3. We show that the graphical model Icarned by our method helps to achieve better generalization and is robust to distribution shifts.
4. The proposed method offers a means to derive an optimal representation from a pro-trained gen- crative model, which enables samplo-efficient domain adaptation, while providing (partial) inter- pretability.
2

---

Published in Transactions on Machine Learning Research (08/2024)
2 Background and related work
2.1 Disentangled/Invan /Invariant generative models
Several prior works (Higgins et al. 2017 Makhzani et al. 2015 Kumar et al. 2017 Kim & Mnih 2018
have considered learning generative models with a disentangled latent space. Many of them define disentan- glement as the independence of the latent representation, i.e., changing one dimension/ Limension/factor ffacoor of the latent representation does not affect others. This involves additional constraints the latent distribution during training opposed to our method.
More generally, disentanglement requires that a primitive transformation in the observed data space (such as translation or rotation of images) results in a sparse change in the latent representation. In practice, the sparse changes can need not be atomic, i.e., a change in the observed space can lead to small but correlated latent factors. Additionally, it has been shown that it maybe impossible to achieve disentanglement without proper inductive biases or form of supervision (Locatello et. al. 2019 In Träuble al. (2021) the authors use small number of ground truth latent factor labels and impose desired correlational structure to the latent space. Similarly, we extract a graphical model between latent representations of the generative model and the attributes that are available in the form of metadata.
Discovering the causal structure the latent factors is an even more challenging but worthwhile pursuit compared to discovering only the correlated latent factors as allows us to ask causal questions such as ones arising from interventions or counterfactuals. Causal GAN (Kocaoglu et al. 2017 assumes that the the causal graph of the generative factors is known priori and learns the functional relations by learning good generative model. CausalVAE (Yang et al.) 2021 shows that under certain assumptions, a linear structural causal model (SCM) on the latent space can be identified while training the generative model. In contrast, here, we learn the graphical structure hidden in the pre-trained generative model in a post-hoc manner which docs not affect the training of the gencrative model. Further, while the dcrived model shows domain generalization, indicating learning of invariant structure across distributions, does not focus on capturing true causality, Therefore, our work is different from Besserve al. (2019). as our goal is learn graphical model in the latent space instead of the causal structure hidden in the decoder of the generative model. Leeb et al. (2023) proposes a latent response framework, where the intervention in the latent variables reveals causal structure of the learned generative process well as the relations between latent variables. The present work is distinct, as it aims to capture the response in the meta attribute space upon intervention on a latent dimension.
2.2 Invariant representation learning
Among conceptually related works, Arjovsky et al. (219 focuses on learning invariant correlations across multiple training distributions to explore causal structures underlying the data and allow OOD gencraliza- tion. To achieve this, they propose to find a representation & such that the optimal classifier given e is invariant across training environments. We leverage the fact that the pre-training encompasses data from wide set of subdomains however our method does not require multiple explicit training environments. Dubois et al. (2020) propose a decodable information bottleneck objective to find an optimal set of predic- tors that are maximally informative of the label of interest, while containing minimal additional information about the datasct under consideration to avoid overfitting. Our work instcad aims to derive an interpretable representation of the data using fixed altributes and guided by a learned graphical model.
2.3 Generative autoencoders for molecular representation learning
In recent years, generative autocncoders have emerged as popular and successful approach for modeling both small and macromolecules (e.g., peptides) (Gómez Bombarelli et al. 2018 Das el al.) 2021 hamarak- shan et al.) 2020: Hoffman et al.) 2021). Often, those generative models are coupled with search or sampling methods to enforce design of molccules with certain attributes. Inspired by the advances in text generation
(uu et al.) 2017) and the widely used text annotations of molecules, many of those frameworks imposed structure in the latent space by semi-supervised training or discriminator learning with metadata /labels. A
3

---

Published in Transactions on Machine Learning Research (08/2024)
number of studies also have provided wet lab testing results of the machine-designed molecules derived from those deep generalive foundation models, confirming the validity of the proposed designs (Nagarajan et al.
2017 Das et al. 2021 Shin et al. 2021 Chenthamarakshan et al. 2023). This demonstrates the effectiveness of generative models for molecular representations and implics that they have learned important insights about the domain that may be extracted.
2.4
Interpretability in molecular learning
While machine learning models, including deep generative models, have been successful in deriving an in- formative representation of different classes of molecules, it remains non-trivial and largely unxxplredd to infer the relationship between different, physicochemical and functional attributes of the data. Though laws of chemistry and physics offer broad knowiedge on that relationship, those are not enough to establish the (causal) mechanisms active in the system. For that reason, most experimental studies deal with partially known causal relationships, while confounding and observational bias factors (e.g., different experimental conditions such as temperature, assays, solution buffer) are abundant. In this direction, recently Zaatdinov et al. (2020) have established the pairwise causal directions between a set of data descriptors of the micro- scopie images of molecular surfaces. Interestingly, they found that the causal relationships are consistent across range of molecular composition.
To our knowledge, this is the first work that infers graphical model between data attributes from the latents of a generative model of molecules and shows the generalizability of the inferred model across different data distributions.
3
Algorithm
3.1
Problem setup
Consider a generative model consisting of a decoder Dec that takes input latent code E gax1 and generates data point Dec(z) x and an encoder Ene that takes a data point and embeds it in the latent space Enc(z). We assume that it has been pre-trained using some training method (c.g., Kingma & Welling 2014 Makhzani et al.] (2015)) on the data distribution PP())) sampled from the domain X. We further have access to a vector of attributes a(x) E g|4/x1 as metadata along with datapoints x. In most cases, these metadata attributes can be directly computed on I, i.c., az(z) x + R. For those attributes we cannot directly compute on æ, we assume that we also have access to an attribute estimator trained on the data from Pa. dRdx1 + producing a regression estimate for attribute a,(x) = Pa, (Enc(x)))
The key problem we would like to solve is to find the structure of the graphical model implied by the generative model and attribute classifier combination, where latent codes z act as the exogenous variables while axx) act as the observed variables. In other words, we hypothesize structural equation model as follows:
axxx) S((apa() (x), =k.); Vi E [1:
(1)
Here, apa()) is a subset of the attributes a(x) = (1(x) ... aja (x) that form the parents of (zx)). Also, note that the 'exogenous' variables (i.e. 2k; E (2I...2g) are actually the latent representations/cove
at the input of the generator. We call 2k, the latent variable associated with attribute a+(-). Together, the map
from the space of latents z to a(x)) can be called the probabilistic mechanism a(x) = (() as represented by the graphical model given by
We can define DAG G(A, E) where (i.5) E E if Pa(j) in the structural equation model. Because it is dependent on the latent representation z, we call this DAG the "structure of the graphical model implied by the generative model," " or simply, the attribute graph. This means it may actually improve as state-of- the-art generative models improve. Our principal aim is to learn an attribute graph that consistent with perturbation experiments on the latent codes 2, The attribute graph should thus be viewed as a summary of how changes to the learnt latent components translate to changes in the predicted/estimant attributes.
4

---

Published in Transactions on Machine Learning Research (08/2024)
W
22
23
-0.20
2, (-
x~X
Enc
Dec
X
2
0.62
-0.40
aj (x), if possible
Qg
0.33
0.22
Pa (Enc else
Aa x)
=a(x) a(x)
a4
a(x)
ao
aj
AA
0.1
(a)
(b)
21
f (az(x'), ((,(('),,,xx)) (x)))
a2
x'"X'
Enc
2x(x))
Z2
a3
a; dx(x', if possible
Pay (Enc(x')), else
23
(c)
(d)
Figure 1: PerturbLearn overview illustrating the: (a) perturbation procedure involving pre-trained generative model Enc(-), Dec(-) and attributes a(x) resulting Aa(x,20 Ø), (b) weight matrix heatmap derivation from perturbations (sparsity threshold of interest and blue indicates the Markov blanket), (d) application of function f learned on Markov blanket features to OOD samples from X".
0.1), (c) DAG building from sparse weights (red indicates the attribute
assuming these changes are mostly mediated through influences among the attributes, rather than through a direct influence from latents on altributes.
We then investigate how the attribute graph implied by the generative autoencoder can help train predictor for given attribute from other attributes, which generalize OOD data. Supposc want to predict the property a,(x), then, using the DAG, take the Markov blanket of a,, i.e., the set of attributes MB(a;) C. A JUZ Z which makes it independent of the rest of the variables. This consists of all parents, children, and co-parents in a DAG. Since this subset contains only the features relevant to predicting da we should end up with a robust model of minimal size. Furthermore, this feature selection is important for explainability since justified by Bayesian reasoning. We assume the structure the graph is valid across both domains only the edge weights might change which we can easily fine-tune with minimal samples necessary from the new domain. We visualize PerturbLearn, solution to this problem, in Figure
3.2 PerturbLearn algorithm overview
Our key idea is the following obscrvation, we take 2μ. associated with attribute a,xx), perturbing εg. would actually affect a, and all its descendants the attribute graph. So, we first obtain the sparse perturbation map between each latent 2j and the subset of attributes A₃ C A that it influences upon
5

---

Published in Transactions on Machine Learning Research (08/2024)
perturbation. Then, apply peeling algorithm that would actually find the attribute that is associated with a specific latent. In other words, we would find that attribute that occurs last in the ancestral order that is influenced by that latent variable. We use this algorithm to find a DAG with minimal edges that is consistent with the attribute sets for every latent variable. These steps are described in more detail below.
3.2.1 Perturbation procedure
Given the pre-trained generative model Ene(-), Dec(.) and a set of property predictors a(-), we:
Encode the sequence into the latent code: a through the encoder Enc(-).
2. Choose any single dimension j in the latent space and change 2j to âj that is uniformly sampled in [-b, bj (approximate range of the latent variables when encoded from the data points),
3. Obtain ā = Dec(E),
Estimate the value of all attributes a()). Note: any attribute which relies on a predictor Pa, will use the matching Enc(ä) Enc(Dec())). The attribute space is effectively discrete only makes sense to discuss attributes of sequences (or encodings thereof) not every perturbed 2 will refer to new sequence (and therefore, attributes).
5. Obtain the net influence vector Aa(z, = a(x) a(t). Influence values are then standardized (zero mean, unit variance) to scale the influences since attributes may have vastly different ranges. This in turn leads to weights which can be compared against each other.
For the purposes of this work, assume linear relationship between latents and attributes although the strength need not be precise, explained later. Therefore, we then learn linear model (OLS) to predict the change in attributes Aa(x, i) from the latent perturbations 2j 2 for all data points æ. This is repeated for every latent dimonsion j. We obtain a woight matrix W E [|| relating attributes in the rows to latent variables in the columns. If we pick the elements whose absolute weights are above a specific threshold 8, we obtain a sparse matrix. In theory, any sparse linear coefficient learning method would work here (e.g., LASSO [Tibshirani] I996)) however, in practice, due to the number of samples necessary, the computational load is intractable without the ability parallelize the computation per latent dimension as in OLS. At this point, the values are simply binary - influenced not. This would represent the attribute subsets each latent dimension influences. This is summarized in Algorithm I Now, we need to associate a latent to one or more attributes such that all attributes influenced by this latent appear later in the ancestral order.
3.2.2 Building the DAG
PerturbLearn constructs the attribute graph iteratively starting from sink nodes (Algorithm 2 Suppose one could find a latent variable that influences only one attribute and suppose that the graph is a DAG, then that attribute should have no children. Therefore that node added to the graph (and the corresponding latent is associated with it.) (Line 10 in Algorithm 2 and removed from the weight matrix (Lines 17 Algorithm 2. Sometimes, during this recursion, there may be no latent variable that may be found to affect single remaining attribute. We find subset with smallest influence (Lines 3 and add a confounding arrow between those attributes (Lines II 13 and add it to the graph. Now, if this latent affects any other downstream nodes previously removed, then we draw an edge from the current set of attributes to them In other words, we assume influences can be explained via mediation through altributes rather than direct influence from the latents. After the recursive procedure is performed, we obtain the transitive roduction (Aho et af. 1972) of the resulting DAG (Line 20 It the minimal unique edge subgraph that preserves all ancestral relations. This would be the minimal graphical model that would still preserve all latent-attribute influence relationships. See Figure g for visualization of this process on toy example.
(Lines 14 15

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm PerturbLearn
Perturbations to influence weights
Require: data X, attribute estimators a, sample size number of perturbations sampling bounds b
sparsity 0 < <s$1
1: AZ,4A=11
b empty arrays
2: for j E 1....Z. do
for 2; sample(X,n) do
2 2 Enc((;)
loop p times
2 (((bb,) b)
7:
i = Dec(2)
8:
append As = = to AZ
D net latent vector net influence vector
9: append Aa = a(x) - aā)) to AA
10: end loop
11: end for
12: end for
13: AA slandardize(AA)
center columns to zero mean and scale to unit variance b least squares weights predicting Aa from Azj b sparsify weights
14: W = =OLS(ZZ, AA)
15: W(abs(Wi.)) _ 0
16: return W
Algorithm 2 PerturbLearn
Sparse weights to attribute graph
Require: attributes A, latent features Z, weights W E M|x
= {2, : {a; Ei E ||] if Wa, # #0 Vj E [2]}
D influence sets for each latent feature store original influence sets for reference
Lorig =
G =mmtty directed graph
E =-[[]
empty list to store confounded pairs of attributes
5: while I is not empty do
6:
N -=[]>
Vzj E
D count number of attributes influenced by each latent fealure
7:
S {2y 1[,y
Vzz E arg Ni
select subsets of minimal influence
for 2, E do
P= S[zp]
parents
10:
add node(a;) to G Va: P if P > 1 then
11:
nodes are confounded save all pairs of nodes in P for later
12:
append permutations(P) (P,2) to E end if
13:
Torig [2p] T[2p] add edge(i, G WEE.EEeC 16: end for
14:
> children b add edges from parents to children
15:
17: drop 2p from 1 Vzp ES
remove latents from further steps remove attributes from further steps
18: drop ag from 1221 Vzp SS S(zp]: Vz E
19: end while
20: G =transttive ==trasstive_reduucioonn
21: add edge(i,j) Vi,j e E
> add confounded edges
22: return G
4 Data and models
4.1
Datasets
We conduct experiments on nine pharmacokinetic property prediction problems, an arca of science where data limitations are prevalent and label acquisition is costly and time-consuming, Acquiring additional data often involves performing physical synthesis and testing of the molecule detailed computer simulations of

---

Published in Transactions on Machine Learning Research (08/2024)
21 [An, Ag, As
21 (A), Az. Ast
21
(A
22 (A,, Ar
Ru { {42. Aot
2g {Aa}
21
A1
L
Stage 1: Find latent of min influence (Here, attribute). Remove the latent and the attribute.
Stage 2: Find of minimum
Stage 3: latent has previously removed attribute descendants, add edges from attributes to them.
influence (Here, attributes), Remove the latent and the attributes.
Underlying Graph
Figure Musttatinn various stages of PerturbLearn (Algorithm
using toy underlying ground truth Algorithm Stage corresponds to
altribute model. Stages and 2 correspond to lines 10 and
lines 14
15
Table 1: Datasets used in experiments along with total size, evaluation metric, and original reference. All datasets were loaded and split according to scaffolds using the Therapentics Data Commons API.
Name
Size
Metric
Split
Property
Reference
FreeSolv
600
MAE
Scaffold
Absorption
Mobley & Guthric
2014
Half-life
667
Spearman
Scaffold
Excretion
Obach
2008
Caco-2
906
MAE
Scaffold
Absorption
Wang et al.
2016)
Hepatocyte
1020
Spearman
Scaffold
Excretion
Wenlock & Tomkinson
201 5
Microsome
1102
Spearman
Scaffold
Excretion
Wenlock & Tomkinson
(0015
VDss
1130
Spearman
Scaffold
Distribution
Lombardo & Jing (2016)
PPBR
1797
MAE
Scaffold
Distribution
Wenlock Tomkinson
2015
Lipophilicity
4200
MAE
Scaffold
Absorption
Wenlock Tomkinson
2015
Solubility 9982
MAE
Scaffold
Absorption
Sorkun al. 20199
the system, which demand time, computational resources, and domain experts. Using existing data more effectively from similar domains is one way to lower barrier costs to solving problems in these areas.
Therapeutics Data Commons (TDC) is a platform for Al-powered drug discovery which contains number of datasets with prediction tasks for drug-relevant properties (Huang et al., 2021 We use all of the regres- sion tasks in the pharmacokinetics domain i.e., drug absorption, distribution, metabolism, and excretion (ADME). These summarized in Table 1 All datasets were divided into training, validation, and testing splits according to ratio of 70%, 10%, and 20%, respectively, based scaffold (the core structure of the molecule) in order to separate structurally distinct compounds. Because the structures differ greatly, unlike in a random split, scaffold splitting approximates realistic distribution shift and is widely used to measure robust generalization in ADME prediction (Huang et al.) 2021). All datasets were also filtered to include only organic molecules (atoms (B, C, N, O, F, P. S, CI, Br, т}) with additional salt ions removed and pre-processed to use a canonical form of SMILES without isomeric information.
4.2 Data attributes
We use a of molecular descriptors from RDKit as features for the experiments. full list of the 207 descriptors can be found in the Table I2 (note: Ipc was not used due numerical instabilities in the implementation).
8

---

Published in Transactions on Machine Learning Research (08/2024)
4.3 Generative autoencoder details
For small molecule representation, we use the VAE from Chenthamaralshan et al. (020). This model was trained primarily on the MOSES dataset (Polykovskiy et al. 2020) which is subset of the ZINC Clean Leads dataset (Trwin et al., 2012). These are 1.6M molecules which are considered lead-like for drug development with benign functionality. The encoder uses a gated recurrent unit (GRU) with a linear output layer while the decoder is a 3-yyyer GRU with dropout. The encoded latent distribution is constrained to be close to a diagonal Gaussian Distribution, i.c., qa(2|2)) = [(2; ((xx), E(z)) where E(x) = (angfcxp(llg()))()))) The model was first trained unsupervised, followed by training jointly with set of attribule regressors to encourage disentanglement in the learnt space. We use the original model as-is, without any fine-tuning
5 Methods
In order to validate the attribute graphs learned on the latent features of the generative autoencoder, we devise a series of experiments for learning with limited data. If the graph is valid, the Markov blanket of a given node should include all the features needed to predict that attribute and only those features. We hypothesize that this will lead to improved generalization and robustness when learning on only a few data points. We assume that the structure of the attribute graph will not change when the domain shifts, although the relative strength of the links may change. However, these functions should be easy to learn given the minimal set of independent variables and therefore only a small datasct is needed.
To this end, we consider a scenario in which we wish to train a regressor to predict a certain value of interest in a niche domain. There is limited data available for training in this domain, however, data from similar domains are more plentiful. This a realistic scenario for many real-world applications where the cost of acquiring more data is high and throughput is low due to the involvement of physical experimentation, slow virtual simulation, or laborous user feedback. Therefore, we opt to leverage the larger dataset for pre-training training the regressor and the smaller dataset for fine-tuning. With our proposed method, also use the larger dataset to learn an attribute graph from which we can extract the Markov blanket features.
For every experiment in this work, we take 10 perturbations of 2500 samples (sampled with replacement from the training data) for each dimension in the latent space. We also train a target attribute estimator using the latent representations of the data which is used to estimate the influence of perturbed samples. We then apply PerturbLearn using these inputs along with the pre-trained generative autoencoder to obtain our attribute graph modcling the original domain.
We test our method against five baseline sets of features. Each baseline uses the same model described in the next paragraph. The first baseline uses the full latent vector, 2, as input features to the predictor model. The second baseline uses all available attributes (i.c., all possible nodes in the attribute graph). The third baseline is simply a concalenation of the first two (all autributes +-).. We also compare our model trained on the features from the Markov blanket of the attribute of interest to the blankets from causal graphs derived using Greedy Equivalence Search (GES) (Chickering] 2002) and LINGAM (Shimizu et al. 2006 as well as the feature sets chosen by Sparse Group Lasso (Simon et al. 2013 and Predictive Permutation Feature Selection (PPFS) (Hassan et al. 2021) learned on the training data split. GES and LiNGAM are causal discovery methods which learn based on observational (labeled) data. Group Lasso is a form of structured sparsity regularization which aims to learn a reduced set of imput features used for prediction while PPFS is a Markov blanket- feature-selection algorithm. Finally, we demonstrate that these methods outperform a naive bascline which always predicts the training mcan (dummy).
For initial (base model) training, we use a multilayer perceptron (MLP) model architecture for each feature set for comparison. The MLP hyperparameters are tuned using 5-fold cross-validation on the training set where the scarch space is grid of combinations of: hidden layer size 64, 128, or 512 (2 hidden layers chosen independently); dropout rale 0.25 or 0.5; and training duration 100 or 500 epochs. All models use a mean squared error (MSE) loss with a batch size of 256 (or the size of the dataset, if smaller), rectified linear unit (RcLU) activations, and Adam optimization with a learning rate of 0.001. Data is also scalod to zero mean and unit variance independently for each feature. For the Half-life and VDss datasets, we scale the targets in the training set by taking the natural logarithm before learning the outputs are exponentiated
9

---

Published in Transactions on Machine Learning Research (08/2024)
Table 2: Mean absolute error (MAE # o; lower is better) of different predictors (with input dimension "size") on the FreeSolv benchmark with varying numbers of test samples included in fine-tuning (n). Best results in bold; second-best is underlined.
size
n = 0
11 = 7
11 =
10
11 = 25
7 = 50
n 100
dummy 2
0
2.9210.00
2.7510.03
20.0
2.65=0.01
2.6210.01
2.6210.02
128
10.04
1.70
1.50
1.22
1.10.1..05
0.99
all attributes
208
1. 0.22
1.49
1.29
1.04=0.05
0.92+0.04
0.80
all 2
336
1.36
1.60
1.36=0.01
1.09=0.03
0.96=t0.01
0.82=+0001
GES blanket
154
1.13
1.50
1.307
1.06
0.93 10.04
0.84.10.03
LINGAM blanket
174
1.16+:.05
1.48
1..04
1. 1.04=0.08
0.
0.81+0.03
Group Lasso blanket
103
1.1540.04
1.3640.03
1.18
1.04
0.83
0.73
PPFS blanket
121
1.04+0.03
1.39
1.19+0.03
0.95=0.02
0.85+0.02
0. 0.76+
PL blanket
174
0.99
1.33+0.07
1.1644.06
0.95=0.04
0.85++0.03
0.75+0.02
before measuring performance metrics. Note, the base model for the z-baseline is effectively equivalent to the target attribute estimator used in the PerturbLearn step to make predictions for generated samples.
Fine-tuning on the second domain uses the MLP base model as a feature extractor by freezing the weights and using the outputs from the last hidden layer. These weights are then fed into Gaussian Process (GP) regressor with a kernel consisting of sum of two radial basis function (RBF) kernels and a white noise kernel. We optimize the kernel parameters with the Broyden Broyden-Fletccher Goldfarb Shanno (L-BFGS-B) method with 50 restarts. In order to get robust estimate of the performance of the model, we each fine-tuning experiment 10 times on randomly drawn subsets and take the mean. We also repeat the entire procedure 8 times (retraining the base model) to obtain the mean and standard deviation values in Tables 2 11
For all the experiments, we utilize the provided scaffold splits from TDC. For each the tasks, we perform graph learning with PerturbLearn and initial regressor training using the "train" split and report the fine- tuned results on the "test" set. We treat the "validation" and "test" sets as effectively two independent shifted domains. In all cases samples to set sizes.
the full split when training (and graph learning) but restrict the testing
When applying our method, PerturbLearn, after learning the linear weights from the perturbation data, we must choose a sparsity threshold before converting the data to a DAG. We tune this hyperparameter by choosing the best performing threshold on the validation set for the smallest fine-tuning size, n. The sparsity threshold is chosen from the set, (0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.25].
For all experiments, we use machines with Intel Xeon Gold 6258R CPUs and NVIDIA Tesla V100 GPUs with up to 32 GB of R.A..
6 Results
Table shows the results of experiments using the FreeSolv dataset. Each row shows mean absolute error (MAF) results (mean standard deviation (o)) averaged over runs of fine-tuning the MLP base model (learned from the "train" split) with GP regressor using either the PerturbLcarn Markov blanket (also learned from the "train" split) or various baselines while each column shows different number of fine-tuning samples, n, from the "test" split. The fine-tuning, on the "test" split. The best performing model on the "validation" split with n 7 is used to choose PerturbLearn hyperparameters..
column shows the performance of the base model, without
The performance of our method ("PL blanket") is best for small fine-tuning sizes (n = 0, 7, 10) and second- best for larger ones (n 25, 50, 100) for the FroeSolv dataset. The performance of our method on the Half-life dataset (Table g shows improvement over all baselines other than PPFS all scenarios. However, the PPFS algorithm involves aggregating à Markov blankets independently derived from separate folds of the
10

---

Published in Transactions on Machine Learning Research (08/2024)
Table 3: Spearman correlation (p = o; higher is better) of different predictors (with imput dimension "size") on the Half-life benchmark with varying numbers of test samples included in fine-tuning (n). Best results in bold; second-best underlined. Note: since the dummy regressor output is constant, the Spearman correlation is undefined for that row.
size
n 0
n 7
71
10
71
25
n1 50
n 100
dummy
0
2
128
0.27+0.04
0.05++..2
0.08=0.02
0.13=0.02
0.21 #0.03
0.28++.04
all attributes
208
0.49
0.17#0.02
0.21
0.32=0.02
0.38
0.42:+0.03
all attributcs+2
336
1.402
0.15
0.1710.03
0.28 0..04
0.04
0.3710..04
GES blanket
11
0.30
0.04+:0.02
0.05
0.07=0.02
0.11+0.08
0.22:10.07
LiNGAM blanket
194
0.49
0.17+0.01
0.20#0.02
0.31 =0.02
0.37
0.41
Group Lasso blanket
178
0.51
0.16+0.02
0.22=0.01
0.33=0.05
0.40=#.001
0.44+0.03
PPFS blanket
58
0.5611.02
0.21 10.02
0.26110.02
0.40
0.47 10.02
0.49 10.02
PL blanket
109
0.53+0.02
0.20+0.03
0.24+
0.35=0.03
0.43+v.08
0.47#0.03
PL blanket (FS)
174
0.48
0.15
0.19
0.29
0.39
0.4610.02
Table 4: Fine-tuning results for TDC datasets = 25. FreeSolv, Caco-2, and PPBR use MAE (1 a) while Half-life, VDss, and the Clearance datasets (Hepato. and Micro.) use Spearman correlation (1: e) as the performance metric. Best result in bold; next-best is underlined.
FreeSolv (4)
Half-life (t)
Caco-2 (1)
Hepato. (t)
Micro. (f)
VDas (1)
PPBR (4)
dummny
2.6510.01
0.58:10.00
11.03:1.00
1.22+0.05
0.13+0.02
0.45++0.01
0.1520.03
0.23+00.3
0.32+0.04
10.9840.1 13
all attributes
1.0440.00
0.32:+0.02
0.42440.01
0.22:10.2
0.47+00.22
0.60+0.02
10.2740.12
all attributes+2
1.09::.03
0.28:10.04
0.41 =0.01
0.21
0.47700.01
0.55400.22
10.26 14
GES blanket
1.06=0.05
0.07+0.02
0.42++001
0.20+0.02
0.45+0.01
0.60+
10. 10.88++0. 12
LiNGAM blanket
1.04+0.03
0.31
0.41=0.00
0.21+
0.49+0001
0.60+0001
10. 10.28+0011
Group Lasso blanket
0.94.10.03
0.33:10.05
0.45:10.01
0.1810.01
0.49
0..02
10.45:10.19
PPFS blanket
0.95
0.40
0.40:10.02
0.21
0.4510.02
0.6010.00
10.56:10.17
PL blanket
0.35:10.0z
0.41
0.19
0.47:0.01
0.5910.02
10.36:10.17
0.95
PL blanket (FS)
0.29:10.0r
0.41:0.01
0.2010.02
0.48+0.01
0.62:10.1
10.16:10.16
data to handle intrinsic sample inefficiency, whereas PertubLearn operates on a single MB, and therefore is less complex.
These are also the smallest of all the datasets used in our experiments, each under 700 total samples. For the rest of the datasets (Tables nearly indist inguishaie,, especially ly with few fine-tuning samples, however, in general, there is little difference between the performance of the feature sets with sufficient training split size.
our method is competitive with the best baselines, often outperforming
In some cases, the performance at n = 0 is better than when we include additional fine-tuning samples (n 0). In these cases, the models may benefit from a different method of fine-tuning which causes less forgetting than learning a new model on top of a pre-trained feature encoder, however, in the interest of consistency across datasets we keep the fine-tuning procedure the same for all experiments. Since the baseline methods also experience a proportional drop in performance in these cases, the performance of our method under the fine-tuning procedure presented here still follows the same trends noted above. Other methods of fine-tuning should be explored in future work.
Given the impressive performance applying this graph to other tasks. In this case, since the tasks are distinct but within the same family (ADME), we can think of this as a more extreme distribution shift problem. These results can be seen in the *PL blanket. (FS)" row of Tables
our attribute graph from the FreeSoly experiment, we can also try
11 In few cases, such as VDss and PPBR, this leads to a noticeable
11

---

Published in Transactions on Machine Learning Research (08/2024)
improvement, surpassing the best baselines. In the rest, there is little difference compared to the task-specific
PerturbLeearn blanket.
Table M shows fine-tuning results at n = 25 where we observe in every case, except Hepatocyte clearance, either the task-specific or transferred PL blanket performs best or second-best. We do not witness a significant decrease in performance margin until data size increases past this range (> 1800 total samples). For larger datasets, such as Lipophilicity and Solubility (Tables I0 most predictors achieve equivalent performance.
7 Discussion
In this study, we proposed a simple framework to extract graphical model relating different data attributes from the latents of a pre-trained generative model. The influence of a latent on attributes is used to construct the graph. In this way, our method leverages the unsupervised learning techniques which produce powerful models trained on vastly more data than is available and labeled for any individual task. This also means our method can be extended to use better models in the future, given they encode information in a continuous latent space.
We have provided here domain generalization method that utilizes latent-mediated influences among molec- ular altributes and have empirically shown its efficacy to predict molecular attributes by learning from limited nut-of-distribution data. We showcased the performance of the proposed method, along with other exist- ing feature selection methods, in predicting molecular propertics in a low-data regime. Thus, this study provides a comprehensive evaluation of feature selection methods for molecular learning tasks, as well as offers insights into the derived statistical relations between molecular attributes. Taken together, the results show the use of derived Bayesian information improves over the baseline feature sets, especially when data from the training and/or test domains is limited, demonstrating that extracting the dependency information underlying pre-trained generative model leads to more robust and generalizable models. These models also use only the meaningful and relevant properties as features meaning the models will be both smaller and more interpretable than their baseline counterparts.
One potential limitation of our framework is the dependency on a pre-determined set of data altributes, which may not be comprehensive (or may already involve experts applying implicit causal models Icarned from experience to choose attributes). Further investigation on the accordance of the derived attribute graph with domain priors, using only low-level and cheaply computable attributes, and/or with expert knowledge would be addressed in future.
References
Alfred V. Aho, Michael R Garey, and Jeffrey D. Ullman. The transitive reduction of a directed graph. SIAM Journal on Computing, 122):131 137, 1972.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXix:1907.022 02803, 2019.
Michel Besserve, Arash Mehrjou, Rémy Sun, and Bernhard Schölkopf. Counteraactuls uncover the modular structure of deep generative models, 2019.
Krzysatof Chalupka, Pietro Perona, and Frederick Eberhardt. Visual causal feature learning, arXiv preprinl arXiv:1412.288: 2309, 2014.
Vijil Chenthamarakshan, Payel Das, Samuel Hoffman, Hendrik Strobelt, Inkit Padhi, Kar Wai Lim, Benjamin Hoover, Matteo Manica, Jannis Born, Teodoro Laino, and Aleksandra Mojsilovic. Cogmol: Target-specific and selective drug design for covid-19 using deep generative models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 4320 4332. Curran Associates, Inc., 2020. URL https://proceeddngss wwweeeeerrippppe wwwrips.cc/pappe.e.. cc/paper files/ hptp://0001111221144444440006662
pdf
Vijil Chenthamarakshan, Samuel C. Hoffman, C. David Owen, Petra Lukacik, Claire Strain-Damerell, Daren Fearon, Tika R. Malla, Anthony Tumber, Christopher J. Schofield, Helen M.F. Duyvesteyn, Wanwisa
12

---

Published in Transactions on Machine Learning Research (08/2024)
Dejnirattisai, Loic Carrique, Thomas S. Walter, Gavin R. Screaton, Tetiana Matvinkk, Aleksandra Mo- jsilovic, Jason Crain, Martin A. Walsh, David I. Stuart, and Payel Das. Accelerating drug target in- hibitor discovery with a deep generative foundation model. Science Advances, (25) 2023. doi: 10.1126 10.1126/sciadv.dn wrrriadv.dd77880 URL https://www. /ttt::::/ww..cceenc.. wtwpeee..rr/////b///////// 1126/ s1ia/sciaad.. sciadv....vv. adg7865
David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 3(Nov):507-553 -5440 2002.
Max Chickering. Statistically efficient greedy equivalence search. In Conference Uncertainty in Artificial Intelligence, pp. 241-249) PMLR, 2020.
Payel Das, Tom Sercu, Kahini Wadhawan, Inkit Padhi, Sebastian Gehrmann, Flaviu Cipcigan, Vijil Chen- thamarakshan, Hendrik Strobelt, Cicero Dos Santos, Pin-Yu Chen, et al. Accelerated antimicrobial dis- covery via deep generative models and molecular dynamics simulations. Nature Biomedical Engineering, 5(6):613 623, 2021.
Yann Dubois, Douwe Kiela, David J Schwab, and Ramakrishna Vedantam. Learning optimal representations with the decodable information bottleneck. Advances in Neural Information Processing Systems, 33:18674 18690, 2020.
Rafael Gómez- Bombarll, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez- Lengelin,,,,,,,,,,,, Dennis Sheberla, Jorge Aguilera- Iparraguire, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using data-driven continuous representation of molecules. ACS central science, 4(2):268 276, 2018.
Atif Hassan, Jiaul H Paik, Swanand Khare, and Syed Asif Hassan. Ppfs: predictive permutation feature selection. arxiv preprint NT::211001071:: 2021.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthow Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. International Conference on Learning Representations, 2017.
Samuel C Hoffman, Vijil Chenthamarakshan, Kahini Wadhawan, Yu Chen, and Payel Das. Optimizing molecules using efficient queries from property evaluations. Nature Machine Intelligence, pp. 11, 2021.
Zhiting Hu, Zichao Yang, Xiaodan Liang, Ruslan Salakhutdinov, and Eric P Xing. Toward controlled 2017.
generation of text. In Internalional Conference on Machine Learning, pp.
Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. Proceedings of Neural Information Processing Systems, NeurIPS Datasets and Benchmorks, 2021.
John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman. Zine: free tool to discover chemistry for biology. Journal of chemical information and modeling, 52(7))1757-1178, 2012.
Markus Kalisch and Peter Bühlman. Estimating high-dimensional directed acyclic graphs with the pe- algorithm. Journal of Machine Learning Research, 8(3), 2007.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning, pp.
PMLR, 2018.
Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http://arxiv. wwtp://./...a....ooo///o/
Murat Kocaoglu, Christopher Snyder, Alexandros G Dimakis, and Sriram Vishwanath. Causalgan: Learning causal implicit generative models with adversarial training. arXiv: 1rXiu:709.00000 709.02023, 2017.
13

---

Published in Transactions on Machine Learning Research (08/2024)
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. arXie: 15:::71110000 100048,, 2017.
Felix Leeb, Stefan Bauer, Michel Besserve, and Bernhard Schölkopf. Exploring the latent space of autoen- coders with interventional assays, 2023.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervisod learning of disentangled represen- tations. In international conference on machine learning, pp. 4114 4124. PMLR, 2019.
Franco Lombardo and Yankang Jing. In silico prediction of volume of distribution in humans. extensive data set and the exploration of linear and nonlinear methods coupled with molecular interaction fields descriptors. Journal of chemical information and modeling, (2(10)):244222)) 2016.
Alireza Makhzani, Jonathon Shiens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoen- coders. arXiv preprint arXiv:1511.0664. 2015.
David I, Mobley and J Peter Guthrie. Freesolv: database of experimental and calculated hydration free energies, with input files. Journal of computer-aided molecular design, 28:711 720, 2014.
Deepesh Nagarajan, Tushar Nagarajan, Natasha Roy, Omkar Kulkarni, Sathyabaarathi Ravichandran, Mad- hulika Mishra, Dipshikha Chakravortty, and Nagasuna Chandra. Computational antimicrobial peptide design and evaluation against multitrug-resistan clinical isolates of bacteria. Journal of Biological Chem- istry, pp- jbc M117, 2017.
R Scott Obach, Franco Lombardo, and Nigel J Waters. Trend analysis of database of intravenous phar- macokinetic parameters in humans for 670 drug compounds. Drug Melabolism and Disposition, 36(7): 1385 1405, 2008.
Judea Pearl. Causality. Cambridge university press, 2009.
Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengeling Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Alekscy Artamonov, Vladimir Aladinskiy, Mark Veselov, Artur Kadurin, Simon Johansson, Hongming Chen, Sergey Nikolenko, Alán Aspuru-Guzik, and Alex Zha- voronkov. Molecular sets (moses): A benchmarking platform for molecular generation models. Fron- tiers in Pharmacology, 11. 2020. ISSN 1663-9812. 7/ww. frontiersin org/ http//rricleess//.. 3389/fphar 2020.
doi: 10.3389 /aa00//fha.22200 2020..5422..6 URL https:
Shohei Shimizu, Patrik O Hoyer, Aapo Hyvärinen, Antti Kerminen, and Michael Jordan. A linear non- gaussian acyclic model for causal discovery. Journal of Machine Learning Research, 7(10), 2006.
Jung-Eun Shin, Adam Riesselman, Aaron W Kollasch, Conor McMahon, Elana Simon, Chris Sander, Aashish Manglik, Andrew C Kruse, and Debora Marks. Protein design and variant prediction using autoregressive generative models. Nature communications, 12(1):1-11, 2021.
Noah Simon, Jerome Friedman, Trevor Hastie, and Robert Tibshirani. A sparse-group lasso. Journal of computational and graphical statistics, 22(2):231-24) -245, 2013.
Murat Cihan Sorkun, Abhishek Khetan, and Süleyman Er. Aqsoldb, curated reference set of aqueous solubility and 2d descriptors for a diverse set of compounds. Scientific data, 6(1):143, 2019.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 5861)2267-288 1996.
Frederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bern- hard Schölkopf, and Stefan Bauer. On disentangled representations learned from correlated data. In International Conference on Machine Learning, pp. 10401- 10412. PMLR, 2021.
Joannis Tsamardinos, Laura E Brown, and Constantin F Aliferis. The max-min hill-climbing bayesian network structure learning algorithm. Machine learning, 65(1):31-74 2006.
14

---

Published in Transactions on Machine Learning Research (08/2024)
Ning-Ning Wang, Jie Dong, Yin-Hua Deng, Min-Feng Zhu, Ming Wen, Zhi-Jiang Yao, Ai-Ping Lu, Jian- Bing Wang, and Dong-Sheng Cao. Adme properties evaluation in drug discovery: prediction of caco-2 cell permeability using a combination of nsga-ii and boosting. Journal of chemical information and modeling, 56(4):763-773, 2016.
Mark Wenlock and Nicholas Tomkinson. Experimental in vitro dmpk and physicochemical data set of publicly disclosed compounds. Technical report, AstraZeneca, February 2015. URL https://doi. org/ 10.
Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, and Jun Wang. Causalvae: Disentangled representation learning via neural structural causal models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 9593 9602, 2021.
Maxim Ziatdinov, Christopher T Nelson, Xiaohang Zhang, Rama K Vasudevan, Eugene Elisoev, Anna N Morozovska, Ichiro Takeuchi, and Sergei V Kalinin. Causal analysis of competing atomistic mechanisms in ferroelectric materials from high-resolution scanning transmission electron microscopy data. npj Com- putational Materials, 6(1):1 9, 2020.
15


Reviews:
[{"review_id": "Fp1FRTUt2N", "reviewer": "Reviewer_3F1X", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "---", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "*Context:* I reviewed an earlier version of this work (https://openreview.net/forum?id=Vyw437epFz), which was rejected with the suggestion for a major revision. Having re-read the revised paper, the below review is focused on the delta to the previous version. \n\nThe main changes made for the revision appear to be relatively minor, and focused on changing individual wordings to remove causal terminology, e.g.,\n- causal graph / causal model --> attribute graph / graphical model\n- causal mechanism ---> probabilistic mechanism\n- causal reasoning / information --> Bayesian reasoning / information\n- causal relations --> statistical relations\n\nWhile this is a step in the right direction, it arguably falls somewhat short of the requested \"adjustment of the framing\". Quoting from the previous AE decision:\n> In some sense, instead of asking \"how was this data generated?\", the paper asks \"what is this pre-trained generative model doing?\".\n\nThis is a spot-on summary, but something that---in my opinion---is still is not sufficiently clear in the revised version. \n\nI still get the impression that the learnt \"attribute graph\" is also meant/presented/interpreted as conveying something about true relations between attributes. However, one aspect that deserves more emphasis/attention is the fact that the learnt \"attribute graph\" also crucially depends on the pre-trained attribute estimators, and not just on the generative model (Enc & Dec), as the title might suggest. The attribute graph should thus be viewed as a summary of how changes to the *learnt* latent components translate to changes in the *predicted/estimated* attributes, assuming these changes are mostly mediated through influences among the attributes, rather than through a direct influence from latents on attributes.\n\nIf the main goal is to understand the interplay between the generative/latent variable model and the learnt attribute predictors, why not simply learn, say, a bipartite graph consisting only of arrows $z_i \\rightarrow a_j$? \n\nAnother issue that deserves further discussion is the reliability of the fixed, learnt components under perturbations to individual latents in light of out-of-support issues: even though the perturbed $\\tilde z_i$ are sampled such that they are marginally in-support, the resulting joint vector $\\bf \\tilde z$ may be out-of-support. As a result, both the Dec/Enc and the attribute predictors may be completely unreliable on the perturbed  $\\bf \\tilde z$ and $\\bf \\tilde x$. \n\nAll in all, it seems that many choices are somewhat arbitrary but presented in a way that could suggest that they follow in a principled way. However, the main evidence presented in favour of the proposed method is empirical. Thus, I would favour a framing along the lines of \"this is just one way of approaching this problem, and there is no reason why this should work, or should be the best way to go about this problem, but it turns out that this method is actually effective and yields good results on real data on a relevant problem\". \n\nI list some further questions below: \n_____\n\n- In the perturbation procedure (Alg. 2, which should be moved to the main paper), why do you opt for (unregularized) OLS + post-hoc thresholding of the weights, instead of a sparse regression approach such as LASSO/L1 regularization?\n- How is the standardization in Step 5 of Sec. 3.2.1 done? I could not find these details in Alg. 2. Why is standardization necessary / how does it impact the behaviour in combination with the subsequent thresholding?\n- In Steps 3 & 4 of Sec. 3.2.1, why are the attributes estimated after first decoding and then re-encoding the perturbed latent? That is, why apply the attribute estimator to $Enc(\\tilde x)=Enc(Dec(\\tilde z)$ rather than to $\\tilde z$ directly? If Enc inverts Dec, the two should be equivalent, but this is generally not the case even in-support, and certainly not out-of-support. \n- In Stage 3 of Figure 2, why add both edges from $A_2,A_3$ to $A_4$. Wouldn't one be enough? If sparsity is desired, adding the single edge $z_2 \\rightarrow A_4$ could also explain the dependence. Why not opt for this? (Related to my point about learning a bipartite graph)", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "----", "broader_impact_concerns": "---", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "RyfbwSGvLY", "reviewer": "Reviewer_BBhC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe model and ideas are clearly written, and the core idea makes intuitive sense.  The performance seems competitive across the board with a reasonable number of relevant comparisons.\n\nWeaknesses:\n\nThe fine-tuning performance is surprising, as the models often get worst across the board when fine-tuning.  The fine-tuning method is not described in sufficient detail to evaluate whether it is reasonable.  Please give a more precise mathematical and algorithmic description.  Furthermore, it is mentioned that the hyperparameters are chosen based upon the n=7 scenario with a validation set.  As such, it seems like performance should not appreciably decrease from the n=0 scenario, as it does in both examples given in the main paper and several in the appendix.\n\nIt appears a reasonable-sized validation split is used for model selection, which seems to contradict the message that this method is most appropriate when there are few labeled samples available for model training and selection.  Please describe how model selection would be done in practice in the limited sample case.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript proposes an approach to learn a graphical model on attributes by using pertained generative models (PerturbLearn), and then use that graphical model for feature selection in prediction tasks.  This idea can improve generalization in semi-supervised or transfer learning tasks by leveraging available generative modeling information to reduce the complexity of predictive models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Needed:\n\nUncertainty metrics in tables need to be defined.\n\nGive full details on fine-tuning procedure, and explain the decrease in performance that frequently occurs from fine-tuning.\n\nDescribe a model selection approach that could be used in practice on the motivating problems.", "broader_impact_concerns": "N/A,", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AJSGDdS3BS", "reviewer": "Reviewer_8DbN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n1.\tThe writing of this paper is much improved compared to the previous version.\n2.\tChanging the concept from causality to correlation to describe the attribute graph seems more reasonable.\n\n**Weaknesses:**\n\n1.\tThe figures still need some revisions to be improved.\n2.\tThe proposed method seems to work well for molecular attribute prediction tasks. However, the title seems quite broad and covers all generative models. To avoid misleading readers, it would be good to narrow down the scope to molecules. I am happy to discuss the title with the author and other reviewers.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an algorithm to use perturbation experiments on latent codes of a pre-trained generative model to discover a causal graph, which can be used for generalizable prediction on out-of-distribution testing. Specifically, a generative model is pre-trained on a large amount of unlabeled data, and then the domain-speciﬁc causal graph is extracted using a smaller set of data with easily obtained labels. After that, the authors do a sensitivity analysis of each of the data attributes by perturbing one latent dimension at a time and using that sensitivity information to construct the causal graph of attributes. The proposed method is evaluated on molecule generation. Results show that their model is more robust to distribution shifts when transferred or ﬁne-tuned with a few samples from the new distribution.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.\tFigure 1 needs to be revised. For example, the font size of texts should be larger considering there is enough space. The order of (b)(c)(d) should be changed. (d) misses a lot of information. How does the train and test process obtain a_2 and a_3? What is the task of this application?\n2.\tFor Figure 2, it would be better if the authors could annotate the corresponding lines of Algorithm 1 in each stage.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a resubmission of a previous TMLR submission (https://openreview.net/forum?id=Vyw437epFz). In the revised submission, the authors have carefully clarified their results and situated them within the context of causal machine learning. Although the proposed attribute graph is non-causal in nature, it is nevertheless useful for downstream tasks as illustrated in the experiments. After revision, all three reviewers are in favour of acceptance.
Decision: Accept (poster)