A probability contrastive learning framework for 3D molecular representation learning
Authors: Jiayu Qin
Jian Chen
Rohan Sharma
Jingchen Sun
Changyou Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: contrastive learning; molecular property prediction; expectation maximization
Primary area: machine_learning_for_healthcare

Abstract: Contrastive Learning (CL) plays a crucial role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. It has inspired various applications in molecular property prediction and drug design.
However, existing molecular representation learning methods often introduce potential false positive and false negative pairs through conventional graph augmentations like node masking and subgraph removal. The issue can lead to suboptimal performance when applying standard contrastive learning techniques to molecular datasets. 
To address the issue of false positive and negative pairs in molecular representation learning, we propose a novel probability-based contrastive learning (CL) framework. Unlike conventional methods, our approach introduces a learnable weight distribution via Bayesian modeling to automatically identify and mitigate false positive and negative pairs. This method is particularly effective because it dynamically adjusts to the data, improving the accuracy of the learned representations. Our model is learned by a stochastic expectation-maximization process, which optimizes the model by iteratively refining the probability estimates of sample weights and updating the model parameters.
Experimental results indicate that our method outperforms existing approaches in 13 out of 15 molecular property prediction benchmarks in MoleculeNet dataset and 8 out of 12 benchmarks in the QM9 benchmark, achieving new state-of-the-art results on average.

A Probability Contrastive Learning Framework for 3D Molecular Representation Learning
Jiayu Qin University at Buffalo jiayuqin laayuuii@buffll..cem buffalo.edu
Jian Chen University at Buffalo
Rohan Sharma
University at Buffalo
(alen378@bufflloo
rohanjag rohanjag@buffall @buff@lo..di
Jingehen Sun University at Buffalo jsun39@ @ buffalo.edu
Changyou Chen
University at Buffalo
changyou@buffalo.co @ufffllo.edu
Abstract
Contrastive Learning (CL) plays crucial role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. It has inspired various applications molecular property prediction and drug de- sign. However, existing molecular representation learning methods often introduce potential false positive and false negative pairs through conventional graph augmen- tations like node masking and subgraph removal. The issue can lead to suboptimal performance when applying standard contrastive learning techniques to molecular datasets. To address the issue of false positive and negative pairs in molecular repre- sentation learning, we propose novel probability- contrastive learning (CL) framework. Unlike conventional methods, weight distribution via Bayesian modeling to automatically identify and mitigate false positive and negative pairs. This method is particularly effective because dy- namically adjusts to the data, improving the accuracy of the learned representations. Our model is learned by stochastic expectation-maxtt maximization process, which optimizes the model by iteratively refining the probability estimates of sample weights and updating the model parameters. Experimental results indicate that our method outperforms existing approaches in 13 out 15 molecular property prediction benchmarks in MoleculeNet dataset and 8 out of 12 benchmarks in the QM9 benchmark, achieving new state-of-the-an results on average.
approach introduces a learnable
Introduction
We investigate the problem of learning representations from molecules, a field known as molecular representation learning (MRL). MRL has gained significant attention due to its critical role in enabling learning from limited supervised data for applications such as molecular property prediction [29] [3
5) and drug design [14 20 22). Molecular representation learning involves creating models that can derive meaningful and generalizable representations of molecules, which can then be used to enhance various downstream applications. Among the most common methods in MRL contrastive learning (CL), which leverages large-scale unlabeled molecular datasets to learn robust representations. CL works by contrasting different augmentations of the same molecule ensure that the model learns to recognize the essential features of the molccule, thereby improving performance on tasks such as molecular property prediction and drug design.
With the success of contrastive learning methods in computer vision and multi-modalin pretraining 7 27]. various contrastive learning approaches have been proposed for molecular representation learning. MolCLR 33 introduces a contrastive learning framework specifically for molecular representation
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

learning. It employs atom masking and edge removal as data augmentations, which enhances the performance of Graph Neural Network (GNN) models on a variety of downstream molecular property prediction benchmarks. In contrast, GraphMVP 20] incorporates both 2D topology and 3D geometry during pre-training, though its downstream tasks primarily utilize 2D topology. These methods highlight different strategies for applying contrastive learning to molecular data, focusing on unique aspects of molecular structures to improve learning efficacy.
Although existing works have demonstrated the success of contrastive learning in molecular property predictions, they still face significant drawback: the reliability of "positive" and "negative" labels in augmented molecule pairs. For example, MolCLR 33 uses augmentations like atom masking and edge removal, which can lead to false negative pairs when molecules with similar structures and chemical properties are labeled as negatives. Similarly, GraphMVP (20), which incorporates both 2D topology and 3D geometry, can also mislabel structurally similar augmented molecules as negatives due to its augmentation processes. These augmentations often remove parts of the molecular graph, such as nodes, edges, and subgraphs, resulting in potentially incorrect pairings. This issue is exacerbated by the large volume and extensive augmentations applied to molecular datasets, naturally leading to numerous falsely aligned pairs.
The fundamental problem lies in the random nature of these augmentations. Existing molecular contrastive learning methods assign hard positive and negatives to molecule pairs and do not account for the probabilistic relationships between molecules. Figure 3 provides an example of false positives and negatives resulting from graph augmentations MoICL ,where two distinct graph augmen- tations are applied to enhance two different molecules. The augmented molecule pair originating from the same molecule is categorized as positive, while other molecule pairs within the same batch are considered negative. However, as illustrated in the figure, the correct contrastive learning setup should consider molecules with structural similarities as positive pairs, even when they originates from different molecules. In contrast, the same molecule subjected to different augmentation methods may also be considered negative due to structural dissimilarities. Existing methods like MoICLR 33 fail to maintain this distinction, where augmented pairs from the same molecule are always treated as positive, while pairs from different molecules within the same batch are always treated negative, regardless of their structural similarity. This mislabeling results in false positives and negatives, undermining the effectiveness of the contrastive learning process.
0.97
0.96
0.67
0.65
SH
SH
Node removal
Edge removal
SH
SH
SH
Node removal
Edge
removal
Figure Existing problem in molecular contrastive learning. Adopt node removal and edge removal for molecular contrastive learning can lead to false positive and false negative problems. Blue lines indicate positive pairs and yellowing lines indicate negative pairs. The numbers on each line indicate the chemical similarity between the augmented pair of molecules. In this case, positive pairs indeed have lower similarity than negative pairs.
To overcome the aforementioned issue, we introduce a generalization of existing contrastive learn- ing frameworks for molecular representation learning with probabilistic modeling. Our approach introduces data-pair weights as additional random variables, and dynamically infers optimal weights to account for false positive and false negative pairs, which can effectively address the mislabeling problem in previous methods. By incorporating a probability framework, we can effectively manage the uncertainty in data pair assignments. Specifically, we introduce a novel Bayesian inference
2

---

methods with Bayesian data augmentation to automatically infer these weights through posterior sampling. This allows us to optimize the model parameters efficiently using stochastic expectation
maximization.
It is worth mentioning that while MoICLR authors introduced i-MoICLR 32 to address similar issues by penalizing faulty negatives with a fingerprint-based similarity metric and a motif-level data augmentation called fragment contrast, our method offers distinet advantages. Unlike i-MoICLR which relies on direct fingerprint similarity, our approach introduces a novel probabilistic contrastive learning framework. This framework dynamically infers weight distributions and optimizes through stochastic expectation maximization, eliminating the need for explicit Tanimoto similarity calculations. Our method addresses the issue of false negative pairs more fundamentally and efficiently, providing a more robust solution for molecular contrastive learning.
In addition, method is flexible and can be applied different molecular representation learning framework. In this paper, we first integrate our method into MolCLR (33] series model and henchmark the performance on 2D non-charality MoleculeNet [35] dataset. We then integrated our method into Uni-Mol [42] and evaluate its performance on M3leculeNet(35]. We also trained and evaluated our model on the QM9 [28] dataset, following Equiformer [17]. With molecular property prediction tasks, we aim to test our model's 's ability in extracting useful features from molecular. Extensive experiments show that our method outperforms all other molecular representation learning baselines, including contrastive and non- con-ootrrattiee methods.
The contributions of this paper can be summarized as follows:
To tackle the challenges posed by false positive and negative pairs, we introduce probability method for molecular contrastive learning. By introducing different weights as random variables various false positive and negative pairs, we effectively mitigate the impact of these erroneous pairs on the learning process.
To optimize our probabilistic contrastive learning framework, we propose a novel and effec- tive optimization algorithm based on Bayesian data augmentation and stochastic expectation maximization, to simultancously perform posterior inference and model optimization.
Through extensive and large-scale experiments, demonstrate enhanced performance across multiple public benchmarks for molecular representation learning, validating the effectiveness of our proposed method.
2 Methods
2.1 Learning Representations from Molecular Graphs
We begin by elucidating the foundational setup and notation in molecular contrastive learning. Molecules can represented as 2D or 3D graphs depending on datasets. 2D molecule graphs have atoms nodes and bond as edges. 3D molecule graphs additionally adds spacial positions of the atoms. For simplicity, we adopt static atom positions this paper.
In molecular representation learning, as illustrated in Figure of N molecules. Each molecule, represented as x4, undergoes stochastic augmentation strategies to generate two augmented versions, denoted as xx; These augmentations involve methods such as atom masking, edge perturbation, and subgraph removal, transforming the original molecular structure while preserving its core characteristics. Among the resulting 2N augmented molecules, each pair (x,,*' x',) is treated
we start by randomly sampling a batch
positive pair, while the remaining 2(N - 1) augmented molecules within the same batch are considered negative samples. This setup allows to utilize contrastive learning effectively by distinguishing between similar and dissimilar molecular structures. neural network encoder [(x; 0), parameterized by 0, is employed to extract representation vectors from the augmented molccular samples. In this paper, we utilize three different types of encoders in various experiments, depicted in Figure B, C, and D. These encoders include Graph Neural Networks (GNNs) and Transformers, each providing unique advantages for capturing the intricate features of molecular structures.
Let si+ A ('m(z,,z')) represent the similarity score between the positive pair (x,, x'() after the encoder, and 8 A sim (Zi, 2x) signifies the similarity score between the negative pair (x;, X1.), and sim(:,-) represents any positive-valued similarity metric. In this paper, adopt the commonly
3

---

Imput 3D Graph
sfm-
5n1-
MatMuu
SB+
Embeddinng
Layer Norm
Transtormer feature extractor
GNN Transtormer feature extractor
Math/ul
Equivariant Gmph Attention
Linear
Unner
Linear
Atom Representaaion
Pair Regresentation
Layer Norm
Network
x1
Xn!
x
Layer Norn
Forward Netwoek
Output
Molecule x1
Molecuiee Xn
Head
Figure 2: (A) Molecular contrastive learning Molecules are represented as 2D or 3D molecule graphs. Two stochastic augmentation strategies are applied to each graph, resulting in aug- mentations. A feature extractor is used to extract features and contrastive loss is used to maximize the similarity positive pairs and minimize the similarity of negative pairs B,C,D: Different ar- chitectures used as feature extractors in different experiments. (B) Uni-Mol [42 architecture used
MolcculeNet (35) Dataset experiment. (C) GCN (2T) architecture from MolCLR (33) used in Non- Non-Chirality MoleculeNet [35] experiment. (D) Equilormer (17] architecture used in QM9 [28
dataset experiment.
used exponential cosine similarity, defined as (2, z2) >2////// I||aall, where t denotes temperature parameter.
2.2 Probability Weighted Contrastive Learning
We describe the proposed probability framework for molecular contrastive learning. In standard contrastive learning, tries to encode data samples to latent space such that positive pairs stay close to each other while negative pairs are pushed away. The contrastive loss function is:
N
Σ ((22 - 1k) + f(2k, 2k - 1)], with {(i.j) - -log 5,+ +
5;-
(1)
N
-2N
2N Siml 1
1=1
31,6-
As mentioned, one issue of directly applying the contrastive learning into molecular representation learning is the potential false positive and negative molecular pairs, as discussed in the introduction. This could confuse the learning, ending up with sub-optimal representations. Is there way to automatically identify and differentiate these pair data? In the following, we propose Bayesian approach address this issue that allows the algorithm for automatic inference of the degree of positiveness and negativeness of data pairs, involving enhancing the standard contrastive loss by incorporating learnable stochastic weights for all data pairs. To more specific, we introduce local learnable weights, denoted as w, for each positive pair and w.J. for each negative pair. We then define a weighted contrastive loss based on these introduced weights. This modification aims to mitigate the issues by automatically assigning relatively lower weights (or weights) to false positive and false negative pairs;
L. N Σ ((2 1, 1,2k) + E(2k, 2k 1)], E(i.j) = =llog
8
(2)
w. 8
Σ
2N
W
8sk-
One problem with this formulation, however, is that not realistic to compute and store all the weights in the learning process. This precaution arises from the quadratic growth in the number of

---

weights to be calculated as the training data size increases. Furthermore, the random nature of our augmentation method further adds complexity to the pre-calculation and storage of these weights.
straightforward baseline for calculating these weights can be envisioned as follows: we can consider these weights in binary fashion, with all weights initialized to one. In the learning process, if for some positive pairs the similarity score falls below a specified threshold, we set the corresponding weights to zero, marking these positive pairs as false positives. Conversely, if for some negative pairs the similarity score exceeds a threshold, we set the associated weights to zero, indicating false negatives. challenge associated with this baseline method, however, lies in the establishment of rigid similarity threshold to create binary division of weights between zero and one. This approach proves less suitable for our molecular contrastive task as these heuristically chosen thresholds might not be optimal.
To address this challenge, propose principled Bayesian approach that allows adaptively inferring the optimal weights by Bayesian inference. Specifically, we treat the weights to be random variables and assign appropriate priors them. We consider two types of priors: a Bernoulli prior to model weights as binary random variables and Gamma prior to represent them as positive values. For simplicity, we model positive weights using the Gamma distribution and negative weights using either the Gamma distribution or the Bernoulli distribution, as expressed by the following formulas:
Option 1
Gamma priors for continuous weighting:
w
Gamma(a+ b4), wiz Gamma(a_,bb).
Option 2 Bernoulli priors for selective weighting:
T Gammala_,ba b+), Wik ~ Bernoulli((.).
here, a.,, b+. a_ and bre are shape and rate parameters for Gamma distribution and d_ is the probability parameter for Bernoulli distribution.
With our reformulation, we can define a joint distribution over the global model parameter and local random weight variables Wi and Wag. as:
P ((ww)) (wi)) ,0; x II
W 85t
p(( Hp((w,))))(0
(3)
x;ED WT Styt Ek=
One problem with the above formulation, however, is that posterior inference of the weights is challenging, due to the lack of convenience posterior distributions.
Fortunately, inspired by (2], we can introduce an augmented random variable u, that is associated to data point x,. Consequently, we can define an augmented joint posterior distribution of the random variables 0, u, w. denoted as [F]]] (w
[
to be
p0 u. u.
II EX.ED
w 8; w.w +ct II
Mil P ((wt)) ((w7)) p(0),
where 2 { U1, u2...... Uipi and - {w f w }. It worth noting that this joint distribution is equivalent to the original distribution (3) because (3) is recovered if one marginalize out the auxiliary random variables in (4). In other words, optimization thought { equivalent to optimization over (3). Consequently, we can perform learning and inference based on the augmented posterior of p(0, u, w D), which preserves a much convenient form for posterior inference. In the following, we propose an efficient algorithm based on stochastic expectation maximization (stochastic EM) to alternatively infer the local random variables w and optimize the global model parameter
2.3 Efficient Inference and Learning with Stocastic Expectation Maximization
We propose stochastic EM algorithm for efficient inference and learning of our model. Stochastic EM (24) is stochastic variant of the EM algorithm, which is an itcrative method for finding the
In the sense that marginalizing over the augmented random variables {w: andd {wij Jn (0, U, (wit) }. [wī] P) gives back to the original p ([wi) {wī) , D). Thus, learning and infer- ences on the two forms equivalent.
5

---

maximum likelihood of model parameters in statistical models when data is only partially, or when model depends on unobserved latent variables (4T).
In our setting, the objective of stocastic EM is to maximize the posterior in equation The basic idea alternatively 1) optimizing model parameter 0 with fixed (u,, w) and 2) sampling (u, w) with fixed O. To this end, we follow standard procedures in stochastic EM to divide the learning into three steps: Simulation, Stochastic Expectation, and Maximization. Specifically, simulation corresponds sampling local random variables u and w for a batch of data; stochastic expectation then uses the sampled auxiliary random variables update the model parameter by maximizing a stochastic objective Q(0), defined as: (0+1(0)) = () + ) (log p
D) Q+())) at iteration t + 1, where {Ae} is sequence of decreasing weights. And maximization corresponds to maximizing the stochastic objective constructed in the previous step. In the following, we detail the three steps.
Simulation Given the joint posterior distribution in equation and the current batch of data, the posterior distributions of the local random variables u and w can be directly read out, which simply follow Gamma or Bornoulli distributions the following forms:
Uy { wt Wa 8 Gamma Cus by w. 85t + E (i Vi, and
w {u, 0} ~ Gamma (1 + a+, U,31+ + b+). and
(5)
Option 1: W/z {u, 0 Gamma (a_, U4 1188;- + b-), Vi
a_e 41804
Option 2: Wik
{u, 0} Bernoulli
a + Pu
We use the Gamma prior because naturally lends itself to conjugacy in the posterior, which significantly eases the posterior sampling procedure. Also, it is known for its flexibility in shape and scale to model positive continuous variables, which suitable for sample weights in our setting.
Stochastic Expectation We then proceed to calculate the stochastic expectation based on the sim- ulated local random variables above. For notation simplicity, we define Qn(A) = 0. Then we can reformulate (+) by decomposing the recursion, resulting in
Q0+1(0) Σ 7r log (0, ur wr Dr), where A, A Ar II
(1-Ae),
(6)
M=T+1
where 7 indexes the minibatch and the corresponding local random variables at the current time т.
Maximization The stochastic
expectation objective (6) pro- vides convenient form for stochastic optimization over
Algorithm 1 Contrastive Learning with Stochastic EM
1: Initialize 0; set f =
2: for batch of molecules in loader do
time, similar to online optimiza- tion (Bent & Van Hentenryck, 2005). Specifically, each time t, we can initialize the param- eter 0 from the last step, and update it by stochastic gradient ascent the logelikelihood, logp(0 u,, w, (D,)) calcu- lated from the current batch of data. To reduce variance, we propose to optimize marginal version by integrating out u. from p (0, u,, w, (D+), which essentially reduces to our original weighted contrastive loss equation 2 With the above steps, it is ready to optimize the model by stochastic EM. The detailed steps are described in the Algorithm [
3: Augment each molecule x; into a pair (X,, x';)
4: Calculate positive/negative similarity scores s and
for all the molecule pairs
5: Initialize all the weights w and be one
6: for k = 1 to iter [4 in practice] do
7:
Sample and w according distributions end for
8:
9:
Calculate the weighted contrastive loss equation with the sampled w the current batch of data
10:
Update the model parameter by stochastic gradient de- scent with the calculated weighted contrastive loss
11:
t t=t
12:
end for
3 Related works

---

Contrastive Learning As popular self-supervised learning paradigm, contrastive learning focuses on learning semantically informative representations for downstream tasks [16] 39 91. The most widely used loss function is InfoNCE [25] which pulls in the representations between positive sample
pairs while pushing away that between negative sample pairs.
Molecular Representation Learning Representation learning on large-scale unlabeled molecules attracts much attention recently SMILES-BERT (31 is pretrained on SMILES strings of molecules using BERT. Subsequent works are mostly pretraining on 2D molecular topological graphs [15] 29]
MolCLR 33] applies data augmentation to molecular graphs at both node and graph levels, using
self-supervised contrastive learning strategy to learn molecular representations. I-MoICLR (32) is
improved version of MoICLR that uses new data augmentation and introduces weighted contrastive learning for mitigating false pair problem. Further, several recent works try to leverage the 3D spatial information of molecules, and focus on contrastive or transfer learning between 2D topology and 3D) geometry of molecules. For example, GraphMVP [20] proposes a contrastive learning GNN-based framework between 2DD topology and 3D geometry. GEM (5) uses bond angles and bond length as additional edge attributes to enhance 3D information. Uni-Mol 42 is universal 3D molecular pretraining framework that significantly enlarges the representation ability and application scope in drug design.
Noisy Pairs in Contrastive Learning Noisy data pair problem have been found and studied in contrastive learning community. NLIP (II) enforces pairs with larger noise to be less similar in embedding space improve model training. Jpplly noise estimation component to adjust the consistency between different modalities for action recognition task. RINCE (8) uses ranked ordering of positive samples to improve InfoNCE loss. (3) introduces a new debiased contrastive learning loss function by transforming the distribution of negative samples. Matchdrop (34) designed a new graph augmentation method to alleviate the false positive sampling problem by retaining the most critical parts of the graph and augmenting the unimportant parts.
Stochastic Expectation Maximization Stochastic EM (24) stands a pivotal algorithm in machine learning and probabilistic modeling for large-scale Bayesian inference. Building upon the foundations of the classical Expectation- peeettation-Maxmizztttee Maximization (EM) algorithm (18). Stochastic EM offers an efficient solution for parameter estimation in situations involving vast datasets or latent variables, c.g.. to maximize the log-likelihood of p(z, 'D 0), where is the dataset, z is the local random variable and 0 is the global model parameter. By leveraging the power of mini-batch sampling, Stochastic EM strikes a balance between computational scalability and estimation accuracy. It has found widespread utility in various domains, including clustering (I), topic modeling (40), and latent variable modeling (4I), making it an indispensable tool to cope with complex probabilistic models and extensive data and natural fit to our problem.
4 Experiments
We evaluate method molecular property prediction tasks. Our approach is designed to be versatile component that can be seamlessly integrated with various molecular property prediction datasets and models. In this study, integrate our model into three different existing models: Uni-Mol (22) I. I-MoICLR (32), Equiformer [17] and assess its performance on three distinct datasets: MoleculeNet (35), MoleculeNet without chirality, and the QM9 (28) dataset. For all experiments, we provide detailed experiment settings in Appendix
4.1 The MoleculeNet Dataset
MoleculeNet (35) is popular benchmark for molecular property prediction, including datasets focusing on different molecular properties, from quantum mechanics and physical chemistry to biophysies and physiology. For a fair comparison, we integrated our method into Uni-Mol 42] framework. We applied both the Gamma and Bernoulli versions of our method, as shown in Table I. In our contrastive learning framework, we used the representation of the [CLS] token as the final encoded representation, representing the entire molecule. Additionally, we incorporated the original dinee-dimensional recovery loss as an extra loss function. The model was trained on the same large-scale dataset, including 19 million molecules and 209 million conformations, as in the original paper. We used the same evaluation metrics: ROC_AUC for classification tasks and RMSE and MAE for regression tasks.
7

---

As shown in Table[and 2 our method outperforms Uni-Mol(4] and GEM (5], the current state- of-the-art methods. with an average gain of 1.3 percent classification tasks and 7.6 percent in regression tasks. This substantiates that our approach facilitates more flexible training with higher tolerance for false positive and false negative data pairs, thereby enhancing the model's performance in molecular representation learning.
Table 1: ROC AUC on molecular property prediction classification tasks (Higher is better)
Dalasets
ВВBP
BACE
ClinTox
Tox21
ToxCast
SIDER
HIV
PCBA
MUV
# Molecules
2039
1513
1478
7831
8575
1427
41127
437929
93078
# Tasks
1
2
12
617
27
128
17
D-MPNN [37
71.0
80.9
90.6
75.9
65.5
57.0
77.1
86.2
78.6
Attentive
FP
36
64.3
78.4
84.7
76.1
63.7
60.6
75.7
80.1
76.6
N-Gramgr
19
69.7
77.9
77.5
74.3
66.8
77.2
76.9
N-GramxGB
19
69.1
75.8
65.5
78.7
74.8
PretrainGNN
T0
68.7
84.5
72.6
78.1
65.7
62.7
79.9
86.0
81.3
GraphMVP
20
72.4
81.2
79.1
75.9
63.1
63.9
77.0
77.7
GEM 55
72.4
85.6
90.1
78.1
69.2
67.2
80.6
86.6
81.7
MolCLR
33
72.2
82.4
91.2
75.0
58.9
78.1
79.6
Uni moi-Molae
72.9
85.7
91.9
79.6
69.6
65.9
80.8
88.5
82.1
Ours (Gamma)
76.7
88.2
80.1
69.9
63.6
83.0
89.6
79.0
Ours (Bernoulli)
73.7
84.3
79.8
68.8
64.9
80.8
89.3
82.9
Table 2: Performance molecular property prediction regression tasks (Lower is better)
Datascts
ESOL
FreeSolv
Lipo 4200
QM7
QM8
QM9
MEAN (RMSE)
MEAN (MAE)
Molecules
1128
642
6830
21786
133885
Metric
RMSET
MAET
D-MPNN 37
1.050
2,082
0.683
103.5
0.019U
0.00814
1.272
34,509
GROVER
29
0.895
2.272
0.823
92.0
0.0224
0.00986
1.33
30.67
MolCLR 33
1.271
2.594
0.691
66.8
0.0178
1.519
GraphMVP
20
1.029
0.681
GEM
0.798
1,877
0.660
58.9
0.0171
0.00746
1.112
19.642
Uni-Mol
42
1,480
0.603
0.0156
0.00467
13.940
Ours (Gamma)
0,590
0,0142
0.00395
12,839
Ours (Bemoulli)
0.664
1,358
0.626
55.6
0.0154
0.0056
0.883
18.541
4.2
Non-Chirality version MoleculeNet
also integrated our method into MoICLR 33 framework. MolCLR and I-MoICLR are 2D based methods, their experiments are conducted different version of MoleculeNet dataset that does not consider chirality. We adopted the same dataset, augmentation, GNN- based encoder and other settings. As shown Table (3] our method outperforms I-MoICLR 7 out of 9 downstream tasks and got average of 2 points increase non-chirality MoleculeNet classification datasets.
In order make fair comparison with I-MoICLR [32].
Table Comparison against i-MolCLR on non-chirality MoleculeNet dataset
Without Chirality
ВВBР
BАCE
Clin'Tox
Tox21
SIDER
HIV
MUV
MEAN
I-MOLCLR [32]
76.4
88.5
95.4
79.9
69.9
80.8
90.8
83.1
Our Method
78.3
94.8
91.4
84.9
72.7
85.5
88.0
85.1
4.3 QM9 Dataset
The QM9 dataset 28 is another popular dataset in molecular property prediction, it consists of 134k small molecules, and the goal is to predict their quantum properties. For this dataset, we choose equiformer [17] as baseline method. The data partition training, validation and testing sets. We use both our contrastive loss function and original minimize mean absolute error(MAE) training objectives.
has 1 10kk,1kk.nd 11k molecules in
As shown in we get state of the art result in 8 out of 12 baselines. The increase is relatively subtle compared with other dataset, argue that this is due to the fact that QM9 is relatively small

---

regarding number of molecules in training set, and also the saturation on performance achieved by different methods.
Table 4: Experiment results on QM9 dataset
Methods
AE
homo
E Jumo
Cv
H
R22
pul
ZPVE
GraphCL
39
0.066
45.5
26.3
22.9
0.027
0.028
10.2
9.6
0.095
9.7
9.6
1.42
JOAOv2
27.8
22.2
9.8
3D-MGP
0.057
37.1
21.3
18.2
0.020
0.026
9.3
8.7
0.092
8.6
8.6
1.38
Transformer-Uu
0.041
27.4
17.5
16.2
0.037
0.022
9.63
9.39
0.075
9.41
9.37
1.18
Equiformer
0.046
30
0.023
6.74
Ours
0.037
24.2
21.1
13.7
0.022
0.022
6.2
6.31
0.082
7.22
9.40
1,09
4.4
Ablation Study
Distribution of similarity scores Our method is largely motivated by the observation that previous MCL. approaches neglect potential semantic dissimilarity between positíve samples and that account- ing for this phenomenon can improve learned molecule representations.
FigureA|Sse See Appendix A). we plot the distribution of similarity scores for both positive and negative samples. Figure left reveals that our method yields larger similarity scores with lower variance for positive pairs compared to MolCLR (33 baseline which uses standard contrastive learning method. Figure Figure/AJright reveals that our method also mitigates the false negative problem standard CL. also shows that our method sometimes assigns lower similarity scores to positive pairs. While it may seem counter intuitive to assign lower similarity scores to positive samples, we argue that doing so the very reason our method captures dissimilarity between positive pairs. By allowing some degree of alignment between the right set negative examples, our method able to minimize the inconsistencies between shared context of related positives and negatives. This in turn allows representation space, resulting increased robustness and downstream performance.
learn overall more coherent
Comparisons with the Standard Contrastive Learning We conducted an ablation study to show- case that our method of probablistic framework of contrastive learning has already achieved strong emperical results and demonstrate the improvement brought by adding the 3D-aware loss functions on MoleculeNet (35] classification dataset. We first examined the effect of adding the probabilistic framework to the standard contrastive loss, and the 3D-aware loss functions as implemented in Uni-Mol (2)
Table 5: Ablation Study MoleculeNet Classification Datasets
BBBP
ClinTox
Tox21
ToxCast
SIDER
PCBA
Standard CL
69.3
81.5
84,1
75.5
63.4
58.9
78.3
84.1
72.5
75.2
CL 3D Loss
75.1
86.8
87.9
78.9
68.5
62.8
81.8
88.0
77.1
78.1
Probabilistic Framework
74.1
86.3
88.2
79.5
68.2
63.1
82.5
88.4
77.1
78.6
Both
76.7
88.2
89.4
80.1
63.6
83.0
89.6
79.0
80.1
Table 5 presents the results of our ablation study. Incorporating the probabilistic framework resulted 3.4-point increase in ROC-AUC, significantly enhances the model's
in a great improvement performance. On the other hand, introducing the additional loss component led to increase in ROC-AUC by 2.9 points, demonstrating its secondary role in enhancing the model's performance. When adopt both of them, can get the final ROC-AUC of 80.1 average MoleculeNet classification datasets.
Hyperparameters We also conducted an ablation study to determine the optimal hyperparameters (e.g., a.,, a_) MoleculeNet classification datasets. We selected 04, a_, b4, and b_ from the range [, 5, 10]. Table inddictss that our method achieves the best performance with a = 5 and a b_ b_ 1. Tuning different hyperparameters affects performance, with increase in a. from to 5 leading to a 1.6 percent performance gain.
Conclusion
In this paper, we investigate important yet unnoticed limitation of molecular contrastive learning, where augmented graph data come with false positive and false negative data pairs. As remedy, we propose principled solution to molecular contrastive learning by reformulating it into a probability

---

Table 6: Abalation studies on hyperparamelers for MoleculeNet classification tasks
4+
5
10
5
5
5
a-
1
10
b+ b_
5
10
5
5
10
Avg. ROC-AUC (%)
78.8
80.4
79.6
79.3
80.0
79.4
79.3
framework and introducing random weights for data pairs. With a Bayesian data augmentation technique, the random weights can be efficiently inferred via sampling, and the model parameter can be efficiently optimized via stochastic expectation maximization.
The effectiveness of our innovative approach has been proven through rigorous evaluations multiple molecular property prediction benchmarks. The results also showcase the wide-ranging applicability and improved robustness of our proposed method over existing methods for learning molecular representations.
We believe our method is valuable addition to the literature on molecular contrastive representation learning, which can further boost the performance of state-of-the-art molecular representation learning models for drug design.
Acknowledgement
This work is partially supported by NSF AI an Amazon research award, and Adobe gift fund. Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation, the Institute of Education Sciences, the U.S. Department of Education.
NSF RI-2223292, NSF IIS I14-174761
References
[1] Stéphanie Allassonnière and Julien Chevallier. "A new class of stochastic EM algorithms. Escaping local maxima and handling intractable sampling". In: Computational statistics and data analysis 159 (2021), p. 107159.
[2] Changyou Chen et al. "Why do We Need Large Batchsizes in Contrastive Learning? A Gradient- Maddient-Bias Perspective". In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., 2022, pp. 33860-33875.
[3] Ching- Yao Chuang, Joshua Robinson, Yen-Cheng Lin, et al. "Debiased contrastive learning". In: Advances in Neural Information Processing Systems. Vol. 33. 2020, pp. 8765-8777
[4] Jacob Eberhardt et al. "AutoDock Vina 1.2.0: New docking methods, expanded force field, and python bindings". In: Journal of chemical information and modeling 61.8 (2021), pp. 3891-
3898.
[5] Xiantao Fang, Ling Liu, Juncheng Lei, et al. "Geometry-zehanced molecular representation learning for property prediction". Nature Machine Intelligence 4.2 (2022), pp. 127-134. Han et al. "Noise-tolerant learning for audio-visual action recognition". In: IEEE Transac- tions Multimedia (2024).
[6]
[7] Kaiming He, Haoqi Fan, Yuxin Wu, al. "Momentum contrast for unsupervised visual representation learning". In: Proceedings of the IEEF/CVF conference on computer vision and pattern recognition. 2020, pp. 9729-9738.
[8] T. Hoffmann et al. "Ranking info noise contrastive estimation: Boosting contrastive learning via ranked positives". In: Proceedings of the AAAI Conference Artificial Intelligence. Vol. 36.
1. 2022, pp. 897-905.
[9] Pengfei Hu. Hao Zhu, Jian Lin, et al. "Unsupervised contrastive cross-modal hashing". In: IEEE Transactions on Pattern Analysis and Machine Intelligence 45.3 (2022). pp. 3877-3889,
[10] Weihua Hu, Bowen Liu, Joseph Gomes, et "Strategies for pre-training graph neural net- works", In: arXiv preprint 1905v:1995..2200 (2019).
10

---

[11] R. Huang et al. "Nlip: Noise-robust language-iman pre-training". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. 2023, pp. 926-934.
[12] Runxuan Jiao, Jingrui Han, Wen Huang, et al. "Energy-motivane equivariant pretraining for 3d molecular graphs" In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37.
7. 2023, pp. -104-
[13] David R Koes, Mark P Baumgartner, and Carlos J Camacho. "Lessons learned in empirical scoring with smina from the CSAR 2011 benchmarking exercise". In: Journal of chemical information and modeling 53.8 (2013), 3), pp. 1893-1904.
[14] Panagiotis Koukos, Li C Xue, and Alexandre MJJ Bonvin. "Protein-ligand pose and affinity prediction: Lessons from D3R Grand Challenge 3". In: Journal of computer-aided molecular design 33 (2019), pp. 83-91.
[15] Peng Li, Jun Wang, Yu Qiao, et al. "An effective self-supervised framework for learning expressive molecular global representations to drug discovery". In: Briefings in Bioinformatics 22.6 (2021), bbab109.
[16] Yiqun Li, Mengya Yang, Dong Peng, et al. "Twin contrastive learning for online clustering".
In: International Journal of Computer Vision 130.9 (2022), pp. 2205-2221.
[17] Yuyang Liao and Tess Smidt. "Equiformer: Equivariant graph attention transformer for 3d atomistic graphs". In: arXiv preprini arXiv:2206. / 1990 (2022).
[18] D. Lin. An Introduction to Sspectation-Maximizzae
2011.
[19] Shengchao Liu, M Fatih Demirel, and Yatao Liang. "N-gram graph: Simple unsupervised
representation for graphs, with applications to molecules". In: Advances in Neural Information Processing Systems 32 (2019).
[20] Shitong Liu, Hui Wang, Wenhao Liu, et al. "Pre-training molecular graph representation with 3d geometry". In: arXiv preprini arXiv:21 20.22000.77222 (2021).
[21] Shitong Luo, Tianlong Chen, Yuyang Xu, et al. "One transformer can understand both 2d & 3d molecular data". In: The Eleventh International Conference on Learning Representations.
2022.
[22] Oscar Méndez-Lucio Muhammad Ahmad, Emilio A del Rio-Chanona, et al. "A geometric deep learning approach to predict binding conformations of bioactive molecules". In: Nature Machine Intelligence 3.12 (2021), pp. 1033-1039.
[23] Garrett M Morris, Ruth Huey, William Lindstrom, et al. "AutoDock4 and AutoDock Automated docking with selective receptor flexibility". In: Journal of computational chemistry 30.16 16 (2009), pp. 2785-2791.
[24] S. F. Nielsen. "The stochastic EM algorithm: estimation and asymptotic results", In: Bernoulli (2000), pp. 457-489.
[25) Aaron van den Oord, Yazhe Li, and Oriol Vinyals. "Representation learning with contrastive predictive coding". In: arXiv preprint 1rXivv1807.03748 (2018).
[26] Rodrigo Quiroga and Miguel A Villarreal. "Vinardo: A scoring function based on autodock vina improves scoring, docking, and virtual screening". In: PloS one 11.5 (2016), e0155183.
[27] Alec Radford, Jong Wook Kim, Chris Hallacy, et al. "Learning transferable visual models from natural language supervision". In: International conference on machine learning. PMLR, 2021, pp. 8748 8763.
[28] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, et al. "Quantum chemistry struc- tures and properties of 134 kilo molecules". In: Scientific data 1.1 (2014), pp. 1-7.
[29] Yu Rong, Yatao Bian, Tingyang Xu, et al. "Self-supervised graph transformer on large-scale molecular data". In: Advances in neural information processing systems 33 (2020), pp. 12559- 12571.
[30] Oleg Trott and Arthur J Olson. "AutoDock Vina: improving the speed and accuracy of docking with new scoring function, efficient optimization, and multithreading". In: Journal of computational chemistry 31.2 (2010), pp. 455-461.
[31] Sheng Wang, Yiwen Guo, Yanming Wang, et al. "Smiles-bert: large scale unsupervised pre- training for molecular property prediction". In: Proceedings of the 10th ACM international conference bioinformatics, computational biology and health informatics. 2019, pp. 429- 436.
11

---

[32] Yuyang Wang, Rahul Magar, Chen Liang. et al. "Improving molecular contrastive learning via faulty negative mitigation and decomposed fragment contrast". In: Journal of Chemical Information and Modeling 62.11 (2022), pp. 2713-2725.
[33] Yuzhou Wang, Jiaqi Wang, Zhi Cao, et al. "Molecular contrastive learning of representations via graph neural networks". " In: Nature Machine Intelligence 4.3 (2022), pp. 279-287.
[34] Fan Wu, Shuoguo Li, Xin Jin, et al. "Rethinking explaining graph neural networks via non- parametric subgraph matching". In: International Conference on Machine Learning. Vol. 375.
Proceedings of Machine Learning Research. PMLR. 2023, pp. 11-23.
[35] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, et al. "MoleculeNet: benchmark for molecular machine learning". In: Chemical Science 9.2 (2018), pp. 513-530.
[36] Zhaocheng Xiong, De Wang, Xia Liu, et al. "Pushing the boundaries of molecular repre- sentation for drug discovery with the graph attention mechanism". In: Journal of Medicinal Chemistry 63.16 (2019), pp. 8749-8760.
[37] Kevin Yang. Kyle Swanson, Wengong Jin, et al. "Analyzing learned molecular representations for property prediction". In: Journal of Chemical Information and Modeling 59.8 (2019), pp. 3370-3388.
[38] Yuning You, Tianlong Chen, Yang Shen, et al. "Graph contrastive learning automated". In:
International Conference on Machine Learning. PMLR. 2021, pp. 12121-12132.
[39] Yuning You, Tianlong Chen, Yuxuan Sui, et al. "Graph contrastive learning with augmen- tations". In: Advances in Neural Information Processing Systems. Vol. 33, 2020, pp. 5812- 5823.
[40]
Manzil Zaheer et al. "Exponential stochastic cellular automata for massively parallel inference".
In: Artificial Intelligence and Statistics. PMLR. 2016, pp. 966-975.
[41] Shujian Zhang and Yuguo Chen. "Computation for latent variable model estimation: A unified stochastic proximal framework". In: Psychometrika 87.4 (2022), pp. 1473-1502.
[42] Guangyong Zhou, Zhen Gao, Qi Ding, et al. "Uni-mol: universal 3d molecular representation learning framework". In: International Conference on Learning Representations (2023).
12

---

A Similarity Score Distribution
Figure 3: Similarity Scores Similarity scores distribution for negative pairs in joint space after pre-training with original MolCLR loss and our proposed loss is provided. Compared to Using pretrained MolCLR model, our method yields similarity scores with lower mean and lower variance for negative pairs. While MoICLR have two peaks of negatives similarity scores around and 2.7, our method concentrates them at only one peak of 1.Our method yields similarity scores with higher mean and lower variance for positive pairs. Our method concentrates higher levels as it allows for some degree of semantic dissimilar between positives. The similarity scores are dot similarity, they not normalized to enhance the difference for visual purposes.
B Limitations
In this section, we discuss the limitations of our proposed EM-based algorithm for molecular contrastive learning.
B.1 Assumptions and Robustness
Our approach relies on several strong assumptions, such as the independence of molecular features and the noisiness nature of the input data. In practice, these assumptions may be violated, potentially affecting the performance and robustness the model. For instance, correlated features could lead to biased estimates of weights, while unnoisy data might degrade the necessity apply method in learning representations. Future work could explore methods to relax these assumptions and enhance the model's robustness such violations.
B.2
Scope of Claims
The empirical results presented in this paper are based on experiments conducted on a specific set of datasets: MoleculeNet and QM9. While these datasets are commonly used in molecular machine learning research, they may not fully represent all possible application domains. Consequently. the generalizability of our findings to other datasets real-world scenarios might be limited. Further validation on a broader range of datasets is necessary to confirm the wide applicability our approach.
Also, one limitation our method that the performance gains brought by the proposed architectural improvements depend datasets and tasks. For small datasets like QM9, the performance gain not significant.
B.3
Privacy and Fairness
While our work does not specifically address issues of privacy and fairness, these are important considerations for any machine learning model, especially those used sensitive domains such as healthcare. The potential for bias in molecular datasets, as well as privacy concerns related to molecular data, are areas that require further exploration. Ensuring that our model adheres to ethical standards and mitigates bias is an avenue for future work.
13

---

Table 7: hyperparameter search space for MoleculeNet dataset
Hyperparameter
Small
Large
HIV
Learning rate
5e 8e
5, le
4e
4, 5e
4
[2c 5, 1e
[2e 5, 5e
Batch size
[32, 64, 128, 256]
[128, 256]
[128, 256]
Epochs
[40, 60, 80, 100]
[20, 40
[2, 5, 10
Pooler dropout
[0.0, 0.2, 0.5]
10.0, 0.1]
0.0,
Warmup ratio
[0.0, 0.06, 0.1]
(0.0, 0.06]
0.1
C
Training details for experiments
C.1 MoleculeNet dataset
Following Unimol, we report the detailed hyperparameters setup during pretraining in Molecular pretraining runs on 4 A6000 GPUs, and the training time is about 48 hours. We split all the datasets with scaffold split, which splits molecules according to their molecular substructure.
C.2 MoleculeNet Non-Chirality Analysis
We basically follow Mol-CLR on experiment settings. During pre-training, the GNN encoder maps each molecular graph 512-dimensional embedding h. projection head, modeled as an MLP with single hidden layer, transforms h into a -i6-dimensi latent vector z. ReLU is utilized as the non-linear activation function. The model undergoes pre-training over 50 epochs with batch size of 512, optimized via the Adam optimizer with an initial learning rate of 5 x 10 and weight decay rate x 10 5, cosine learning rate decay schedule is applied throughout pre-training. For most datasets, we use scaffold-based data split; however, the QM9 subtask follows random split line with the Mol-CLR methodology.
In the fine-tuning phase, the projection head is replaced with randomly initialized MLP that maps h the target property prediction, while the pre-trained GNN encoder remains fixed. The fine-tuning process spans 100 epochs per benchmark task, with hyperparameters tuned via random search on validation sets, and results reported on test sets. Each benchmark is evaluated over three independent runs, with average performance reported. The model implementation is based on Py Torch Geometric.
C.3 QM9 Dataset Experiments
Data partitioning for the QM9 tasks follows the scheme utilized in the Equiformer. For tasks predicting Or, SHOMO, SLUMO, Ae, and C., our configuration includes a batch size of 64, 300 training epochs, learning rate of 10-1, and Gaussian radial basis functions with 128 bases. The architecture comprises six Transformer blocks, a weight decay of 5 x 10-3, and a dropout rate of 0.2. Mixed-precision training employed for these tasks.
For the Rt task, the setup consists of a batch size of 48, 300 epochs, a learning rate of 1.5 10 10-4 Gaussian radial basis functions with 128 bases, five Transformer blocks, weight decay of x 10 and dropout rate of 0.1, executed in single precision.
The ZPVE task employs batch size of 48, 300 epochs, learning rate of 1.5 x 10-4, Gaussian radial basis functions with 128 bases, five Transformer blocks, a weight decay of 10-3 and a dropout rate of 0.2, with single-precision training.
For the G, H, U, and Uo tasks, the setup includes a batch size of 48, 300 epochs. a learning rate of 1.5 x 10-4 Gaussian radial basis functions with 128 bases, five Transformer blocks, with both weight decay and dropout omitted, utilizing single precision.
All models were trained on a single A6000 GPU, with mixed- tasks requiring 81 GPU-hours and Annlle-reccssio tasks requiring 151 GPU-hours. Model complexity includes 11.20 million parameters for configurations with six blocks and 9.35 million parameters for configurations with five blocks.
The data partitioning approach adheres to the random splitting strategy outlined the Equiformer paper.
14

---

D Comparison against other weight calculation methods
Here we show the comparison against using simpler approach based on similarity scores. To thoroughly investigate this, we designed ablation experiments using the chiral version of MoleculeNet classification tasks and compared three different methods: Bayesian Inference (Our Method); In this approach, we calculate the weights using Bayesian inference, as described in our paper. 2. Fingerprint-based Similarity: This method calculates the weights based on the similarity scores derived from molecular fingerprints, similar to the approach used in I-MoICLR [32]. 3. Encoder- based Similarity: Here, we first extract features of data pairs using encoders and then calculate their similarity scores. These scores are then regularized to the [0, 1] range.
For methods 2 and 3, we compute the weights using the following formulas:
W; W = A x Sim (2,,
Sim (Zt, xx)
In our experiments, we set (A
1).
MoleculeNor Experiments
Tox21
ToxCast
SIDER
ChinTox
BACE
MUV
PCBA
Bayesian Interence
76.7 1 2.0
80.1 1 1.0
69.91225
64.9 3.3
89.4
88.2 = L.3
82.9 1 3.1
83.0 1 L.7
82.9 I
0.5
Fingerprint Similarity
77.9
80.0
64.9
83.6
78.7
83.8
Encoder- bused Similarity
75.2
79.6
67.8
58.5
90.4
82.8
80.5
81.0
74.5
From these results, we observe the following:
1. While fingerprint-based similarity showed improvements in out of 9 tasks compared to our original method, but did not perform as well overall. This indicates that they may not be flexible enough to fully capture the complexities of molecular representations required for robust performance across diverse tasks.
2. Encoder-based Similarity performed worse than both the Bayesian inference method and the fingerprint-base similarity approach, further suggesting that using a direct similarity- -baeed method does not necessarily yield better results.
These findings suggest that while simpler methods may work in some cases, they do not outperform our proposed Bayesian inference method which can dynamically adapt and provide better alignment of positive and negative pairs. Thus, our approach is essential for achieving state-of-the-art performance across various molecular property prediction benchmarks.
E Protein-ligand binding task
We also conducted the protein-ligand binding pose prediction task. This is one of the most important tasks in structure based drug design. The task is to predict the complex structure of protein binding site and a molecular ligand. We need to consider how ligand lays in the pocket, that is, the 6 degrees (3 rotations and 3 translations) of freedom of a rigid movement.
Following Uni-Mol, the molecular representation and pocket representation are firstly obtained from their own pretraining models by their own conformations; then, their representations concatenated as the input of an additional 4-layer Transformer decoder, which is finetuned to learn the pair distances of all heavy atoms molecule and pocket. Then, with the predicted pair-distance matrix as scoring function, we first randomly place the ligand and then optimize the coordinates of its atoms by directly back proppaatatttnonnn the loss between current pair-distance and predicted pair-distance.
We evaluate our method using the metric binding pose accuracy. Specifically, we keep the pocket conformation fixed, while the ligand conformation is fully flexible. We evaluate the RMSD(root mean squared distance) between the prediction and the ground truth. Following previous works, we the percentage of results below predefined RMSD thresholds as metrics.
The binding pose accuracy results are shown in Table. Not surprisingly, our model again outperforms all the baseline methods, achieving ofte-of-the-an results with our Gamma-prion version model.
15

---

Table 8: Performance on binding pose prediction.
Methods
1.0
1.5
2.0
3.0
5.0
Autodock Vina 30
44.21
57.54
64.56
73.68
84.56
Vinardo 26
41.75
57.54
62.81
69.82
76.84
Smina [13
47.37
59.65
65.26
74.39
82.11
A 23
21.75
31.58
35.44
47.02
64.56
Uni 42
43.16
68.42
80.35
87.02
94.04
Ours (Bernoulli)
48.77
70.18
78.95
85.26
94.04
Ours GGamma)
45.61
69.47
80.70
88.42
96.84
Table 9: ablation study on au and by.
Ou
5
10
10
5
10
bn
1
5
5
5
10
10
10
Avg ROC AUC%
80.4
80.2
80.1
80.6
80.7
80.2
80.3
80.1
80.5
F More ablation on au and by
Here we conducted an ablation study on the choice of an and by- Generally speaking, the choice of au and by will not influence the experiment results to large margin, the top performance is au = -b= = 5.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state the key contributions of the paper, including the development and evaluation of the proposed EM-based algorithm for molecular contrastive learning, as well as the comparison with baseline methods.
Justification:
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
The paper includes dedicated "Limitations" section where it discusses the strong assumptions made in the model, such as the independence assumptions and the potential impact of noisy data. It also reflects on the scope the claims, emphasizing that the results are based on experiments conducted on specific datasets.
Justification:
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper includes detailed presentation of all theoretical results, with each theorem and lemma clearly numbered and cross-referenced. All assumptions are explicitly stated within the statements of the theorems, and complete proofs are included in the main text.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper provides comprehensive details on the experimental setup, including dataset descriptions, preprocessing steps. model architectures, hyperparameter settings, and evaluation metrics. Additionally, the paper outlines the exact procedures followed during experimentation and includes detailed pseudocode in the appendix. These disclosures ensure that the main experimental results can be reproduced and verified independently.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The paper provides open access to both the data and the code. Detailed instructions are included in the supplemental material, covering the exact commands and environment settings needed reproduce the results. The data access instructions encompass steps for obtaining the raw data, preprocessing methods, and generating the necessary datasets. Scripts to reproduce all experimental results, including those for the proposed method and baselines, are provided. Any deviations or omitted experiments are clearly stated with justifications.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [TODO]
17

---

The paper provides comprehensive details on the experimental setup, including specific data splits, hyperparameters. and their selection pro-cess. It also covers the type of optimizer used and all other relevant parameters. These details are presented clearly in the appendix of the paper.
Justification:
7. Experiment Statistical Significance
Question: Does the paper report enror bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The paper does report error bars, confidence intervals, or statistical signifi- cance tests for the experimental results. Instead, all the experiments are carried out for three times and the average reported.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Answer: [Yes]
The paper provides detailed information on the computer resources required for each experiment. It specifies the type of compute workers (GPU), memory usage, and the time of execution for each experimental run.
Justification:
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p/p////ooooo..o/o//o////aaa
Answer: [Yes]
Justification: There is no consider of ethics according to the NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: There is no societal impact discussed in the paper.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not pose such risks.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The authors have cited the original, paper that produced the code package or dataset, included the license and provided the copyright and terms of services.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Dastification:Tee authors have cited the original paper that produced the code package or dataset, stated its" versions, included the license and provided the copyright and terms of services.
18

---

14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsoureing nor research with human subjects 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/teview based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve studying participants
19


Reviews:
[{"review_id": "l36h7taZSp", "reviewer": "Reviewer_T3Hf", "summary": "Commonly used data augmentation methods may produce false positives or negatives in learning molecular representations. This study proposes a novel probability-based contrastive learning framework to tackle false positive and negative pairs in molecular representation learning. Bayesian modeling is used to learn a weight distribution on sample pairs, automatically identifying and mitigating false pairs. This dynamic adjustment enhances representation accuracy. The model is optimized through a stochastic expectation-maximization process that iteratively refines sample weight probabilities and updates model parameters.", "questions": "* Figure 1, which node(s) and edge(s) are removed? How is similarity calculated? Does augmentation consider structural validity and stability?\n* Why do you choose Gamma or Bernoulli distribution over other distribution?\n* The authors pointed out that i-MolCLR did something similar, however, claimed that the proposed approach offered a better solution. Why not compare with i-MolCLR in all experiments?\n* The paper shows that the best setting is positive pair weights following  a+ =5 and b+ = 1, and negative pairs following a- = 1 and b- = 1. The significant difference between these two distributions is approximately close to the threshold for decision-making. When the distributions are similar, such as positive pair weights following a+ =1 and b+ = 1 while negative pair weights following a- =1 and b- = 1, performance is worse according to the ablation study. Besides, the author should discuss about the choices of au and bu.", "limitations": "None.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* Present a good analysis on the impact of false positive and negative pairs\n* The proposed probability-weighted CL is novel\n* The overall structure of the manuscript is clear and easy to follow (make the fonts in the figures bigger, they are hard to read).\n* The experimental results show the average performance. How about variances, an indication of robustness?", "weaknesses": "* Why not design the probability distribution based on the similarity score distribution, a simpler approach? The best hyperparameters (Table 6) indicate that the positive and negative pairs have quite different distributions. Maybe a even simpler approach using a similarity thresholding would work. \n* What are the appropriate values for a_u and b_u.\n* Figure 3 is confusing. Is it used to show that, before learning, some positive samples are highly similar to negative samples, which is no longer the case after learning?\n* Recent and better performance on the MoleculeNet datasets should be considered.\n\n[1] Fang, Yin, et al. \"Knowledge graph-enhanced molecular contrastive learning with functional prompt.\" Nature Machine Intelligence 5.5 (2023): 542-553.\n[2] Hajiabolhassan, Hossein, et al. \"FunQG: Molecular representation learning via quotient graphs.\" Journal of chemical information and modeling 63.11 (2023): 3275-3287.\n[3] Ren, Gao-Peng, Ke-Jun Wu, and Yuchen He. \"Enhancing molecular representations via graph transformation layers.\" Journal of Chemical Information and Modeling 63.9 (2023): 2679-2688.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for recognizing the presentation and novelty of our work.\n\nW1: Why not design the probability distribution based on the similarity score distribution, a simpler approach?\n\nWe agree that using a similarity thresholding is a simpler approach. However, as the experiments showed, even if we use fingerprint-based similarity calculation, the issue cannot be completely resolved, and the performance gaps between this simple approach and our methods are in fact still reasonably large. Please refer to our response to all the reviewers for more detailed analysis. \n\nW2: What are the appropriate values for a_u and b_u.\n\nIn code, we set a_u=b_u=1, here we conducted an ablation study on the choice of a_u and b_u. Generally speaking, the choice of a_u and b_u will not influence the experiment results to a large margin, the top performance is at a_u=b_u=5.\n\n| a_u          | 1    | 5    | 10   | 1    | 5        | 10   | 1    | 5    | 10   |\n|--------------|------|------|------|------|----------|------|------|------|------|\n| b_u          | 1    | 1    | 1    | 5    | 5        | 5    | 10   | 10   | 10   |\n| Avg ROC_AUC% | 80.4 | 80.2 | 80.1 | 80.6 | **80.7** | 80.2 | 80.3 | 80.1 | 80.5 |\n|              |      |      |      |      |          |      |      |      |      |\n\nW3: Figure 3 is confusing. Is it used to show that, before learning, some positive samples are highly similar to negative samples, which is no longer the case after learning?\n\nFigure 3 is not to illustrate the similarity of positive samples to negative samples before and after learning. Instead, it compares the distributions of similarity scores for positive and negative pairs after pre-training with the original MolCLR loss and the proposed loss.\n\nLeft plot (Positive samples): It shows the distribution of similarity scores for positive samples. The proposed method (blue) produces higher similarity scores with a higher mean and lower variance compared to the original MolCLR method (orange). This indicates that the positive samples are more consistently similar to each other with the proposed method.\n\nRight plot (Negative samples): It shows the distribution of similarity scores for negative samples. The proposed method (blue) yields lower mean and lower variance in similarity scores compared to MolCLR. MolCLR shows two peaks (around 1 and 2.7), suggesting that it incorrectly assigns higher similarity to some negative pairs. The proposed method concentrates the negative similarity scores around 1, indicating better separation of negative pairs from positive pairs.\n\nW4 Recent and better performance on the MoleculeNet datasets should be considered.\n\nWe appreciate the reviewer's suggestion. We will cite these papers in our related work section. In fact, we can apply our method on top of these methods, we will update the experiments in final version of our paper.\n\nHowever, we would like to emphasize that direct comparisons with these methods may not be fair due to differences in additional information sources and data splitting methods:\n\nKANO: This method utilizes additional knowledge graph as an information source, which provides a significant advantage. Our method does not use such external information.\n\nLineEvo: This method employs random data splitting, which can yield artificially higher performance.\n\nFunQG: This method uses a molecular graph coarsening framework that significantly reduces the graph's complexity, losing 3D information of each individual atom. Our approach retains the original graph structure.\n\nHere, we present another fairer comparison with a recent baseline Instruct-Bio [1]. The following analysis shows that our method outperforms InstructBio across almost all subtasks in the MoleculeNet classification and regression experiments.\n\n|               |  BBBP        | Tox21        | ToxCast      | SIDER         | ClinTox       | BACE          | ESOL             | FreeSolv        | Lipo             |\n|---------------|--------------|--------------|--------------|---------------|---------------|---------------|------------------|-----------------|------------------|\n| Insturct-bio[1] | 67.4±0.5     | 75.6±0.3     | 65.1±1.5     | 61.5±1.9      | 78.0±0.6      | 78.5±0.8      | 1.771±0.015      | **0.832±0.020** | 0.752±0.017      |\n| Ours          | **76.7±2.0** | **80.1±1.0** | **69.9±2.5** | **64.9± 3.3** | **89.4± 0.1** | **88.2± 1.3** | **0.664± 0.024** | 1.358± 0.452    | **0.590± 0.003** |\n|               |              |              |              |               |               |               |                  |                 |                  |\n\n[1] Wu F, Qin H, Li S, et al. Instructbio: A large-scale semi-supervised learning paradigm for biochemical problems[J]. arXiv preprint arXiv:2304.03906, 2023.", "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": "eQW6Lxxs2q", "reviewer": "Reviewer_Jokw", "summary": "This paper proposes a probability-based contrastive learning framework for 3D molecular representation learning. It addresses the issue of false positive and negative pairs in existing methods. Experiments show its effectiveness, outperforming other baselines. The approach has wide applicability and can boost the performance of molecular representation learning models.", "questions": "1. What is the augmentation strategy for QM9 and MoleculeNet? Are only MoleculeNet and Non-Chirality version MoleculeNet pretrained while QM9 is trained directly?\n2. If we calculate a similarity based on the molecular structure using chemical software like RDKit to use as a weight, should we achieve better performance? What do you think?", "limitations": "See **Weaknesses** Section.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. To the best of my knowledge, this paper raises an intriguing question and warrants attention.\n2. The approach taken to address the problem is also innovative.\n3. The experiments were comprehensive and yielded significant gains.", "weaknesses": "There are no major flaws in this article overall, but it requires attention to detail in the writing. For example:\n1. Eq.3 is supposed to be $e^{w_i^+s_i}$, not $w_i^+s_i$. Or am I misunderstanding?\n2. Line 129 contains a duplicated \"positive.\"\n3. The authors should explain the function of an augmented random variable $u_i$ in the appendix, in case readers are not familiar with the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for recognizing the innovation of our work.\n\nW1: Eq.3 is supposed to be $e^{w_i^+s_i}$, not $w_i^+s_i$. Or am I misunderstanding?\n\nThis might be a misunderstanding, Equation 3 is correct as written. This equation is part of our Bayesian augmentation method, designed to create a joint distribution that maintains conjugacy between the priors and posteriors of the weights.\n\nW2: Line 129 contains a duplicated \"positive.\"\n\nWe appreciate your careful review and will remove the duplicated.\n\nW3: The authors should explain the function of an augmented random variable ui in the appendix, in case readers are not familiar with the method.\n\nWe will add a detailed explanation of the function of the augmented random variable u_i​ in the appendix, providing context and clarity on its role in our method.\n\nQ1: What is the augmentation strategy for QM9 and MoleculeNet? Are only MoleculeNet and Non-Chirality version MoleculeNet pretrained while QM9 is trained directly?\n\nA1: We will add this in the appendix.\n\nFor non-chiral MoleculeNet, we apply a combination of node removal and edge removal as augmentation strategies, as shown in Figure 1.\n\nFor chiral MoleculeNet and QM9, we adopt two different augmentations, following Uni-Mol:\n\nAtom Masking Augmentation: This augmentation randomly masks a certain percentage of atoms.\n\nCoordinate Masking: This augmentation adds random gaussian noise to atom coordinates.\n\nAlso, QM9 has been trained directly without the pretraining on a larger dataset.\n\nQ2: If we calculate a similarity based on the molecular structure using chemical software like RDKit to use as a weight, should we achieve better performance? What do you think?\n\nA2: \nAs discussed in our common response to all reviewers, we have conducted experiments comparing different weighting strategies, including similarity-based weights derived from molecular fingerprints based on RDkit.\n\nOur experiments demonstrate that while this similarity-based approach shows some improvement in certain tasks, it does not outperform our proposed method in the majority of the tasks. This suggests that the added complexity and adaptability of our method provides better alignment with the diverse tasks in molecular property prediction.", "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": "s1IwHgS3zU", "reviewer": "Reviewer_NjGA", "summary": "This paper introduces a probability-based contrastive learning framework. It regards learnable weights as variables with different distributions and automatically identify and mitigate false positive and negative pairs via Bayesian modeling. The author verify the effectiveness of their method in 13 out of 15 property prediction tasks in MoleculeNet and QM9, which are standard benchmarks.", "questions": "(1) Why 0.251 for Equiformer in Table 4 is in bold? \n\n(2) I remember that Equiformer is purely a backbone architecture without pretraining. Therefore, it may be unfair to compare it directly in QM9. Moreover, it has an improved version Equiformer V2 [A]. \n\n[A] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. \n\n(3) Similarly, the strongest baseline listed in the paper, Uni-Mol, also has an improved version, Uni-Mol+ [A]. However, it is a pity that it does not release its performance on QM9 and MoleculeNet. \n\n[A] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+. \n  \n(4) I would recommend the author add InstructBio[A], another interesting work as an extra baseline in Table 1 and 2, which also reports performance on MoleculeNet. \n\n[A] InstructBio: A Large-scale Semi-supervised Learning Paradigm for Biochemical Problems.", "limitations": "Yes, the authors addressed limitations adequately.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The probability weighted contrastive learning mechanism seems novel to me. It is new to treat model weights as random variables and sample optimal weights using Bayesian inference. \n\n(2) The performance is brilliant and experiments are comprehensive, ranging from MoleculeNet to QM9. It also outperforms strong baselines such as Equiformer, Transformer-X, etc. \n\n(3) The method is presented in a clean and readable way, easy to understand.", "weaknesses": "(1) Algorithm A in lines 199-216 is filled with text. Can the author present the algorithm more elegantly? \n\n(2) Some related works should be mentioned in the related work and be compared in the experiments. For example, MatchDrop [A] also noticed the false positive sampling issue in the graph data augmentation technique. They propose to keep the crucial subgraph unchanged and only move the less informative part. Can this MatchDrop augmentation be considered in your contrastive learning framework? \n\n[A] Rethinking explaining graph neural networks via non-parametric subgraph matching, ICML 2023\n\n(3) MoleculeNet is somehow of small sizes. Some large datasets should be considered such as PCQM4MV2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for recognizing the presentation and novelty of our work.\n\nWeaknesses:\n\n(1) Algorithm A in lines 199-216 is filled with text. Can the author present the algorithm more elegantly?\n\nA1:We appreciate the suggestion. We will revise the manuscript by converting the text into functions and formulas, hopefully providing a clearer and more structured presentation of the algorithm.\n\n(2) Some related works should be mentioned in the related work and be compared in the experiments. For example, MatchDrop [A] also noticed the false positive sampling issue in the graph data augmentation technique. They propose to keep the crucial subgraph unchanged and only move the less informative part. Can this MatchDrop augmentation be considered in your contrastive learning framework?\n\n[A] Rethinking explaining graph neural networks via non-parametric subgraph matching, ICML 2023\n\nThanks for the suggestion, we will add this paper in the related works part. However, MatchDrop did not include results on Molecule datasets, we will try to include this data augmentation method and the results will be included in the final version of this paper.\n\n(3) MoleculeNet is somehow small in size. Some large datasets should be considered such as PCQM4MV2.\n\nA3: Thanks for your suggestion, We added the following experiments on PCQM4Mv2 dataset: We start from the checkpoint of pretrained Uni-Mol+ small model, we adopt Atom Masking and Coordinate Masking augmentations and substitute the original 3D coordinate prediction loss function with our weighted contrastive learning loss function. The model is trained for 150K steps. We get 0.0690 valid MAE, marginally better than 0.0696 of the original model. Using larger backbones and longer training time might help to achieve better performance.\n\nQuestions:\n\n(1) Why 0.251 for Equiformer in Table 4 is in bold?\n\nThank you for catching this mistake. We will correct this in the manuscript.\n\n(2) I remember that Equiformer is purely a backbone architecture without pretraining. Therefore, it may be unfair to compare it directly in QM9. Moreover, it has an improved version Equiformer V2 [A].\n\n[A] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations.\n\nWe acknowledge the point. While Equiformer is indeed a backbone architecture without pretraining, our experiments on QM9 were also \nconducted without the use of external data. The contrastive learning was performed directly on the QM9 dataset. \nPlease refer to the following table, our method outperforms EquiformerV2 in 6 out of 12 tasks on QM9, despite using significantly fewer transformer layers (2 vs 6). This suggests that our approach is both efficient and effective.\n\n|               | α    | ΔE   | E_homo | E_lumo | μ     | C_v   | G    | H    | R^2   | μ     | μ0    | ZPVE |\n|---------------|------|------|--------|--------|-------|-------|------|------|-------|-------|-------|------|\n| EquiformerV2  | 0.050 | 29   | **14**     | **13**     | **0.010**  | 0.023  | 7.57 | **6.22** | 0.186  | **6.49**  | **6.17**  | 1.47 |\n| Equiformer |0 .046 | 30   | 15     | 14     | 0.011  | 0.023  | 7.63 | 6.63 | 0.251  | 6.74  | 6.59  | 1.26 |\n| Ours          | **0.037** | **24.2** | 21.1   | 13.7   | 0.022  | **0.022**  | **6.2**  | 6.31 | **0.082**  | 7.22  | 9.40  | **1.09** |\n\n(3) Similarly, the strongest baseline listed in the paper, Uni-Mol, also has an improved version, Uni-Mol+ [A]. However, it is a pity that it does not release its performance on QM9 and MoleculeNet.\n\n[A] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+.\n\nWe appreciate this. We will cite the Uni-Mol+ paper in related works. \n\nPlease refer to our answer on weakness 3 for the experiment to validate our framework’s effectiveness based on Uni-Mol+. Experiments show that our method get 0.0690 valid MAE, marginally better than 0.0696 of the original Uni-Mol+ model. \n\n(4) I would recommend the author add InstructBio[A], another interesting work as an extra baseline in Table 1 and 2, which also reports performance on MoleculeNet.\n\n[A] InstructBio: A Large-scale Semi-supervised Learning Paradigm for Biochemical Problems.\n\nThanks for recommendation. We will add this paper to our experiment section. The following analysis shows that our method outperforms InstructBio across all subtasks in the MoleculeNet classification and regression experiments.\n\n|            |  BBBP    | Tox21    | ToxCast  | SIDER     | ClinTox   | BACE      | ESOL         | FreeSolv     | Lipo         |\n|------------|----------|----------|----------|-----------|-----------|-----------|--------------|--------------|--------------|\n| Instrctbio | 67.4±0.5 | 75.6±0.3 | 65.1±1.5 | 61.5±1.9  | 78.0±0.6  | 78.5±0.8  | 1.771±0.015  | 0.832±0.020  | 0.752±0.017  |\n| Ours       | 76.7±2.0 | 80.1±1.0 | 69.9±2.5 | 64.9± 3.3 | 89.4± 0.1 | 88.2± 1.3 | 0.664± 0.024 | 1.358± 0.452 | 0.590± 0.003 |\n|            |          |          |          |           |           |           |              |              |              |", "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": "yVO72lYGbp", "reviewer": "Reviewer_oFgE", "summary": "To address the problem of potential false positives and false negatives in contrastive learning of molecules, this paper proposed a learnable weighted contrastive learning approach for molecular representation learning. The effectiveness of the proposed method is tested on the MoleculeNet and QM9 datasets.", "questions": "See weaknesses.", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The false pair label is an issue for contrastive learning. \n2. The experimental result shows improved performance. \n3. The proposed method is clear and easy to follow.", "weaknesses": "1. There are existing works about the false labels in contrastive learning [1,2] and weighted contrastive learning [3,4,5]. None of these are discussed in the related works while they are quite related to the issue that the paper tries to address. This limits the novelty of this paper. \n2. In the experiment, how to do data splitting is not clear. I did not find it in the paper. In molecular tasks, data splitting may have a big effect on performance.\n3. In the title and the ablation, the author mentioned 3D structures. However, it seems that 3D is not a key component of the paper. And the 3D loss is taken from Uni-mol which is not a contribution of this paper.\n4. The reason for using Gamma prior should be justified. It is an essential step in the proposed approach and related to the technical soundness. \n5. In MolCLR, the contrastive is done between two augmentations Xi_1 and Xi_2, which is the same as Figure 1. But in the code, the contrastive is done between the original sample Xi and one augmented sample Xi_1. Maybe the latter case is less affected by the false labels.  \n6. The major part of Figure 2 is copied from the cited papers (MolCLR, Uni-mol, Equiformer). Is this permitted?\n7. Some typos: e.g. line 285, missing \"Table\" before the reference.\n\nOverall, the application of advanced contrastive learning to molecular representation learning is a promising idea. But the current paper needs to be improved in several aspects for publication. \n\n[1] Debiased Contrastive Learning, NeurIPS 2020  \n[2] A Theoretical Analysis of Contrastive Unsupervised Representation Learning, ICML 2019  \n[3] CWCL: Cross-Modal Transfer with Continuously Weighted Contrastive Loss, NeurIPS 2023  \n[4] Camera Alignment and Weighted Contrastive Learning for Domain Adaptation in Video Person ReID, WACV 2023  \n[5] Weighted Contrastive Learning With False Negative Control to Help Long-tailed Product Classification, ACL 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for recognizing the promising potential of our application.\n\nW1: There are existing works about the false labels in contrastive learning [1,2] and weighted contrastive learning [3,4,5]. None of these are discussed in the related works while they are quite related to the issue that the paper tries to address. This limits the novelty of this paper.\n\n[1] Debiased Contrastive Learning, NeurIPS 2020\n\n[2] A Theoretical Analysis of Contrastive Unsupervised Representation Learning, ICML 2019\n\n[3] CWCL: Cross-Modal Transfer with Continuously Weighted Contrastive Loss, NeurIPS 2023\n\n[4] Camera Alignment and Weighted Contrastive Learning for Domain Adaptation in Video Person ReID, WACV 2023\n\n[5] Weighted Contrastive Learning With False Negative Control to Help Long-tailed Product Classification, ACL 2023\n\nA1: We appreciate the reviewer's suggestion and will update our manuscript to include and discuss these references in the related works section. Our contribution lies in adapting the weighted contrastive loss to molecular data with a new optimization algorithm that allows for simultaneous posterior inference and model optimization, coupled with extensive evaluation from multiple aspects. \n\nSpecifically, our method diverges from existing works such as [1] and [2] by focusing on a Bayesian inference framework tailored for weighted contrastive learning. Additionally, [3], [4] and [5] address weighted contrastive loss in various domains, our approach uniquely integrates 3D molecular structures and graph-based representations. By referencing these works, we aim to frame our approach within the broader context of addressing false negatives in contrastive learning and clearly delineate our novel contributions.\n\nW2: In the experiment, how to do data splitting is not clear. I did not find it in the paper. In molecular tasks, data splitting may have a big effect on performance.\n\nA2: We apologize for the confusion for data splitting methodology in the paper. Here, we provide a comprehensive explanation. \nChiral MoleculeNet Experiments: We split all the datasets with scaffold split, which splits molecules according to their molecular substructure.\n\nNon-Chiral MoleculeNet Experiments:We created a scaffold split for most datasets. Except that the subtask QM9 adopts a random splitting setting following Mol-CLR,\n\nQM9 Experiments: For the QM9 dataset (which have different downstream regression tasks compared with QM9 subtask in MoleculeNet dataset), we adhered to the random data splitting strategy employed in the Equiformer paper.\n\nWe will clarify it in our paper.\n\nW3: In the title and the ablation, the author mentioned 3D structures. However, it seems that 3D is not a key component of the paper. And the 3D loss is taken from Uni-mol which is not a contribution of this paper.\n\nA3: We appreciate the reviewer’s observation and acknowledge that the 3D loss is derived from Uni-mol. However, our contribution lies in exploring the combined use of this 3D loss with a weighted contrastive learning loss, which our ablation study in table 5 demonstrates enhances performance. The 3D information is vital in our work for one more main reason: The calculation of our weighted contrastive loss depends on the 3D structural information, providing spatial context that complements the molecular graph structure, which enables performing protein-ligand binding task(Table 8, Appendix D) and potentially many other real-world drug design tasks that other contrastive learning based method, like MolCLR, don’t work. We will ensure these points are clearly articulated in the revised manuscript.\n\nW4: The reason for using Gamma prior should be justified. It is an essential step in the proposed approach and related to the technical soundness.\n\nA4: We use the Gamma prior because it naturally lends itself to conjugacy in the posterior, which significantly eases the posterior sampling procedure. Also, it is known for its flexibility in shape and scale to model positive continuous variables, which is suitable for sample weights in our setting. Additionally, we compared it with a baseline using binary sample weights with the Bernoulli prior in Table 1. Our results show that the Gamma prior offers superior robustness and accuracy. This advantage is likely due to the broader search space available to continuous weights, with binary weights being a subset of this continuous space.\n\nW5: In MolCLR, the contrastive is done between two augmentations Xi_1 and Xi_2, which is the same as Figure 1. But in the code, the contrastive is done between the original sample Xi and one augmented sample Xi_1. Maybe the latter case is less affected by the false labels.\n\nA5: Thank you for pointing out the difference. We agree that performing contrastive learning between the original sample \\(X_i\\) and one augmented sample \\(X_{i1}\\) can potentially reduce the impact of false labels, as it introduces less randomness and maintains higher consistency. We will add a discussion in the paper to highlight the benefits of our trick. Additionally, we will include visualizations to empirically demonstrate how our method effectively identifies and mitigates the impact of false labels.\n\nW6: The major part of Figure 2 is copied from the cited papers (MolCLR, Uni-mol, Equiformer). Is this permitted?\n\nA6: We have contacted the original authors for permission to use these images and have received approval. In addition, we will improve the quality of the figure with a similar design to enhance the presentation in the camera-ready version.\n\nW7: Some typos: e.g. line 285, missing \"Table\" before the reference.\n\nA7: Thank you for pointing this out. We will correct it in our paper.", "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:
We thank the reviewers for their valuable comments. We are happy that the reviewers find our work innovative and promising in general. We notice there are also some questions and concerns from several perspectives. A common issue raised by the reviewers is the comparison against directly using similarity scores threshold to decide weight distributions, which we will address below. For other specific questions raised by each reviewer, we will post our responses separately. We also have revised the manuscript by incorporating some extra results and explanations based on the reviews. We will incorporate more detailed revisions into the camera-ready version according to responses and further discussions. We hope the reviewers can read our response and reconsider your decision if necessary. Thanks again for helping us to make this work better.

Q: Why not design the probability distribution based on the similarity score distribution?

We appreciate the reviewers’ insightful suggestion to use a simpler approach based on similarity scores. To thoroughly investigate this, we designed ablation experiments using the chiral version of MoleculeNet classification tasks and compared three different methods:

1. Bayesian Inference (Original Method): In this approach, we calculate the weights using Bayesian inference, as described in our paper.

2. Fingerprint-based Similarity: This method calculates the weights based on the similarity scores derived from molecular fingerprints, similar to the approach used in I-MolCLR.

3. Encoder-based Similarity: Here, we first extract features of data pairs using encoders and then calculate their similarity scores. These scores are then regularized to the [0,1] range. 

For methods 2 and 3, we compute the weights using the following formulas:

$w_i^- = 1 - \lambda \times Sim(x_i, x_k)$,

$w_i^+ = \lambda \times Sim(x_i, x_k)$

In our experiments, we set $\(\lambda = 1\)$.

|  MoleculeNet Experiments |  BBBP     | Tox21     | ToxCast   | SIDER     | ClinTox   | BACE      | MUV       | HIV       | PCBA      |
|--------------------------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|
| Bayesian Inference       | 76.7± 2.0 | **80.1± 1.0** | **69.9± 2.5** | **64.9± 3.3** | 89.4± 0.1 | **88.2± 1.3** | **82.9± 3.1** | **83.0± 1.7** | 82.9± 0.5 |
| Fingerprint Similarity   | **77.9**      | 80.0      | 68.9      | 64.9      | 87.6      | 83.6      | 78.7      | 79.8      | **83.8**      |
| Encoder-based Similarity | 75.2      | 79.6      | 67.8      | 58.5      | **90.4**      | 82.8      | 80.5      | 81.0      | 74.5      |

From these results, we observe the following:

1. Fingerprint-based Similarity: While this method showed improvements in 2 out of 9 tasks compared to our original method, but it did not perform as well overall. This indicates that while similarity-based methods are simple, they may not be flexible enough to fully capture the complexities of molecular representations required for robust performance across diverse tasks.

2. Encoder-based Similarity: This approach performed worse than both the Bayesian inference method and the fingerprint-based similarity approach, further suggesting that using a direct similarity-based method does not necessarily yield better results.

These findings suggest that while simpler methods like similarity thresholding may work in some cases, they do not outperform our proposed Bayesian inference method which is designed to dynamically adapt and provide better alignment of positive and negative pairs of molecular representation learning. Thus, our Bayesian inference-based approach is essential for achieving state-of-the-art performance across various molecular property prediction benchmarks.

We hope this detailed comparison addresses the reviewer’s concerns and provides clarity on the advantages of our proposed method.

Meta-review: This paper proposes a contrastive learning framework for molecular representation learning. The main idea is to remove the "false" pairs of data augmentation from severely modifying the graph structure. The authors propose a Bayesian learning framework to solve this problem. 

Most of the reviewers strongly support the paper. One reviewer raised concerns on (1) justification behind the Gamma prior, (2) Figure copyrights, and (3) missing discussion of related works. However, after carefully reading through the author rebuttal, I believe the author rebuttal is persuasive enough to alleviate the concerns.

Therefore, I recommend acceptance for this work.
Decision: Accept (poster)