Learning Molecular Representation in a Cell
Authors: Gang Liu
Srijit Seal
John Arevalo
Zhenwen Liang
Anne E Carpenter
Meng Jiang
Shantanu Singh
Venue: ICLR 2025 Poster (2025)
Keywords: Molecular Representation Learning
Drug Discovery
Cell Morphology
Gene Expression
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose InfoAlign, a method that learns minimal sufficient molecular representations from cellular response data, including cell morphology and gene expressions.

Abstract: Predicting drug efficacy and safety in vivo requires information on biological responses (e.g., cell morphology and gene expression) to small molecule perturbations. However, current molecular representation learning methods do not provide a comprehensive view of cell states under these perturbations and struggle to remove noise, hindering model generalization. We introduce the Information Alignment (InfoAlign) approach to learn molecular representations through the information bottleneck method in cells. We integrate molecules and cellular response data as nodes into a context graph, connecting them with weighted edges based on chemical, biological, and computational criteria. For each molecule in a training batch, InfoAlign optimizes the encoder's latent representation with a minimality objective to discard redundant structural information. A sufficiency objective decodes the representation to align with different feature spaces from the molecule's neighborhood in the context graph. We demonstrate that the proposed sufficiency objective for alignment is tighter than existing encoder-based contrastive methods. Empirically, we validate representations from InfoAlign in two downstream applications: molecular property prediction against up to 27 baseline methods across four datasets, plus zero-shot molecule-morphology matching. The code and model are available at https://github.com/liugangcode/InfoAlign.

Published as a conference paper at ICLR 2025
LEARNING MOLECULAR REPRESENTATION IN A CELL
Gang Liu', Srijit Seal?, John Arevalo?, Zhenwen Liang'
Anne E. Carpenter Meng Jiang', Shantanu Singh²
University of Notre Dame Broad Institute of MIT and Harvard {gliu7, zliang6, mjiang2 @nd. edu (seal, jarevalo, anne, chantanul aaraatannnnoroddnsssiiiiioo /bbroaddsssitttteerrr.
ABSTRACT
Predicting drug efficacy and safety in vivo requires information on biological responses (e.g., cell morphology and gene expression) to small molecule perturba- tions. However, current molecular representation learning methods do not provide comprehensive view of cell states under these perturbations and struggle to remove noise, hindering model generalization. We introduce the Information Alignment (InfoAlign) approach to learn molecular representations through the information bottleneck method in cells. We integrate molecules and cellular response data as nodes into context graph, connecting them with weighted edges based on chemical, biological. and computational criteria. For each molecule a training batch, InfoAlign optimizes the encoder's ' latent representation with a minimality objective to discard redundant structural information. sufficiency objective de- codes the representation to align with different feature spaces from the molecule's s neighborhood in the context graph. We demonstrate that the proposed sufficiency objective for alignment is tighter than existing encoder-baad based contrastive methods. Empirically. we validate representations from InfoAlign in two downstream appli- cations: molecular property prediction against up to 27 baseline methods across four datasets, plus zero-shot molecule-morphology matching. The code and model are available at htt.p:://github... http:/tt..oooo/oggggugooeoeeeee. ToaAll
INTRODUCTION
Drug properties, e.g., toxicity and adverse effects (Liu et al. 2023a), are induced by molecular initiating events-interaacie between molecule and biological system- that first impact the cellular level and ultimately influence tissue or organ functions (Mast et al. 2014). However, a chemical molecule's structure alone is insufficient information to predict its impact on cells: each chemical interacts with multiple cells and genes and induces complex changes in gene expression and cell morphology, making predictions of downstream responses challenging (Carpenter al. 2006 Moshkov et al. (2023). Hence, molecular representation learning should make use of information about cellular response, enhancing the representation of the mode of action and thereby improving predictions for downstream bioactivity tasks (Liu et al. 2023a: Wang et al. 2023a).
There is lack of exploration for holistic molecular representations from molecular structure, cell morphology, and gene expression (Hu et al. 2020a You et al. 2020 Liu et al.) 2002 Wang et al. 2023a; Sanchez-Fernandez al. 2023). For example, graph sell-supervised methods only manipulate molecular structures to perturb or mask molecular graphs using contrastive or predictive losses (Hu al., 2020al You al. 2020; Inae et al. 2023). Moshkov al.] (2023) explored the ability of different data modalities, taken independently. to predict molecules' assay activity in a diverse set of assays (tasks). They found (from (Moshkov al. 2023)'s Fig.2) that molecular structure supports highly accurate prediction (AUC> 90%) in 31% (16/52) of tasks, gene expression in 37% (19/52) and cell morphology in 54% (28/52). Similarly, in our experiments (Figure 3). we observe that molecular structure is not one-size-fitse ail solution.
Cells can be perturbed by treating them with chemicals or genetic reagents that disrupt particular gene pathway. These chemical and genetic perturbations in vitro naturally bridge molecules with cell morphology and gene expression, as illustrated in Figure (b). However, multi-modal
contrastive methods such CLOOMEE (Sanchez-Fermandzzeee et al. 2023) and InfoCORE (Wang et al..

---

Published as a conference paper at ICLR 2025
(a) Existing: Contrastive
(b) Proposed: InfoAlign with Information Bottleneck
BDruB
Call Encoder
Gene Encoder
implement
Cell Decoder
Drug Decoder
Gene Decoder
Drug Encoder
Not Shared Drug
Cell
Gene
Drug
Encoder
Minimal Sufficlent Objectives
Encoder
Chemical Perturbation
Min [ 0.5 (aj
Redundant Information
Bottleneck Information
MI: Infrrsitnoomtttinn
Figure 1: Comparison of Representation Learning Methods: (a) Existing contrastive methods use two encoders one for molecules and another for cell morphology or gene expression features-witte sharing the molecule encoders for different alignment targets. (b) InfoAlign remove redundant information from molecules, cell morphology, and gene expressions based the information bottleneck, resulting in more concise yet predictive molecular representations (Alemi et al. 2016).
2023a), depicted in Figure (a), focus primarily aligning molecular representations with cell
morphology (Sanchez-Fernandez Fernandez et al. 2023 Wang et al. 2023a) or gene expression (Wang et al. 2023ab. These approaches fall short in ways. (1) They do not remove redundant information, grey-colored area in Figure [ (b), that may harm representation generalization. The presence of redundant information (Wang al. 2023a) may induce spurious correlations, adversely affecting the generalization molecular representations. For example, in small molecule perturbations (Bray ctt al., 22166] Chandrasekaran al. 2023), batch identifiers can signify confounding technical factors, creating misleading associations between molecular structures and cell morphology (Wang et al.
2023a). (2) They treat molecules the sole connectors between gene expression and cell morphology, ignoring the potential for genetic perturbations (Chandrasekarau) et al. 2023) to directly establish connections between these modalities. Genetic perturbations (Chandrasekarail et al.] (2023) not only enrich the feature space of gene expression and cell morphology but also enhance the navigation of molecular representation learning towards the overlapped (bottleneck) area in Figure (b)).
To address the aforementioned challenges, we conceptualize the cellular response processes context graph, capturing a more complete set of interactions among molecules, gene expression, and cell morphology. We identify the neighborhood of the molecule on the context graph and apply the information bottleneck (Tishby et al. 2000) to optimize molecular representations, which aligns them with neighboring biological variables remove redundant information and improve generalization.
We propose the Information Alignment (InfoAlign) approach, as presented in Figure (b). InfoAlign one encoder and multiple decoders with information bottleneck for minimal sufficient statistics in representation learning. The minimality objective optimizes the encoder to learn the minimal informative representation from molecular structures by discarding redundant information. The suffi- ciency objective ensures the encoder retains sufficient information, allowing decoders to reconstruct features for biological variables in neighborhood areas of the context graph. We construct the context graph based on molecule and genetic perturbations (Bray et al. 2017] Chandrasekaran et al. 2023 Subramanian al. 2017) and introduce more biological (gene-gem gene interaction (Himmelstein et al. 2017)) and computational (cosine similarity) criteria to increase edge connectivity. We conduct random walks on the context graph, beginning with the molecule the training batch, to identify its neighborhood. Cumulative edge weights indicate similarity between the molecule and variables along the path. The molecule is encoded, and its latent representation decoded to align with features identified in the random walk. Encoders and decoders jointly optimized using upper bound for the minimality objective and a lower bound for the sufficiency objective.
The sufficiency objective introduces decoder-basse based bound for multi-modal alignment. We show its theoretical advantages by demonstrating that it provides a tighter bound than the encoder- based approaches used in previous contrastive methods (Oord al. 2018 Radford et al. 2021), discussed
2

---

Published as a conference paper at ICLR 2025
in Section 4.3 In experiments, InfoAlign outperforms up 27 baselines across three classification and one regression dataset, covering 685 tasks, with average improvements of up to 6.4%. InfoAlign also demonstrates strong zero-shot multi-moda matching on two wolecule-morpholgge datasets.
2 RELATED WORK
Representation Learning on Molecular Structure: Representation learning approaches for
molecules can be categorized into sequential-based (Krenn et al. 2022 Ross et al. (022) or graph-
based models (Hu et al. 2020a; You et al. 2020; Zhang et al., 202 Liu et al. 2023b). Sequential
models, utilizing string formats of molecules like SMILES and SELFIES (Krenn et al.
2022). have
evolved from Recurrent Neural Networks (RNNs) to Transformers (Chithrananda al. 2020 Ross
eet al.) 2022). These models typically follow specific pretraining strategies similar to language models such as BERT (Devlin et 2018). RoBERTa (Liu et al., 2019 Chithrananda al.) 2020) and
GPT (Radford et al. 2019). The pretraining targets are thus often the next token predictions or mask language modeling (Devlin et al. 2018, Chithrananda et al. 2020) on SMILES or SELFIES se- quences (Radford et al. 2019). Graph Neural Networks (GNNs) are the architectures for graph-based approaches (Hu al., [2200l] You et al.,
2020; Zhang al. 2021 Liu al. 2024b), where methods to pretrain GNNs often perturb mask the atoms, edges, substructures of molecular graphs with contrastive (Hu al., 2020a; You et al. 2020) and predictive losses (Zhang et al. [02T] Inae et al.) 2023). Recent evidence highlights the challenges of developing universal molecular representations based solely on molecular structures without integrating domain knowledge (Bray et al. 2016 Seal
al. 2022 [unn et al., 2022 Seal et al. 2023 Liu et al. (0024a). Although using motifs is common method to incorporate such knowledge (Rong et al.] [020]] Inae et al.) (023). the incorporation of information about molecules' biological impacts much less explored. We aim to enhance molecular representation learning by incorporating domain knowledge from cellular response data.
Representation Learning with Cellular Response Data: A primary goal of molecular representation learning is to predict molecular bioactivity. Likewise, emerging gene expression Subramanian et al. (2017) and morphological profiling approaches Carpenter al. (2006); (Seal al. (2024) that describe perturbed genetic or cellular states in cell cultures can also be used to predict bioactivity. In some datasets, molecules are the perturbations, and the perturbed cell states measured are gene expression values for a thousand or more genes (Subramanian et al.) 2017) and/or microscopy Cell Painting images, which can be represented as thousand or more morphology features Cimini et al. (2023). Recently created large-scale perturbation datasets (Subramanian et al., 2017: Chandrasekaran et al., 2023) could enrich molecular representation learning approaches. CLOOME
Semnchez-Ferandez
al. (023).. MIGA (Zheng et al. 22024) MoCoP (Nguyen et al. (2023), and MolPhenix (Fradkin et al. 2024) contrast cellular images with molecules. InfoCORE (Wang et al.] 2023a) contrasts molecule with either morphological profiling (Bray al. 2017) or gene expression (Wang al., 2023a). InfoCORE Wang et al. (2023a) mitigates confounding batch identifiers using a batch classifier, which may not be practical when batch identifiers are unavailable during training.
3 PROBLEM DEFINITION
We denote x E X as the molecule from the space X. An encoder with parameters pe(z x)) maps I to a D-dimensional latent representation z E RDD. One may implement a Graph Neural Network (GNN) (Xu et al. 2019) as the encoder. The GNN first updates node representations and then performs a readout operation (e.g., summation) over the nodes to obtain the latent representation.
Existing research has extensively used structural features to pretrain the GNN encoder (Hu al.. 2020a Inae et al. 2003). However, incorporating more expressive features from the cellular context, such cell morphology and gene expression, remains largely unexplored for improving molecular representations. In this work, we use these features as targets to optimize molecular representations.
MULTI-MODAL ALIGNMENT WITH INFOALIGN
We present the overall representation learning framework in Figure 2 In Section 4.1
we construct
the context graph for cellular response data. In Section 4.2 we introduce representation learning
3

---

Published as a conference paper at ICLR 2025
Path P
X1
Perturb. Interact
Perturb.
X3
Interact
Pertura.
Encoder
P
x1
Perturo
C1
Decoders m
X1
Interact
Perturb.
eg
X2
Internct
C2
C1
Perturb.
Shared ...
Shared
Path P Px4
Parturb.
Internct
E
CB
Xg
Casine
Perturb
X5
Interast
Parturh
Perturi
C3
Encoder
Xs
P
x4
Decoders
E5
Cosine
C3
(a) Random Walk on Context Graph
(b) Representation Learning over Walk Paths
Figure 2: Molecular Representation Learning Using the Context Graph: (a) In Section 4.1
construct the graph with various interaction, perturbation, and cosine similarities among molecules æ, cell morphology profiles and genes Given a training batch of molecules, including x1 and T.+ random walk extracts paths, for instance, of length four. (b) In Section 4.2 we aim to learn molecular representations based on the information bottleneck, preserving minimal information from the input molecule while ensuring sufficient information for decoding the target along the walk path P..
methods based on the principle minimal sufficiency for molecules and their related modalities. In Section 4.3 we demonstrate the theoretical advantages of the proposed method.
4.1 RANDOM WALKS ON CELLULAR CONTEXT GRAPH
Node Construction: We model the interactions of the molecule with other molecules, the cell and genes e using the context graph. They are nodes with different features Node features for
molecules binary vectors obtained using fingerprints (Rogers & Hahn (010). Cell morphology features are derived from CellProfiler (Carpenter et al.) (006) applied to Painting microscopy images. Gene nodes have expression values obtained using L1000 methods (Subramanian al.)
2017). We rescale the cell morphology and gene expression features to a range between 0 and 1.
Edge Construction: We link nodes using chemical, biological, and computational criteria. For example, molecules can perturb cultured human cells, inducing changes in cell morphology (Chan 2017). thus linking molecules to cell morphology and gene expression nodes. Genes could also perturb cells, inducing links between genes and cell morphology (Chandrasekaran et al. 2023). Additionally, We calculate cosine similarity for nodes of the same type and use biological criteria, such gene-gene interactions (Himmelstein (et al., (201)), to enrich edges. Each edge is assigned a weight w ranging from 0 to 1. For example, edges derived from computational criteria between molecule nodes are assigned weights based on the assumption that structurally similar molecules may exhibit similar biological effects, a concept widely used in drug discovery, such as lead optimization. We construct the context graph as detailed in Section 5 and appendix C.T with an example shown in Figure (a).
drasekaran et al. 2023) and gene expression (Subramanian
The context graph identifies related cellular response patterns for input molecules in representation learning. Given an input molecule x, we extract its neighborhood through random walks starting from æ. Specifically, we employ degree-based transition probabili-
Random Walk Path Extraction:
ties (Perozzi et al. 2014) and denote the walking path as Pa 11, 112
u2 w4, UL, where 122
is direct neighbor of z. To quantify the similarity between and node U2 (2 S i S It) on Pz, we
compute the cumulative product of edge weights a(v; P.) I1 Wy.
4.2 OPTIMIZATION FOR REPRESENTATION WITH INFORMATION BOTTLENECK
The information bottleneck (IB) Tishby et al. 2000 Alemi et al. 2016) is an appealing method for defining concise representations with strong predictive power. For molecular representation,
4

---

Published as a conference paper at ICLR 2025
extract minimal sufficient information from the random variable X of molecules. This is achieved by aligning the molecular representations Z with the targets Y. derived from node features along the walk path P. The IB has two principles based on mutual information (MI): (1) the minimality principle, which minimizes MI hetween molecules and their latent representations as T T(;; 2), and (2) the sufficiency principle, which decodes latent representations maximally reconstruct feature spaces for variables along the walk path / (Z; Together, these form the optimization objectives:
min [-1 (2;Y) + ()); p(a(=)
where 3 controls the trade-off between minimality and sufficiency. The exact computation of (Z: Y) and T(X; : 2) is intractable due to the unknown conditional distribution p(ylz) and the marginal p(z). We introduce the variational approximations (yylz) and q(z) for them, respectively. This results in a lower bound estimation for the first decoding term IDLB and upper bound for the second encoding
term LUU (Poole et al. (201)).
I(ZY)) 2 Eptz,) [log ( 2)] + HYY IDLB
(2)
T(X; 2) S ppt)) [KL (p(z
x) [)))
A IEUB
HYY) is the differential entropy. Proofs are in appendix B.1 Together, JDLB and TEUB upper bound Eq. ①. forming a tractable objective - DLLBB + TEUB to optimize the encoder. For the target Y. the IDLB objective requires decoders rather than encoders, as typically used in prior work (Sanchez, Fernandez et al. 2023). We use distinct decoders, denoted as da with parameters Φ, for various targets, including molecular fingerprints, gene expressions, and cell morphology features.
After ignoring the constant terms, one could formulate the loss function according to Eq. (2) for the molecule sample r, its latent representation and the targets yeu from Pa:
L C= -- - Σ a(r(Px) |- log (1o 22) B KT. (po(z x:) [W 7)), L
(3)
PEP=
where the first term aligns the representation with other features, and KL is the Kullback-Leiblr divergence used for regularization. (N, 7) is the Gaussian prior. In this formulation, the encoder models a distribution instead of a single representation z, learning the mean and variance μt, o E RR One may use parameterization tricks sample from the distribution (Alemi et al. 2016). The decoder then reconstructs Ve. the features of the neighboring node v E Ps from the walk path.
InfoAlign uses multiple decoders for qs to align multi-modal features, while prior work relies on encoders with CLIP-like losses to align the latent space (Radford et al. [021 Girdhar et al. 2023 Wang et al. 2023a decoder-based alignment alongside the empirical advantages in Section
al., 2023). Next, we provide the theoretical benefits of
4.3 THEORETICAL MOTIVATION FOR DECODER BASED ALIGNMENT
InfoNCE (Oord et al. 2018 is the contrastive loss used for most CLIP-like methods (Radford et al., 2021 Wang al. 2023a In this work, we show that the MI lower bound based on InfoAlign is tighter than that based on InfoNCE.
Proposition 4.1. For the molecular representation 7. and target Y (from cell morphology, gene expressions, or molecular fingerprints), the encoder-based MI lower bound TELB for InfoNCE can be derived by incorporating K additional samples, denoted as y2.K, build the Monte Carlo estimate ml()) of the partition function (Nguyen et al. 2010 Poole et al. 2019):
((22))
Maav
LELB = 1 1+ (p(2.))(
log
Ep(z)p(va (plv)
m(z; y, V2.k)
where 1((z,)) v) is the neural network parameterized critic for density approximation with the energy- based variational family. The decoder-based lower bound I is defined in Eq. ②2, then we have that IDLB is tighter than IELB, i.e., I(Z;Y) 2 IDLB(Z:Y) IELB(Z;Y).
Proofs are in appendix B.2 The result aligns with empirical observations in previous studies such as DALL-E 2 (Ramesh et al. 2022). where a prior model was introduced to improve representations
from CLIP (Radford et al. (021)) before decoding to another modality. In this work, we learn decodable latent representations from molecules to align with different biological features.
5

---

Published as a conference paper at ICLR 2025
5 IMPLEMENTATION OF CONTEXT GRAPH AND PRETRAINING SETTING
Data Source of Context Graph: We create the context graph based on (1) two Cell Painting datasets (Bray al. 2017; Chandrasekaran al. (023). containing around 140K molecule per- turbations (molecule and cell morphology pairs) and 15K genetic perturbations (gene and cell morphology pairs) across 1.6 billion human cells; (2) Hetionet (Himmelstein et al. 2017), which captures gene-gene and gene-molecule relationships from millions of biomedical studies; and (3)
dataset reporting differential gene expression values for 978 landmark genes (Wang et al. 2016) for chemical perturbations (molecule and gene expression pairs) (Subramanian al. 2017).
Node Features: Different profiling methods provide node features different ways. Morgan fingerprints (Rogers & Hahn 2010) are feature vectors extracted from each molecule's structure, CellProfiler features (Carpenter et al. 2006) are computed from the image of each cell and represent cell morphology. and L1000 profiles (Subramanian et al. (017)) capture gene expression values 978 landmark genes from cells treated with a chemical perturbation. Here are two practical considerations for the context graphs: (1) Chandrasekaran et al. (2023) provided one dataset that measured the cell morphology impacts of perturbing individual genes. The 15K genetic perturbations (Chandrasckaran al. 2023) provide gene-cell morphology pairs but lack corresponding gene expression profiles. Still, we keep the gene nodes from this dataset to account for potential gene-gene interactions and incorporate cell morphology features into them. (2) All 978 landmark genes have expression values linked to the molecules used in (Wang al. 2016). We update new gene expression nodes with 978-dimensional feature vectors. These vectors summarize all molecule-gene expression connections for small molecule perturbation. This approach efficiently reduces dense connections between landmark genes and molecules. We select the top 1% of gene-molecule expression values as new edges to enrich the context graph's 's connectivity. We scale cell morphology and gene expression features to range of 0 to 1 using the Min-Max scaler along each dimension.
For edges based chemical perturbations (Chandrasekaran et al. 2023 we assign the edge weight of 1. We also compute cosine similarity for nodes if they are in the same feature space (such as two cell morphology/gene expression profiles, or Morgan fingerprints). To avoid noisy edges from computations, we (1) apply a 0.8 threshold for cosine similarity, and additionally (2) explicitly enforce 99.5% sparsity by selecting top similar edges.
Edge Weights:
All together, this results in a context graph of 276,855 nodes (129,592 molecules, 4533 genes + 13,795 gene expressions, and 128,935 cell morphology) and 366,384 edges.
Encoder and Decoder: We use a five-layer Graph Isomorphism Network (GIN) (u et al. 2019)
with sum as the readout function as the molecule encoder. All molecules on the context graph are used to pretrain the encoder. Since we extract feature vectors as decoding targets in different modalities, we efficiently use a Multi-Layer Perceptron (MLP) as modality decoders. We set the hidden dimension to 300, 3 = 4, and T, = 4. In each training batch, random walks start from the molecule node to extract the walk path. The decoders are then pretrained reconstruct the corresponding node features along the path. Further details can be found in Section appendix C
EXPERIMENTS
We demonstrate the effectiveness of InfoAlign'' representation in (1) molecular property prediction, (2) molecule-morpholee matching, and (3) analyze the performance of InfoAlign. These lead to three research questions (RQs).
6.1 RQI: MOLECULAR PROPERTY PREDICTION
Better molecular representations should improve prediction performance. We train MLPs on different representations to predict molecular properties in both classification and regression tasks.
6.
EXPERIMENTAL SETTING
Dataset and Evaluation: We select datasets for important tasks in drug discovery, including activity classification for various assays in ChEMBL2K (Gaulton et al.] 2012) and Broad6K (Moshkov et al.) 2023), drug toxicity classification using ToxCast (Richard et al. 2016), and absorption, distribution, metabolism, and excretion (ADME) regression using Biogen3K (Fang et al. 2023). The dataset

---

Published as a conference paper at ICLR 2025
Table 1: Results on ChEMBL2K and Broad6K. We report average AUC (Avg.), as well as the percentage of tasks achieving AUC above 80%, 85%, and 90%. We highlight the best and second best mean. We also highlight the row of the best method in each category.
Dataset
ChEMBL2K (AUC f)
Broad6k (AUC f)
(# Molecule/ # Task)
(2355/4))
(6567 /22)
Method
Avg.
>80%
>85%
>90%
Avg.
>80% >85%
>90%
Morgan Fingerprints
MLP
76.8+2.2
48.8+3.9
34.6+6.3
21.9=5.7
63.3+0.3
6.3+0.0
1.4+1.7
3.1+0.0
RF
54.7=0.7
0.0=0.0
0.0+0.0
0.0+0.0
55.5+0.1
0.0=0.0
0.0+0.0
0.0+0.0
GP
51.0=0.0
0.0=0.0
0.0+0.0
0.0+0.0
50.6+0.0
0.0:0.0
0.0+0.0
0.0+0.0
Pretrained GNN
AttrMask (Hu al.
2020u
73.9=0.5
46.8+2.7
31.244.4
14.6=1.7
59.8+0.2
3.1=0.0
3.1=0.0
3.1+0.0
ContextPred (Hu et al.
2020a
77.0=0.5
55.1+1.3
34.124.6
14.6=1.7
60.0+0.2
7.5
3.1=0.0
3.1+0.0
EdgePred (Hu ct al.
2020a
75.6=0.5
54.2+4.0
34.6+7.2
12.2+2.4
59.9.0.2
3.1=0.0
3.1=0.0
3.1+0.0
GraphCL You et al.
2020
75.621.6
46.827.6
32.246.8
18.0=3.7
67.240.5
15.6+3.1
3.1+0.0
3.1+0.0
GROVER RRng al.
2020
73.3=1.4
38.5+2.0
22.443.6
14.6+2.4
66.240.1
15.6+0.0
3.+1.4
3.1+0.0
JOAO You ct al. 2020
75.1=1.0
47.845.1
33.742.0
19.0=3.2
67.340.4
12.5+0.0
38+1.4
3,1+0.0
MGSSL (Zhang al.
2021
75.1=1.1
39.0+4.6
29.343.0
10.3+3.2
66.9+0.5
13.8+2.8
3.1=0.0
3.1+0.0
GraphLoG Xu
73.5=0.7
41.9+2.0
29.3+3.4
15.6+2.8
62.9.0.4
4.4+1.7
0.0+0.0
0.0=0.0
GraphMAE Hou et al. 2022
74.7=0.1
33.2+1.3
27.8+1.3
12.2+1.7
66.8+0.3
14.411.7
3.1+0.0
3.1+0.0
DSLA Kim et al. 2022
69.3#1.0
23.9+4.7
14.625.5
6.8+1.1
63.340.3
6.3+0.0
3.1+0.0
3,1+0.0
UniMol Zhou et al. 2023
76.8=0.4
46.8+2.0
33.7+1.1
24.9=2.0
65.4+0.1
7,5=1.7
3.1+0.0
3.1+0.0
Pretrained Chemical Language Models
Roberta
Mary et al. 2024
74.7=1.9
46.343.4
35.1
22.9+1.3
59.8.0.7
5.0=1.7
3.1+0.0
3.1+0.0
GPT2 Mary et al. 2024
71.0=3.4
31.2=11.2
20.0+9.4
7.346.9
60.640.3
7.52177
1.91.7
1.921.7
MolT5 wdwards al.
2022
69.9=0.8
32.242.0
21.0:4.1
8.8+1.3
56.420.8
3.8+1.4
2.5+14
2.5+1.4
ChemGPT Frey et al.
2023
65.0=1.1
16.142.8
11.243.3
5,411.1
55.140.9
3.140.0
3.1+0.0
1.3+1.7
Cell Morphology
MLP
64.3+2.4
15.646.6
8.343.7
4.943.9
51.9&1.0
0.0+0.0
0.0+0.0
0.0+0.0
RF
55.9
3.9+1.3
3.941.3
2.4+0.0
55.3+0.1
0.0g0.0
0.0+0.0
0.0+0.0
GP
50.1
0.0+0.0
0.0+0.0
0.0+0.0
54.7+0.0
0.0+0.0
0.0+0.0
0.0+0.0
Gene Expression
MLP
56.1=1.1
5.1=1.4
3.4+1.3
3,4+1.3
56.9+1.4
1.9=1.7
1.9+1.7
1.91..7
RF
52.8=0.3
0.0=0.0
0.0+0.0
0.0+0.0
55.240.2
0.0=0.0
0.0+0.0
0.0+0.0
GP
Run out of time
50.1
0.0=0.0
0.0+0.0
0.0+0.0
Multi-modal Alignment
CLOOME
66.7+1.8
26.814.6
16.1+3.7
10.7=5.1
61.7+0.4
3.1=0.0
3.1+0.0
0.0=0.0
InfoCore (GE)
79.3=0.9
62.442.8
46,343.0
30.3+2.2
60.240.2
3.140.0
0.0=0.0
0.0+0.0
InfoCore (CP)
73.8=2.0
37.629.2
26.344.7
10.7+4.1
61.120.2
6.3e0.0
3.1+0.0
0.0+0.0
InfoAlign (Ours)
81.3=0.6
66.3+2.7
49.3+2.7
35.1=3.7
70.0=0.1
18.8+2.2
3.1=0.0
3.1=0.0
We apply scaffold-splitingg for all datasets. We follow a 0.6:0.15:0.25 ratio for training, validation, and test sets for all datasets. We use the Area under the curve (AUC) for classification and mean absolute error (MAE) for regression. Mean and standard deviations are reported from ten runs.
covers 685 tasks with details in Table[]and appendix D.I
Baseline: We include 27 basclines across six categories: (1) three molecular fingerprint (FP)-based methods (Rogers & Hahn 2010); (2) eleven pretrained GNNs; (3) four pretrained chemical language models; (4,5) six methods based on cell morphology and gene expression values from cells treated with each molecule; (6) CLOOME (Sanchez-Fernander et al.] (203)) and InfoCORE (Wang et al.) 2023ab for multi-modal alignment using structure, morphology. and gene expression data. We use MLPs, Random Forests (RF), and Gaussian Processes (GP) for methods in categories (1,4,5). Setting details and all results are in appendices) D.I
6.2.2 RESULTS AND ANALYSIS
We present results across various assays in Tables | and and and Figure 3
Key observations include:
(1) Molecular structures are superior compared to cell morphology and gene expression features for predicting various molecular assays. This is likely because the datasets and tasks selected
7

---

Published as a conference paper at ICLR 2025
Table 2: Results on ToxCast and Biogen3K. We report the average AUC and the percentage of AUC above 80% on ToxCast, and regression MAE (scaled by 100) for Biogen3K. We highlight the best and second best mean. We also highlight the row of the best method in each category.
Dataset
ToxCast (AUC 1)
Biogen3K (MAE x
x100 4)
(# Molecule/ Task)
(8576/ /17)
(3521/6)
Mcthod
Avg.
>80 %
Avg.
hPPB
rPPB
RLM
HLM
ER
Solubility
Morgan Fingerprints
MLP
57.6+1.0 1.6
66.2+2.4 66. 66.142.66
56.8+2.3
56,5+4.2
74.616.2
73.747.3
69.5+3.0
1 0.2+0.1
52.8+0.2
44,2=0.1 60.0x0.0
44.240.1
42,0+0.2
67.740.7
66.940.9
51.6+0.1
GP
Run out of Time
51.3=0.0 59.5+0.0
49.7+0.0
68.8+0.0
69.3x0.0
61.620.0
Pretrained GNN
AttrMask et
2020a
63.1
3.2+1.2 1.2
67.3+0.3
82.4#1.1
49.8+0.7
51.7+1.0
57.9
62.640.5
99.12
ContextPreS (Hu crain
2020a
63.0+0.6
1.3
68.5
85.0
49.7+0.4
55.1+2.7
614+1.
63.1+0.5
96.5+3.7
EdgePred (Hi
63.5+1.1
4.8+3.0
81.2=10.2
48.0+0.5
53.5+2.8
62.2+1.8
62.9+0.7
GraphCL Youl et
2020
52.2#0.2
0.5+0.3
53.940.6
43.8=0.3
45.420.6
40.6#0.5
76.7+1.0
67.142.2
49.620.3
GROVER Rone
2020
53.120.4
0.5+0.1
54,9116
44.5=0.4
46.520.7
41.7=0.6
73.2+5.7
7L024.3
52.6+0.3
JOAO You or d.. 2021
52.3+0.2
0.420.1
55.0#0.8
44.5=0.5
47.620.5
40.6+0.2
74.3+2.8
71.542.
51.420.6
MGSSL Zhane
2021
64.2+0.2
4.0=0.4
53.240.3
44.8=0.6
49.720.3
41.5+022
65.621.8
64.640.5
52.7=0.5
GruphLoG (Xu èt all
2021
58.6#0.4
2.5+0.3
56.9=0.4
49.3=0.3
54.8
42,6+0.3
66.8+1.7
69.0+1.3
58 8+0.5
GraphMAE
(Hou (Hou
2022
53.320.1 0.6+0.1
82.8
43.3=0.9
51.240.8
40.9+0.3
64.422.7
65.943.8
50.9=14
DSLA (Kim air 2029
57.8+0.5
0.7+0.1
57.9+0.7
50.4.0.7
53.6+1.7
43.3+0.9
68.6
708+2.0
60.90.66
UniMol Zhou a. 2026
64.620.2
4.8+10
55.8 50.1
49.915.6
43.6211
65.424.9
65.841.2
59.9
Pretruined Chemical Language Models
Robcrta Mary al. 2024
64.220.8 3.1+1.8
69.022..6
71.4=14.5 65.1 63.7+24.6
67.5+5.2
69.944.9
76.7+13.2
GPT2 Mary elal a 2032)
61.5#1.1 2.420.6
74.0+8.5 15.9
73.1 54.1=12.9 83.2+21.5
161.8
8:25,5
(xrwards 2 2022
64.7+0.9 3.6+1.1
140.5
76.7+2.1
55.9+1.1
49,2+1.0
70.340.8
73.1+1.0
65.3+1.7
ChemGPT (Prey 01 2023
Token Error
75.7+8.5
59.5=7.3 88.8+32.3
16.1=1188
84.0=20.6
77.248.5
68.627.1
Multi-modal Alignment
CLOOME
54.220.9 0.9+0.2
64.310.4
65.2=1.5
56.920.8
44.220.8
70.740.4
73.640.8
75.0=2.1
InfoCORE (GE)
65.3+0.2 5.4+1.7
69.9+1.2
79.9=3.6 51.62188
51.3+2.1
78.6
77.8119
80.3+0.9
InfoCORE (CP)
62,4+0.4 1.3+0.5
71 71.0+0.6
74.5+4.9
53.5+0.7
53.6+2.1
80.8+1.5
79.443.4
84.4+1.
InfoAlign (Ours)
66.4#1.1
6.6+1.6
49.4+0.2
39.7=0.4
39.2+0.3
40.5+0.6
66.7+1.7
62.0+1.5
48.4=0.6
Single Rep.
Single
78.0
7.3 14.6
90.6
6.231
Molecular Structure
InfoAlign InfoCORE CLOOME
Molecular Structure
InfoAlign InfoCORE CLOOME
Cell Morphalogy
Cell Morphology
Gene Expression
Gene Expression
Aligned Rep.
Aligned Rep.
61.0
34.1
4,9
81.2
80
100
Tasks (%) with Best Performance
Tasks (%) with Best Performance
(a) ChEMBL2K with Tasks.
(b) Broad6K with Tasks.
Figure 3: Percentage of Tasks Where Representations Excel: The top bar compares the best baselines using single-moda representations (Single Rep.) across representation categories. The bottom bar compares three aligned representations (Aligned Rep.): InfoAlign, CLOOME, and InfoCORE.
fundamentally involve predicting the binding affinity of a molecule to a protein Gaulton et al. 2002 22) furthermore, in these datasets, molecules with activity in a given assay tend have highly re related structures, rather than representing two more structurally distinct classes of molecules with activity; together this implies that molecular structure alone will tend yield strong results. When comparing the three popular structure-based representation approaches, no single method outperforms the others across all four datasets. Pretrained GNNs generally perform better than fingerprint-base methods and pretrained chemical language models, thanks recent advancements. However, continued efforts in universal structural representation are still necessary.
(2) Cell morphology and gene expression features may complement molecular structures, yield- ing more generalizable representations. As shown in Figure cell morphology and gene expression outperform molecular structure in approximately 20% and 10% of tasks the ChEMBL2K and Broad6K datasets, respectively. This suggests that incorporating cell context into representation learning would be beneficial. That said, existing multi-modal baselines (InfoCORE, CLOOMME)
8

---

Published as a conference paper at ICLR 2025
Table 3: Retrieval results on ChEMBL2K and Broad6K: Ranking metrics for top candidates.
NDCG & (t)
HIT (f)
NDCG & (t)
HIT % (t)
ChEMBL2K
Broad6K
lop-1
top-10
lop-1
top-10
top-1
top-10
top-1
top-10
CLOOME
0
2.0
CLOOME
0.5
0.9
0.5
1.5
InfoCORE
4.5
11.3
InfoCORE
1.0
2.5
1.0
4.6
InfoAlign
1.3
5.7
12.5
InfoAlign
0.5
2.3
0.5
5.1
CLOOME
InfoCORE
InfoAlign
30-
CLOOME
InfoCORE
InfoAlign
Frequency
Frequency
5
10
20
40
60
80
50
100
150
200
Ranking Position of Correct Match
Ranking Position of Correct Match
Figure 4: Histogram of rankings for the correct matching on ChEMBL2K (left) and Broad6K (right).
only outperform molecular structure-based approaches on ChEMBL2K and ToxCast, as they do not construct molecular representations holistically by using all cell-related modalities.
(3) InfoAlign achieves the best average performance on all tasks compared 27 baselines. The improvements from InfoAlign range from 2.5% to 6.4% on average across four datasets compared to the second-best method. These gains are more significant when using the 80% AUC threshold on classification datasets. While InfoCORE (GE) performs best among baselines on the ChEMBL.2K and ToxCast datasets, struggles to align molecular representations with more than two modalities and sometimes leads to negative transfer, as seen in Broad6K and Biogen3K.
6.2 RQ2: MOLECULE-MOR MORPHOLLOGY CROSS-MODAL MATCHINGG
Molecular representations are aligned with cell morphology. The zero-shot matching performance of queried molecule to cell morphology features evaluates the alignment between the two modalities.
6.2.1 EXPERIMENTAL SETTING
We compared InfoAlign with the CLOOME and InfoCORE (CP) for retrieving cell morphology from molecular representations. We calculate the cosine similarity between the molecular representation and all cell morphology candidates, rank these candidates, and compute Normalized Discounted Cumulative Gain (NDCG) and HIT scores for the top-1 and top-10 candidates as metrics. For fair evaluation zero-shot matching, we exclude the cell morphology data for molecules that were used train the baselines. Consequently, we have 80 molecule-cell morphology pairs from ChEMBL2K and 196 pairs from Broad6K. All the morphology data are used candidates for matching.
For InfoAlign, we use the pretrained decoder from Section 5 to extract the morphology features of the encoded molecule and then calculate the likelihood of these decoded features against the candidate morphology data. We then rank the candidates in the decoding space based on their likelihood scores.
6.2.2 RESULTS AND ANALYSIS
Cross-modal matching results with ranking metrics for the top candidates are in Table InfoAlign outperforms InfoCORE on ChEMBL2K and is comparable on Broad6K, with both surpassing CLOOME. Additionally, present histogram of ranking positions for correct matching pairs in Figure compare overall retrieval performance. The results show that InfoAlign and InfoCORE perform similarly, while CLOOME tends ranks correct pairs lower.
6.3
RQ3: PERFORMANCE ANALYSIS
6.3.1
ABLATION STUDIES
We perform ablation studies on Eq. B by pretraining encoders with different targets removed: (1) molecule-related, (2) cell rellhelogy-related, and (3) gene expression-related features. The results

---

Published as a conference paper at ICLR 2025
Table 4:
Ablation studies on the pretraining loss. Different node features are removed from the context graph to assess their impact on downstream tasks. Avg. AUC is reported.
ChEMBL AUC
Broad6K AUC
ToxCast AUC f
Biogen3K MAE 4 100)
Default as Eq. 3
81.3 +: 0.6
70.0 0.1 66.4 = 1.1
49.4 0.2
0.2
w/o Cell Morphology
80.7 - 0.6
68.6
0.1
66.5-11
51.7
Gene Expressions
78.3
68.6
6
50.3
Molecular Features
79.1 + 0.2
67.1
0.4
65.8 233
5.6
and
1.5
Pretraining Loss
Default
1.4
79.9
ChEMBLOK BONBBLKKKA AUC
1e-9: 81.3
1.3
1e-12:
1e-14: 80.1
Baseline: InfoCORE (GE)
79-
1.1
10 12
300
Random Walk Path Length
Epoch
(b) AUC for varying T. (Default: 4). The error bar
(a) Losses and AUC varying 3 (Default: Ic-9).
represents one standard deviation from ten runs.
Figure 5: Analysis on the hyperparameters: strength prior 3 and random walk length L. AUC is
computed the test set of ChEMBL.2K.
in Table g cover all datasets. We observe that both cell morphology and gene expression features are crucial for achieving the best performance. Different biological targets have varying impacts across datasets: molecular structure has more influence on Broad6K and Biogen3K, while gene expression more important for ChEMBL.2K and ToxCast. Further studies on the cell morphology and gene expression are in appendix D.4
6.3.2 HYPERPARAMETER ANALYSIS
Lastly, we perform analysis for the hyperparameters: strength of the regularization to the prior Gaussian distribution 9 and the length of the random walk paths L. Results presented in Figure 5 We observe a trade-off between the principles of minimality and sufficiency in Figure Sal a too-high value (minimal information) makes it challenging for the representation to be sufficiently expressive for molecular, gene expression, and cell morphology features, potentially degrading downstream performance. Conversely, too-low value weakens minimality and may impair generalization. The convergence of the pretraining loss could serve as a good indicator to balance these aspects. For the hyperparameter L, we observe in Figure 5b that downstream performance ChEMBL2K is relatively robust across wide range of walk lengths. Further analysis of the random walk sampling provided in appendix D.5 DD.
CONCLUSION
In this work, we proposed learning molecular representations in cell context with three modalities: molecular structure, gene expression, and cell morphology. We introduced the information bottleneck approach, InfoAlign, using a molecular graph encoder and multiple MI.P decoders. InfoAlign learned minimal sufficient molecular representations extracted by reconstructing features in the random walk path on cellular context graph. This context graph incorporated molecules, cell morphology, and gene expression information defined in scalar vector spaces to construct nodes, and used various chemical, biological, and computational criteria to define their weighted edges. We demonstrated the theoretical and empirical advantages of the proposed method. InfoAlign outperformed other representation learning methods in various molecular property prediction and zero-shot molecule- morphology matching tasks.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This study was supported by National Institutes of Health (R35 GM122547 AEC) and internship funded by the Massachusetts Life Sciences Center (to GL). This work was supported by NSF IIS-2142827, IIS-2146761, IIS-2234058, CBET-2332270, and ONR N00014-22-1-2507.
REFERENCES
Alexander A Alemi, lan Fischer, Joshua Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410, 2016.
John Arevalo, Ellen Su, Robert van Dijk, Anne E Carpenter, and Shantanu Singh. Evaluating batch correction methods for image-based cell profiling. bioRxiv, 2023.
Mark-Anthony Bray, Shantanu Singh, Han Han, Chadwick T Davis, Blake Borgeson, Cathy Hartland, Maria Kost-Alimova, Sigrun M Gustafsdottir, Christopher C Gibson, and Anne E Carpenter. Cell painting, high-content mmage-based assay for morphological profiling using multiplexed fluorescent dyes. Nature protocols, (1(9):1757-11 2016.
Mark-Anthony Bray, Sigrun M Gustafsdottir, Mohammad H Rohban, Shantanu Singh, Vebjorn Ljosa, Katherine L Sokolnicki, Joshua A Bittker, Nicole Bodycombe, Vlado Dančík, Thomas P Hasaka, et al. A dataset of images and morphological profiles of 30 000 small-molccule treatments using the cell painting assay. Gigascience, 6(12)giw011, 2017.
Anne E Carpenter, Thouis R Jones, Michael R Lamprecht, Colin Clarke, In Han Kang, Ola Friman, David A Guertin, Joo Han Chang, Robert A Lindquist, Jason Moffat, et al. Cellprofiter: image analysis software for identifying and quantifying cell phenotypes. Genome biology, 7:1-11, 2006.
Srinivas Niranj Chandrasekaran, Jeanelle Ackerman, Eric Alix, D Michael Ando, John Arevalo, Melissa Bennion, Nicolas Boisseau, Adriana Borowa, Justin D Boyd. Laurent Brino, et al. Jump cell painting dataset: morphological impact of 136,000 chemical and genetic perturbations. Biorxiv, pp. 2023-03, 03, 2023.
Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: large-scale self- supervised pretraining for molecular property prediction. arXiv preprint arXiv:2010.09885, 2020.
Beth Cimini, Srinivas Niranj Chandrasekaran, Maria Alsm-Ali Lisa Miller, Amy Goodale, Briana Fritchman, Patrick Bymne Sakshi Garg. Nasim Jamali, David J Logan, et al. Optimizing the cell painting assay for image-based profiling. Nature protocols, 18(7):1981-20) 2023.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint / 1810.888000 2018.
Carl Edwards, Tuan Lai, Kevin Ros, Garrett Honke, Kyunghyun Cho, and Heng Ji. Translation between molecules and natural language. arXiv preprint arXiv:2204.11817 2022.
Cheng Fang, Ye Wang, Richard Grater, Sudarshan Kapadnis, Cheryl Black, Patrick Trapa, and Simone Sciabola. Prospective validation of machine learning algorithms for absorption, distribution, metabolism, and excretion prediction: An industrial perspective. Journal of Chemical Information and Modeling, (3(11):323333274 2023.
Philip Fradkin, Puria Azadi, Karush Suri, Frederik Wenkel, Ali Bashashati, Maciej Sypetkowski. and Dominique Beaini. How molecules impact cells: Unlocking contrastive phenomolecular retrieval. arXiv preprint urXiv:2409.088 000020 2024.
Nathan C Frey, Ryan Soklaski, Simon Axelrod, Siddharth Samsi, Rafael Gomez-Bombarell Con- nor W Coley, and Vijay Gadepally. Neural scaling of deep chemical models. Nature Machine Intelligence, 5(1 5(11):1297-1305 2023.
Anna Gaulton, Louisa J Bellis. A Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey, David Michalovich, Bissan Al-Lazikani, et al. Chembl: a large-scale bioactivity database for drug discovery. Nucleic acids research, 40(D1):D1 100-D1107, 2012.
11

---

Published as a conference paper at ICLR 2025
Rohit Girdhar, Alaacldin El-Nouby, Zhuang Liu, Mannat Singh, Kalyan Vasudev Alwala, Armand Joulin, and Ishan Misra. Imagebind: One embedding space to bind them all. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 15180-15190, 2023.
Gokul Gowri, Xiao-Kang Lun, Allon M Klein, and Peng Yin. Approximating mutual information of high-dimensional variables using learned representations. arXiv preprint arXiv:2409.002755 02732, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Daniel Scott Himmelstein, Antoine Lizee, Christine Hessler, Leo Brueggeman, Sabrina L Chen, Dexter Hadley, Ari Green, Pouya Khankhanian, and Sergio Baranzini. Systematic integration of biomedical knowledge prioritizes drugs for repurposing. Elife, 6xe26726, 2017.
Zhenyu Hou, Xiao Liu, Yukuo Cen, Yuxiao Dong, Hongxia Yang. Chunjje Wang, and Jie Tang. Graphmae: Self-supervised masked graph autoencoders. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 594-604, 2022.
W Hu, B Liu, Gomes, M Zitnik, Liang, V Pande, and J Leskovee. Strategies for pre-training graph neural networks. In International Conference on Learning Representations (ICLR), 2020a.
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:22 18-22133, 2020b.
Kexin Huang, Cao Xiao, Lucas M Glass, and Jimeng Sun. Moltrans: molecular interaction transformer for druget interaction prediction. Bioinformatics, 37(6):830-836, 2021.
Eric Inae, Gang Liu, and Meng Jiang. Motif-aware attribute masking for molecular graph pre-training, arXiv preprint arXiv:2309.0455 2023.
Bowen Jin, Gang Liu, Chi Han, Meng Jiang, Heng Ji, and Jiawei Han. Large language models on graphs: A comprehensive survey. arXiv preprint arXiv:23/2.02788, 2023.
Dongki Kim, Jinheon Back, and Sung Ju Hwang. Graph self-supervised learning with accurate discrepancy learning. Advances in Neural Information Processing Systems, 35:14085-14098, 2022.
Mario Krenn, Qianxiang Ai, Senja Barthel, Nessa Carson, Angelo Frei, Nathan C Frey, Pascal Friederich, Théophile Gaudin, Alberto Alexander Gayle, Kevin Maik Jablonka, et al. Selfies and the future of molecular string representations. Patterns, 3(10), 2022.
Anika Liu, Srijit Seal, Hongbin Yang, and Andreas Bender. Using chemical and biological data to predict drug toxicity. SLAS Discovery, 28(3):53-6)) 64, 2023a.
Gang Liu, Tong Zhao, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph rationalization with environment-based augmentations. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1079-1078, 2022.
Gang Liu, Tong Zhao, Eric Inae, Tengfei Luo, and Meng Jiang. Semi-supervised graph imbalanced regression. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1453-1465, 2023b.
Gang Liu, Eric Inae, Tong Zhao, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Data-centric learning from unlabeled graphs with diffusion model. Advances in neural information processing systems, 36, 2024a.
Gang Liu, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph diffusion transformer for multi-conditional molecular generation. arXiv preprint wrXiv:2401.13882 2024b.
Yinhan Liu. Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. / 2019.
12

---

Published as a conference paper at ICLR 2025
Hadrien Mary, Emmanuel Noutahi, DomInvivo, Lu Zhu, Michel Moreau, Steven Pak. Desmond Gilmour, Shawn Whitfield, t, Valence-JonnyHsu, Honoré Hounwanou, Ishan Kumar, Saurav Maheshkar, Shuya Nakata, Kyle M. Kovary, Cas Wognum, Michael Craig, and DeepSource Bot. datamol-io/datammll: 0.12.3, January 2024. URL httth:tt///oo org 5281 /zenodo [0535844]
Fred D) Mast, Alexander V Ratushny, and John D Aitchison. Systems cell biology. Journal of Cell Biology, 206(6):695-700 2014.
Nikita Moshkov, Tim Becker, Kevin Yang, Peter Horvath, Vlado Dancik, Bridget K Wagner, Paul A Clemons, Shantanu Singh, Anne E Carpenter, and Juan C Caicedo. Predicting compound activity from phenotypic profiles and chemical structures. Nature Communications, 14(1):1967, 2023. 2023.
Cuong Q Nguyen, Dante Pertusi, and Kim M Branson. Molecale-morrhology contrastive pretraining for transferable molecular representation. bioRxiv, pp. 2023 05, 2023.
XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56 (6(11):5847-580) 2010.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint (.XXv:8807.0374 2018.
Bryan Perozzi, Rami Al-Rfou, and Steven Skiena. Deepwalk: Online learning of social representa- tions. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 701-710, 2014.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pp. 5171-5180. PMLR, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint 22Xi:2204400625,, 1(2):3 2022.
Ann M Richard, Richard S Judson, Keith A Houck, Christopher M Grulke, Patra Volarath, Inthirany Thillainadarajah, Chihae Yang, James Rathman, Matthew Martin, John F Wambaugh, et al. Toxcast chemical landscape: paving the road to 21 century toxicology. Chemical research in toxicology, 29(8): 2016.
David Rogers and Mathew Hahn. Extended-cconeettvvss fingerprints. Journal of chemical information and modeling, 50(5):742-754 2010.
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer large-scale molecular data. Advances in neural information processing systems, 33: 15559-12577, 2020.
Jerret Ross, Brian Belgodere, Vijil Chenthamarakshan. Inkit Padhi, Youssef Mroueh, and Payel Das. Large-scale chemical language representations capture molecular structure and properties. Nature Machine Intelligence, 155-11264 2022.
Ana Lanchez-Fernandez Elisabeth Rumetshofer, Sepp Hochreiter, and Günter Klambauer. Cloome: contrastive learning unlocks bioimaging databases for queries with chemical structures. Nature Communications, 14(1):7339, 2023.
13

---

Published as a conference paper at ICLR 2025
Srijit Seal, Jordi Harreras-Puigvers Maria-Anne Trapotsi, Hongbin Yang, Ola Spjuth, and Andreas Bender. Integrating cell morphology with gene expression and chemical structure to aid mitochon- drial toxicity detection. Communications Biology, 5(1):858, 2022.
Srijit Seal, Hongbin Yang, Maria-Anna Trapotsi, Satvik Singh, Jordi Carreras-Puigwer Ola Spjuth, and Andreas Bender. Merging bioactivity predictions from cell morphology and chemical finger- print models using similarity to training data. Journal of Cheminformatics, 15(1):56, 2023.
Srijit Seal, Maria-Anna Trapotsi, Ola Spjuth, Shantanu Singh, Jordi Carreras-Puigvert Nigel Greene, Andreas Bender, and Anne E Carpenter. A decade in systematic review: The evolution and impact of cell painting. bioRxiv, pp. 2024-05, 2024.
Aravind Subramanian, Rajiv Narayan, Steven M Corsello, David D Peck, Ted E Natoli, Xiaodong Lu, Joshua Gould, John F Davis, Andrew A Tubelli, Jacob K Asiedu, et al. A next generation connectivity map: L1000 platform and the first 1,000,000 profiles. Cell, 171(6): 1337-1452, 2017.
Ruoxi Sun, Hanjun Dai, and Adams Wei Yu. Does gnn pretraining help molecular representation?
Advances in Neural Information Processing Systems, 35:1 35:12096-13 12109, 2022.
Naftali Tishby, Fernando Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/00040057 2000.
Chenyu Wang. Sharut Gupta, Caroline Uhler, and Tommi S Jaakkola. Removing biases from molecular representations via information maximization. In The Twelfth International Conference on Learming Representations, 2023a.
Hanchen Wang, Jean Kaddour, Shengchao Liu, Jian Tang. Joan Lasenby, and Qi Liu. Evaluating self- supervised learning for molecular graph embeddings. Advances in Neural Information Processing Systems, 36, 2024.
Zichen Wang, Neil R Clark, and Avi Ma' ayan. Drug-induced adverse events prediction with the lincs 11000 data. Bioinformatics, 32(15):2338 2345, 2016.
Zifeng Wang, Zichen Wang, Balasubramaniam Srinivasan, Vassilis N Ioannidis, Huzefa Rangwala, and Rishita Anubhai. Biobridge; Bridging biomedical foundation models via knowledge graph. arXiv preprini arXiv:2310.03320 2023b.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Minghao Xu, Hang Wang. Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supe graph-level representation learning with local and global structure. In International Conference on Machine Learning, pp. 11548 1 1558. PMLR, 2021.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in neural information processing systems, 33: 5812-5823, 2020.
Yuning You, Tianlong Chen, Yang Shen, and Zhangyang Wang. Graph contrastive learning automated. In International Conference on Machine Learning, pp. 12121-12132. PMLR, 2021.
Zaixi Zhang, Qi Liu, Hao Wang, Chengqiang Lu, and Chee- Kong Lee. Motif-based graph self- supervised learning for molecular property prediction. Advances in Neural Information Processing Systems, 34: 14:05880-,,882, 2021.
Shuangjia Zheng, Jiahua Rao, Jixian Zhang, Lianyu Zhou, Jiancong Xie, Ethan Cohen, Wei Lu, Chengtao Li, and Yuedong Yang. Cross-modal graph contrastive learning with cellular images. Advanced Science, 11(32):4408455, 2024.
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, 2023. URL https 7/openreview.o net/ www.mmoddde22eeeeooor.o.
14

---

Published as a conference paper at ICLR 2025
A
MORE RELATED WORK ON REPRESENTATION LEARNING WITH DIFFERENT MODALITIES
Existing methods multimodal alignment, such CLIP (Radfford et al. 2021), primarily address pairwise relationships between texts and images and use methods like InfoNCE. (Oord et al. 2018 Wang et al. 2023a Sanchez-Femmandzz al. 2023). These approaches use separate encoders for different modalities to compute contrastive loss, which is upper bounded by the number of negative examples (Poole al. 2019). Subsequent research on molecules similarly focuses on pairwise alignment between molecules and cell images
et al. [2023] Wang et al. 2023a
2021), molecules and protein sequences (Huang [ and molecules and text (Edwards al. 2022
Jin et al. (023). Although BioBridge (Wang et al. 2023b) handles multiple modalities, it leverages knowledge graph for transforming representations between modalities rather than optimizing molecular representations.
B
PROOF DETAILS
B.1 PROOF OF EQ.
Approximating the mutual information of high-dimensional variables is challenging task (Gowni et al., 2024). For the input, latent, and target variables X, Z, and Y, the exact computation of the mutual information (MI) Y) and T(X; Z) is intractable. For the molecule x, its latent representation and any biological target from cellular responses y, we introduce variational approximations ((y|z)) to obtain a lower bound on (2;Y):
log p(2,8)g(ul2) p(U)p(z)g(ulz))
I(Z:Y)= = Ep(x.)
log {)y[=)
(5)
= Ep(z.))
ply) + Ep(z) [KL (p(vl=)
2)) 2 0. We introduce the variational approximations ((2) for 2 Bp(z,)
[log ((v((())))))))))))))))) HYY) !DLB
This is because that KL (p(y)=)]l<())) a upper bound on 1(;;2):
log 2)=)))((=) ] p(x))(2)g(())
T(X; 2) E Ep(x,))
log p(2|x)]
(6)
Ep(,,))
q(z) Ep(x) [(p())]la()))
< E p(x) KT. (p(2)=))l<<=)) e TEuB
B.2 PROOF OF PROPOSITION
For the molecule æ, latent representation 2, and any biological target from cellular responses y- we use the neural network parameterized critic h(2, v) with the energy-based variational family for density approximation (Poole al.. 2019);
p(v) M(.,))
<(yl=)
Ep(w) [ch(2.))
Thus, we can rewrite IDLB based on the unnormalized distribution of q(yl2):
IDLB Ep(z.)) [log [(>|=)] (YY)
p(y) Eppsi [chlzwn] (ha.ee )
Ep(x.w) log
Ep(i) [log p(y)],
Ep(z.w) [1(2,, y)] Ep(z.w) [Epcr
= Eptea) [M(,, Pp(e) (()))
15

---

Published as a conference paper at ICLR 2025
where 2(2) Ee is the partition function.
Note that the log partition function is intractable. Poole et al. (201)) introduced new variational parameter al-) to upper bound (z), deriving a tractable lower bound for IDLB:
E (ppo))e
JDLB 2 Pepea [
pp)))
+ log(a()))-11
(8)
a(2)
This is because Vz, a > the inequality log (z) S + log a holds, which can be applied to
the second term of Eq. The INw bound Nguyen et al. 2010) is special case where a(z)
Epty) [e (2.0))
I INwJ E Ep(z,n) ()) - Ep(x)
+ log(e) -
E(zyy [h(2, y)1 E
I has high variance due to the estimation of the upper bound on the log partition function. Based on [ and multiple examples, one can derive the encoder-based lower bound IELB for InfoNCE.
Suppose there are K additional examples independently and identically sampled and denoted as AAle.w)
K2.K, and the critic is configured with parameters a(.) as 1 + log a(2:g.02:x) Then, we can rewrite J for its multi-sample version:
JNwJ (p(e.))(va..))
1 log
(2,,))
e- Bp(o))(s)((nux)))
e 1tlog a(srvin2:k)
chco.vl
a(2: ys (2:K)
(10)
(=..))
((x,))
11 (p(x.))(vaacc)
log
a(z; y. (2:K)
a(2; ys 92.
Multiple samples be utilized for the Monte Carlo method m(2; V2.k) to estimate the upper bound on the partition function a(z; 92:K)::
a(2:y. B2.K) m( V2.K) K
(h(2.))) + Σ ((2.)))))
i=2
where K - independent samples are drawn from 11, P(U4) and one sample from p(z, v) for the
term Ep((aa))(vax or samples from 1 p(y:) (we set 11 = y) for p(z) sample in the
Ep())ramn) H term. Therefore, we can derive TELB INCE:
((2M))
b(2.))
JLLB A INCE = 1+ Bp(x.))(vakk) log m(2; 9, (2.K)
Eptule(uan)e(e))
m 92.K)
M(299)
+ (2)))2)))))))
log
K Ei=1 c((2))))
Diei K
K
Ep(as) (M(2.))] - Ep(a)
log K
(11)
Note that for E Epichp(ygn [-]=) we average the bound over K replicates as well to ensure that the last term in Eq. (TO) the constant 1. Now, TELB or INCE is upper bounded by log K, rather than a(-). Hence, the difference between IDLB and IELB is
IDLB JELB = En(z)
log Σ ManaE) K
1
Epp Z(2)) 2 0.
(12)
When K is sufficiently large to estimate the partition function, we have Ep(x) [log (Ep(s) [c ((
for the left term, indicating that IDLB IELB O. Since INCE is upper bounded by log K (ord ct al. 2018), smaller values K may result in a less tight LELBS causing I IELB 2 to always hold. In particular, T(Z;Y) log K implies that the bound TELB will be loose.
16

---

Published as a conference paper at ICLR 2025
Pentra
Perturt
Iniewrncs
Intersc
Penun
Penuth
C
Display all landmark genes (in X-peturbbb relations)
GBg
Initial
Ponura
68
Cosinu
Interna
Perturb:
Interact
Pertinh
Pertin
Perturs
Interecs
Interia
62
Gi2
Pertar
Cosina
Pertunt
Perunn
Interent
Interoo
e
Interast
Perturty
Perturh
Cooting
Used for pretraining
Perturb.
Xs
Pertue
Fertrr
GE
Gene expression feature vectors
Cosnna
Genes with cell morphologyy feature
INTERBOL
X4
PartInA
Pettuz
Implementation with practical considerations
Cownne
Node Type: Feature Type
Molecula: Vector
Gene merged with cell morphology: Vector
Gene expression: Vector
Cell morphology
Gene: No feature
Landmark gene: Scalar
Vector
Figure 6: From the initial idea in Section 4 to the practical implementation of the context graph, first display relations between molecules and all the landmark genes from Wang al. (2016) for the X1 E3 and X3 - Es relationships. Es and Eg are landmark genes involved in small molecule perturbations and cell morphology perturbation; we display them separately for clarity, Next, we merge all landmark genes into new gene expression nodes and integrate genes from genetic perturbations in the JUMP dataset (Chandrasekaran et al. (023) with cell morphology features. Practical considerations are detailed in Section 5 and appendix
C CONTEXT GRAPH AND MODEL DETAILS
C.I EDGE CONSTRUCTION
Edges represent similarity relationships between molecules, genes and cells. According to chemical or biological criteria, we have following types of edges:
1. Molecule- Cell Morphology Edges: These edges are introduced through molecule perturba- tion experiments from cell painting datasets created by Bray et al.) (2016) and the JUMP dataset (Chandrasekaran al.) 2023). links molecule nodes with cell morphology nodes.
We use the edge weight all these edges.
2. Edge-Cell Morphology Edges: These edges should be introduced by genetic perturbation from the JUMP dataset Chandrasekaran al. (2023). The perturbations either based on gene overexpression (ORF) or gene knockout techniques (CRISPR). They link the gene nodes and the cell morphology nodes. However, the genes introduced by the genetic perturbations lack gene expression profiling from (Subramanian et al., 2017) as node features. We did not implement gene-cell morphology edges from (Chandrasckaran et al. (023)) due to the absence of differential gene expression profiling values (Subramanian et al. (017). Instead, we merged the gene nodes from (Chandrasekaran morphology nodes, creating single nodes. This approach enables more efficient context graph, incorporating some gene nodes with cell morphology features.
al., 2023) with their linked cell
3. Molecule- Gene Edges: These edges could represent molecule- sonnecul-ggene gene binding and regulation relationships, linking molecules to genes (Himmelstein et al.] (017). Some links be sourced from (Himmelstein et al.) 2017). and we also retrieve gene-molecule links from (Wang al. 2016) by selecting the top 5% absolute differential expression values.
17

---

Published as a conference paper at ICLR 2025
4. Gene-Gene Edges: These edges denote the relationships of gene-gent covariance and interaction and we use the links from (Himmelstein et al. 2017).
We enrich the edges in the context graph by incorporating computational similarity edges, where cosine similarity is computed among within nodes having the same type and feature vectors. We note that the cell morphology features from (Bray et al., (016) and (Chandrasckaran et al., 2023) have different dimensions since the latter has applied batch correction techniques Arevalo et al. (2023) on the CellProfiler features (Carpenter et al. 2006). Thus, we cannot compute the similarity between these two subsets of cell morphology nodes. We use (1) 0.8 similarity threshold and (2) minimal sparsity of 99.5% % by selecting top 0.5 % similar edges to avoid excessive noise in computational similar edges.
C.2 DATASET SOURCES OF NODES
Here the datasets used create different types of nodes the context graph:
Molecule nodes: Molecular nodes are sourced from two cell painting datasets: one by Bray et 2017 7) and the other from the recently released JUMP dataset (Chandrasekaran et al. 20233, and the third source from Wang et al. (2016), which are used to study adverse drug reactions.
Gene nodes: Gene nodes are from the landmark genes used by Wang et al. (2016) in creating the LINCS L1000 profiling of drugs. Other gene nodes come from genetic perturbations in the JUMP dataset (Chandrasekaran et al. 2023). The gene nodes from (Chandrasekaran 2023) have cell morphology features as deseribed in appendix C.1 The landmark gene nodes from (Wang et al. (016)) have scalar gene expression profiles, but these values are updated in the new gene expression nodes.
et al.
Cell morphology nodes: Cell nodes sourced from the two cell painting datasets (Bray et al. 2016 Chandrasekaran et al. 2023).
Gene expression nodes: Based on landmark genes from (Wang et al., 2016), each gene expression node summarizes all gene expression profiles into vectors from small molecule perturbation. Since Wang et al. (2016) measured the same landmark genes for a set of molecules, we update new gene expression nodes with feature vectors for all these land- mark genes. This approach efficiently constructs decoding targets from molecules to gene expression profiles and prevents redundant gene-molecule connections.
We present example of the cellular context graph in Figure 6
C.3 MODEL DETAILS
We use live-layer GIN (Xu et al. 2019), following previous GNN pre-training work (Hu et al. 2020a). The hidden dimension is set to 300, with a batch normalization layer and summation used for readout of the graph representation from node-level features.
For the MLP decoders, they have three layers. The input dimension 300, with a hidden dimension of 1200. The output dimension matches the corresponding node features from the context graph in pretraining. For the newly introduced hyperparameters 3 = 1e - 9 and I. = 4, details on their studies can be found in Section 6.3.2
In downstream tasks, the GNN encoder is frozen, and new MLP decoders, which output task predictions, are applied the molecular representation from the GNN encoder.
D EXPERIMENT DETAILS
D. PREDICTION DATASETS
single 32G V100. Molecules interact with multiple cells and genes to generate cell morphology and gene expression features. These molecules act as perturbations, with the resulting cell states measured as gene expression values for thousands of genes (Subramanian
All experiments
18

---

Published as a conference paper at ICLR 2025
Table 5: Datasets and task information. Classf. denotes classification and Regr. denotes regression.
# Atoms Avg. /Mxx
# Edges Avg./Max
# Available Cell Morphology
Available Gene
Dataset
Type
Task
Molecules
Expressions
ChEMBL2K
Classf.
2355
23.7/61
25.66 25.6/68
2353
Broad6K
Classf.
6567
34.1/74
36.8/82
2673
1138
ToxCast
Classf.
8576
18.8/124
19.3/134
N.A.
N.A.
Biogen3K
Regr.
6
3521
23.2/78
25.3/84
N.A.
N.A.
SLEGAA CHRMISIGHARIHTIOC
ADRAZB
HTRZA
0R04
HTRI
SLC6A
GGF
2N81
InfoAlign
Molecular Structure
0.4
0.6
Cell Morphology
Gene Expression
CHRM
ABL1
C#22C
CYPLA2
MCL2
CYP2D6 PTGS]
Figure 7: An overview of the representation's 's predictive performance on all 41 bioactivity prediction tasks in ChEMBL2K. Results for molecular structure are obtained from the best method ContextPred. Results for cell morphology and gene expression come from the best method based on MLPs.
et al.
2017 and/or as morphology features from Cell Painting microscopy images (Chandrasekaran
al. 2023 Cimini al. (023). Due the high cost and complexity of these experiments,
all downstream datasets include morphology gene expression data for cach molecule. Notably, ToxCast and Biogen3K lack these features. Prediction dataset statisties detailed Table 5
ChEMBL2K (Gaulton et al. 2012): The dataset is subset of the ChEMBL dataset (Gaulton ct al. 2012), overlapping with the JUMP CP (Chandrasckaran et al. 2023) datasets. We determined activity using the "activity comment" provided by ChEMBL. If not, applicd a threshold of 6.5, labeling compounds with pChEMBL > 6.5 as active. We exclude all molecules in the dataset from the pretraining to avoid data leakage. There are total of 41 tasks related to protein binding affinity, which are converted to binary activity values. We filter the dataset ensure that each task has least one positive and five negative examples.
Broad6K (Moskkov et al. 2023): The original version provided by Moshkov et al. (023)
is a collection of 16,170 molecules tested in 270 assays, resulting in a total of 585, 585,439 4399 readouts. However, there
large number of missing values, with 153 assays having missing value percentage above 99%. To mitigate bias the conclusions, we extract subsets where the percentage is less than 50%.
19

---

Published as a conference paper at ICLR 2025
Cell
Vira
IntsAlign
Moleculer Structure
Cell Morphology
0.4
sterial
Gene Expression
Blochem
Figure An overview of the representation's predictive performance across five major task categories on Broad6K. Results for molecular structure are obtained from the best method based on JOAO. Results for cell morphology come from the best method based on RF. Results for gene expression are derived from the best method based on MLP.
ToxCast (Richard et al. 20166; The toxicology data is collected from the "Toxicology in the 21st Century" initiative, widely utilized in many graph machine learning models (Hu et al. 2020b). The dataset comprises 8,576 molecules and 617 binary classification tasks.
Biogen3K (Fang et al.) 2023): The dataset includes properties that describe the disposition of a drug in the body, including absorption, distribution, metabolism, and excretion (ADME). It is collected from 120 Biogen datasets across six ADME in vitro endpoints over 20 time points spanning about years. The endpoints include human liver microsomal (HLM) stability reported as intrinsic clearance (Clint, mL/min/kg), MDR1 1-DCCK efflux ratio (ER), Solubility at pH 6.8 (ug/mL.,, rat liver microsomal (RLM) stability reported as intrinsic clearance (Clint, mL/min/kg), human plasma protein binding (hPPB) percent unbound, and rat plasma protein binding (rPPB) percent unbound.
We utilize scaffold-splitinng with a ratio of curve (AUC) score for classification and mean absolute error (MAE) for regression. We report the mean and standard deviations from ten runs.
for all datasets. We use the Area under the
D.2 IMPLEMENTATION AND BASELINE
We consider baselines from three representation sources: molecular structures, cell morphology, and gene expressions. Moreover, we have three different ways to represent molecular structures, including fingerprints based on domain knowledge, GNNs based on the graph structure of molecules, and chemical language models (ChemLM) based on SMILES sequence structure of molecules.
1. Molecular descriitoor//ingeeprnnts (Rogers & Hahn 2010) (Structure only): We train MLPs, Random Forests (RF), and Gaussian Processes (GP) these representations.
2. Pretrained GNN representations (Hu et al. 2020a) (Structure only): We consider AttrMask,
ContextPred, and EdgePred with supervised pretraining (Hu et al.) 2020a). We also include
GraphCL Youu et al.) 2020), GROVER (Rong et al.) 2020). JOAO (You et al.) 2021).
MGSSL Zhang et 202T), GraphLoG (Xu al.12021).. GraphMAE (Hou al. 2022).
DSLA (Kim et al. 2022), and UniMol (Zhou et al. (023). We implement GraphCT, GROVER, and JOAO based on (Wang et 2024). Fine-tuned ML.ss are applied on top of the pretrained representations.
3. Pretrained ChemLM representations (Frey et al. 2023 (Structure only): We consider pretrained models such 102M Roberta and 87M GPT2 implemented by (Mary al.
20

---

Published as a conference paper at ICLR 2025
Table 6: Ablation studies on the pretraining loss. Different node features are removed from the context graph to assess their impact on downstream tasks. Avg. AUC is reported.
ChEMBL.2K AUC
Broad6K AUC
ToxCast AUC ^
Biogen3K
AUC
InfoAlign
81.33 + 0.62 69.95 0.09
06.36 + 1.05 49.42
0.18
w/o Cell-Related Nodes
79.57 + 0.58 68.41 11 0.31
65.11
0.82
51.21
+ 0.17
w/o Nodes
77.97 + 0.33
67 10
0.17
64.98
+ 0.96
51.57
2 0.46
Cell-Relater
80.70 +: 0.60
68 60
65.50
1.10
51.70 1.10
w/o Bene-Related Loss
78.30 + 0.50
68.00
0.20
64.70
+ 1.00
50.30
+ 0.50
2024
We also include MolT5 (Edwards et al. 2022) and 19M ChemGPT (Frey et al. We apply fine-tuned MLPs on top of these pretrained representations.
2023
4. Cell Morphology (Rogers & Hahn 2010) (Cell Structure only): Cell morphology features are available in for part of molecules in the ChEMBL2K and Broad6K datasets. We train MLPs, RF, and GP on these representations. Note that not all molecules have corresponding cell morphology feature vectors; in such cases, we replace the predictions on the missing feature with ML predictions on the structure.
5. Gene Expression (Rogers & Hahn, 2010) (Gene Structure only): Differential gene expres- sion values are available for part of molecules the ChEMBL.2K and Broad6K datasets. We train MIPs, RF, and GP on these representations. Note that not all molecules have corresponding gene expression vectors over landmark genes; such cases, we replace the predictions on the missing feature with ML predictions the structure.
6. CLOOME "Sanchez-Fernandeee et al. 2023) and InfoCORE (Wang et al. 2023a (Structure- Cell Structure- Gene aligned): CLOOME utilizes ResNet (He et al. 2016) and descriptor- based MLP to align representation from cell morphology images with the molecular structure representation. We use their pretrained MLP to obtain molecular representations and fine- tune another MLP on top of these representations. InfoCORE has two versions, InfoCORE- CP and InfoCORE-GE which align the molecular graph representation with cell morphology features or differential gene expression features, respectively. We use both versions as baselines and fine-tune another MLP on top these representations.
D.3 MORE RESULTS FOR MOLECULAR PROPERTY PREDICTION
We present additional comparisons on the ChEMBL2K dataset between basic representation ap- proaches and InfoAlign aeross all task dimensions in Figure ?| Similarly, results for the Broad6K dataset, comparing basic representations across five major task dimensions (Cell, Yeast, Viral, Biochem, and Bacterial-related targets), are shown in Figure Combined with Tables I and these detailed results lead to further observations:
(1) Different structure-based based molecular representations vary in sensitivity to model architecture. Dramatic performance drops occur with Morgan FP when replacing the MLP architecture with RF or GP in the ChEMBL2K and Broad6K datasets. Conversely, in the Biogen3K dataset, RF and GP significantly outperform MLP. In contrast, pretrained GNN and ChemLM representations maintain more consistent performance across various datasets.
(2) Learning universal molecular representations solely from molecular structures remains chal- lenging, even within the representation category. For pretrained GNN representations, ContextPred outperforms others on the ChEMBI.25 .2K dataset. JOAO excels the Broad6K dataset. UniMol and GraphMAF are the best pretrained GNN representations on ToxCast and Biogen3K datasets, respectively. For ChemLM representations, MolT5 excels over other sequential-basee models in the ToxCast and Biogen3K datasets, but this is not the case with the ChEMBL2K and Broad6K datasets. Different datasets may emphasize varied aspects of bioactivity classification regression and pose generalization challenges for molecular representation learning.
(3) InfoAlign shows strong generalization for the targets of non-human cells, shown in Figure g Although the context graph primarily uses data from small molecule and genetic perturbation datasets (Bray et al. 2016; Chandrasekaran al. 2023) focused human cell cultures, InfoAlign
21

---

Published as a conference paper at ICLR 2025
Unique Nodes in Extracted Paths
250
Mean Unique Nodes with one standard deviation
200
150
100
2
3
4
5
6
8
10 12
Random Walk Path Length
Figure 9: The mean, with one standard deviation, of the number of unique nodes extracted from random walk paths. We vary the walk length L from 2 to 12, showing that random walks can sample diverse neighbors, with the number exceeding the walk length.
also exhibits robust generalization to bacterial and viral targets compared to basic representation approaches.
Table 7: Performance comparison of fully fine-tuned UniMol and InfoAlign models with their representations across various datasets. "Others" refers to cell morphology and gene expression features.
Fine-tuning Method
ChEMBL2K
Broad6K
ToxCast
Biogen3K
Representatio from UniMol
76.8+0.4
65.440.1
64.6+0.2
55.8+2.8
Representation from UniMol and Others
77.5+0.1
66.440.5
NA
NA
Fully-tuned UniMol
78.9+0.2
65.1+1.0
71.3+0.6
43.621.3
Representation from InfoAlign
81.3+0.6
70.0+0.1
66.4+1.1
49 19.4+0.2
Fully-tuned InfoAlign
80. 10.9
69.220.7
72.0+0.5
42.8
Additionally, we set up further comparisons, including: (1) the concatenated representation from UniMol (Zhou et al., 2023) using cell morphology and gene expression on the ChEMBL2K and Broad6K datasets, and (2) fully fine-tuned InfoAlign and UniMol, including both the representation encoder and MLP predictors.
Results are presented in Table First, we observe that concatenating UniMol representations with cell morphology and gene expression features improves performance in prediction tasks, though still does not match the performance of InfoAlign. InfoAlign achieves the best results by aligning molecular representations with these features during pretraining, rather than in the downstream stage.
Second, we observe that fully fine-tuning benefits both UniMol and InfoAlign on ToxCast and Biogen3K datasets. While fully fine-tuning improves UniMol representation on ChEMBL. InfoAlign, with only MLP decoder tuning, achieves the best performance. On Broad6K, fully fine-tuning is less effective for both Info.liign and UniMol compared tuning only the MLP. These results suggest that, if resources allow, fully fine-tuning should be preferred for better overall performance, especially UniMol, which requires more time and resources due to the use of 3D molecular structures. resources are limited, InfoAlign's representation provides a strong alternative without the need for full fine-tuning.
D.4 MORE ABLATION STUDIES ON CELL MORPHOLOGY AND GENE EXPRESSION
We also conducted ablation studies to assess the importance of cell morphology and gene expression in constructing the context graph. We removed either all cell morphology-related nodes or gene
22

---

Published as a conference paper at ICLR 2025
100
L=2
100.0 92.5 92.0
91.8
91.7
91.5 91.4 91.3
99
92.5 100.0 13..093.2 93.0
92.8
92.6 92.5 92.3
98
Iaccand J d Similarity (%)
92.0
93.2 100.0
93.3 93.1 92.9 92.8
97
96
L=5
91.8
93.0
114.0.0..33.. 93.3 93.2 93.1
95
L=6
91.7
92.8 93.3
93.4
93.4
100.0
93.4 93.3
94
L=8
-9.55
92.6
93.1
93.3 93.5
100.093.6 93.6
93
L=10
91.4
92.5
92.9
93.2 93.4
93.6
100.0093333
92
91
L=12
91.3
92.3
92.8
93.1 93.3 93.6 93.7
100.0
90
L=2
L=6
L=8
L=10
L=12
Figure 10: Jaccard Similarity of Neighborhoods Across Different Walk Lengths: We caleulate the Jaccard similarity neighborhoods from different walk lengths. We compare pairwise similarities for the molecule across varying walk lengths L in the pretraining and report the average similarity for the set.
expression-relatee nodes from the context graph and pretrained InfoAlign. Combined with results from Table (without cell/gene-related loss terms), additional results where nodes are removed from the context graph are presented in Table 6
First, confirm the importance of cell morphology and gene expression nodes in the context graph for pretraining. Second, we observe a greater performance drop when removing nodes compared to excluding the loss terms. This highlights the importance of including diverse data types as nodes, even without features, and suggests a promising direction for improving InfoAlign's 's pretraining by enhancing the context graph with virtual nodes.
D.5 MORE PERFORMANCE ANALYSIS ON THE RANDOM WALK
We conduct more experiments to analyze how random walk helps InfoAlign pretraining with diverse neighbor sampling. We cache the random walk results for 100 epochs and study the number of unique nodes at varying walk lengths. In Figure we report the mean and standard deviation of the number of unique nodes for all molecules in the pretraining set each walk length. We observe that the number of unique nodes is larger and varies compared to the corresponding walk length, indicating that diverse neighbors are sampled in different training steps.
We further explore the Jaccard similarity of neighborhoods extracted for the same molecule across varying walk paths, averaging similarity scores over all pretraining molecules. The pairwise similari- ties for different walk lengths are shown in Figure 10 We observe that similarity decreases as the difference in walk lengths increases, but remains above 90%. This suggests that while random walks sample diverse neighbors, varying the walk length does not significantly affect the neighborhood, which may explain the stable performance of InfoAlign in Figure These results also show that even with walk length of diverse neighbors can be obtained, likely due high-degree nodes in the context graph.
23


Reviews:
[{"review_id": "oT9Qo4UF1b", "reviewer": "Reviewer_hxs4", "summary": "The paper introduces a novel approach called Information Alignment (InfoAlign) for learning molecular representations by integrating molecular structure, cell morphology, and gene expression data. The method leverages the information bottleneck principle to optimize a molecular graph encoder and multiple MLP decoders, aiming to achieve minimal yet sufficient molecular representations.\n\nThe authors demonstrate the effectiveness of InfoAlign through extensive experiments on molecular property prediction and zero-shot molecule-morphology matching, showing superior performance compared to 27 baseline methods across four datasets.\n\nI find the paper of good quality in general. Besides, the problem it is tackling is meaningful but less explored. I suggest an acceptance to advocate this direction.", "questions": "- What are the minimum data requirements for cell morphology and gene expression to effectively apply InfoAlign? How does the method perform when these data are sparse or incomplete?\n- Given the promising performance of structure-based pretrained GNNs, have the authors considered using representations from these pretrained models instead of Morgan fingerprints as molecular features? This might leverage the strengths of both approaches.\n- How does the computational complexity of InfoAlign compare to existing methods, and what are the practical limitations in terms of dataset size and computational resources?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The approach is well-motivated. Applying information bottleneck to this problem for learning minimal yet sufficient molecular representations seems a valid match of theory and real-world problem.\n- The approach is comprehensively evaluated by comparisons across different methods and even paradigms. InfoAlign demonstrates improved accuracy over up to 27 baseline models across four datasets.\n- The paper is well-organized and clearly presented.\n- In addition to empirical evidence, the paper provides theoretical proofs to support the advantages of the proposed method.\n- The provided supplementary material contains code, dataset and checkpoints. This suggests good reproducibility of the results in the paper.", "weaknesses": "- As I mentioned in the summary, the problem is meaningful yet less tackled in the AI community. To my eyes, it is mostly due to the missing prerequisites of biological knowledge. The paper provides some explanation of the problem in the introduction, but it would be much more helpful if more context could be provided (maybe in the appendix).\n- Line 215: The motivation for computing edge weights on random walk paths is unclear, and no empirical evidence is provided to support it. Since the context graph incorporates data from three modalities, the edges likely exhibit strong heterogeneity. Is there evidence suggesting that a cumulative product of edge weights effectively captures dependency or similarity between nodes?\n- Line 303: The approach for avoiding noisy edges in computations lacks motivation and an ablation study. Providing more detail here would offer valuable insights for researchers interested in extending InfoAlign to new contexts.", "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": "4awvGMrXaW", "reviewer": "Reviewer_ALhV", "summary": "The authors propose a method called InfoAlign for predicting molecular properties by integrating three different modalities: molecular structures, gene expression, and phenomics embeddings. To learn useful representations, they construct a weighted connected graph over cell morphology profiles, related molecules, and gene expression values. They train an encoder-decoder architecture where, for a molecule of interest, they encode its representation and decode both itself and all other nodes encountered during a random walk on a pre-specified graph. This approach results in mutual information maximization between the compound of interest $x_i$ and related entities discovered through the random walk on the pre-encoded graph. The authors test their method on a variety of chemical property prediction datasets, demonstrating that they outperform various baselines, including pre-trained Graph Neural Networks (GNNs), chemical language models, uni-modal models such as cell morphology or gene expression, and some multi-modal alignment models.", "questions": "- Are the authors surprised by the relatively minor drop in ROC AUC values when omitting individual modalities?\n- Some additional relevant literature that could enhance the discussion includes:\n    [0] Cross-Modal Graph Contrastive Learning with Cellular Images\n    [1] How Molecules Impact Cells: Unlocking Contrastive PhenoMolecular Retrieval (a work evaluating zero-shot classification)\n    [2] Approximating Mutual Information of High-Dimensional Variables Using Learned Representations (proposes a scalable approach for approximating mutual information of high-dimensional objects)\n- One of the conclusions from the work is emphasizing the importance of molecular features. Dot he authors have an explanation for why the absence of molecular features in the ablation results in a ToxCast AUC that overlaps the non-ablated model performance? Without seeing the results I would expect the performance in the absence of an ablated molecular feature reconstruction loss to be a lot worse.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The authors perform a comprehensive evaluation against baseline models across a variety of datasets.\n- They convincingly demonstrate that including additional modalities improves performance, as evidenced by thorough evaluations and ablation studies.\n- The incorporation of a graph is an interesting way to introduce prior knowledge into the learning representations for a particular molecule.", "weaknesses": "- The validity and robustness of the pre-specified graph are not thoroughly explored. It would be informative to assess how sensitive the method is to the quality of the graph. For example, one experiment could involve removing 50% of valid connections and replacing them with random pairs of nodes; another could involve using a completely random graph.\n- The second gap identified by the authors is slightly misformulated: \"They treat molecules as the sole connectors between gene expression and cell morphology, ignoring the potential for genetic perturbations.\"\n    - Essentially, the authors are arguing that incorporating genetic perturbation data can further improve predictive capacity. However, there is no ablation study where this information is omitted to directly validate its impact on empirical performance.\n- Regarding the ToxCast dataset, the authors report a performance of 0.72 ROC AUC using GROVER. Did the authors use a different partitioning of the dataset than previous works?\n- The ablation loss is only regarding removal of the losses as far as I understand the data itself is still input into the training. Can the authors perform an ablation where a full data modality is not added as part of training?", "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": "KsR5yjMqGt", "reviewer": "Reviewer_VNuN", "summary": "Thank you for a really interesting read, I agree that the most important next step in molecular representation is combining different modalities and current methods leave much to be desired, and I really liked the novel way shown here to combine data in a constricted way. It is well mathematically motivated and the work done clearly extensive. \n\nSummary:\n* The authors present a new multi-domain method to learn the representation between drugs, genes and cells. In contrast to typical contrastive training setups the InfoAlign method removes the redundant information through a bottleneck using a well formulated mutual information method. \n* They show a method to create walks over the cellular context  graph representing the interactions between the modalities, and use this to populate the compute graph for their representation learning. \n* This representation framework uses Morgan fingerprints as molecular node features, CellProfiler features for the cell node features, and L1000 features for the gene node features, and the connections are based on chemical perturbations and cosine similarities. \n* The authors show results across a range of benchmarks, demonstrating good performance, achieving top ranking scores in all but 2 / 15 of the key datasets/criteria evaluated in Tables 1 and 2.", "questions": "Questions / suggestions:\n* I would ask for some more clarity on the results presented in table 1, specifically for the Broad6k results. Many columns have a +/- 0 error on a reoccurring 3.1 result. I suspect this means that there are a few tasks that are always identified with a high AUC, but given the lack of discussion of these results it’s hard to interpret. It might be worth using a different set of thresholds for the Broad6K to get more resolution? \n* It would also help in section 6 for each of the research questions to have a couple of sentences explaining what each question is and why we are asking them, otherwise the results feel disjointed to the uninitiated reader and hard to connect. \n* I found the explanation of the exact model architecture / parameters used to be lacking, and while pointed in the main text to the appendix I did not find enough information there on the size of/ structure of the MLPs / training regime to feel like the result could be replicated. This is of particular importance in Table 1 where comparing to other methods I can only guess at things like parameter efficiency etc. \n\n* In general I feel like the weighting of the paper is very theoretical, no bad thing I really like the inclusion of the mutual information, but I wonder if more details could be moved to the appendix to free up more space for the experimental method, the model construction, and discussion of the results. In the current format of the paper I find a really interesting set of ideas, that I find very hard to understand how to weight the importance of / a method that I could follow to replicate certain components. \n\n\nSpecific formatting suggestions:\n* Fig. 3. I don’t really understand what this graph is conveying? The top bar shows the proportion of single representation tasks, but the bottom splits along the model types? I either need more detailed explanation, or maybe consider a better type of graph to convert this information?\n* Table 3. (right) (This should be a separate figure - I understand it might make formatting harder but it is hard to reference) I personally don’t like these KDE type plots as they imply smooth functions from what is usually limited data, and make it almost impossible to draw quantified conclusions from, I would be much happier with a histogram if this is important information to convey. Additionally x-axis labels. \n* Fig 4 (a) - this plot is almost unreadable with the overlapping lines / y-scale, and the meaning I think is being conveyed with the LR spans such a number of magnitudes in range that I would rather see both more granularity and perhaps these results presented in a table? \n* Fig 4 (b) - I can guess what the plot means but without a key / description detailing elements like the error band (? 1 std dev I assume, based on what variation is unclear though), the points are connected, but given the discrete data this is misleading, and the huge discrepancy at length 8 suggests to me either the error band is under-estimating the variance, or there are properties of a random walk length 8 that are not discussed. Again this is a really interesting result as it's looking at the way that different parts of the compute graph contribute to the final result, but I'm left asking more questions with this figure than it answers. \n    * While I thin this statement on line 521 is correct, “we observe in Figure 4b that downstream performance on ChEMBL2K is relatively robust across a wide range of walk lengths.” The inclusion of this plot raises questions that are not answered.  \n* On a similar point to Fig 4b, I would be very interested to see a description / plot of what the typical construction of the random walk graphs contain. \n\n\n\nThank you again for a really interesting read, I like the approach with InfoAlign, and appreciate the large quantity of effort put into this work. \nI'd request a slight look again at the weighting of different sections in this paper, as I found it slightly unbalanced, and think as a reader I would find benefit from more detail in the method so I could reproduce if I wanted, and more detail in the discussion of the results to understand how to place the results in better context. i.e. which parts of the method are the most important.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The novel construction of the problem mitigates a significant issue with current implementations using a standard contrastive learning approach, by combining the modalities into a single method there is no decoupling between the different representations, and the MI approach to filter the data seems an effective way to reduce the data required to train over, potentially increasing data quality and improving time to train. \n* The visual presentation is good, with tables are figures clearly designed to convey meaning. \n* The breadth of models evaluated against gives a really clear status against both several standard approaches one might take, and also against SOTA prior work. This comparison is really nice to see.", "weaknesses": "* Model clarity - I found it hard reading this paper to extract exactly how the model was constructed to run these experiments. The method is clearly novel, and an interesting approach to the multi-domain problem, but after being strongly theoretically motivated the method / implementation details, or even a description of model size was lacking. If these details could be expanded on it would help place the model / method in the appropriate context. \n* Impact of the MI bottleneck - The mutual information bottleneck is well theoretically motivated, however I didn’t feel the power of the bottleneck was evaluated in the paper? A different approach would be just using all the data and relying on data size / model size to outweigh data quality. I like the idea with MI, but would have liked to see an ablation with this feature included / not as that would help me evaluate the importance of the bottleneck vs the joined training setup?\n* Impact of context graph - Similarly the impact of the context graph and different ways of including the context felt under-explored. Fig 4(b) I think shows that the random walk length had little impact, but how this is impacted by composition of the random walk path (are the walks just reconstructing the most likely combinations that would form a triplet anyway?) and what happens if instead of walking the relevant combinations are just grouped together. This feels like an important baseline. \n* The presentation of the research questions felt more like a report than a paper, I would have liked to see some more motivation / explanation of each rather than assuming knowledge on the readers part.", "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": "F25Mg26Q8H", "reviewer": "Reviewer_Sxqu", "summary": "This article proposes a pre-training method where the authors utilize genes, cells, and molecular information as nodes, and construct a context graph using the interactive or similar relationships between different components as edges. The authors then extract paths on the graph using random walks to construct training data, design training objectives from the perspective of mutual information, and integrate cellular and genetic information into molecular representations. After pre-training, the authors fine-tune the model on downstream tasks across four datasets to validate its performance.", "questions": "- **Methodology**\n  - How is the neighboring node of $x$ defined in line 250? Is it the neighbor of $x$ on the context graph or all the $v_i$ on the sampled path?\n  - As shown in Fig. 2, different decoders are used to reconstruct the features of different types of nodes. However, if there are multiple neighbors $v_1, v_2, \\ldots, v_n$ of a certain molecule $x$ sharing the same type, how can the decoder decode various $y_{v_i}$ from the same encoded latent  $z$ of $x$ without any additional information about $v_i$ being provided?\n\n- **Experiments**\n  - The ChEMBL dataset provides information on whether a certain molecule exhibits activity against a specific target. In this manuscript, a task is defined as predicting whether a molecule can interact with a given target. However, it is unclear why the molecule would be characterized by **Cell Morphology** and **Gene Expression**. Such information should be related to the task (target) itself rather than the input molecule. Could the authors explain how **Gene Expression** and **Cell Morphology** data are generated for molecules?\n  - To substantiate the claim that the proposed multimodal alignment method more effectively models molecular properties within a cellular context, the authors must include a baseline comparison using a simple concatenation approach for alignment. Specifically, the authors should employ a pretrained GNN, such as Uni-Mol[1], to extract molecular features and concatenate them with Cell Morphology and Gene Expression representations as inputs for prediction tasks on the ChEMBL2k and Broad6k datasets.\n  - During the fine-tuning for downstream tasks, is the encoder frozen or is a full-fine-tune performed with the encoder also being trained?\n\n**If the authors can address all my major concerns, I would be pleased to raise the score.**\n\n[1] Zhou G, Gao Z, Ding Q, et al. Uni-mol: A universal 3d molecular representation learning framework[J]. 2023.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This article pioneeringly incorporates cellular modality information in molecular representation learning, which enables the learned molecular representations to achieve better results in biochemistry-related tasks.\n- The article conducts sufficient experiments to validate the effectiveness of the model.\n- The paper is well-motivated; adding cellular modality data to enhance the model's performance in biochemical aspects is very reasonable.", "weaknesses": "- The writing requires improvement; the methodology is difficult to follow, and many variables are not clearly defined. See Questions for details.\n- The description of the dataset and downstream tasks is not sufficiently detailed, with some parts being confusing. See Questions for details.", "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)