Title: A probability contrastive learning framework for 3D molecular representation learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: contrastive learning; molecular property prediction; expectation maximization
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

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.

1 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, 33, 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 is contrastive learning (CL), which leverages large-scale unlabeled molecular datasets to learn robust representations. CL works by contrasting different augmentations of the same molecule to ensure that the model learns to recognize the essential features of the molecule, 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-modality 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

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 a 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 in MolCL[33] ,where two distinct graph augmentations 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 MolCLR [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 as negative, regardless of their structural similarity. This mislabeling results in false positives and negatives, undermining the effectiveness of the contrastive learning process.

0

.

67

0

.

65

0

.

96

0

.

97

Node

removal

Edge

removal

Node

removal

Edge

removal

SH

SH

SH

F

SH

F

SH

F

SH

To overcome the aforementioned issue, we introduce a generalization of existing contrastive learning 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

<latexi sh

1_b

64="9A

W3fJT+nE

oFw5I

DKQ

>

B7

cdV

SgN

O

jXGL

v

RP

0

y2

M

k

p

C

8

q

/

z

H

r

u

Z

m

Y

U

<latexi sh

1_b

64="

IO

q53pE

KvrWnCw

Tm

>A

B9X

c

DJSgN

j

GL

Q

RPY

k

H

Z

d0

M+

8

2/

fF

z

o

u

<latexi sh

1_b

64="k nUFIOycBS23/C

rA

>

8

V

N

q

9

vg

T

7Ho

Z

+Q

m

zY

Gd

W

5

X

fJ

u

j

D

pM

K

P

R

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

<latexi sh

1_b

64="9A

W3fJT+nE

oFw5I

DKQ

>

B7

cdV

SgN

O

jXGL

v

RP

0

y2

M

k

p

C

8

q

/

z

H

r

u

Z

m

Y

U

<latexi sh

1_b

64="FWRuB+c

H80

vf

z

U3E

q

M

>A

V

NS

J

r

/

9L

oQ

Y

g

C

I

m

X

w

p

5

KZ

7

T

j

2

d

D

y

k

O

P

G

n

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

<latexi sh

1_b

64="LS82wRumB

rv

Z

3

EJ/U

0

>A

c

V

N

q

9

d

W

Q

5

F

k

O

n

H

p

Y

7

f

I

og

z

G

CD

P

y

T

M

+

<latexi sh

1_b

64="9A

W3fJT+nE

oFw5I

DKQ

>

B7

cdV

SgN

O

jXGL

v

RP

0

y2

M

k

p

C

8

q

/

z

H

r

u

Z

m

Y

U

<latexi sh

1_b

64="

IO

q53pE

KvrWnCw

Tm

>A

B9X

c

DJSgN

j

GL

Q

RPY

k

H

Z

d0

M+

8

2/

fF

z

o

u

<latexi sh

1_b

64="Q0g29vukpd

GDw8

y

FH7

>A

n

c

V

NS

E

3Ur q/

I

T

o

K

P

5

+

m

J

X

L

f

W

Y

C

R

O

M

Z

z

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

<latexi sh

1_b

64="

IO

q53pE

KvrWnCw

Tm

>A

B9X

c

DJSgN

j

GL

Q

RPY

k

H

Z

d0

M+

8

2/

fF

z

o

u

<latexi sh

1_b

64="k nUFIOycBS23/C

rA

>

8

V

N

q

9

vg

T

7Ho

Z

+Q

m

zY

Gd

W

5

X

fJ

u

j

D

pM

K

P

R

<latexi sh

1_b

64="HGY9j

CKVuEArB

n

o

+m

>

3

c

NS8

J

U

q/

L

0

X

W

Q

Z

5

w

F

k

R

2

O

v

p

7

f

d

I

g

z

D

P

y

T

M

<latexi sh

1_b

64="cj vLFAH

PI

N

D9dV

S

k

T

>

B

3

8

EJ

Ur q/

0

o

W

Q

+

Y

g

C

mz

X

w

p

5

KZ

7

2

f

u

R

y

M

O

G

n

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

<latexi sh

1_b

64="

IO

q53pE

KvrWnCw

Tm

>A

B9X

c

DJSgN

j

GL

Q

RPY

k

H

Z

d0

M+

8

2/

fF

z

o

u

<latexi sh

1_b

64="Q0g29vukpd

GDw8

y

FH7

>A

n

c

V

NS

E

3Ur q/

I

T

o

K

P

5

+

m

J

X

L

f

W

Y

C

R

O

M

Z

z

<latexi sh

1_b

64="

r

j

I

oA+C

8

HySmnX

E

GU

>

B

cdVDJ

gN

O2

W

K

T

3

5pFk

0

q

L

P

w

f

Q

v

M

9

zu

Y

R

7

/

Z

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 MolCLR[33] authors introduced i-MolCLR[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 distinct advantages. Unlike i-MolCLR 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, our method is flexible and can be applied to different molecular representation learning framework. In this paper, we first integrate our method into MolCLR [33] series model and benchmark the performance on 2D non-charality MoleculeNet[35] dataset. We then integrated our method into Uni-Mol[42] and evaluate its performance on MoleculeNet[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 ability in extracting useful features from molecular. Extensive experiments show that our method outperforms all other molecular representation learning baselines, including contrastive and non-contrastive methods.

The contributions of this paper can be summarized as follows:

To tackle the challenges posed by false positive and negative pairs, we introduce a probability method for molecular contrastive learning. By introducing different weights as random variables to 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 effective optimization algorithm based on Bayesian data augmentation and stochastic expectation maximization, to simultaneously perform posterior inference and model optimization.

Through extensive and large-scale experiments, we 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 be represented as 2D or 3D graphs depending on datasets. 2D molecule graphs have atoms as nodes and bond as edges. 3D molecule graphs additionally adds spacial positions of the atoms. For simplicity, we adopt static atom positions in this paper.

In molecular representation learning, as illustrated in Figure 2, we start by randomly sampling a batch of N molecules. Each molecule, represented as x i , undergoes stochastic augmentation strategies to generate two augmented versions, denoted as ( x i , x ′ i ) . 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 2 N augmented molecules, each pair ( x i , x ′ i ) is treated as a positive pair, while the remaining 2( N -1) augmented molecules within the same batch are considered negative samples. This setup allows us to utilize contrastive learning effectively by distinguishing between similar and dissimilar molecular structures. A neural network encoder f ( x ; θ ) , parameterized by θ , is employed to extract representation vectors z from the augmented molecular samples. In this paper, we utilize three different types of encoders in various experiments, as depicted in Figure 2 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 s i + ≜ sim( z i , z ′ i ) represent the similarity score between the positive pair ( x i , x ′ i ) after the encoder, and s ik -≜ sim ( z i , z k ) signifies the similarity score between the negative pair ( x i , x k ) , and sim( · , · ) represents any positive-valued similarity metric. In this paper, we adopt the commonly

A

ContrastiveLoss

B

D

FFN

sn1-

↑

MatMul

S1+

sn+

Softmax

Embedding

Z1

z1'

zn

zn'

Scale

LayerNorm

GNNorTransformer

GNNorTransformer

MatMul

Equivariant

featureextractor

featureextractor

↑

↑

GraphAttention

Linear

Linear

Linear

K

Atom Representation

Pair Representation

LayerNorm

FeedForward

Network

x1

X1

xn

xn'

×N

LayerNorm

FeedForward

Network

Output

Moleculex1

Molecule xn

Head

used exponential cosine similarity, defined as sim( z 1 , z 2 ) ≜ e z T 1 z 2 / ∥ z 1 ∥∥ z 2 ∥ τ , where τ denotes a temperature parameter.

2.2 Probability Weighted Contrastive Learning

We describe the proposed probability framework for molecular contrastive learning. In standard contrastive learning, one tries to encode data samples to a latent space such that positive pairs stay close to each other while negative pairs are pushed away. The contrastive loss function is:

̸

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 a way to automatically identify and differentiate these pair data? In the following, we propose a Bayesian approach to 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 be more specific, we introduce local learnable weights, denoted as w + i for each positive pair and w -ik 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 no weights) to false positive and false negative pairs;

̸

One problem with this formulation, however, is that it is 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.

Astraightforward baseline for calculating these weights can be envisioned as follows: we can consider these weights in a 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. A challenge associated with this baseline method, however, lies in the establishment of a rigid similarity threshold to create a 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, we propose a 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 to them. We consider two types of priors: a Bernoulli prior to model weights as binary random variables and a 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:

Option 2 - Bernoulli priors for selective weighting:

here, a + , b + , a -and b -are shape and rate parameters for Gamma distribution and ¯ a -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 w + i and w -ik , as:

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 i that is associated to data point x i . Consequently, we can define an augmented joint posterior distribution of the random variables θ , u , w , denoted as p ({ w + i } , { w -ik } , θ | D ) 1 , to be

where u ≜ { u 1 , u 2 , · · · , u |D| } and w ≜ { w + i } ∪ { w -ik } . It is 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 u in (4). In other words, optimization thought (4) is equivalent to optimization over (3). Consequently, we can perform learning and inference based on the augmented posterior of p ( θ , 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 a stochastic EM algorithm for efficient inference and learning of our model. Stochastic EM [24] is a stochastic variant of the EM algorithm, which is an iterative method for finding the

maximum likelihood of model parameters in statistical models when data is only partially, or when model depends on unobserved latent variables [41].

In our setting, the objective of stocastic EM is to maximize the posterior in equation 4. The basic idea is to alternatively 1) optimizing model parameter θ with fixed ( u , w ) and 2) sampling ( u , w ) with fixed θ . 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 to sampling local random variables u and w for a batch of data; stochastic expectation then uses the sampled auxiliary random variables to update the model parameter θ by maximizing a stochastic objective Q ( θ ) , defined as: Q t +1 ( θ ) = Q t ( θ ) + λ t (log p ( θ , u , w | D ) -Q t ( θ )) at iteration t +1 , where { λ t } is a 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 3 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 of the following forms:

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.

Stochastic Expectation We then proceed to calculate the stochastic expectation based on the simulated local random variables above. For notation simplicity, we define Q 0 ( θ ) = 0 . Then we can reformulate Q t +1 ( θ ) by decomposing the recursion, resulting in

where τ indexes the minibatch and the corresponding local random variables at the current time τ .

Maximization The stochastic expectation objective (6) provides a convenient form for stochastic optimization over time, similar to online optimization (Bent & Van Hentenryck, 2005). Specifically, at each time t , we can initialize the parameter θ from the last step, and update it by stochastic gradient ascent on the log-likelihood, log p ( θ , u τ , w τ | D τ ) calculated from the current batch of data. To reduce variance, we propose to optimize a marginal version by integrating out u τ from p ( θ , u τ , w τ | D τ ) , which essentially reduces to our

Algorithm 1 Contrastive Learning with Stochastic EM

1: Initialize θ ; set t = 1

2: for a batch of molecules in loader do

3: Augment each molecule x i into a pair ( x i , x ′ i )

4: Calculate positive/negative similarity scores s + and s -for all the molecule pairs

5: Initialize all the weights w + and w -to be one

6: for k = 1 to iter [4 in practice] do

7: Sample u and w according to distributions

8: end for

9: Calculate the weighted contrastive loss in equation 2 with the sampled w on the current batch of data

10: Update the model parameter by stochastic gradient descent with the calculated weighted contrastive loss

11: t = t +1

12: end for

original weighted contrastive loss in equation (2). With the above steps, it is ready to optimize the model by stochastic EM. The detailed steps are described in the Algorithm 1.

3 Related works

Contrastive Learning As a popular self-supervised learning paradigm, contrastive learning focuses on learning semantically informative representations for downstream tasks [16, 3, 39, 9]. 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 a self-supervised contrastive learning strategy to learn molecular representations. I-MolCLR [32] is a improved version of MolCLR 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 2D topology and 3D geometry. GEM [5] uses bond angles and bond length as additional edge attributes to enhance 3D information. Uni-Mol[42] is a 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 [11] enforces pairs with larger noise to be less similar in embedding space to improve model training. [6]apply noise estimation component to adjust the consistency between different modalities for action recognition task. RINCE [8] uses a 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 as a pivotal algorithm in machine learning and probabilistic modeling for large-scale Bayesian inference. Building upon the foundations of the classical Expectation-Maximization (EM) algorithm [18], Stochastic EM offers an efficient solution for parameter estimation in situations involving vast datasets or latent variables, e.g., to maximize the log -likelihood of p ( z , D | θ ) , where D is the dataset, z is the local random variable and θ 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 [1], topic modeling [40], and latent variable modeling [41], making it an indispensable tool to cope with complex probabilistic models and extensive data and a natural fit to our problem.

4 Experiments

We evaluate our method on molecular property prediction tasks. Our approach is designed to be a versatile component that can be seamlessly integrated with various molecular property prediction datasets and models. In this study, we integrate our model into three different existing models: Uni-Mol[42]], I-MolCLR [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 C.

4.1 The MoleculeNet Dataset

MoleculeNet [35] is a popular benchmark for molecular property prediction, including datasets focusing on different molecular properties, from quantum mechanics and physical chemistry to biophysics 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 1. 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 three-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.

As shown in Table 1 and 2, our method outperforms Uni-Mol[42] and GEM [5], the current stateof-the-art methods, with an average gain of 1.3 percent in classification tasks and 7.6 percent in regression tasks. This substantiates that our approach facilitates more flexible training with a higher tolerance for false positive and false negative data pairs, thereby enhancing the model's performance in molecular representation learning.

4.2 Non-Chirality version MoleculeNet

In order to make a fair comparison with I-MolCLR [32], we also integrated our method into MolCLR [33] framework. MolCLR and I-MolCLR are 2D based methods, their experiments are conducted on different version of MoleculeNet dataset that does not consider chirality. We adopted the same dataset, augmentation, GNN-based encoder and other settings. As shown in Table 3, our method outperforms I-MolCLR on 7 out of 9 downstream tasks and got an average of 2 points increase on non-chirality MoleculeNet classification datasets.

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 a baseline method. The data partition we use has 110k,10k,and 11k molecules in training, validation and testing sets. We use both our contrastive loss function and original minimize mean absolute error(MAE) as training objectives.

As shown in 4, we get state of the art result in 8 out of 12 baselines. The increase is relatively subtle compared with other dataset, we 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.

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 positive samples and that accounting for this phenomenon can improve learned molecule representations. In Figure A(See Appendix A), we plot the distribution of similarity scores for both positive and negative samples. Figure A 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 A right reveals that our method also mitigates the false negative problem in standard CL. It 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 is the very reason our method captures dissimilarity between positive pairs. By allowing some degree of alignment between the right set of negative examples, our method is able to minimize the inconsistencies between shared context of related positives and negatives. This in turn allows us to learn an overall more coherent representation space, resulting in increased robustness and downstream performance.

Comparisons with the Standard Contrastive Learning We conducted an ablation study to showcase 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[42].

Table 5 presents the results of our ablation study. Incorporating the probabilistic framework resulted in a great improvement of 3.4-point increase in ROC-AUC, significantly enhances the model's performance. On the other hand, introducing the additional loss component led to an increase in ROC-AUC by 2.9 points, demonstrating its secondary role in enhancing the model's performance. When we adopt both of them, we can get the final ROC-AUC of 80.1 average on MoleculeNet classification datasets.

Hyperparameters We also conducted an ablation study to determine the optimal hyperparameters (e.g., a + , a -) on MoleculeNet classification datasets. We selected a + , a -, b + , and b -from the range [1 , 5 , 10] . Table 6 indicates that our method achieves the best performance with a + = 5 and a -= b + = b -= 1 . Tuning different hyperparameters affects performance, with an increase in a + from 1 to 5 leading to a 1.6 percent performance gain.

5 Conclusion

In this paper, we investigate an important yet unnoticed limitation of molecular contrastive learning, where augmented graph data come with false positive and false negative data pairs. As a remedy, we propose a principled solution to molecular contrastive learning by reformulating it into a probability

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 on 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 a 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.

References

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.

Changyou Chen et al. 'Why do We Need Large Batchsizes in Contrastive Learning? A Gradient-Bias Perspective'. In: Advances in Neural Information Processing Systems . Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., 2022, pp. 33860-33875.

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-8775.

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. 38913898.

Xiantao Fang, Ling Liu, Juncheng Lei, et al. 'Geometry-enhanced molecular representation learning for property prediction'. In: Nature Machine Intelligence 4.2 (2022), pp. 127-134.

H. Han et al. 'Noise-tolerant learning for audio-visual action recognition'. In: IEEE Transactions on Multimedia (2024).

Kaiming He, Haoqi Fan, Yuxin Wu, et al. 'Momentum contrast for unsupervised visual representation learning'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2020, pp. 9729-9738.

D. T. Hoffmann et al. 'Ranking info noise contrastive estimation: Boosting contrastive learning via ranked positives'. In: Proceedings of the AAAI Conference on Artificial Intelligence . V ol. 36. 1. 2022, pp. 897-905.

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.

Weihua Hu, Bowen Liu, Joseph Gomes, et al. 'Strategies for pre-training graph neural networks'. In: arXiv preprint arXiv:1905.12265 (2019).

R. Huang et al. 'Nlip: Noise-robust language-image pre-training'. In: Proceedings of the AAAI Conference on Artificial Intelligence . Vol. 37. 1. 2023, pp. 926-934.

Runxuan Jiao, Jingrui Han, Wen Huang, et al. 'Energy-motivated equivariant pretraining for 3d molecular graphs'. In: Proceedings of the AAAI Conference on Artificial Intelligence . Vol. 37. 7. 2023, pp. 8096-8104.

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), pp. 1893-1904.

Panagiotis I 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.

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.

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.

Yuyang Liao and Tess Smidt. 'Equiformer: Equivariant graph attention transformer for 3d atomistic graphs'. In: arXiv preprint arXiv:2206.11990 (2022).

D. Lin. An Introduction to Expectation-Maximization . 2011.

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).

Shitong Liu, Hui Wang, Wenhao Liu, et al. 'Pre-training molecular graph representation with 3d geometry'. In: arXiv preprint arXiv:2110.07728 (2021).

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.

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.

Garrett M Morris, Ruth Huey, William Lindstrom, et al. 'AutoDock4 and AutoDockTools4: Automated docking with selective receptor flexibility'. In: Journal of computational chemistry 30.16 (2009), pp. 2785-2791.

S. F. Nielsen. 'The stochastic EM algorithm: estimation and asymptotic results'. In: Bernoulli (2000), pp. 457-489.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. 'Representation learning with contrastive predictive coding'. In: arXiv preprint arXiv:1807.03748 (2018).

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.

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.

Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, et al. 'Quantum chemistry structures and properties of 134 kilo molecules'. In: Scientific data 1.1 (2014), pp. 1-7.

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. 1255912571.

Oleg Trott and Arthur J Olson. 'AutoDock Vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading'. In: Journal of computational chemistry 31.2 (2010), pp. 455-461.

Sheng Wang, Yiwen Guo, Yanming Wang, et al. 'Smiles-bert: large scale unsupervised pretraining for molecular property prediction'. In: Proceedings of the 10th ACM international conference on bioinformatics, computational biology and health informatics . 2019, pp. 429436.

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.

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.

Fan Wu, Shuoguo Li, Xin Jin, et al. 'Rethinking explaining graph neural networks via nonparametric subgraph matching'. In: International Conference on Machine Learning . Vol. 375. Proceedings of Machine Learning Research. PMLR. 2023, pp. 11-23.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, et al. 'MoleculeNet: a benchmark for molecular machine learning'. In: Chemical Science 9.2 (2018), pp. 513-530.

Zhaocheng Xiong, De Wang, Xia Liu, et al. 'Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism'. In: Journal of Medicinal Chemistry 63.16 (2019), pp. 8749-8760.

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.

Yuning You, Tianlong Chen, Yang Shen, et al. 'Graph contrastive learning automated'. In: International Conference on Machine Learning . PMLR. 2021, pp. 12121-12132.

Yuning You, Tianlong Chen, Yuxuan Sui, et al. 'Graph contrastive learning with augmentations'. In: Advances in Neural Information Processing Systems . Vol. 33. 2020, pp. 58125823.

Manzil Zaheer et al. 'Exponential stochastic cellular automata for massively parallel inference'. In: Artificial Intelligence and Statistics . PMLR. 2016, pp. 966-975.

Shujian Zhang and Yuguo Chen. 'Computation for latent variable model estimation: A unified stochastic proximal framework'. In: Psychometrika 87.4 (2022), pp. 1473-1502.

Guangyong Zhou, Zhen Gao, Qi Ding, et al. 'Uni-mol: A universal 3d molecular representation learning framework'. In: International Conference on Learning Representations (2023).

A Similarity Score Distribution

Distribution of Scores for positive samples

Ours

MolCLR

4

0.5

1.0

1.5

Score

2.0

2.5

Distribution of Scores for negative samples

Ours

MolCLR

0.5

1.0

1.5

2.0

2.5

Score

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 of the model. For instance, correlated features could lead to biased estimates of weights, while unnoisy data might degrade the necessity to apply our method in learning representations. Future work could explore methods to relax these assumptions and enhance the model's robustness to 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 or real-world scenarios might be limited. Further validation on a broader range of datasets is necessary to confirm the wide applicability of our approach.

Also, one limitation of our method is that the performance gains brought by the proposed architectural improvements can depend on datasets and tasks. For small datasets like QM9, the performance gain is 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 in 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.

C Training details for experiments

C.1 MoleculeNet dataset

Following Unimol, we report the detailed hyperparameters setup of during pretraining in 7. 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 to a 512-dimensional embedding h . A projection head, modeled as an MLP with a single hidden layer, transforms h into a 256-dimensional latent vector z . ReLU is utilized as the non-linear activation function. The model undergoes pre-training over 50 epochs with a batch size of 512, optimized via the Adam optimizer with an initial learning rate of 5 × 10 -4 and a weight decay rate of 1 × 10 -5 . A cosine learning rate decay schedule is applied throughout pre-training. For most datasets, we use a scaffold-based data split; however, the QM9 subtask follows a random split in line with the Mol-CLR methodology.

In the fine-tuning phase, the projection head is replaced with a randomly initialized MLP that maps h to 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 PyTorch Geometric.

C.3 QM9 Dataset Experiments

Data partitioning for the QM9 tasks follows the scheme utilized in the Equiformer. For tasks predicting µ, α, ε HOMO , ε LUMO , ∆ ε , and C ν , our configuration includes a batch size of 64, 300 training epochs, a learning rate of 5 × 10 -4 , and Gaussian radial basis functions with 128 bases. The architecture comprises six Transformer blocks, a weight decay of 5 × 10 -3 , and a dropout rate of 0.2. Mixed-precision training is employed for these tasks.

For the R 2 task, the setup consists of a batch size of 48, 300 epochs, a learning rate of 1 . 5 × 10 -4 , Gaussian radial basis functions with 128 bases, five Transformer blocks, a weight decay of 5 × 10 -3 , and a dropout rate of 0.1, executed in single precision.

The ZPVE task employs a batch size of 48, 300 epochs, a learning rate of 1 . 5 × 10 -4 , Gaussian radial basis functions with 128 bases, five Transformer blocks, a weight decay of 5 × 10 -3 , and a dropout rate of 0.2, with single-precision training.

For the G,H,U , and U 0 tasks, the setup includes a batch size of 48, 300 epochs, a learning rate of 1 . 5 × 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-precision tasks requiring 81 GPU-hours and single-precision 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 in the Equiformer paper.

D Comparison against other weight calculation methods

Here we show the comparison against using 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 (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-MolCLR [32]. 3. Encoderbased 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:

In our experiments, we set ( λ = 1) .

From these results, we observe the following:

While fingerprint-based similarity showed improvements in 2 out of 9 tasks compared to our original method, but it 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.

Encoder-based Similarity 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 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 a 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 are concatenated as the input of an additional 4-layer Transformer decoder, which is finetuned to learn the pair distances of all heavy atoms in molecule and pocket. Then, with the predicted pair-distance matrix as a scoring function, we first randomly place the ligand and then optimize the coordinates of its atoms by directly back-propagation 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 use 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 state-of-the-art results with our Gamma-prior version model.

F More ablation on a u and b u

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 .

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]

Justification: 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.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The paper includes a 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 of the claims, emphasizing that the results are based on experiments conducted on specific datasets.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The paper includes a 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of 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 instructions 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 to 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 (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer:

[TODO]

Justification: 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.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The paper does not report error bars, confidence intervals, or statistical significance tests for the experimental results. Instead, all the experiments are carried out for three times and the average is reported.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: 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.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper does not pose such risks.

12. Licenses for existing assets

Question: 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?

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]

Justification:The 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.

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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: the paper does not involve crowdsourcing 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/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve studying participants

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n(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.\n\nWEAKNESSES:\n(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.\n\nQUESTIONS:\n(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.\n\nLIMITATIONS:\nYes, the authors addressed limitations adequately.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

comments: 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.

STRENGTHS:
1. To the best of my knowledge, this paper raises an intriguing question and warrants attention.
2. The approach taken to address the problem is also innovative.
3. 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:
1. Eq.3 is supposed to be $e^{w_i^+s_i}$, not $w_i^+s_i$. Or am I misunderstanding?
2. Line 129 contains a duplicated "positive."
3. 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.

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?
2. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nCommonly 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.\n\nSTRENGTHS:\n* 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?\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}