Removing Biases from Molecular Representations via Information Maximization
Authors: Chenyu Wang
Sharut Gupta
Caroline Uhler
Tommi S. Jaakkola
Venue: ICLR 2024 poster (2024)
Keywords: Molecular Representation
Batch Effect
Contrastive Learning
Information Maximization
Drug Discovery
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose InfoCORE to deal with batch effects in drug screens and obtain refined molecular representations. It is established on a variational lower bound of the conditional mutual information between latent representations given a batch identifier.

Abstract: High-throughput drug screening -- using cell imaging or gene expression measurements as readouts of drug effect -- is a critical tool in biotechnology to assess and understand the relationship between the chemical structure and biological activity of a drug. Since large-scale screens have to be divided into multiple experiments, a key difficulty is dealing with batch effects, which can introduce systematic errors and non-biological associations in the data. We propose InfoCORE, an Information maximization approach for COnfounder REmoval, to effectively deal with batch effects and obtain refined molecular representations. InfoCORE establishes a variational lower bound on the conditional mutual information of the latent representations given a batch identifier. It adaptively reweights samples to equalize their implied batch distribution. Extensive experiments on drug screening data reveal InfoCORE's superior performance in a multitude of tasks including molecular property prediction and molecule-phenotype retrieval. Additionally, we show results for how InfoCORE offers a versatile framework and resolves general distribution shifts and issues of data fairness by minimizing correlation with spurious features or removing sensitive attributes.

Published as a conference paper at ICLR 2024
REMOVING BIASES FROM MOLECULAR REPRESENTA- TIONS VIA INFORMATION MAXIMIZATION
Chenyu W Sharut Gupta', Caroline Uhler2.3, Tommi Jaakkola'
Computer Science and Artificial Intelligence Laboratory, Massachusetts Institute of Technology Eric and Wendy Schmidt Center, Broad Institute of MIT and Harvard Laboratory for Information and Decision Systems, Massachusetts Institute of Technology
ABSTRACT
High-throughput drug screening using cell imaging or gene expression mea- surements readouts of drug effect - is a critical tool in biotechnology to as- sess and understand the relationship between the chemical structure and biologi- cal activity of drug. Since large-scale screens have to be divided into multiple experiments, a key difficulty is dealing with batch effects, which can introduce systematic errors and non-biological associations in the data. We propose In- foCORE, an Information maximization approach for COnfounder REmoval, to effectively deal with batch effects and obtain refined molecular representations. InfoCORE establishes a variational lower bound on the conditional mutual infor- mation of the latent representations given batch identifier. adaptively reweighs samples to equalize their implied batch distribution. Extensive experiments on drug screening data reveal InfoCORE's superior performance in a multitude of tasks including molecular property prediction and poolecuee-phnntyye Additionally, we show results for how InfoCORE offers a versatile framework and resolves general distribution shifts and issues of data fairness by minimizing correlation with spurious features or removing sensitive attributes. The code available https: //githb/ com/ htt//ubless uulerlall eeteerrrbbbbbbb Iab/IRECCOR..
retrieval.
INTRODUCTION
Representation learning (Bengio et al., 2013) has become pivotal in drug discovery (Wu al., 2018) and understanding biological systems (Yang et al., 2021b). It serves as pillar for recognizing drug mechanisms, predicting a drug's activity and toxicity, and identifying disease-associated chemical structures. A central challenge in this context is to accurately capture the nuanced relationship be- tween the chemical structure of small molecule and its biological or physical attributes. Most molecular representation learning methods only encode a molecule's 's chemical identity and hence provide unimodal representations (Wang et al., 2022; Xu et al., 2021b). A limitation of such tech- niques is that molecules with similar structures can have very different effects in the cellular context.
High-content drug screens have been developed that output post- post-eerturbation (i.e., after the applica- tion of drug) cellular images and gene expression (Chandrasekaran et al., 2021). Such datasets pro- vide a unique opportunity to improve our understanding of the biological effect of a compound and help refine the representation of molecules. Given the huge chemical space of over 1060 molecules of possible interest for drug discovery (Reymond et al., 2010), the full space cannot be explored experimentally. Thus, computational methods are needed to obtain biologically informed molecular representations that can be generalized to untested molecules.
Recent works in this direction (Nguyen et al., 2023; Zheng al., 2022) focused on training models to map 2D molecular structures to high-content cell microscopy images (Bray et al., 2017; Chan- drasekaran et al., 2023) through multimodal contrastive learning (Radford et al., 2021). Generalizing across molecular structures has been challenging due to the variability in the training data as drug responses vary under different conditions. In particular, batch effects are pervasive, stemming from the fact that large-scale drug screens have to be divided into multiple experiments over many years.
correspondence wangchy@mi edu

---

Published as a conference paper at ICLR 2024
D - PCR (D)
Observed Batch Identiter
Drug Encoder
Daselier
(D)
Drug
DN
GrD:
GrD:
Gr
B,
Image/ Geno Encoder
...
B:
Cell Imazinn
:
G+D:
G
Classifer
M2 G. a)) -Dils
Ø = GGB(G)
The model takes drug structure and screening data (gene expression or cellular imaging) as input and encodes it into latent representations D; and Gg. Two batch classifiers are trained to obtain the posterior batch distribution G; = p (B) and D, = p,) The two encoders are jointly trained to predict the correct pairings of molecule-phenotyae training samples. The loss is function of two matrices: the pairwise cosine similarity of the representations (left) and the weight matrix for each pair, where the weight Miy is the weighted (with scaler a)) average posterior probability of the latent representations being in batch B,, i.e. G [BB]] and [B:] (right). Sample pairs with more similar posterior batch distributions are emphasized more in L. The figure illustrates the term in the loss L using G as an anchor; using D as an anchor can be done analogously.
Figure 1:
Batch effects refer to non-biological associations introduced through the measurement process. They can systematically distort the data and make it challenging to isolate the true biological signal.
While various statistical approaches have been proposed to correct for batch effects (Johnson et al., 2007: Korsunsky et al., 2019), these generally consist of preprocessing step that independent of the downstream training task. We provide more effective batch correction method by directly integrating it with the learning algorithm. framing the task similar removing the impact from sensitive attribute (batch number). While various techniques have been developed for this problem, including in computer vision, they do not directly address the challenges drug screening. For example, when observations of batch-drug combinations are limited in number and coverage, using contrastive samples with the same value of the sensitive attribute as in Ma et al. (2021) can lead to poor generalization. Zhang al. (2022) rely models to generate unbiased image training data, but generative models for high-content screening data are in their infancy (Yang et al., 2021a).
In this paper, we introduce novel method, InfoCORE, designed mitigate confounding factors in multimodal contrastive learning. While we describe our method in the context of batch effect removal for molecular representation learning, framework, such as for the removal of sensitive information to enhance fairness. InfoCORE formu- lates an intuitive and casy-to-optimize objective based on variational lower bound on conditional mutual information. In essence, sample pairs with more similar batch distributions are emphasized more in the InfoNCE loss function (Oord et al., 2018). This weighting scheme enables the model to adapt its learning strategy for each sample based on the batch-related information present in the latent representations. The proposed framework is illustrated in Figure 1.
also show its effectiveness as a general-purpose
We conduct experiments with two common readouts of high-content drug screens: LINCS gene expression profiles (Subramanian al., 2017) and cell imaging profiles (Bray et al., 2017). We show that InfoCORE consistently outperforms the baseline models across range of downstream tasks, including molecule- phrnoolpeppeettpeeeeeeee retrieval and molecular property prediction. Furthermore, our empirical evaluations demonstrate InfoCORE's broad applicability well beyond drug screening. In particular, we demonstrate that InfoCORE obtains improved representations with respect to various fairess measures over existing baselines across multiple datasets, including UCI Adult (Asuncion & Newman, 2007), Law School (Wightman, 1998), and Compas (Angwin et al., 2022).
To summarize, the main contributions of work are:
2

---

Published as a conference paper at ICLR 2024
We propose InfoCORE, a framework for multimodal molecular representation learning, capable integrating diverse high-content drug screens with chemical structures.
Theoretically, we show that InfoCORE maximizes the variational lower bound on the conditional mutual information of the representation given the batch identifier. It empirically outperforms var- ious baselines on tasks such as molecular property prediction and molecule-phenee
retrieval. The information maximization principle of InfoCORE extends beyond drug discovery. We empir- ically demonstrate its efficacy in eliminating sensitive information for representation fairness.
METHOD
In this section, we introduce InfoCORE, provide the intuition for how counteracts biases from irrelevant attributes, and derive the corresponding training objective. In Section 2.1, we describe the underlying graphical model as well as the main learning objective, which is based on conditional mutual information. In Section 2.2, we establish tractable lower bound for this objective, which we show can be interpreted as the InfoNCE loss (Oord et al., 2018), but with unequally weighted negative samples. See Appendix E for short review and additional backgrounds on InfoNCE.
2.1 CONDITIONAL MUTUAL INFORMATION MAXIMIZATION TO REMOVE BIASES FROM IRRELEVANT ATTRIBUTES.
Figure depicts the graphical model for both the observations and the learned representations. In this model, (Xa, X.. Xx) represent observed variables Xa signifies the molecular structure of drug. X, indicates the shift in gene expression other phenotype induced by applying the drug, and X6 represents an irrelevant attribute such as the experimental batch in the context of molecular representation learning. Illustrated by dashed lines in Figure 2, (Zd Zg) are representations gener- ated from the trained encoders: Za = Enca( Xa; 0a). Zg = Encg [X 0g), where Enca(.; 0a) is the encoder of the drug structure and Enc, 02) is the encoder for the screening readout.
The graphical model illustrates that X, is confounding factor affecting the learned representations. The observed phenotype X, is influenced by the drug's effect, derived from its molecular structure, but also by external non-biological factors (i.e. batch effects) Xs. Since batch assignment in biological experiments is not always (often) fully randomized, the batch number can affect which drug is applied, explaining the arrow from X, to X,- The presence of a backdoor path from Xa to Xg through X,, means that naive estimates of the effect of Xa on Xg will be biased.
X
Observed
X
Fncoder
Z
Representation
We adopt conditional mutual information objective to address this challenge, conditioned on the irrelevant attributes Xg. This results in learned latent representations Za. Zg. where variations associated with X5 are effectively excluded. Intuitively, this because the features tied X, do not improve the conditional mutual information objective, and because the latent representation has a finite dimension (implying rep- resentational resource constraint). This claim further justified Ma et al. (2021) and Proposition 2 of Robinson et al. (2021). More precisely, based on the graphical model given above, specify a conditional mutual information criterion for learning latent representations of X and Xg that do not contain the effect of X,. This criterion satisfies the following bound!:
Figure 2: Graphical Model
max 0g.0, 2
(T(Za; X,Ix;; + I(Zg: Nx: Z max T(Za; Zgl Xx; 02, 0g).
82.9,
As noted earlier, this objective function emphasizes drug's 's bioactivity by focusing on shared features of the two modalities that are unrelated to batch.
2.2 REWEIGHTED INFONCE OBJECTIVE AS TRACTABLE LOWER BOUND.
Direct estimation of conditional mutual information is computationally intractable for high- dimensional continuous random variables. The InfoNCE objective in contrastive learning (Poole
The incquality is based the Markov relationship in the graphical model and data processing inequality.
3

---

Published as a conference paper at ICLR 2024
et al., 2019; Oord et al., 2018; Tian et al., 2020) has been introduced tractable alternative and was extended to the multimodal case in CLIP (Radford et al., 2021). optimizes the represen- tation to bring each "anchor" data point close to its paired "positive" and away from "negative" samples in the latent space (details in Appendix E). It was further extended to the conditional case by Ma et al. (2021), where the relationship between conditional mutual information and conditional contrastive learning (CCL) was built. In CCL, the negative samples are drawn from the conditional marginal distributions of the representations zd and Zg. (() and p(2y| [mb), conditioned on the an- chor's batch number zt. The downside is that when there are limited observations for particular value of 2h (limited number drugs), the small number of possible negative samples will lead to poor gen- cralization to new molecular structures.
Table I: Comparison of various methods such CCL and CLIP with InfoCORE
Method
CLIP
CCL
Ours
Debiasing
/ /
Large Supply of Negatives
X
/
To address this challenge, we introduce InfoCORE, which reweighs the negative samples in the InfoNCE objective differentially based on the posterior batch distribution. This acts as a practi- cal lower bound for the conditional mutual information I [(Za: Zg Xx) and dynamically adjusts the weight for each machor-negative pair. A comparison of various methods is provided in Table 1.
2.2. SINGLE-SAMPLE LOWER BOUND OF CONDITIONAL MUTUAL INFORMATION.
InfoNCE (Oord al., 2018) is one of the most popular objectives in contrastive learning. Poole al. (2019) showed that the InfoNCE objective is lower bound on mutual information. We obtain the bound in Proposition by extending the variational lower bound from the mergy-based variational family in Poole et al. (2019) to conditional mutual information. The proof is given in Appendix A. Proposition 1. Given random variables Zg. Za as representations Iwo data modalities and Xn as the irrelevant attribute, the conditional mutual information 1 a(Z,; Xx) between representations conditioned on the irrelevant attribute has the following lower bound:
I ZZz (x) Bo(sa=") an) [M(aa Zg. [6)] - EEptes Ep(ey ((sasFgota)
I(Za; Xo).
which results from using an energy-basee based variational family (Z2n volza) to approximate the distri- bution p(z, Iib zz))
a(2g) xnl (za) = p(2g, mn) M(saigga)) where ((za) = (ppts [ [chteatgmes] is the partition partitionfuuncionn function,
Z (zaa
with equality holding when the critic? satisfies h²" (zd: Zg, Tb) xb)= = + log p(z,.ealza) p(=g.m))
This bound provides a tractable estimator - where the first two terms can be estimated given samples from the joint distribution and product marginal distribution, and I(Za; can be optimized via adversarial training as shown by Guo et al. (2023). However, as explained in Poole et al. (2019), this bound may exhibit high variance due its reliance on the upper bounds of the log partition function, i.e., log (za) S e -ZZz),,), which introduces high variance in its sample approximations.
2.2.2 INFOCORE AS A MULTI-SAMPLE LOWER BOUND CONDITIONAL MUTUAL INFORMATION.
To reduce the variance of the single sample lower bound discussed above, we develop a method that makes use of additional samples from the data. An intuitive approach is to extend the un- conditional mutual information proposed in Oord et al. (2018) and Poole et al. (2019). However, applying this strategy directly in our use case would mean an adversarial optimization of I I(Za: X₂) (see Equation 2), which is disadvantageous from computational perspective. Instead, propose to decompose the distribution p(=g: my/ =a) into the product of a conditional distribution p(zylza) which is independent of batch .i and the posterior batch distribution based the latent represen- tation p(zy|2a 2d, 2.). We show that by separately approximating these distributions, maximizing the conditional mutual information maintains connections to the widely-used InfoNCE objective.
In the following propositions, we build on the decomposition of p(zg. Ty|za) to derive a multi-sample lower bound of the conditional mutual information objective shown in Equation 3. The proof is given
We follow the terminoloyy in Poole et al. (2019) refer to h(2g: Zg. 26) as critic.
4

---

Published as a conference paper at ICLR 2024
in Appendix B. Variables with superscript 1 (i.e. Z3 X.) denote the corresponding variables for the unchor-positive pair. Those with superscript other than 1 (i.e. ZJ-K, zz denote the negative samples in the multi-sample case; K is the number of negative samples. A lookup table of the definition of all variables is provided in Appendix J.
Proposition 2. Given samples (zd, zg. mh) drawn from the joint distribution (Za, Zg X.)) pf z4, 2g. In) and 2.K drawn ii.d. from the marginal distribution Za p(za) for i = 2, ... K, then the conditional mutual information I(Z'; Zg xx)) has the following lower bound:
1(22: (x,) 2 Leur LLur +HHH())).
(3)
where Leus
B of
log
2
R
+ Bptehege//bn000m
log
DN Mabiit) Pa(xrllag =b)
Law 2 Bete [DEL (p (abled) 11 (251 + Bnltt DRL (p (x>lig) 11 (=dl=)))
1 C 2 (ptaeeaatt log Po ( zt) (a =
This holds for any choice of critic f(z,, =d) and variational distribution (g zd). Pa |zg0 zd), with equality holding when f (za, 2g) = log p(zylza) and (zolzg, za) = (2(2112 za) = p( (2gs za).
p(sy)
We note that LCLIP resembles the symmetrical InfoNCE loss in multimodal contrastive learning (Radford et al., 2021), but with the negative samples reweighted according to the posterior batch distributions; i.e., the negative samples z/ that are more likely to be the anchor's batch zi (having higher value Py (ar 22 2 are weighted more. LCLF denotes the average classification loss achieved when training a classifier to predict the batch number from given latent representation (2g or za).
Posterior Batch Distribution Estimation. Estimating the lower bound in Proposition 2 requires estimating the reweighting factors P. (z) 22)) and pa(z)) [g))))) This presents several challenges, 1) The corresponding empirical observations are absent, leading to a lack of data; 2) Since most inputs zd and 2g are unpaired, directly training classifier with the paired input concat =3] could result in poor nut-of-distribution generalization. Addi- tionally, this approach would require the computationally intensive step of rerunning the classifier for each pair separately.
especially for negative samples when i & 1.:
Given that p(=,|2),, z²') represents the posterior batch distribution informed by both latent represen- tations, both modalities offer valuable but possibly overlapping information about the batch number x/.. Hence, we can estimate p(r) (..l)..23 =d)) using the weighted arithmetic average of the posteriors given each individual latent, p(x))= ) and p(x))=4), which measures the remaining batch-related information in each latent representation. This offers an intuitive and computationally inexpensive estimate for p(ar! 1/ ={)):
Py (x) 12g = = a p( + (1 a1 ((al))))))) pa(a) sg. =3) = a ((=(())))) + (1 (1-a) (L-a))p(allt) p(alle)).
Based on this, the denominator in LCLiP can be viewed as a combination of uniformly weighted neg- ative samples $(*) 1 Σi= ((())=))) and a ((xb)=))- t/2 Σx (((2)))))) as well confounder biased weighted negative samples (1 - a) ((ab)=)))))))) . (=) and (1 a) / Σ ((=b1=)) ( serves tuning parameter, determining the degree of reliance on the posterior esti- mated from the common anchor (=) or z/). It presents a tradeoff between correcting for irrelevant attributes and enhancing model generalization by incorporating larger supply of negative samples. When & = 1, LCLP collapses to the unweighted InfoNCE loss in CLIP (Radford et al., 2021). The following proposition shows that under these assumptions, C' can be lower bounded by 0, thus simplifying the lower bound further. The proof is given in Appendix C.
Proposition 3. When estimating Po (=g)zz =²) as the weighted average of ((ttl=)) and ((r)lz)), and analogously for Pa(x)) sy, =d), the term C defined in Proposition 2 is lower bounded by zero.
We note that p(2ol=z) and p( za) cancel out in - -eeu C, but we spell out these terms for ease of estimating C later.
5

---

Published as a conference paper at ICLR 2024
Based on Proposition 3, the lower bound in Proposition 2 can be further reduced to [(Z): -3(X;) 2
- Lecp - LCIF H. Since (X;)) is a constant, maximization of the conditional mutual infor- mation lower bound gives rise to the minimization of our InfoCORE loss function:
LIBOCOOEE = LLLP + LELF,
(4)
where LELi and LeLF are defined in Proposition 2.
Estimating the batch distribution as opposed to directly utilizzig the observed values of X, for the negative samples in Ma et al. (2021) and Tsai et al. (2021) has advantages: Since experimental batches may contain molecules that share scaffolds, while others may possess randomly assigned ones, the batch confounding effect can vary from batch batch. InfoCORE fact that positive and negative samples may be affected differently by the batch confounder. More- over, since batch distribution is estimated using the latent representations and not the original data, once the batch effect has been mitigated, InfoCORE implicitly adjusts during training and ceases to reweight the negative samples.
deal with the
2.2.3 COMPUTATIONAL CONSIDERATIONS.
We iteratively optimize the loss function by updating the encoders based LELIP and then the classifiers based on LaFF- Note that although we freeze the classifiers when optimizing ICLP,, p( (2)) and p(z depend on the encoder parameters, which introduces competing objectives. Precisely, the gradient of LCLPP w.r.t. the representations z² and zy can decomposed into two components: a standard component as in CLIP (Radford et al., 2021), involving (. (f..) 220/ and 0f(s), =5)/0zá, and competing component involving =
and 00(x)) (2)) l=a)/ (see de- tails in Appendix D). Note that both gradient components drive the representations towards reduced batch effect: In the standard component, samples with similar batch distributions
wp-weighted and thus the latent variables are driven to be less informative of batch number. In the other com- ponent, when
1, the derivatives ==(x))=))//0 and 0p(z))=]) =/2 update the encoder to make p(at) 22)) and p(x))==)) smaller, making the latent representations less informative of batch num- ber. When
1, according to the confounder-biased weighted part of LCuP's 's denominator, i.c. (1 11-aa + E ((x)) (z.) = we can infer that ((x]12)) is adjusted similar in response to e((s)=+t) as ((())=t)) is adjusted the basis of ((xr)(=() Therefore, the gradient updating schema of ((x))=²)) in terms of the underlying representation 2/z is, in broad terms, consistent with that of (1))) This motivates our approach to use either gradient or stop gradient on á and treat them as constant weights. In practice, we hyperparameter À to control the gradient update schedule.
EXPERIMENTS
In this section, present comprehensive experiments demonstrating the effectiveness of InfoCORE to remove biases in representation learning. Across experiments, we compare InfoCORE with the unconditional multi-modal contrastive learning method CLIP (Radford et al., 2021; Nguyen et al., 2023) and the recent conditional contrastive learning method CCL (Ma et al., 2021).
3.1 SIMULATION STUDY
To offer insights into InfoCORE's capability discern drug's biological effect from noisy data influenced by batch confounders, we conduct low-dimensional simulation experiment that repli- cates our data generation process. Here, each drug. associated with a real effect from one of five categories, is randomly assigned to one of 25 batches, with total of 50 drugs per batch. The ob- servational data Xg. corresponding to screen output, and X corresponding to drug structure, are 10-dimensional vectors. These vectors are generated by a randomized function of the input which concatenates biological effect, batch effect, and Gaussian noise. The dataset is split, with half uti- lized for training and the remaining half held out for evaluation. Additional details pertaining to data generation and experiments can be found in Appendix H.
The 2-dimensional latent representation for the drug screens. Zg. learned by different methods for the held-out data and several quantitative metrics (Xu et al., 2021a) are depicted in Figure 3; the corresponding figure for drug structure and detailed explanation of these metries provided in

---

Published as a conference paper at ICLR 2024
Inf:
Prodiction Accuracy
KNN
By Neal affrect
50-
Entropy of batch mixing
Betreh Eiflect
GLP
IntaD0m2
CLIP
CCL
InfoCORE
15
Figure 3: 2D representation visualizations and quantitative metries for each method: top row uses real effect category for coloring and calculating prediction accuracy; bottom row uses batch identi- fier. The rightmost column shows KNN purity and entropy of batch mixing across different k values.
Appendix I.
An ideal representation should distinctly separate real effect categories without being influenced by batch identifier. This is characterized by high accuracy in predicting real effects, low accuracy for batch identifiers, and elevated values of KNN purity and entropy of batch mixing. CLIP is affected by batch confounders and has difficulty identifying real effect categories, while CCL's embeddings, despite fusing batch identifiers, falter in distinguishing biological effects due to limited negative samples. Conversely, InfoCORE achieves a balance between debiasing and expressivity, offering superior representations that discern real effects and effectively mitigate batch effects.
3.2 REPRESENTATION LEARNING OF SMALL MOLECULES.
Dataset Description. We use two high-content drug screening datasets, L1000 gene expression profiles (Subramanian et al., 2017) (GE) and cell imaging profiles obtained from the Cell Painting assay (Bray al., 2017) (CP). In GE, we select data from the nine core cell lines, resulting in 17,753 drugs and 82,914 drug-cell line pairs. In CP, 30,204 small molecules are screened in one cell line (U2OS). We use the hand-crafted image features obtained by the popular CellProfiler method (McQuin et al., 2018). The chemical structures are featurized using Mol2vec (Jaeger et al., 2018).
Molecule-Phenotype Retrieval. The drug repurposing or drug discovery task can be viewed as follows: identify molecules (c.g., from drug repurposing library) that are most likely to induce given desired phenotypic change (i.e., gene expression change from diseased towards normal). To mimic this process, similar to the setting in Zheng et al. (2022), we randomly split both datasets into a training set consisting of 80% of the molecules and hold out the remaining molecules for testing. We use top N accuracy (N=1, 5, 10) as an evaluation metric and analyze two retrieval libraries: (a) all molecules in the held-out set (whole), and (b) held-out set molecules that are in the same experimental batch as the retrieving target (batch). Accuracies over both libraries whole reflect the model's s ability of molecule- patecuull-phnootpee retrieval for drug discovery and drug repurposing.
Table 2: Retrieving accuracy of different methods for gene expression and cell imaging screens.
Dataset
Gene Expression (GE)
Cell Painting (CP)
Retrieval Library
whole
batch
whole
batch
Top N Ace (%)
N=1
N=5
N=10
N=1
N=5
N=10
N=1
N=5
N=10
N=1
N=5
N=10
Random
0.03
0.13
0.27
1.58
7.90
15.81
0.02
0.08
0.17
1.59
7.97
15.94
CLIP
5.96
18.59
27.17
12.23
30.29
42.63
7.23
20.95
28.89
13.20
37.78
52.72
CCL
1.93
5.85
8.37
12.76
32.39
45.77
1.31
4.93
7.38
13.20
37.99
53.13
InfoCORE
6.39
18.99
27.18
14.03
33.63
46.78
6.93
20.65
28.22
13.26
38.50
53.13
Since CLIP and CCL are not specifically designed for cross-modal molecular representation learn- ing, we made several modifications to the vanilla algorithms for this task; see details in Appendix F. As shown in Table 2, InfoCORE is the only model that has strong performance in both libraries, especially in GE, which agrees with the belief that gene expression data is more affected by batch
7

---

Published as a conference paper at ICLR 2024
effects. Batch-related features dominate the CLIP representation, leading the model's poor perfor- mance in batch. For example, InfoCORE outperforms CLIP in GE over batch with 15% increase in top 1 accuracy and 11% in top 5 accuracy. CCL's reliance on the few hundred negative samples within the same experimental batches hinders its out-of-hatch generalization, leading to its poor performance in whole. Standard deviations over 3 random seeds are provided in Appendix I.
Transfer Learning for Property Prediction. Our trained model can be fine-tuned for various downstream tasks. Given that InfoCORE is pre-trained using drug screening data that contains in- formation on the biological effect of a drug, transfer learning is expected to do well on the prediction of bioactivity-related properties. We test this by analyzing the following classification and regression tasks: For classification, we select 7 bioactivity-related benchmarks from MoleculeNet (Wu et al., 2018) and follow the standard scaffold splitting procedure as suggested by Hu al. (2020). Since the regression tasks in MoleculeNet (i.e. ESOL, Lipo. FreeSolv) are not directly bioactivity-relate we use post-perturbation cell viability in the PRISM dataset (Corsello al., 2020) for the regression task using the same scaffold splitting procedure. Further details are provided in Appendix H.
We analyze InfoCORE when pretraining using the two different drug screening datasets and compare the performance to CLIP and CCL. We report the performance of Mol2vec (Jaeger et al., 2018) with randomly initialized MLP of the same neural network architecture as a bascline for no no-pretraining. Mean AUC-ROC is reported for the classification tasks and R² for the regression tasks, together with standard deviations based on 3 random seeds. The results are shown in Table 3. Compared to Mol2vec, the results clearly demonstrate the benefit of pretraining using drug screening data: signif- icant gains are obtained in all properties except SIDER. Moreover, among the different pretraining strategies, InfoCORE is highly competitive, achieving the best performance in almost all properties when using GE, especially in ClinTox, HIV, and BACE, and most properties with CP.
Table 3: Performance of different methods molecular property prediction task.
Classification (ROC-AUC %)
Reg (R² %) 1
Datasets
BBBP
BACE
ClinTox
Tox21
ToxCast
SIDER
IIIV
Avg.
PRISM
Molecules
2039
1513
1478
7831
8575
1427
41127
3172
Tasks
2
12
617
5
Mol2vee
70.7(0.4)
82.9(0.7)
84.9(0.3)
76.0(0.1)
74.4(0.5)
64.9(0.3)
77.7(0.))
75.9
8.5(0.7)
CT.IP
73.5(0.4)
(0.4()
19.6(2.1)
77.3(0.0)
75.7(0.6)
63.7(0.6)
77.7(0.6)
77.6
(3.9(0.4)
GE
CCL
73,0(0.8)
85.9(0.6)
90.5(1.0)
77.0(0.2)
75.8(0.2)
63.4(0.5)
77.5(0.9)
77.6
16.0(0.5)
InfoCORE
73,5(0.3)
86.6(0.3)
91.9(1.9)
77.4(0.4)
75.7(0.2)
64.8(0.6)
78.5(0.2)
78.3
14.8(0.1)
CLIP
73.4(0.8)
85.2(0.4)
87.3(0.1)
76.4(0.1)
76.7(0.1)
64.8(0.6)
78.(00.4)
77.4
16.2(0.2)
CP
CCL
73.7(0.5)
84.9(0.9)
87.7(1.)
75.9(0.3)
75.7(0.4)
65.2(0.4)
79.3(0.3)
77.5
14.7(0.3)
InfoCORE
74.0(0.8)
85.0(0.2)
89.3(0.5)
76.6(0.1)
76.9(0.1)
65.2(0.1)
78.7(0.1)
78.0
16.2(0.3)
3.3
REPRESENTATION FAIRNESS.
Dataset and Fairness Criteria. In the following, we demonstrate InfoCORE's broad applicability by conducting experiments in representation fairness using three fairness datasets: UCI Adult (Asun-
Table 4: Performance of various methods on representation fairness task.
UCI Adult
Law School
Compas
Method
Acct
EOJ
BOPP
Acct
BOG
EOPP
Accf
EO4
EOPP
CLIP
85.1(0.1)
20.7(1.8)
15.2(1.7)
83.1(0.2)
30.9(1.4)
7.9(0.8)
60.8(2.3)
18.4(2.4)
11.7(1.9)
CCL
85.1(0.2)
19.0(.3)
13.3(2.8)
83.0(0.3)
27.8(1.5)
6.7(0.9)
59.5(2.3)
17.1(3.4)
10.1(3.0)
InfoCORE
85.2(0.1))
14.9(1.1)
9.7(0.8)
82.7(0.4) 25.4(3.6)
6.0(1.4)
60.1(2.1)
15.3(2.5)
9.3(0.8)
cion & Newman, 2007), Law School (Wightman, 1998), and Compas (Angwin et al., 2022). Race and gender are used as protected attributes, separating the data into four subgroups. Considering that minority groups often encounter data limitation issues, subsample the training data to mimic an unbalanced population. To assess representation fairness, we follow the setup in Lahoti et al. (2020) and Ma al. (2021) and analyze three common fairness criteria (Feldman al., 2015; Hardt
8

---

Published as a conference paper at ICLR 2024
et al., 2016): equalized odds (EO), equality of opportunity (EOPP), and demographic parity (DP), with definitions in Appendix G and DP results in Appendix I. We also calculate prediction accuracy (Acc) to account for the utiliity faiimess trade-off in the representations (Zhao & Gordon, 2022).
Results. As shown in Table 4, while achieving similar levels of prediction accuracy, InfoCORE consistently obtains more fair representations across datasets in terms of the different criteria, espe- cially in UCI Adult, CCL improves representation fairness over CLIP by constraining the negative sampling to be within the same sensitive attribute subgroup. However, given the limited sample size of the minority groups, it performs inferior to InfoCORE.
4 RELATED WORK
Representation Learning for Molecules. Traditional unimodal techniques for molecular represen- tation learning (Rogers & Hahn, 2010; Durant et al., 2002: Wang et al., 2022; Xu et al., 202 lb; Wang et al., 2019) often falter because molecules with similar structures can have very different effects in the cellular context. Given such limitations, researchers are increasingly turning to multi- modal methods- - incorporating additional modalities like 3D structure (Stärk et al., 2022; Zhou et al., 2022) as well as high-throughput cell imaging. For instance, Nguyen et al. (2023) utilize the CLIP model (Radford et al., 2021) to learn multimodal molecular and cell image representations. Zheng et al. (2022) incorporate masked graph modeling and generative graph-image matching objectives to elevate the quality of the learned representations. However, they have difficulties generalizing across molecular structures due to the variability and batch effect in the drug screens as the training data, which can introduce confounding variables and bias the representations.
Batch Effect Removal. Various approaches have been proposed to correct for batch effects, in- cluding a linear method (Johnson et al., 2007), a mixture-model based method (Korsunsky al., 2019), neighbor-based methods (Hie et al., 2019; Haghverdi et al., 2018), and variationallinferrnnce based methods (Lopez et al., 2018: Li et al., 2020). These methods, typically used independent preprocessing steps designed for specific biology experimental techniques, are challenging to be ap- plied to more general contexts. In contrast, our approach treats batch effects as sensitive attributes to remove, providing a flexible framework applicable across various domains.
Representation Learning with Sensitive Attributes. Wu et al. (2022); Chen et al. (2022); Yang et al. (2022); Miao et al. (2022); Fan et al. (2022) learn invariant features for better out-of distribu- tion generalization in graph classification or interpretability. However, they focus on the unimodal supervised learning setting. In the unsupervised setting, to learn fair representations, Song et al. (2019) use variational and adversarial objectives as a tractable lower bound on conditional mutual information, which requires challenging tri-level optimization. Ma et al. (2021) relate conditional mutual information with a conditional contrastive learning objective, where the negative pairs are drawn from the conditional marginal distributions. Tsai et al. (2021) further extend it to handle continuous conditioning variables by leveraging similarity kernels. However, this does not resolve the confounder issue in drug screening data, since observations per batch are limited and similarity between batches as a categorical identifier is hard to measure. Zhang et al. (2022) rely on recent advances in image generation obtain balanced training data prior the representation learning task, but generative models for high-content screening data are in their infancy (Yang et al., 2021a).
5 CONCLUSION
We present InfoCORE, a general framework for multimodal molecular representation learning. In- foCORE effectively integrates diverse high-content drug screens with the chemical structure in the presence of confounders like batch effects. Theoretically, we show that InfoCORE maximizes the variational lower bound on the conditional mutual information of the representations given the batch identifier. Empirically, we demonstrate that InfoCORE outperforms baselines on two drug screen- ing datasets (gene expression and single-cell imaging) on two tasks, molecular property prediction and molecule-phenotype retrieval. Finally, we show that the information maximization principle underlying InfoCORE extends well beyond drug discovery and can be viewed as a general-purpor framework. In particular, we empirically demonstrate its efficacy in eliminating sensitive informa- tion, focusing on fairness applications across three major datasets and various fairness criteria.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Adiryanarayanan Radhakrishnan, Joshua Robinson, Yonglong Tian, Yilun Xu, Shangyuan Tong, Wengong Jin, Hannes Stärk, Boyuan Chen, Zongyu Lin, and Mengfei Xia for helpful discus- sions and comments on the manuscript.
CW and TJ acknowledge support from the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium. SG acknowledges fundiing from the Office of Naval Research grant 9000014-2-1-2200 (MURI ML-SCOPE) and NSF award CCF-21 12665 (TILOS Al Institute). CU acknowledges support by NCCIH/NIH (IDP2AT012344 ONR (N00014-22) 1200014-22--11 1-4--22--2210- 16), the MIT- IBM Watson Al Lab, AstraZeneca, the Eric and Wendy Schmidt Center at the Broad Institute, and Simons Investigator Award.
REFERENCES
Julia Angwin, Jeff Larson, Surya Mattu, and Lauren Kirchner. Machine bias. In Ethics of data and analytics, pp. 254-264. Auerbach Publications, 2022.
Arthur Asuncion and David Newman. Uci machine learning repository, 2007.
Philip Bachman, Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. JEEE transactions on pattern analysis and machine intelligence, 35(8):1798-15)) 2013.
Mark-Anthony Bray, Sigrun Gustafsdottir, Mohammad Rohban, Shantanu Singh, Vebjorn Ljosa, Katherine L Sokolnicki, Joshua A Bittker, Nicole E Bodycombe, Vlado Dančík, Thomas P Hasaka, et al. dataset of images and morphological profiles of 30 000 small-molecule treat- ments using the cell painting assay. Gigascience, 6 6(12)giw004, 2017.
Srinivas Niranj Chandrasekaran, Hugo Ceulemans, Justin D Boyd, and Anne Carpenter. Image- based profiling for drug discovery: due for machine-learning upgrade? Nature Reviews Drug Discovery, 20(2):145 159, 2021.
Srinivas Niranj Chandrasekaran, Jeanelle Ackerman, Eric Alix, D Michael Ando, John Arevalo, Melissa Bennion, Nicolas Boisseau, Adriana Borowa, Justin D Boyd, Laurent Brino, al. Jump cell painting dataset: morphological impact of 136,000 chemical and genetic perturbations. bioRxiv, pp. 2023-03, 2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, MA Kaili, Binghui Xie, Tongliang Liu, Han, and James Cheng. Learning causally invariant representations for distr--iistiitttoo generalization on graphs. Advances in Neural Information Processing Systems, 35:22131-22148. 2022.
Steven M Corsello, Rohith T Nagari, Ryan D Spangler, Jordan Rossen, Mustafa Kocak, Jordan G Bryan, Ranad Humeidi. David Peck. Xiaoyum Wu, Andrew A Tang, al. Discovering the an- ticancer potential of non-oncology drugs by systematic viability profiling. Nature cancer, 1(2): 235-248 2020.
Joseph L Durant, Burton A Leland, Douglas R Henry, and James G Nourse. Reoptimization of mdl keys for use in drug discovery. Journal of chemical information and computer sciences, 42(6): 1273-1280, 2002.
Shaohua Fan, Xiao Wang, Yanhu Mo, Chuan Shi, and Jian Tang. Debiasing graph neural networks via learning disentangled causal substructure. Advances in Neural Information Processing Sys- fems, 35:24934-24999 24946, 2022.
10

---

Published as a conference paper at ICLR 2024
Michael Feldman, Sorelle A Friedler, John Moeller, Carlos Scheidegger, and Suresh Venkatasubra- manian. Certifying and removing disparate impact. In proceedings of the 21th ACM SIGKDD international conference knowledge discovery and data mining, pp. 159-268, 2015.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning, pp. 1180-1189. PMLR, 2015.
Xingzhuo Guo, Yuchen Zhang, Jianmin Wang, and Mingsheng Long. Estimating heterogeneous
treatment effects: Mutual information bounds and learning algorithms. 2023.
Laleh Haghverdi, Aaron TI. Lun, Michael D Morgan, and John C Marioni. Batch effects single- cell ma-sequencing data are corrected by matching mutual nearest neighbors. Nature biotechnol- ogy, 366)):41--427, 2018.
Moritz Hardt, Fric Price, and Nati Srebro. Equality of opportunity in supervised learning. Advances in neural information processing systems, 29, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Brian Hie, Bryan Bryson, and Bonnie Berger. Efficient integration of heterogeneous single-cell transcriptomes using scanorama. Nature biotechnology, 37(6):685-691, 2019.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118 18-22133, 2020.
Sabrina Jaeger, Simone Fulle, and Samo Turk. Mol2vec: unsupervised machine learning approach with chemical intuition. Journal of chemical information and modeling, 58(1):27-35, 2018.
Gwanghoon Jang, Sungjoon Park. Sanghoon Lee, Sunkyu Kim, Sejeong Park, and Jacwoo Kang. Predicting mechanism of action of novel compounds using compound structure and transcrip- tomic signature coembedding. Bioinformatics, 37 (
2021.
W Evan Johnson, Cheng Li, and Ariel Rabinovic. Adjusting batch effects in microarray expression data using empirical bayes methods. Biostatistics, 8(1):188-1)) 112, 2007.
Ilya Korsunsky, Nghia Millard, Jean Fan, Kamil Slowikowski, Fan Zhang, Kevin Wei, Yuriy Baglaenko, Michael Brenner, Po-ru Loh, and Soumya Raychaudhuri. Fast, sensitive and accurate 2019.
integration of single-cell data with harmony. Nature methods,
Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang, and Ed Chi. Fairness without demographics through adversarially reweighted learning. Advances in neural information processing systems, 33:728 740, 2020.
Xiangjie Li, Kui Wang. Yafei Lyu, Huize Pan, Jingxiao Zhang, Dwight Stambolian, Katalin Susztak, Muredach P Reilly, Gang Hu, and Mingyao Li. Deep learning enables accurate clustering with batch effect removal in single-cell ma-seq analysis. Nature communications, 1111::2338, 2020.
Romain Lopez, Jeffrey Regier, Michael B Cole, Michael Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature methods, 15112):0033-1088 2018.
Martin Q Ma, Yao-Hung Hubert Tsai, Paul Pu Liang, Han Zhao. Kun Zhang, Ruslan Salakhutdinov, and Louis-Philippe Morency. Conditional contrastive learning for improving fairness in self- supervised learning, arXiv preprint arXiv:2106.02866 2021.
Claire McQuin, Allen Goodman, Vasiliy Chernyshev, Lee Kamentsky, Beth A Cimini, Kyle W Karhohs, Minh Doan, Liya Ding, Susanne M Rafelski, Derek Thirstrup, et al. Cellprofiler 3.0: Next-generation image processing for biology. PLoS biology, 16(7):e2005970, 2018.
Siqi Miao, Mia Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic atten- tion mechanism. In International Conference on Machine Learning, pp. 15524-15543. PMLR, 2022.
11

---

Published as a conference paper at ICLR 2024
Cuong Q Nguyen, Dante Pertusi, and Kim M Branson. Molecule-morpholee contrastive pretrain- ing for transferable molecular representation. bioRxiv, pp. 2023-05, 2023.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv:
2018.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds mutual information. In International Conference on Machine Learning, pp. 5171- 5180. PMLR, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- $748-8766 PMLR, 2021.
Jean-Louis Reymond, Ruud Van Deursen, Lorenz C Blum, and Lars Ruddigkeit. Chemical space as source for new drugs. MedChemComm, 1 111::30-38 2010.
Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shortcut solutions? Advances neural information processing systems, 34:4974-4986, 2021.
David Rogers and Mathew Hahn. Extended-commecctvve fingerprints. Journal of chemical informa- tion and modeling, 50(5):742-75, 2010.
Yang Shu. Zhangjie Cao, Chenyu Wang, Jianmin Wang. and Mingsheng Long. Open domain gen- cralization with Nomain-augmentt meta-learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9624-9633, 2021.
Jiaming Song, Pratyusha Kalluri, Aditya Grover, Shengjia Zhao, and Stefano Ermon. Learning controllable fair representations. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2164 2173. PMLR, 2019.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp. 11479-205022 PMLR, 2022.
Aravind Subramanian, Rajiv Narayan, Steven M Corsello, David D Peck, Ted E Natoli, Xiaodong Lu, Joshua Gould, John F Davis, Andrew Tubelli, Jacob Asiedu, et al. A next generation connectivity map: L1000 platform and the first 1,000,000 profiles. Cell, 17166):137-1142, 2017.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In Computer EssionEECCV0020 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part 16, pp. 776-794. Springer, 2020.
Yao-Hung Hubert Tsai, Tianqin Li, Martin Q Ma, Han Zhao, Kun Zhang, Louis-Philippe Morency, and Ruslan Salakhutdinov. Conditional contrastive learning with kernel. International Confer- on Learning Representations, 2021.
Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. Smiles-bert: large scale unsupervised pre-training for molecular property prediction. In Proceedings of the 10th ACM international conference on bioinformatics, computational biology and health informatics, pp. 429 436, 2019.
Yuyang Wang. Jianren Wang. Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks, Nature Machine Intelligence, 4(3):279- 287, 2022.
Linda F Wightman. Lsac national longitudinal bar passage study. Isac research report series. 1998.
Ying- Xin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. Discovering invariant rationales for graph neural networks. arXiv preprint arXiv:2201. /2872, 2022.
12

---

Published as a conference paper at ICLR 2024
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9(2):113-530 2018.
Chenling Xu, Romain Lopez, Edouard Mehlman, Jeffrey Regier, Michael Jordan, and Nir Yosef. Probabilistic harmonization and annotation of single-cell transcriptomies data with deep genera- tive models. Molecular systems biology. 17(1):e9620, 2021a.
Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graph-level representation learning with local and global structure, In International Conference Machine Learning, pp. 11548-11550 1155.. PMLR, 2021b.
Karren Yang. Samuel Goldman, Wengong Jin, Alex X Lu, Regina Barzilay, Tommi Jaakkola, and Caroline Uhler. Mol2image: improved conditional flow models for molecule to image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6688-6699 2021a.
Karren Dai Yang, Anastasiya Belyaeva, Saradha Venkatachalapathy, Karthik Damodaran, Abigail Katcoff, Adityanarayanan Radhakrishnan, GV Shivashankar, and Caroline Uhler. Multi-domain translation between single-cell imaging and sequencing data using autoencoders. Nature commu- nications, 12(1):31, 2021b.
Nianzu Yang, Kaipeng Zeng, Qitian Wu, Xiaosong Jia, and Junchi Yan. Learning substructure in- variance for out-of-distribution molecular representations. Advances Neural Information Pro- cessing Systems, 35:12964-12978 2022.
Fengda Zhang, Kun Kuang, Long Chen, Yuxuan Liu, Chao Wu, and Jun Xiao. Fairness-awaee contrastive learning with partially annotated sensitive attributes. In The Eleventh International Conference on Learning Representations, 2022.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09411 2017.
Han Zhao and Geoffrey J Gordon. Inherent tradeoffs in learning fair representations. The Journal of Machine Learning Research, 23(1)2222--255, 2022.
Shuangjia Zheng, Jiahua Rao, Jixian Zhang, Ethan Cohen, Chengtao Li, and Yuedong Yang. Cross- modal graph contrastive learning with cellular images. bioRxiv, pp. 2022-06, 06, 2022.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2022.
13

---

Published as a conference paper at ICLR 2024
A PROOF OF PROPOSITION 1.
Proposition 1. Given random variables Zgg Za representations of two data modalities and Xa as
the irrelevant attribute, the conditional mutual information T T(Za; Zgl X6) between representations conditioned on the irrelevant attribute has the following lower bound:
[Za; (x) 2 Bpizaseaas) [h(za, 2g v6)] - EEplzn Ep(sg (((as))g)mo) - [(a,;
(5)
which results from using an energy-based variational family g(2g) Tolza) to approximate the distri- bution plz Ib Iza):
((2g: zblza) p(2g. [bb (( where Z(za) Bpp(ggar)) [c [ the partition function,
[(za)
with equality holding when the critic h satisfies h" (Zd, Eg. Tb) 1 log ((2g))o)=d) p(
Proof. By definition of conditional mutual information.
[Za: Xo) Ep(suraams) Doo plas)-p((,, p((s.
Replace the intractable conditional distribution p(2g- Th (ea) with the corresponding variational dis-
tribution (( no (2d), can get lower bound due to the non-negativity of KL divergence:
log p(x6)) (Zg 16 (b(lza)
+ Ep(aa) [DKI. (p (Eg, on (za) ly (Egs an =a))]
p(ra 22) p(2g, Tb)_
log p(xr)) (Zzg I'b (2
2 Ep(zy
p(rz (za) p(zg, 2b)
(6)
We can choose an energy-based variational family that uses critic lz(za, Egz 1b), scaled by the marginal density p(zy: .76) and normalized by the partition function Z ((z)):
<gg nb Ea) p(zg: 2b) Maaaizg)m)) where ((za) TE ())
{ (h
((za)
By plugging this into the lower bound in Equation 6, we obtain
log p(x6)) ( (=
I (Zu; (olX,)= 2 E Ep(=
p( ka) ((2a)
Ep(znr)g [h(za, 2g. xb)] - Bp(za) [log ((xa) - (2ai
The intractable log-pan regpaariitinn function can be upper bounded using the inequality log(x) < +
log(a ()) 1, V x. a > 0, which is tight when x a. Applying the inequality log Z(za) and taking we obtain the following upper bound of the log-partition function:
log (zz) < e Z(a)).
This results in the following single sample lower bound of the conditional mutual information:
T T(Za; Xo) 2 (p( [h(=a, Zg, vb)] e - Bp(ss) Bplenies) (
I(Zd: X6).
Denoting the optimal critic ht (Ea, Rg. Rb). Equation 6 is tight when q Ib (za) = p( I'd (zd),
iee.
htt (za, 2y: r2) = log p(gg lea) c(2g, where c(2d, 21) an arbitrary function solely depend on zu and J7b- Equation 7 is tight when Z(za)= e, i.e.
Plzg. e (#a. deydos =
plzy, ma) xb rbl deydrs:
5g.00
which holds when c(2y, mp) =
Thus, the optimal critic satisfies h² h 2g. 2bb) 1 p(zgsallsu) which completes the proof. - p(2y,2))
14

---

Published as a conference paper at ICLR 2024
B PROOF OF PROPOSITION 2.
Proposition 2. Given samples (2 mb) drawn from the joint distribution (Z Xx)
plza, Rg. Ib, and zá drawn ii.d. from the marginal distribution Z' p(za) for 2, ..., K, then the conditional mutual information I(Za; Zg has the following lower bound:
T(Z xx)x 2 Tarr Tar + C -H(());
(8)
(iggiad)
where Leup
B p "leg
log
2
Σ k el (+g Do (3r =2)
K i=
log
E Mag. pa =g. =2)
Las
(p())) [DKL (p (uNlz)) 11 (<>1la)) Bp(
[KL. (p (1blz²) ll (=))=)))
log Po (xbl 2 (a =
C - E (p(s).)).a))
p(at)=#)) p(..) =2)
The lower bound holds for any choice of critic { 2a) and variational distribution Po (20 (zg za). pa(2olzzg za), with equality holding when 8" (2a, log pizalzal and Py (2ul2y za) P2 (ag zd) p(aa)
p(rolzg 2g, =a).*
Proof. By definition of conditional mutual information,
log p(ab)) p(=d 2/1 an)) p(zd xt)) p(zz xx))
Zx)) = I(2) ( = E Epcet
log p(a)1=2) pl2g- (2d)
p(x!)
= Ep(s)===thee++
pl plzg.00) xn))
Then applying
variational distribution ((2). l/ 1:1 approximate p(z]. 2f.22" 1:K)
p(z). # (2))) yields the following lower bound on the conditional mutual information:
plat) pagle o altg.adle
Z(Z):)z|x!)> 2
log
p())aa))
Observing that -[====/())))) = (( =/), we can approximate the two terms separately. For this, we define the variational distribution as:
[1c]
Po +
g(2) ./) e. e.p(z) 2.))
a(2) ap 1.
2
(Eddi=b)) Do (2 (2)
where Z(z]:K) = Ep(ezazt)
a(2) mn. 2 l:k)
K
a(z). .bb; 2)
(ashit Do (2bb1
K
'We note that p(21o==)) and (( cancel out in - -ecu C, but we spell out these terms for ease of estimating later.
15

---

Published as a conference paper at ICLR 2024
By plugging this term into the inequality above, obtain
p(an) pla,
Po 2
T(2); (g(xt) 2 E pe ((
log
(2
p(a)))
e 2 20) Po (20 12 lzg 2nd)
= + E pp(2)e))2))00
log
pla
Ep(o2in) [log ((zd 1: )]
al a a k)
p(an)
(odi Pog (2r lagg zd))
2 1 + B Bp(sh)e
log
plan 2
BB Bp(=tin) [Z(=2(k]
a mp: )
For the last inequality we used log $ + log a - and took a = e.
By symmetry, we can conclude as follows that the last term equals to the constant 1:
( Po (xt =d)
E k) [Z(2](x]] = e E Bx(
e Ep(s).at)
a(z =
(x22 Pg (20) 2d) a(2) mg 2dd ..ii))
K
e (
po (3/ b2g 2nd) a(2g, :/: 1:k
E pptaallee
K
DE el(tiit)) (1b /gg 2'1) =1.
K ii=
a(z x2:
Therefore,
p(ub) p(agl=s)
(=ti Po ( =d))
= ; 2 #x(=)=))+)))33"))
log
al a(2), xhi;; 21:
with optimal critics (za: 2g) log p(zz (z)) + c(zg); 17 (10l2g zd) = p( Zg: zd)-
Using p(z)) as an estimator of the posterior batch distribution p(zz |zd]) the lower bound can be written as:
p(att) log F(a
- (=g( zd))
p(agleh)
[ZZ; 2 E E =ct
R Ein ((
(= x/)
((
p(all=))
p(a))
((
(= 2/)
p(abl=))
E Eeee
log
Epts). log
k K DE ((E))=)))
(1/11 =1)
p(z)122)
Note that the last term is the expectation of the KL-divergence between p(x)12g) and م = ie..
Ep(eg=0)) log p
p( (
= Ex(2))
Ep(o
E
log
Ex(a)) [KLL (p (abled) ll (=blzd))
p(xr [=]]
16

---

Published as a conference paper at ICLR 2024
Thus the lower bound can be rewritten as follows:
p(rn) el (zgitt) Pa (1gggg1 kg) 2g)) Zz1x)) 2 Bp ag) log
Ma Σ (( Po z1))
-P(t)) [DEL (p (at 11 (xb12d))
e (=
= =f
log
K -i=1
(sg Py (z) 12g 27)
Pa (061
+ =p
log
H((;)) - Ep(zt) [KI.. (r (mbl=d) l1 (cbl=)))]
Note that, anagolously, we have
I(Z); 2,(X,)
I(Z Z (x,)
(agieth)
2 E =f
log R Ex el (3g (u lag 2dd
Pa(= =g: =dd)
+ +B
log
H(x)) EEpca) [KL (p (xbl=)) 11 (=)]=)))]
with optimal critics I" (=a, =g) = (ogg((alzz) c(za): Pā =d) = p( zd)- Incorporating the two losses, we obtain
T(Z); 2 x)) = [(Z] Zg1x, + T(Z); Z [x:)
(+
> B =p00)13)=11000344)
log 1 E (=giee)) (atblay =1)
K
((
+ E p a ())
log
K (+g pa(z) l2g 2d)
R
1
2 [Epc)) [DKL (o (xbl=2)
11 (z) [ Ep(z))
[DKL (p (xbl=)) 11 (=]
1
log (1 lag. =)) p. )
- H(x)).
+ 2 Bp(sdiggia))
p(
(
When e(zg) log p(=g), and c(za) log (2a), the optimal critics agree with each other,
log
i..,, (Zd, = log m(aplaa) = log m(salzp), Pg (20 2gg za) Pa za) = p(zz) [g, 2d), which p(z,)
completes the proof.
Note that the KL-divergence in LeLF can be equivalently written as the sum of an entropy term and a cross-entropy term:
DKL. (p (xbl=d) II (x))=)) H (x(x))))) E [(gg((|)))]]
H (p<b)=))) CE (p(z))=)), ((r))=))),
DKa. (p (1bl=)) l1 (=)l=)))
H (xdl=))) - E [log((b))))
H (p (2))= + CE (p(a))=)) (())))))
Then LELF can be factorized into the cross-entropy loss of the classifiers p(= |x)) and [(]])=]) and a constant term. Therefore, minimizing LeLF can be achieved by optimizing the classifiers via the cross-entropy loss.
17

---

Published as a conference paper at ICLR 2024
C PROOF OF PROPOSITION 3
Proposition 3. When estimating Po zá) as the weighted average of (()b)))) and p(z))z)), and analogously for pa(a) Rgg zd), i.e.
Po Eg. =d) = a pluble)) + (1 a) p(ab (a)). Pa(lab ag. zd) = a p(=bl=a) (1 - a) p(ab
then C' = Bpt)) log Palalle Ba is lower bounded by the constant zero.
Proof. The weighted AM-GM inequality indicates that for non-negative numbers {ax)i== and non- negative weights {w: Jor)en we have the inequality
2 + 22 +
2 x222
M1
an""
where w = W1 + W2 +
+ Wn.
Thus, by applying the weighted AM-GM inequality, we obtain
Py (= 2/) pa(=)/2g 2/)) p( p(211=²)
C' = (f(
log
(a p( bzg (1 a) p(xx) (2)) (a ((a) (1- p(2))1g))
= -Ep(s)esfaa))
log
(p(xbl=g) ((ab)=))--a)
(f(z)) 200)).pagglee (
2 =
log
p 2/) p(x z/
log 0,
which completes the proof.
Additionally, show that this bound can be generalized from the arithmetic average to the weighted geometric average, i.e.,
Py (abll)) = p(ch)=))" ((ah)=a)'' Pa(ah a.j =d) a p(abl=))" ((ahl=))'->
In this case, the estimated probabilities need be normalized:
p(z)1=>)) ((x))2)))
Σx; ((zj p( ((
((al=))a. Eeg p()))=)) - )(
pa(=tleg 2d)
Applying the weighted AM-GM inequality, we obtain
Σ ((ril=)) (1 < Σ a ((an))))) + (1 a)) p(ri)=d)
#
p(zz (2)) +
(1 (1 a) Σ p( (xt =
Analogously, we obtain
Eplelldd) ((rik=))10 < 1.
18

---

Published as a conference paper at ICLR 2024
Plugging this into the definition of C, obtain
P z)) Pa(a zd)
C = (p
log
p zg)) 2(0
p 22)" -(an) 12))"
((abla2)) ((
Σ p( (( E.j
log
b(ttlz)) ((
log Σ p(aik)g) ((akk)) Eolalleda p( (2)
2 0.
Therefore, is lower bounded by the constant zero if p( =d) (and analogously for p(at lay 22)) is estimated as the weighted (either arithmetic or geometric) average of ((0) and p(ag z)).
D GRADIENT OF LCLIP W.R.T. REPRESENTATIONS AND OPTIMIZATION DETAILS.
Denote
LLr -E =(
log
1 ΣK (egiet)) (a ((tll=g + (1- ((x)(=2))
K
LCu -p -r (()(()))
log
I K
k E
(egiee)) (a p(aB + (1 p(
Note that the gradients of LEPP w.r.t. anchor =yg positive = and negatives 2/ (i * 1) have the following form:
Σ (()))))) Ox(>||=!)
aL 0z1
0zh
E ( (a (( + a) p(z)1=())
ER ef( (a p( + (1 a) f(zB)=)) 25())
Of 2n)
Ozg
ER ( (a (( (1 a) p(zb1zt)
( a
0=3
E B(te)=)aa)mt3
dz)
EK ((()))t)) (a ((=b(=))))) + a) p(al(=1)
el(sgied) (a p(zb1=)) + (1 a) p(abl=))) )((
0f(e),2) 2d)
+
dzd
ΣK ef(b)i=)) (a ) (1 a) ((++l2))
E 0L
0=2
+ (1- - at) ΣK (((s)izt) (a +(
(+
(f(x (b) + pi o
027
K -
(( (a pl |z)) + (1 - a) (x
aL 0z
aL and
a1
Similar formulas can be derived for
dzh
19

---

two views of Published as a conference paper at ICLR 2024
The second line of each gradient formula contains (A) and (.((=)=a)/ (and analo- gously (f(=); =02/ and (f()).=))/08 Thus it represents a standard component that optimizes the representations in similar manner to CLIP, i.. making the anchor and positive closer each other and the anchor and negatives farther away in the latent space, but to different extents ac- cording the reweighting factors in InfoCORE. The first line of each gradient formula contains Op(x)) (p(())=))// /02) and 0p(a)]xa)/0zz Thus it is a competing component that optimizes the representa- tions so that their predictive power of the batch identifier becomes lower. Both components drive the representations towards a reduced batch effect.
This allows us to employ a hyperparameter tributes to the gradient update. We follow a similar practice as in the gradient reversal layer by Ganin & Lempitsky (2015). However, instead of reversing the gradient of p( =0)) w.r.t zy (and that of P(ag (2)) w.r.t 2.), we preserve the sign of the gradient while weighting the magnitude by A.
regulate the extent to which each component con-
During training. we update the encoders and the classifiers iteratively. To be more specifie, in each iteration, we first update the parameters of the encoders (i.c., the drug structure encoder (:ca(:;0 fa) and the drug screens encoder Enc,(-: (g))) optimize the latent representations hased Leup, while keeping the classifiers' parameters fixed. Then we fix the encoders and update the parameter of the batch classifiers p(x (z/)) and (z)) by optimizing the batch classification loss LCLF.
E REVIEW OF CONTRASTIVE LEARNING AND THE INFONCE OBJECTIVE.
Contrastive learning (Tian et al., 2020; Oord et al., 2018; Bachman et al., 2019), self-supervised learning approach, aims to learn a representation space of high-dimensional data. Take the uni- modal contrastive learning of images as an example. Two views of the same image are generated to form "positive pair" via data augmentations, one of which serves as the "anchor" and the other serves the "positive". Meanwhile, to avoid representations collapsing to a single point, "negative" samples, which are views of different images, are included to form "negative pairs" with the anchor.
The success of contrastive learning has been connected to maximizing a lower bound of mutual in- formation between the observation X and the representation Z, which is lower bounded by [ZZ Z based on the data processing inequality, where Z and Z are latent representations the same observation X. Since mutual information is the KL divergence between the joint distribu- tion and the product of marginal distributions, as proposed by Oord et al. (2018), we can maximize lower bound of mutual information by optimizing the InfoNCE objective:
(a
LInfONCE Exp(e.s (p((2)0)
log
E (x(
K
where {= K negative samples sampled i.i.d. from the marginal distribution p(Z), i.e. latent representation of random images different from the anchor. Cosine similarity between z and z is usually used for the critic function f, projecting all the data observations into the representation space of a unit hypersphere.
In the case bi-modal contrastive learning, the InfoNCE objective is generalized two analogous terms, with one modality serving as the anchor and the other modality serving as the positive and the negatives. For instance, CLIP (Radford al., 2021) learns representations of paired texts Tj and images I, denoted as Z/11 and / using the following loss:
(x
(+ log R
|- log K Σ (=))es))
E Be(ehaellocep
+ TE PPpce
EK e (zt).
In the presence of a sensitive attribute X,, CCL by Ma et al. (2021) was proposed using the condi- tional contrastive learning objective tolower bound the conditional mutual information and reduce the information of the sensitive attribute from the representations by taking X, as the conditional variable. This leads an objective resembling IIMINNCE but with expectation over the conditional distributions:
f
LeCL = Ep(t,=))) (p(2)
log
2i=z e
K
20

---

Published as a conference paper at ICLR 2024
where F& is the value of the sensitive attribute corresponding to the anchor- settorr-ositive pair (z, 2 In Lea,, negative samples {=
K are sampled i.i.d. from the condition marginal distribution
p(2 Xx mb). Similarly, this can be extended to the bi-modal case as follows:
(xb..))
e (ab.a))
B <p log R Di=1 (ll)).)))
+ *
log
2
K Ei= (oppitt)
F ADAPTATION OF MULTIMODAL CONTRASTIVE LEARNING TO DRUG SCREENING DATA
Most multimodal contrastive learning methods are developed in the computer vision domain. We in- troduce several adaptions to enable application to drug screening data. We apply these modifications to all the benchmark models used in our experiments.
Data Augmentation for Drug Screening. Data augmentations often play a vital role in contrastive learning (Tian et al., 2020; Chen et al., 2020; He et al., 2020). However, data augmentation methods are missing for drug screening data with few replicates for each perturbation condition. Often, the experiment of applying drug on given cell line with certain dosage and perturbation time is repeated only several times (typically 3-5 times). Based on the setup of drug screening experiments, propose three kinds of data augmentations, which we find to effectively improve representation quality in practice:
1) Adding Gaussian Noise. Gaussian noise is added to X, before it is input into the encoder. The level of noise is controlled by a hyperparameter Anoosec We set A'noise 0.5 for GE and O'noisc 0 for CP in our experiments.
2) Dirichlet Mixup. Mixup (Zhang al., 2017) generates a new data point by the weighted sum of existing samples. Given replicates of each experimental condition, we use mixup among the replicates; this helps filling in the support in high-dimensional data. To achieve broader coverage, we utilize Dir-mixup (Shu et al., 2021) to account for the multiple-replicaes scenario. More pre- cisely, we generated augmented samples according to the weighted average of the 3-5 replicates of each experimental condition (drug, cell line, dosage, etc.), and we samples the mixup weights from symmetric Dirichlet distribution with hyperparameter Odiir: w Dirichlet(aai)). We
set Adii = 0.6 for GE and Adir = 0.8 for CP in our experiments.
3) Dropout (Masking). Random dropout is conducted on the input data, which corresponds to masking expression level of certain genes, or values of certain features. Dropout proportions Oldrop are set 0.1 for both datasets.
To Accommodate for Multiple Cell Lines. In GE, each drug is screened on multiple cell lines, and drug effect varies across different cell lines. To learn a universal drug representation that gathers information from all cell lines, make two modifications that directly use the average across cell lines or ignore the cell line labels.
existing methods (Jang et al., 2021)
Firstly, after applying molecular structure encoder that takes molecular structure input gen- erate drug embeddings, we utilize the cell line specific linear projection heads to map the universal embeddings into context-aware latent space corresponding to the phenotype's cell line context. To be more precise, let X denote the drug screening profile in cell line and Z = responding representation. Then the cell line contexualized drug representation Zž is calculated via (Xa)) and both representations are optimized through the objective in Equation 4.
the cor-
Secondly, since drug screening data is typically dominated by cell line effect, if we randomly sample data in the training batch, cell line related features will distract the model from learning drug-specine information. Therefore, we propose to sample training batches from experiments on the same cell line, that the model can focus more on the difference in terms of drug structure. In practice, to achieve a good balance, we use training batches that iterate between a normal randomly sampled batch and batch with all samples from the same cell line.
Additionally, we use the MoCo (He et al., 2020) framework for InfoCORE and CLIP and adapt it to multimodal scenario with multiple cell lines. Two main adaptations as follows: 1) maintaining the queue of momentum encoder outputs separately for each modality, and similarly the queue of batch distributions output by the momentum classifiers in InfoCORE; 2) constructing
21

---

Published as a conference paper at ICLR 2024
cell line specific momentum queues for each of the nine cell lines, and extracting negative samples from cell line specific queues when the training batch has all samples from the same cell line (as discussed above). For CCL, in drug representation experiments, since the MoCo framework requires individual queues for each conditioning variable, i.e. each batch number, and the total batch number large (97 for CP and about 1000 for GE), it is intractable in practice. We the SimCLR (Chen et al., 2020) framework to overcome this. For better comparison, we also report the performance of CLIP and InfoCORE using the SimCLR framework in Appendix I. In the representation fairness experiments, since there are only 4 subgroups, we construct queues for each subgroup and use the MoCo framework directly.
To showcase the effectiveness of these designs, we conducted an ablation study by removing each component from the CLIP model and calculating the retrieving accuracy in the GE dataset. The components include data augmentation (aug). the momentum contrastive learning framework (MoCo), the cell line specifie projection heads (cellproj). and training batches iterating between random batch and cell line specific batch (train bycell). The results are shown in Table For comparison, we also report the performance of the vanilla CLIP model with none these adaptations. Note that all the results of CLIP and CCL reported in Table 2 and Table 3 are based on the modified model instead of the vanilla model, thereby providing a fair comparison reflecting the effectiveness of our InfoCORE algorithm to removing batch-related biases for multimodal molecular representation learning.
Ablation Study.
Table Retrieving accuracy of different methods for drug screens based on gene expression data.
Retrieval Library
whole
batch
Top N Acc (%)
N=1
N=5
N=10
N=1
N=5
N= N=
InfoCORE
6,48
19.13
27.53
14.16
33,93
47.15
CLIP
6.01
18.64
27.16
12.40
30.47
42.77
CLIP (wfo MoCo)
5.90
18.06
26.56
11.87
29.41
42.16
CLIP (w/o aug)
5.25
16.00
23.99
12.25
29.74
41.88
CLIP (wfo cellproj)
5.39
17.36
25.77
11.62
29.58
41.78
CLIP (wlo train bycell)
5.02
16.45
24.90
10.74
28.16
40.66
vanilla CLIP
4.64
14.86
22.16
11.32
28.46
40.55
G FAIRNESS CRITERIA
Following the approach by Ma et al. (2021), we use three different fairness criteria: equalized odds (EO), equality of opportunity (EOPP), and demographic parity (DP). Denote as the label and as the model prediction, which both binary variables. Denote Z as the group identifier of sensitive attributes (also binary, considering the case of one binary value as the sensitive attribute). The definitions are provided in the following list:
1) Equalized odds (EO): EO calculates the sum of the difference (in absolute value) of the true positive rate and the false positive rate of the model predictions between two groups:
EO IP( Z 0, 1) -P( 1 1,1= })|+|P(( 1|2 0, 0) -)-P((-==--------p = 1 1, 0)).
2) Equality of opportunity (BOPP): EOPP calculates the difference (in absolute value) of the true positive rate of the model prediction between the two groups:
EOPP IP(Í 1 0, 1) P 1 = 1,1= 1)).
3) Demographic parity (DP): DP calculates the difference (in absolute value) in model predictions between two groups:
DP 1 0) - P(f 1|2
These criteria can be calculated based on the above definitions for the case with single binary sensitive attribute. We further adapt them to setting where multiple subgroups exist and thus differences can be calculated for each pair of subgroups. Following the approach in Zhang et al. (2022), we calculate EO, EOPP, and DP for each subgroup pair, and then aggregate the pairwise values by taking the average. We observed in our experiments that using the maximum or median leads to similar results.
22

---

Published as a conference paper at ICLR 2024
H DETAILS OF EXPERIMENTAL SETUP
Simulation Study. We generate the simulation imput data according to the graphical model in Fig- ure 2. To be more specific, we randomly assign real effect identifier (1-5) and a batch effect identifier (1-25) to each of the 1250 samples. Each real effect category and batch effect category is represented by a 10-dimensional vector randomly sampled from the multivariate standard normal distribution. To introduce noise to the data, we also generate a 10-dimensional random Gaussian vec- tor for each sample. Concatenating the three vectors into one 30-dimensional vector, and inputting it into two different random 2-layer Multi-Layer Perceptrons (MLP) corresoonding to the two modal- ities, we obtain the simulated observations as outputs of the neural network. The simulated dataset is then split, with half utilized for training and the remaining half reserved for visualization in the 2-dimensional representation space. We use the same encoder architecture, i.e. a 3-layer MI.P with hidden dimension size 128, for all the different models. For InfoCORE, the weighting hyperparam- eter a is set to be 0.09 and the gradient adjustment hyperparameter À is set to be 0.1.
Representation Learning of Small Molecules. We use L1000 gene expression profiles (Subrama- nian et al., 2017) (GE) and cell imaging profiles obtained from the Cell Painting assay (Bray et al., 2017) (CP) as pretraining datasets. In GE, 19,81 I small molecules were screened across 77 cancer cell lines, and the drug effect was measured using 1,1000 profiles (i.e., the expression of 978 land- mark genes was measured), with most data coming from cell lines. We use the data from these nine cell lines. Since for most drugs only one perturbation dosage and perturbation time is available per cell line, for simplicity, we drop the few samples with multiple dosages or perturbation times. This results in 17,753 drugs and 82,914 drug-cell line pairs. We conduct standard batch correction preprocessing step by normalizing the gene expression vectors by the mean over the control groups in the same batch and then use this model input. In CP, 30,204 small molecules are screened in one cell line (U2OS). We use the hand-crafted image features obtained by the popular CellProfiler method (MeQuin et al., 2018), which gives rise 701-dimensional tabular features after dropping missing values. The chemical structures are featurized using Mol2vec (Jaeger et al., 2018), which leads to -i0-dimensional vector features. For both datasets and all the models, we use 3-layer MLPs as both the gene expression cell imaging encoder and the molecular structure encoder, and we use 256-dimensional representation embedding. For InfoCORE, we use 2-layer MLPs as the classifiers and set the hyperparameter a to be 0.33 for GE and 0.83 for CP, and À to be 0 for both datasets.
For the molecule- peenooyyppp retrieval task, we calculate the representation embedding for each data sample in the validation set. Additionally, we calculate the drug structure representation of each molecule in the retrieval library (whole or batch, as discussed in Section 3.2). Then, for each data sample in the validation set, we rank the drugs in the drug library according to the cosine similarity between the drug structure embedding and the embedding of the data sample. Top N accuracy is then calculated based the rank of the ground truth drug structure. In GE, the retrieval accuracy is calculated within each cell line, and we report the average over all cell lines. In CP, it is calculated on the U2OS cell line.
The retrieval library whole contains all molecules in the held-out set; so the accuracy among whole captures the model's general ability pair drugs with their effects. The retrieval library batch contains held-out set molecules that have the same batch identifier as the retrieving target drug; so the accuracy among batch reflects the model's ability to distinguish drugs when the spurious features of the batch confounder are not available. Consider the extreme case: if a model only learns batch- related features, it can still have good accuracy in whole by randomly selecting drug candidates in the correct batch, but its performance in batch will be poor. Therefore, accuracies over both libraries as a whole reflect the model's ability of molecule-phenotype retrieval for drug discovery and drug repurposing.
In the downstream property prediction task, we follow the standard scaffold splitting procedure as suggested by Hu et al. (2020) for all classification tasks, and conduct scaffold splitting for the regression task of the PRISM dataset. For all pretrained models by different methods, we conduct a hyperparameter grid search for the learning rate and training epochs of the finetuning stage, select the ones with the best performance in the validation set (AUC for classification and R2 for regression), and report the mean and standard deviation the test set over random seeds.
Representation Fairness. For this, we use three fairness datasets with tabular features for binary classification: UCI Adult, Law School, and Compas. Race and gender features are binarized and
23

---

Published as a conference paper at ICLR 2024
Accuracy
purity (%)
Roal Elfiet
60-
78
68
Entropy of batch mixing
By Batch Billeot
16
CLI
10-
InfoCOAU
CLIP
CCL
InfoCORE
Figure 4: Visualization and quantitative metrics of the drug structure representations learned by different methods in the simulation experiment. The top row visualizations are colored by real effect category and the accuracy is calculated for the prediction of real effect category, while the bottom row uses batch identifier for both coloring and accuracy calculation. The farthest right column shows the KNN purity and entropy of batch mixing for different values of k.
combined to form the protected attributes, separating the data into four subgroups. Considering that minority groups often encounter data limitation issues, we subsample the training data to mimic an unbalanced population. To be specific, the training data consists of samples from one of the groups (white, male), (white, female), (black, male), (black, female) at the proportion of 10:2:2:1. At the pretraining stage, Gaussian noise is added as data augmentation for the tabular features. We use the MoCo framework for all three models. At the evaluation stage, procedure by Lahoti et al. (2020) and the approach taken by Ma et al. (2021) to freeze the encoder and train small additional network (2-layer MLP) for label classification. Various fairesss criteria are then measured based on the model predictions, with mean and standard deviation reported across 5 seeds.
follow the data splitting
I ADDITIONAL RESULTS.
Simulation Study. Analogously to Figure 3, the latent representations Za for the molecular structure and the corresponding quantitative metrics are shown in Figure 4. Similar to Xu et al. (2021a), we calculated three quantitative metrics evaluate the learned representations, including prediction accuracy, KNN purity, and entropy batch mixing.
To calculate prediction accuracy, linear classifiers trained predict either the real effeett cate- gory or the batch identifier based on the latent representation of drug screens and drug structures. Average accuracy and standard deviation are calculated based 5-fold cross-validation. KNN pu- rity calculates the average ratio of the intersection of the K-nearest neighbors based on the original data and the learned representation in each batch. It measures how well the representation retains the original structure of input data in each batch. Entropy of batch mixing calculates the average entropy of the empirical batch frequencies for the K-nearest neighbors of each drug. measures whether drugs from different batches are well mixed in the representation space.
Table 6: Prediction accuracy (%) of the representation to real effect category and batch identifier.
Representation
Drug Screens
Drug Structure
Target
Real Effect
Batch Identifier
Effect
Batch Identifier
CLIP
44.2(7.0)
8.6(2.6)
61.4(11.7)
14.4(4.5)
CCL
45.1(4.6)
7.8(1.9)
48.0(10.2)
11.7(3.4)
InfoCORE
73.3(5.0)
5.6(3.6)
77.8(6.0)
5.4(2.4)
The same results apply for the simulated molecular structure representations. In the visualization, CLIP representations cannot fully mitigate batch effect, CCL representations fail distinguish real
24

---

Published as a conference paper at ICLR 2024
effect categories, while InfoCORE is able recover the 5 real effect categories well. In terms of quantitative evaluation, InfoCORE outperforms CLIP and CCL in all metrics, especially in predic- tion accuracy of real effect category and KNN purity. The only exception occurs in the entropy of batch mixing for drug structure representation, where CCL has slightly higher entropy than CLIP and InfoCORE. However, the KNN purity of CCL. is much worse. Considering the trade-off of these two metrics, CCL learns better batch-independent representation, but fails to preserve the original data structure (i.e. real effect). The numerical results of prediction accuracy are shown in Table 6.
Drug Representations, As discussed in Appendix F, We use the MoCo (He et al., 2020) framework for CLIP and InfoCORE, and the SimCLR (Chen et al., 2020) framework for CCL (since MoCo requires individual queues for each conditioning variable and is not tractable for the drug screening datasets we experiment with). For better comparison with CCL. using the SimCLR framework, we also run CLIP and InfoCORE with the SimCLR framework.
Table 7: Retrieving accuracy of different methods using either the MoCo or the SimCLR framework for gene expression and cell imaging screens and standard deviations over random seeds.
Dataset
Gene Expression (GE)
Retrieval Library
whole
batch
Top N Acc (%)
N=1
N=5
N=10
N=1
N=5
N=10
Random
0.03
0.13
0.27
1.58
7.90
15.81
CLIP (MoCo)
5.96(0.08)
18.59(0.07)
27.17(0.02)
12.23(0.17)
30.29(0.16)
42.63(0.17)
CLIP (SimCLR)
5.81(0.09)
18.26(0.18)
26.65(0.08)
11.97(0.13)
29.49(0.07)
41.87(0.26)
CCL (SimCLR)
1.93(0.10)
5,85(0.18)
8.37(0.04)
12.76(0.29)
32.39(0.16)
45.77(0.08)
InfoCORE (MoCo)
18.99(0.19) 6.39(0.16)
27.18(0.30)
14.03(0.33)
33.63(0.27)
46.78(0.38)
InfoCORE (SimCLR)
6.34(0.24)
18.80(0.33)
27.05(0.23)
13.88(0.18)
33.25(0.24)
46.38(0.20)
Dataset
Cell Painting (CP)
Retrieval Library
whole
batch
Top N Acc (%)
N=1
N=5
N=10
N=1
N=5
N=10
Random
0.02
0.08
0.17
1.59
7.97
15.94
CLIP (MoCo)
7.23(0.05)
20.95(0.28)
(8.89))334)
13.20(0.10)
37.78(0.36)
52.72(0.16 16)
CLIP (SimCLR)
6.99(0.16)
21.09(0.24)
29.16(0.29)
12.96(0.28)
37.58(0.59)
52.87(0.31)
CCL (SimCLR)
1.31(0.04)
4.93(0.16)
7.38(0.34)
13.20(0.23)
37.99(0.12)
53.13(0.30)
InfoCORE (MoCo)
6.93(0:26)
20.65(0.25)
28.22(0.14)
13.26(0.10)
38.50(0.35)
53.13(0.08)
InfoCORE (SimCLR)
7.29(0.29)
21.15(0.33)
29.16(0.30)
13.16(0.12)
38.24(0.36)
53.00(0.25))
Table 8: Performance of different methods using either the MoCo or the SimCLR framework on the molecular property prediction task.
Classification (ROC-AUC $
Reg (R² %)
Datasets
BBBP
BACE
ClinTox
Tox21
ToxCast
SIDER
HIV
Avg.
PRISM
Molecules
2039
1513
1478
7831
8575
1427
41127
3172
Tasks
12
617
Mol2vec
70.7(0.4)
82.9(0.7)
84.9(0.3)
76.0(0.1) 1)
74.4(0.5)
64.9(0.3)
77.7(0.1)
75.9
8.5(0.7)
(McCo)
73.5(0.4)
86.1(0.4)
89.6(2.1)
77.3(0.0)
75.7(0.6)
63.7(0.6)
77.7(0.6)
13.9(0.4)
CLIP (SimCLR)
73.2(0.8)
85.8(0.5)
88.5(2.7)
77.2(0.2)
76.0(0.1) 1)
64.6(0.1) 1)
78.9(0.7)
77.7
15.7(0.5)
GE
CCL (SimCLR)
73.0(0.8)
85.9(0.6)
90.5(1.0)
77.0(0.2)
75.8(0.2)
63.4(0.5)
77.5(0.9)
77.6
16.000.5)
InfoCORE (MoCo)
73.5(0.3)
86.6(0.3)
91.9(1.9)
77.4(0.4)
75.7(0.2)
64.8(0.6)
78.5(0.2)
78.3
14.8(0.1)
InfoCORE (SimCLR)
73.6(0.3)
85.8(0.8)
91.((2.1)
77.3(0.3)
76.1(0.2)
65.3(0.6)
79.1(0.2)
78.4
14.9(0.1)
CLIP (MoCo)
73.4(0.8)
85.2(0.4)
87.3(0.1)
76.4(0.1)
76.7(0.1)
64.8(0.6)
78.2(0.4)
77.4
16.2(0.2)
CLIP (SimCL.)
74.0(0.6)
84.2(0.4)
88.8(0.4)
77.0(0.))
76.7(0.2)
64.6(0.8)
79.1(0.4)
77.8
15.6(0.2)
CP
CCL (SimCLR)
73.7(0.5)
84.9(0.9)
87.7(18)
75.9(0.3)
75.7(0.4)
65.2(0.4)
79.3(0.3)
77.5
14.7(0.3)
InfoCORE (MoCo)
74.0(0.8)
85.0(0.2)
89,3(0.5)
76.6(0.1)
76.910.1)
65.2(0.1)
78.7(0.1)
78.0
16.2(0.3)
InfoCORE (SimCLR)
74.3(0.5)
84.9(0.2)
88.7(0.8)
77.1(0.3)
77.0(0.1)
65,3(0.3)
80.0(0.8)
78.2
16.3(0.2)
25

---

Published as a conference paper at ICLR 2024
We calculate the standard deviations of the retrieval accuracy for the molecule-phenotee retrieval task based on random seeds. The results of retrieval accuracy for all models are reported in Table 7 and the results of property prediction reported in Table 8.
Representation Fairness. Results of DP with different methods over datasets are shown in Ta- ble 9. InfoCORE outperforms CLIP and CCL in terms of DP.
Table 9: Performance of various methods representation fairness task, reported by DP (in per- centage values).
Method
UCT Adult
Law School
Compas
CLIP
13.1(0.4)
18.4(0.9)
9.5(1.1)
CCL
13.2(1.1)
16.7(0.9)
8.9(1.7)
InfoCORE
12.4(0.5)
15.7(1.8)
8.2(0.8)
J VARIABLES DEFINITION
In this section, We provide a summary of the definitions of all variables used our paper. Variables used in the angle-sample setting in Section 2.2.1 are defined in Table 10. They expanded the salti-sample setting in Section 2.2.2, as defined Table 11. The definition of anchor, positive, and negative samples in contrastive learning is reviewed in Appendix E.
Table 10: Definition of variables in the single-sample setting
Variable Definition
Xa
Na D C Tha is the molecular structure of drug.
Xg E C R is drug screens data reflecting the phenotypic change of cells induced by applying the drug
X,
(eg. gene expression and cell imaging).
X, E B a categorical variable representing the experimental batch identifier of the drug screening experiment.
X
Za Enca( (a) is the representation of the drug structure;
Za
0. is the drug structure encoder parameter.
Zy = Enc, (Xg: 0g) is the representation of the drug screens data; Zg
0, the drug screen encoder parameter.
Table 11: Definition of variables in the multi-sample setting
Variable
Definition
Z
The drug structure representation of the mnchor-poittiv pair,
The drug screens representation of the anchor-positive pair.
X
The experimental batch identifier of the mnchor-positiee pair.
Z (2: K}
The drug structure representation of a negative sample.
Z E (2: K}
The drug screens representation of a negative sample.
26


Reviews:
[{"review_id": "Ob90CHfX6w", "reviewer": "Reviewer_g7wR", "summary": "The authors tackle the problem of removing batch effect in high-throughput drug screening.\nThe proposed method InfoCORE leverages a tractable information theoretic framework by establishing a variational lower bound on the conditional mutual information of some latent representation (obtained by a NN) given a \"batch identifier\" (i.e. the batch effect).\nThe method is rather generic in the sense that it can handle multiple modalities such as gene expression levels, imaging profiles, along with the chemical structure of the drug at hand.\nUnlike many methods that try to learn a generic representation, as a pre-processing task, InfoCORE trains end-to-end for a final task (e.g. molecular property prediction, molecule-phenotype retrieval).\nWhile very closely related to existing frameworks such as CLIP and CCL, InfoCORE's main contribution is to adaptively reweigh negative samples (in a contrastive learning sense) in the objective, so as to make better use of low sample sizes for rare conditions.\nWhile adaptively adjusting the sample weights, the resulting approach is proven to be optimizing a lower bound for the conditional mutual information.\nLast, the method is benchmarked against its very related methods CLIP and CCL on simulated data, representation learning of small molecules and on standard UCI datasets for fairness of the learnt representation.", "questions": "1) Can the authors comment on Proposition 3? Is this bound tight?\n2) At the end of Section 2.2.2, the authors make the interesting comment that \"once the batch effect has been mitigated, InfoCORE implictly adjusts and ceases to reweigh the negative samples\". Is it something that has been observered in practice? It would be great to show some evidence of it.\n3) One important motivation for the derivation of InfoCORE is that it should be more robust to settings where some conditions feature few observed samples. Do you have further experimental results that characterizes how this is the case (against CCL for instance)? For instance, it could be interesting to control the number of samples from one or a few \"rare\" conditions, or at least to show the performance on test samples from such rare conditions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "While building on already very established work, InfoCORE is addressing an important shortcoming of InfoNCE and the high variance of the sample approximation of the log partition function.\nThis shortcoming is particularly problematic in the scenario where samples are coming from different batches and some of them have few observations, which is relevant to many biomedical applications.\nThe proposed approach seems to be well motivated and derived.\nWhile I am not entirely familiar with this literature, the authors seem to have satisfactorily put their contribution in the context of the already abundant literature of extensions of InfoNCE.", "weaknesses": "Overall, the experimental part of the paper is a little bit underwhelming for a paper that claims to be rooted deeply into molecular representation for drug screening data, which is a crowded area.\nThe benchmarks only feature other generic approaches (i.e. CLIP and CCL) and the experimental details are often lacking to ensure a proper reproducibility of the results.\nIn fact, the initial simulation study has very little to do with drug effect, as the data used there is Gaussian noise (from various, controlled mixtures).\nThe results of this first part are purely qualitative (plots of low-dimensional representations of the learnt embedding).\nAs such, I don't find them particularly convincing compared to CLIP or CCL and it would have been nice to have quantitative metrics showing that InfoCORE representations preserve the original structure of the data (e.g. neighbourhoods) within batches but mix them well across batches after integration (see for instance the scANVI paper from Wu et al. 2021).\nThe second part tackles a specific (somewhat unusual maybe?) task of identifying drugs that are likely to induce a given effect (I find that the details describing this task to be lacking even after going through appendices).\nThe quantitative results there do not seem to be significantly better than other methods and the experimental setting seems a bit convoluted as it requires adapting the methods benchmarked against in a non-trivial way.\nIn all tables, most values written in bold do not seem to be significantly better. This is a bit problematic, IMO.\nPerhaps more importantly, the paper avoids addressing what exactly is intended by \"batch effect\" or \"confounding factors\".\nNotably, the treatment of the fact that LINCS data features drug screening of different cell lines is (it is considered a batch effect?) is relegated pretty far into the appendices while I think it is a crucial question.\nIn general, although a lot of other methods operate similarly, I find that methods that predict drug perturbation without modelling explicitly the state of the cells those perturbations are measured on to be of limited impact (impossible to extrapolate to cells that are not coming from the exact cell lines present in LINCS, let alone on real tissues).\nAt the end of the day, I think InfoCORE presents interesting novelties and seems to be sound, but I am not entirely convinced it is a great match for drug screening data.\nPerhaps the interesting analogy with learning fair representation could be the more important application (where the experimental results also seem to be superior)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "tL3ShcrUK1", "reviewer": "Reviewer_JEFo", "summary": "This paper studies multimodal molecular representation learning. The authors analyze the confounding batch effects to the learned molecular representations, and adopt a conditional mutual information maximization objective among the modalities, called the Information maximization approach for COnfounder Removal (InfoCORE). InfoCORE also reweighs the negative samples in the InfoNCE objective differentially based on the posterior batch distribution, to resolve the issue of insufficient negatives.", "questions": "1. The paper is hard to follow:\n- **All of the variables** in the paper have not been defined;\n- The problem is undefined;\n- The bias issue caused by the batching effects seems to be interesting, but does it really exist? Are there any realistic examples for the biases?\n- What does it mean by claiming X_b as a irrelevant attribute? \n- Why does equation (1) “emphasizes the shared features of the two modalities that are unrelated to batch” and “thus emphasizing the bioactivity of a drug”?\n- What is the exact objective of InfoCORE?\n- What is the algorithm of InfoCORE?\n- The information and the caption of Figure 3 is hard to read. It seems the batching identifier has not influence to the learned representations for all methods.\n\n2. The novelty is limited. It seems the main objective InfoCORE comes from CCL. What are the differences between InfoCORE and CCL?\n\n3. The improvements are limited. \n- In Table 2, the improvements of InfoCORE over CLIP or CCL is marginal while no standard deviations are given. Even the vanilla objective CLIP can outperform InfoCORE.\n- In Table 3, InfoCORE has little-to-no improvements over baselines, when considering the standard deviations.\n\n4. Some related works including debiasing or invariant molecular representation learning such as [1,2,3,4,5] have not been compared or discussed.\n\n\n**References**\n\n[1] Discovering invariant rationales for graph neural networks, ICLR 2022.\n\n[2] Learning causally invariant representations for out-of-distribution generalization on graphs, NeurIPS 2022.\n\n[3] Learning substructure invariance for out-of-distribution molecular representations, NeurIPS 2022.\n\n[4] Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism, ICML 2022.\n\n[5] Debiasing Graph Neural Networks via Learning Disentangled Causal Substructure, NeurIPS 2022.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(+) The studied problem is critical;\n\n(+) The proposed method is interesting;", "weaknesses": "(-) The paper is hard to follow;\n\n(-) The novelty is limited;\n\n(-) The improvements are limited;\n\n(-) Some related works have not been compared or discussed;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "fCV9vdJDFh", "reviewer": "Reviewer_FJph", "summary": "The paper presents a way to remove batch effect in contrastive learning of multimodal data while drawing negative samples from its marginal distribution. The weighting adjustment for negative samples is semi-parametric and it ensures that the weighting for unlikely negative samples (w.r.t. a batch number) does not shrink to zero. The model is primarily established in the context of molecular representation learning, but could be generalized to the removal of sensitive information in representation fairness problems. In experiments, the proposed model is examined in both context and overall achieves a better performance compared to the benchmarks.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Paper is well written: review of prior work is given and comparison to prior work is clear, walkthrough of derivation is intuitive and easy to follow. Figures of the graphical model and computational workflow are presented in a clean fashion that aids understanding the framework. The method is adequately novel and addressed the negative sample supply problem in conditional contrastive learning: the framework is capable of conducting training with batch data randomly sampled across different experiments. Experiments are done thorough.", "weaknesses": "1. The computational considerations for $L_{CLF}$ should be illustrated. $L_{CLF}$ by itself is not tractable without the true sampling distribution $p(x_b | z_d)$ and $p(x_b | z_g)$, and is not exactly a classification loss. Without any clear statement that you're replacing $L_{CLF}$ with conditional density learning through classification loss, there isn't any direct justification for factoring $ \\mathbb E_{p(z_d^1, z_g^1, x_b^1)} \\log \\frac{(\\hat{p}(x_b^1 | z_d^1, z_g^1))^2}{p(x_b | z_d^1)p(x_b | z_g^1)}$ into $C - L_{CLF}$.\n\n2. Line 2 and line 3 of Table 1 should be merged, or line 3 should be just removed. Reweighting of samples is an adjustment you have to make for large supply of negative samples to work, it is not a standalone contribution or an advantage. The current formulation of the table is slightly misleading.\n\n3. Figure 1. does not demonstrate the classification workflow clearly. The figure and the yellow coloring on the right grid makes it seem like you're optimizing $M_{ii}$ w.r.t. batch identifiers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "DKkpe6SkzX", "reviewer": "Reviewer_Tyeq", "summary": "* A molecular representation learning method, InfoCORE, is introduced to mitigate the presence of confounders, such as batch effects, in multi-modal contrastive learning.\n* Theoretical results were provided, showing that InfoCORE maximizes the variational lower bound on the conditional mutual information of the representations given the batch identifier.\n* Empirical experiments were conducted to evaluate the proposed method along with baselines on one simulated dataset and two real drug screening datasets for two downstream tasks, including molecule-phenotype retrieval and molecular property prediction.\n* A real-data experiment was presented to demonstrate that InfoCORE can be applied beyond drug discovery, improving fairness metrics when learning representations with sensitive attributes.", "questions": "What is driving the algorithmic adjustment described in the computational considerations, what is the risk of doing so, and how the choice of the hyperparameter will impact the final results?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed method is an extension of CLIP, an unconditional multi-model contrastive learning method, and CCL, a conditional contrastive learning method that does not address the confounder issue, thus tending to overcorrect for biological effects. The use of reweighting the negative samples differentially based on the posterior batch distribution solves the challenge of poor generalization in the case of limited negative samples for CCL. The results show a clear improvement in the metrics for both simulated data and real data tasks.", "weaknesses": "* It is not clear what is driving the algorithmic adjustments described in the Computational Considerations section, and how the choice of hyperparameters will impact the final results. \n* In the simulation data experiment section, it can be relatively difficult to evaluate the quality of the learned representation by examining only the first two principal components since the variance in other dimensions can be concealed. Visualization through manifold learning techniques such as t-SNE or UMAP, or evaluation using various metrics, can provide more reliable insights.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)