A probability contrastive learning framework for 3D molecular representation learning

Authors: jiayu  qin
jian  chen
rohan  sharma
jingchen  sun
changyou  chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:43:39

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.

Main paper: 
 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. Q i V W G K I q V L W n Q X P 0 5 k W G h 9 V S E t l N g M 9 b L X i 7 + 5 / V S E 1 0 F G Z N J a q g k i 0 V R y p G J U f 4 4 G j J F i e F T SQ i V W G K I q V L W n Q X P 0 5 k W G h 9 V S E t l N g M 9 b L X i 7 + 5 / V S E 1 0 F G Z N J a q g k i 0 V R y p G J U f 4 4 G j J F i e F T SF I O y c B S 2 3 / C t s / r A O 2 L 0 E w = " > A A A B 8 n i c b V B N S 8 N A E N 3 U r 1 q / q h 6 9 B I v g q S T i 1 7 H o x Z N U s B + Q h r L Z T N q l m 2 z Y n Y g l 9 G d 4 8 a C I V 3 + N N / + N 2 z Y H r T 4 Y e L w 3 w 8 y 8 I B V c o + N 8 W a W l 5 Z X V t f J 6 Z W N z a 3 u n u r v X 1 j J T D F p M C q m 6 A d U g e A I t 5 C i g m y q g c S C g E 4 y u p 3 7 n A Z T m M r n H c Q p + T A c J j z i j a C S v h / C I + a 0 M Y d K v 1 p y 6 M 4 P 9 l 7 g F q Z E C z X 7 1 s x d K l s W Q I B N U a 8 9 1 U v R z q p A z A Z N K L 9 O Q U j a i A / A M T W g M 2 s 9 n J 0 / s I 6 O E d i S V q Q T t m f p z I q e x 1 u M 4 M J 0 x x a F e 9 K b i f 5 6 X Y X T p 5 z x J M 4 S E z R d F m b B R 2 t P / 7 Z A r Y C j G h l C m u L n V Z k O q K E O T U s W E 4 C 6 + / J e 0 T + r u e f 3 s 7 r T W u C r i K J M D c k i O i U s u S I P c k C Z p E U Y k e S I v 5 N V C 6 9 l 6 s 9 7 n r S W r m N k n v 2 B 9 f A O 9 P J G R < / l a t e x i t > Node < l a t e x i t s h a 1 _ b a s e 6 4 = " a I O h q 5 3 p E 6 s K v r W n C w t T m w s V 7 y U = " > A A A B 9 X i c b V D J S g N B E O 1 x j X G L e v Q y G A R P Y U b c j k E v H i O Y B Z I x 9 H R q k i b d 0 0 N 3 T T Q M + Q 8 v H h T x 6 r 9 4 8 2 / s L A d N f F D w e K + K q n p h I r h B z / t 2 l p Z X V t f W c x v 5 z a 3 t n d 3 C 3 n 7 N q F Q z q D I l l G 6 E 1 I D g M V S R o 4 B G o o H K U E A 9 7 N + M / f o A t O E q v s d h A o G k 3 Z h H n F G 0 0 k M L 4 Q k z D V I N q B i 1 C 0 W v 5 E 3 g L h J / R o p k h k q 7 8 N X q K J Z K i J E J a k z T 9 x I M M q q R M w G j f C s 1 k F D W p 1 1 o W h p T C S b I J l e P 3 G O r d N x I a V s x u h P 1 9 0 R G p T F D G d p O S b F n 5 r 2 x + J / X T D G 6 C j I e J y l C z K a L o l S 4 q N x x B G 6 H a 2 A o h p Z Q p r m 9 1 W U 9 q i l D G 1 T e h u D P v 7 x I a q c l / 6 J 0 f n d W L F / P 4 s i R Q 3 J E T o h P L k m Z 3 J I K q R J G N H k m r + T N e X R e n H f n Y 9 q 6 5 M x m D s g f O J 8 / a i K T H w = = < / l a t e x i t > removal < l a t e x i t s h a 1 _ b a s e 6 4 = " Q 0 g 2 9 v u k p d 2 G D w 8 w y 2 s j B w F H 7 s Q = " > A A A B 8 n i c b V B N S 8 N A E N 3 U r 1 q / q h 6 9 B I v g q S T i 1 7 E o g s c K 9 g P S U D a b S b t 0 k w 2 7 E 7 G E / g w v H h T x 6 q / x 5 r 9 x 2 + a g 1 Q c D j / d m m J k X p I J r d J w v q 7 S 0 v L K 6 V l 6 v b G x u b e 9 U d / f a W m a K Q Y t J I V U 3 o B o E T 6 C F H A V 0 U w U 0 D g R 0 g t H 1 1 O 8 8 g N J c J v c 4 T s G P 6 S D h E W c U j e T 1 E B 4 x v w k H M O l X a 0 7 d m c H + S 9 y C 1 E i B Z r / 6 2 Q s l y 2 J I k A m q t e c 6 K f o 5 V c i Z g E m l l 2 l I K R v R A X i G J j Q G 7 e e z k y f 2 k V F C O 5 L K V I L 2 T P 0 5 k d N Y 6 3 E c m M 6 Y 4 l A v e l P x P 8 / L M L r 0 c 5 6 k G U L C 5 o u i T N g o 7 e n / d s g V M B R j Q y h T 3 N x q s y F V l K F J q W J C c B d f / k v a J 3 X 3 v H 5 2 d 1 p r X B V x l M k B O S T H x C U X p E F u S Z O 0 C C O S P J E X 8 m q h 9 W y 9 W e / z 1 p J V z O y T X 7 A + v g G j O Z G A < / l a t e x i t > Edge < l a t e x i t s h a 1 _ b a s e 6 4 = " a I O h q 5 3 p E 6 s K v r W n C w t T m w s V 7 y U = " > A A A B 9 X i c b V D J S g N B E O 1 x j X G L e v Q y G A R P Y U b c j k E v H i O Y B Z I x 9 H R q k i b d 0 0 N 3 T T Q M + Q 8 v H h T x 6 r 9 4 8 2 / s L A d N f F D w e K + K q n p h I r h B z / t 2 l p Z X V t f W c x v 5 z a 3 t n d 3 C 3 n 7 N q F Q z q D I l l G 6 E 1 I D g M V S R o 4 B G o o H K U E A 9 7 N + M / f o A t O E q v s d h A o G k 3 Z h H n F G 0 0 k M L 4 Q k z D V I N q B i 1 C 0 W v 5 E 3 g L h J / R o p k h k q 7 8 N X q K J Z K i J E J a k z T 9 x I M M q q R M w G j f C s 1 k F D W p 1 1 o W h p T C S b I J l e P 3 G O r d N x I a V s x u h P 1 9 0 R G p T F D G d p O S b F n 5 r 2 x + J / X T D G 6 C j I e J y l C z K a L o l S 4 q N x x B G 6 H a 2 A o h p Z Q p r m 9 1 W U 9 q i l D G 1 T e h u D P v 7 x I a q c l / 6 J 0 f n d W L F / P 4 s i R Q 3 J E T o h P L k m Z 3 J I K q R J G N H k m r + T N e X Re n H f n Y 9 q 6 5 M x m D s g f O J 8 / a i K T H w = = < / l a t e x i t > removal < l a t e x i t s h a 1 _ b a s e 6 4 = " k x 1 n U F I O y c B S 2 3 / C t s / r A O 2 L 0 E w = " > A A A B 8 n i c b V B N S 8 N A E N 3 U r 1 q / q h 6 9 B I v g q S T i 1 7 H o x Z N U s B + Q h r L Z T N q l m 2 z Y n Y g l 9 G d 4 8 a C I V 3 + N N / + N 2 z Y H r T 4 Y e L w 3 w 8 y 8 I B V c o + N 8 W a W l 5 Z X V t f J 6 Z W N z a 3 u n u r v X 1 j J T D F p M C q m 6 A d U g e A I t 5 C i g m y q g c S C g E 4 y u p 3 7 n A Z T m M r n H c Q p + T A c J j z i j a C S v h / C I + a 0 M Y d K v 1 p y 6 M 4 P 9 l 7 g F q Z E C z X 7 1 s x d K l s W Q I B N U a 8 9 1 U v R z q p A z A Z N K L 9 O Q U j a i A / A M T W g M 2 s 9 n J 0 / s I 6 O E d i S V q Q T t m f p z I q e x 1 u M 4 M J 0 x x a F e 9 K b i f 5 6 X Y X T p 5 z x J M 4 S E z R d F m b B R 2 t P / 7 Z A r Y C j G h l C m u L n V Z k O q K E O T U3 p E 6 s K v r W n C w t T m w s V 7 y U = " > A A A B 9 X i c b V D J S g N B E O 1 x j X G L e v Q y G A R P Y U b c j k E v H i O Y B Z I x 9 H R q k i b d 0 0 N 3 T T Q M + Q 8 v H h T x 6 r 9 4 8 2 / s L A d N f F D w e K + K q n p h I r h B z / t 2 l p Z X V t f W c x v 5 z a 3 t n d 3 C 3 n 7 N q F Q z q D I l l G 6 E 1 I D g M V S R o 4 B G o o H K U E A 9 7 N + M / f o A t O E q v s d h A o G k 3 Z h H n F G 0 0 k M L 4 Q k z D V I N q B i 1 C 0 W v 5 E 3 g L h J / R o p k h k q 7 8 N X q K J Z K i J E J a k z T 9 x I M M q q R M w G j f C s 1 k F D W p 1 1 o W h p T C S b I J l e P 3 G O r d N x I a V s x u h P 1 9 0 R G p T F D G d p O S b F n 5 r 2 x + J / X T D G 6 C j I e J y l C z K a L o l S 4 q N x x B G 6 H a 2 A o h p Z Q p r m 9 1 W U 9 q i l D G 1 T e h u D P v 7 x I a q c l / 6 J 0 f n d W L F / P 4 s i R Q 3 J E T o h P L k m Z 3 J I K q R J G N H k m r + T N e X Re n H f n Y 9 q 6 5 M x m D s g f O J 8 / a i K T H w = = < / l a t e x i t > removal < l a t e x i t s h a 1 _ b a s e 6 4 = " Q 0 g 2 9 v u k p d 2 G D w 8 w y 2 s j B w F H 7 s Q = " > A A A B 8 n i c b V B N S 8 N A E N 3 U r 1 q / q h 6 9 B I v g q S T i 1 7 E o g s c K 9 g P S U D a b S b t 0 k w 2 7 E 7 G E / g w v H h T x 6 q / x 5 r 9 x 2 + a g 1 Q c D j / d m m J k X p I J r d J w v q 7 S 0 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 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.v L K 6 V l 6 v b G x u b e 9 U d / f a W m a K Q Y t J I V U 3 o B o E T 6 C F H A V 0 U w U 0 D g R 0 g t H 1 1 O 8 8 g N J c J v c 4 T s G P 6 S D h E W c U j e T 1 E B 4 x v w k H M O l X a 0 7 d m c H + S 9 y C 1 E i B Z r / 6 2 Q s l y 2 J I k A m q t e c 6 K f o 5 V c i Z g E m l l 2 l I K R v R A X i G J j Q G 7V M B R j Q y h T 3 N x q s y F V l K F J q W J C c B d f / k v a J 3 X 3 v H 5 2 d 1 p r X B V x l M k B O S T H x C U X p E F u S Z O 0 C C O S P J E X 8 m q h 9 W y 9 W e / z 1 p J V z O y T X 7 A + v g G j O Z G A < / l a t e x i t > Edge < l a t e x i t s h a 1 _ b a s e 6 4 = " a I O h q 5 3 p E 6 s K v r W n C w t T m w s V 7 y U = " > A A A B 9 X i c b V D J S g N B E O 1 x j X G L e v Q y G A R P Y U b c j k E v H i O Y B Z I x 9 H R q k i b d 0 0 N 3 T T Q M + Q 8 v H h T x 6 r 9 4 8 2 / s L A d N f F D w e K + K q n p h I r h B z / t 2 l p Z X V t f W c x v 5 z a 3 t n d 3 C 3 n 7 N q F Q z q D I l l G 6 E 1 I D g M V S R o 4 B G o o H K U E A 9 7 N + M / f o A t O E q v s d h A o G k 3 Z h H n F G 0 0 k M L 4 Q k z D V I N q B i 1 C 0 W v 5 E 3 g L h J / R o p k h k q 7 8 N X q K J Z K i J E J a k z T 9 x I M M q q R M w G j f C s 1 k F D W p 1 1 o W h p T C S b I J l e P 3 G O r d N x I a V s x u h P 1 9 0 R G p T F D G d p O S b F n 5 r 2 x + J / X T D G 6 C j I e J y l C z K a L o l S 4 q N x x B G 6 H a 2 A o h p Z Q p r m 9 1 W U 9 q i l D G 1 T e h u D P v 7 x I a q c l / 6 J 0 f n d W L F / P 4 s i R Q 3 J E T o h P L k m Z 3 J I K q R J G N H k m r + T N e X R e n H f n Y 9 q 6 5 M x m D s g f O J 8 / a i K T H w = = < / l a t e x i t > removal < l a t e x i t s h a _ b a s e = " x r b j i I o A + C + C H y S m n X C E A G U = " > A A A B H i c d V D J S g N B E O J W x b K O X x i B G m a S T I y o J c c I p F k i H d H q S J j L T V i G P I V X j w o t X P e b f F k E F X Q H i v i q p X i y A s v M D I r q v r G n N N b z u e f v + g p a J E U t a k k Y h k x y O K C R y J n A Q r B N L R g J P s L Y v p z T s m F Y / C G j E z A I M O Q + p w S d N s D d g / p d X a z x c s s + K U z i t k y n W i W H E K C K b h w F t E S j n / v D S K a B C w E K o h S X d u K w U J B E F m + Z i W I x o W M y Z F N Q x I w a b z g f R C s D E d S V w h r n f S E m g C T w d G d A Y K R + e z P x L + b g F U x G C b C Q L h b i c A Q d n e M A l o y A m m h A q u b V x G R h I L O K K d D + P o U / a R d O u m M V u V C W M a R R U f o G J i G h G q q j B m o i i g L g J Q s y G N R + P F e F Z o z l z C H A e P t E O W k O g = < / l a t e x i t > SH < l a t e x i t s h a 1 _ b a s e 6 4 = " x r b j i I 4 o A + C + C 8 1 H y S m n X C 4 E A G U = " > A A A B 8 H i c d V D J S g N B E O 2 J W 4 x b 1 K O X x i B 4 G m a S T I y 3 o J c c I 5 p F k i H 0 d H q S J j 0 L 3 T V i G P I V X j w o 4 t X P 8 e b f 2 F k E F X 1 Q 8 H i v i q p 6 X i y 4 A s v 6 M D I r q 2 v r G 9 n N 3 N b 2 z u 5 e f v + g p a J E U t a k k Y h k x y O K C R 6 y J n A Q r B N L R g J P s L Y 3 v p z 5 7 T s m F Y / C G 5 j E z A 3 I M O Q + p w S 0 d N s D d g / p d X 3 a z x c s s + K U z 6 0 i t k y n W i 2 W H E 1 K C 2 K b 1 h w F t E S j n 3 / v D S K a B C w E K o h S X d u K w U 2 J B E 4 F m + Z 6 i W I x o W M y Z F 1 N Q x I w 5 a b z g 6 f 4 R C s D 7 E d S V w h 4 r n 6 f S E m g 1 C T w d G d A Y K R + e z P x L 6 + b g F 9 1 U x 7 G C b C Q L h b 5 i c A Q 4 d n 3 e M A l o y A m m h A q u b 4 V 0 x G R h I L O K K d D + P o U / 0 9 a R d O u m M 5 V u V C 7 W M a R R U f o G J 0 i G 5 2 h G q q j B m o i i g L 0 g J 7 Q s y G N R + P F e F 2 0 Z o z l z C H 6 A e P t E 4 O W k O g = < / l a t e x i t > SH < l a t e x i t s h a 1 _ b a s e 6 4 = " x r b j i I 4 o A + C + C 8 1 H y S m n X C 4 E A G U = " > A A A B 8 H i c d V D J S g N B E O 2 J W 4 x b 1 K O X x i B 4 G m a S T I y 3 o J c c I 5 p F k i H 0 d H q S J j 0 L 3 T V i G P I V X j w o 4 t X P 8 e b f 2 F k E F X 1 Q 8 H i v i q p 6 X i y 4 A s v 6 M D I r q 2 v r G 9 n N 3 N b 2 z u 5 e f v + g p a J E U t a k k Y h k x y O K C R 6 y J n A Q r B N L R g J P s L Y 3 v p z 5 7 T s m F Y / C G 5 j E z A 3 I M O Q + p w S 0 d N s D d g / p d X 3 a z x c s s + K U z 6 0 i t k y n W i 2 W H E 1 K C 2 K b 1 h w F t E S j n 3 / v D S K a B C w E K o h S X d u K w U 2 J B E 4 F m + Z 6 i W I x o W M y Z F 1 N Q x I w 5 a b z g 6 f 4 R C s D 7 E d S V w h 4 r n 6 f S E m g 1 C T w d G d A Y K R + e z P x L 6 + b g F 9 1 U x 7 G C b C Q L h b 5 i c A Q 4 d n 3 e M A l o y A m m h A q u b 4 V 0 x G R h I L O K K d D + P o U / 0 9 a R d O u m M 5 V u V C 7 W M a R R U f o G J 0 i G 5 2 h G q q j B m o i i g L 0 g J 7 Q s y G N R + P F e F 2 0 Z o z l z C H 6 A e P t E 4 O W k O g = < / l a t e x i t > SH < l a t e x i t s h a 1 _ b a s e 6 4 = " 9 A a W 3 f J T + n E A o F w 5 I J w T 6 D K Q J e 4 = " > A A A B 7 3 i c d V D J S g N B E O 1 x j X G L e v T S G A R P w 0 w y 2 W 5 B Q T x G M A s k Q + j p 9 C R N e h a 7 a 8 Q w 5 C e 8 e F D E q 7 / j z b + x s w g q + q D g 8 V 4 V V f W 8 W H A F l v V h r K y u r W 9 s Z r a y 2 z u 7 e / u 5 g 8 O W i h J J W Z N G I p I d j y g m e M i a w E G w T i w Z C T z B 2 t 7 4 Y u a 3 7 5 h U P A p v Y B I z N y D D k P u c E t B S p w f s H t L L a T + X t 8 x S o e T U L G y Z Z a t S d Y q a O M V a w S 5 j 2 7 T m y K M l G v 3 c e 2 8 Q 0 S R g I V B B l O r a V g x u S i R w K t g 0 2 0 s U i w k d k y H r a h q S g C k 3 n d 8 7 x a d a G W A / k r p C w H P 1 + 0 R K A q U m g a c 7 A w I j 9 d u b i X 9 5 3 Q T 8 q p v y M E 6 A h X S x y E 8 E h g j P n s c D L h k F M d G E U M n 1 r Z i O i C Q U d E R Z H c L X p / h / 0 i q Y d t k s X T v 5 + v k y j g w 6 R i f o D N m o g u r o C j V Q E 1 E k 0 A N 6 Q s / G r f F o v B i v i 9 Y V Y z l z h H 7 A e P s E z u q Q g Q = = < / l a t e x i t > F < l a t e x i t s h a 1 _ b a s e 6 4 = " x r b j i I 4 o A + C + C 8 1 H y S m n X C 4 E A G U = " > A A A B 8 H i c d V D J S g N B E O 2 J W 4 x b 1 K O X x i B 4 G m a S T I y 3 o J c c I 5 p F k i H 0 d H q S J j 0 L 3 T V i G P I V X j w o 4 t X P 8 e b f 2 F k E F X 1 Q 8 H i v i q p 6 X i y 4 A s v 6 M D I r q 2 v r G 9 n N 3 N b 2 z u 5 e f v + g p a J E U t a k k Y h k x y O K C R 6 y J n A Q r B N L R g J P s L Y 3 v p z 5 7 T s m F Y / C G 5 j E z A 3 I M O Q + p w S 0 d N s D d g / p d X 3 a z x c s s + K U z 6 0 i t k y n W i 2 W H E 1 K C 2 K b 1 h w F t E S j n 3 / v D S K a B C w E K o h S X d u K w U 2 J B E 4 F m + Z 6 i W I x o W M y Z F 1 N Q x I w 5 a b z g 6 f 4 R C s D 7 E d S V w h 4 r n 6 f S E m g 1 C T w d G d A Y K R + e z P x L 6 + b g F 9 1 U x 7 G C b C Q L h b 5 i c A Q 4 d n 3 e M A l o y A m m h A q u b 4 V 0 x G R h I L O K K d D + P o U / 0 9 a R d O u m M 5 V u V C 7 W M a R R U f o G J 0 i G 5 2 h G q q j B m o i i g L 0 g J 7 Q s y G N R + P F e F 2 0 Z o z l z C H 6 A e P t E 4 O W k O g = < / l a t e x i t > SH < l a t e x i t s h a 1 _ b a s e 6 4 = " 9 A a W 3 f J T + n E A o F w 5 I J w T 6 D K Q J e 4 = " > A A A B 7 3 i c d V D J S g N B E O 1 x j X G L e v T S G A R P w 0 w y 2 W 5 B Q T x G M A s k Q + j p 9 C R N e h a 7 a 8 Q w 5 C e 8 e F D E q 7 / j z b + x s w g q + q D g 8 V 4 V V f W 8 W H A F l v V h r K y u r W 9 s Z r a y 2 z u 7 e / u 5 g 8 O W i h J J W Z N G I p I d j y g m e M i a w E G w T i w Z C T z B 2 t 7 4 Y u a 3 7 5 h U P A p v Y B I z N y D D k P u c E t B S p w f s H t L L a T + X t 8 x S o e T U L G y Z Z a t S d Y q a O M V a w S 5 j 2 7 T m y K M l G v 3 c e 2 8 Q 0 S R g I V B B l O r a V g x u S i R w K t g 0 2 0 s U i w k d k y H r a h q S g C k 3 n d 8 7 x a d a G W A / k r p C w H P 1 + 0 R K A q U m g a c 7 A w I j 9 d u b i X 9 5 3 Q T 8 q p v y M E 6 A h X S x y E 8 E h g j P n s c D L h k F M d G E U M n 1 r Z i O i C Q U d E R Z H c L X p / h / 0 i q Y d t k s X T v 5 + v k y j g w 6 R i f o D N m o g u r o C j V Q E 1 E k 0 A N 6 Q s / G r f F o v B i v i 9 Y V Y z l z h H 7 A e P s E z u q Q g Q = = < / l a t e x i t > F < l a t e x i t s h a 1 _ b a s e 6 4 = " x r b j i I 4 o A + C + C 8 1 H y S m n X C 4 E A G U = " > A A A B 8 H i c d V D J S g N B E O 2 J W 4 x b 1 K O X x i B 4 G m a S T I y 3 o J c c I 5 p F k i H 0 d H q S J j 0 L 3 T V i G P I V X j w o 4 t X P 8 e b f 2 F k E F X 1 Q 8 H i v i q p 6 X i y 4 A s v 6 M D I r q 2 v r G 9 n N 3 N b 2 z u 5 e f v + g p a J E U t a k k Y h k x y O K C R 6 y J n A Q r B N L R g J P s L Y 3 v p z 5 7 T s m F Y / C G 5 j E z A 3 I M O Q + p w S 0 d N s D d g / p d X 3 a z x c s s + K U z 6 0 i t k y n W i 2 W H E 1 K C 2 K b 1 h w F t E S j n 3 / v D S K a B C w E K o h S X d u K w U 2 J B E 4 F m + Z 6 i W I x o W M y Z F 1 N Q x I w 5 a b z g 6 f 4 R C s D 7 E d S V w h 4 r n 6 f S E m g 1 C T w d G d A Y K R + e z P x L 6 + b g F 9 1 U x 7 G C b C Q L h b 5 i c A Q 4 d n 3 e M A l o y A m m h A q u b 4 V 0 x G R h I L O K K d D + P o U / 0 9 a R d O u m M 5 V u V C 7 W M a R R U f o G J 0 i G 5 2 h G q q j B m o i i g L 0 g J 7 Q s y G N R + P F e F 2 0 Z o z l z C H 6 A e P t E 4 O W k O g = < / l a t e x i t > SH < l a t e x i t s h a 1 _ b a s e 6 4 = " 9 A a W 3 f J T + n E A o F w 5 I J w T 6 D K Q J e 4 = " > A A A B 7 3 i c d V D J S g N B E O 1 x j X G L e v T S G A R P w 0 w y 2 W 5 B Q T x G M A s k Q + j p 9 C R N e h a 7 a 8 Q w 5 C e 8 e F D E q 7 / j z b + x s w g q + q D g 8 V 4 V V f W 8 W H A F l v V h r K y u r W 9 s ZS R g I V B B l O r a V g x u S i R w K t g 0 2 0 s U i w k d k y H r a h q S g C k 3 n d 8 7 x a d a G W A / k r p C w H P 1 + 0 R K A q U m g a c 7 A w I j 9 d u b i X 9 5 3 Q T 8 q p v y M E 6 A h X S x y E 8 E h g j P n s c D L h k F M d G E U M n 1 r Z i O i C Q U d E R Z H c L X p / h / 0 i q Y d t k s X T v 5 + v k y j g w 6 R i f o D N m o g u r o C j V Q E 1 E k 0 A N 6 Q s / G r f F o v B i v i 9 Y V Y z l z h H 7 A e P s E z u q Q g Q = = < / l a t e x i t > F < l a t e x i t s h a 1 _ b a s e 6 4 = " x r b j i I 4 o A + C + C 8 1 H y S m n X C 4 E A G U = " > A A A B 8 H i c d V D J S g N B E O 2 J W 4 x b 1 K O X x i B 4 G m a S T I y 3 o J c c I 5 p F k i H 0 d H q S J j 0 L 3 T V i G P I V X j w o 4 t X P 8 e b f 2 F k E F X 1 Q 8 H i v i q p 6 X i y 4 A s v 6 M D I r q 2 v r G 9 n N 3 N b 2 z u 5 e f v + g p a J E U t a k k Y h k x y O K C R 6 y J n A Q r B N L R g J P s L Y 3 v p z 5 7 T s m F Y / C G 5 j E z A 3 I M O Q + p w S 0 d N s D d g / p d X 3 a z x c s s + K U z 6 0 i t k y n W i 2 W H E 1 K C 2 K b 1 h w F t E S j n 3 / v D S K a B C w E K o h S X d u K w U 2 J B E 4 F m + Z 6 i W I x o W M y Z F 1 N Q x I w 5 a b z g 6 f 4 R C s D 7 E d S V w h 4 r n 6 f S E m g 1 C T w d G d A Y K R + e z P x L 6 + b g F 9 1 U x 7 G C b C Q L h b 5 i c A Q 4 d n 3 e M A l o y A m m h A q u b 4 V 0 x G R h I L O K K d D + P o U / 0 9 a R d O u m M 5 V u V C 7 W M a R R U f o G J 0 i G 5 2 h G q q j B m o i i g L 0 g J 7 Q s y G N R + P F e F 2 0 Z o z l z C H 6 A e P t E 4 O W k O g = < / l a t e x i t > SHIt 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.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. 
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 2N 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 [21] architecture from MolCLR [33] used in Non-Chirality MoleculeNet [35] experiment. (D) Equiformer [17] architecture used in QM9 [28] dataset experiment. used exponential cosine similarity, defined as sim(z 1 , z 2 ) ≜ e z T 1 z2/∥z1∥∥z2∥τ , where τ denotes a temperature parameter.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:EQUATIONAs 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;EQUATIONOne 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.A straightforward 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:w + i ∼ Gamma(a + , b + ), w - ik ∼ Gamma(a -, b -). Option 2 -Bernoulli priors for selective weighting:w + i ∼ Gamma(a + , b + ), w - ik ∼ Bernoulli( ā-). here, a + , b + , a -and b -are shape and rate parameters for Gamma distribution and ā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:EQUATIONOne 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 wEQUATIONwhereu ≜ 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 θ.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 1 In the sense that marginalizing over the augmented random variables w + i and w - ik in p θ, U, w + i , w - ik | D gives back to the original p w + i , w - ik , θ; D . Thus, learning and inferences on the two forms are equivalent. 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:EQUATIONw + i | {u, θ} ∼ Gamma (1 + a + , u i s i+ + b + ) , and Option 1: w - ik | {u, θ} ∼ Gamma (a -, u i s ik -+ b -) , ∀i, k Option 2: w - ik | {u, θ} ∼ Bernoulli a -e -uis ik- 1 -a -+ a -e -uis ik- (5)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.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 inEQUATIONwhere τ indexes the minibatch and the corresponding local random variables at the current time τ .Algorithm 1 Contrastive Learning with Stochastic EM 1: Initialize θ; set t = 1 2: for a batch of molecules in loader doAugment each molecule x i into a pair (x i , x ′ i )Calculate positive/negative similarity scores s + and s - for all the molecule pairs 5:Initialize all the weights w + and w -to be one 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. 
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 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.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. 
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.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 AU C 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. 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. 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. [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.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.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. 
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.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. 
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]  
 
 Appendix: 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 MolCLR have two peaks of negatives similarity scores around 1 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 at higher levels as it allows for some degree of semantic dissimilar between positives. The similarity scores are dot similarity, they are not normalized to enhance the difference for visual purposes.In this section, we discuss the limitations of our proposed EM-based algorithm for molecular contrastive learning.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.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.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.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:w - i = 1 -λ × Sim (x i , x k ) w + i = λ × Sim (x i , x k )In our experiments, we set (λ = 1). From these results, we observe the following:1. 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.2. 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.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. 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. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: * 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. 
* What are the appropriate values for a_u and b_u.
* 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?
* Recent and better performance on the MoleculeNet datasets should be considered.

[1] Fang, Yin, et al. "Knowledge graph-enhanced molecular contrastive learning with functional prompt." Nature Machine Intelligence 5.5 (2023): 542-553.
[2] Hajiabolhassan, Hossein, et al. "FunQG: Molecular representation learning via quotient graphs." Journal of chemical information and modeling 63.11 (2023): 3275-3287.
[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.
questions: questions: * Figure 1, which node(s) and edge(s) are removed? How is similarity calculated? Does augmentation consider structural validity and stability?
* Why do you choose Gamma or Bernoulli distribution over other distribution?
* 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?
* 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: limitations: None.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: See **Weaknesses** Section.
paper_summary: paper_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.
strength_weakness: strength_weakness: (1) Algorithm A in lines 199-216 is filled with text. Can the author present the algorithm more elegantly? 

(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? 

[A] Rethinking explaining graph neural networks via non-parametric subgraph matching, ICML 2023

(3) MoleculeNet is somehow of small sizes. Some large datasets should be considered such as PCQM4MV2.
questions: questions: (1) Why 0.251 for Equiformer in Table 4 is in bold? 

(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]. 

[A] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. 

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

[A] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+. 
  
(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. 

[A] InstructBio: A Large-scale Semi-supervised Learning Paradigm for Biochemical Problems.
limitations: limitations: Yes, the authors addressed limitations adequately.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
2. 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.
3. 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.
4. The reason for using Gamma prior should be justified. It is an essential step in the proposed approach and related to the technical soundness. 
5. 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.  
6. The major part of Figure 2 is copied from the cited papers (MolCLR, Uni-mol, Equiformer). Is this permitted?
7. Some typos: e.g. line 285, missing "Table" before the reference.

Overall, 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. 

[1] Debiased Contrastive Learning, NeurIPS 2020  
[2] A Theoretical Analysis of Contrastive Unsupervised Representation Learning, ICML 2019  
[3] CWCL: Cross-Modal Transfer with Continuously Weighted Contrastive Loss, NeurIPS 2023  
[4] Camera Alignment and Weighted Contrastive Learning for Domain Adaptation in Video Person ReID, WACV 2023  
[5] Weighted Contrastive Learning With False Negative Control to Help Long-tailed Product Classification, ACL 2023
questions: questions: See weaknesses.
limitations: limitations: Yes.

Comments: title: I moved my decision from 5 to 6.
Gratitude for Your Review and Enhanced Rating
Gratitude for Your Review and Enhanced Rating
Follow-up on Rebuttal: Request for Further Review and Feedback
Update to the Response
Official Comment by Reviewer Jokw
Rebuttal Continued
comment: Thank you for your continued feedback. We appreciate your recognition of our idea's potential and would like to address your remaining concerns:

Gamma Prior: We chose the Gamma prior for its conjugacy properties and flexibility in modeling positive continuous variables, which suit our Bayesian framework. We are open to further discussion or providing additional comparative experiments.

Figure Copyright: We have obtained permissions for the figures used. We can either re-design them or provide the permissions documentation in our revised submission.

Related Works: We have now included a detailed discussion of related works on weighted contrastive learning in our updated manuscript, clarifying our unique contributions.

We are committed to making the necessary improvements and welcome any further guidance to ensure our work is ready for publication.

Best regards,

Authors
Thanks for the rebuttal. But I still have concerns about the use of the Gamma prior, the copyright of the figures, and the ignorance of related works about weighted contrastive learning. I do like the idea but the current submission is not ready for publication.
The responses address some of my concerns. It will be interesting to see additional experimental results.
Dear Reviewer,

Thank you for taking the time. We really appreciate your decision. Your support is invaluable in improving our work.

Best regards,

Authors
Dear Reviewer,

Thank you for taking the time. We really appreciate your decision. Your support is invaluable in improving our work.

Best regards,

Authors
Dear Reviewer oFgE,

We hope this message finds you well. We wanted to follow up on our rebuttal and ensure you had the opportunity to review our responses to your valuable feedback. We have carefully addressed each of your concerns and made significant improvements to the manuscript based on your suggestions. We would greatly appreciate it if you could take a moment to review our rebuttal and consider the clarifications and enhancements we've made. Your feedback has been instrumental in refining our work, and we hope it convinces you of the strength of our contributions.

We kindly ask you to re-evaluate our paper, and if you find our improvements satisfactory, we hope you might consider raising the score.

Thank you again for your time and consideration.

Best regards,

Authors
Dear Reviewer T3Hf,

Thank you for your thorough feedback on our work. We have carefully addressed each of your comments and provided additional analyses to clarify the points you raised. We believe our detailed rebuttal directly addresses your concerns and strengthens the validity of our approach. We kindly encourage you to review our responses at your earliest convenience, as your insights are invaluable to improving our manuscript.

We appreciate your time and look forward to any further comments you might have.

Best regards,

Authors
The authors answered all my questions and conducted relevant experiments and comparisons. I have no more problems and highly recommend the AC for acceptance. Thanks for their efforts and I have raised my score accordingly.
Thanks for conducting the experiment and your rebuttal; my concern has been resolved, leading me to adjust my rating from 6 to 7.
Questions:

Q1: Figure 1, which node(s) and edge(s) are removed? How is similarity calculated? Does augmentation consider structural validity and stability?

[A1]  Sorry for making the confusion, we will add this to explain in our manuscript.

In Figure 1, the nodes marked with red dots and the edges marked with red lines are the ones removed. When we “remove” a node, we do not entirely eliminate it from the graph. Instead, we substitute it with a special [MSK] node that does not correspond to a specific element type(like C or O). This approach is employed to avoid any changes in the topology that would occur due to node removal. Structural validity and stability is also preserved.

The similarity is calculated based on the representations from our trained encoders. The reason for not using fingerprint-based similarity is that such methods are specifically designed for unmodified molecules. The introduction of [MSK] nodes leads to discrepancies when using fingerprint-based similarities, as these traditional methods cannot accommodate the presence of masked nodes effectively. Therefore, we rely on our model's learned representations, which are robust to such augmentations, to measure similarity accurately.

Q2: Why do you choose Gamma or Bernoulli distribution over other distribution?

A2: 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 Bernoulli prior which uses binary sample weights in Table 1. Our results show that 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.

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

A3: We appreciate your suggestion. Please refer to the common reply to all reviewers for experiment using fingerprint-based similarity on Chiral MoleculeNet dataset (that is exactly the method proposed in I-MolCLR) and tabel 3 for experiment on non-Chiral MoleculeNet dataset. 

We are currently working on extending our experiments to include QM9 dataset, and we will incorporate these results in a future version of the paper. 

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

A4: We agree that the choice of hyperparameters in the prior distribution plays a crucial role in the performance of our method. The significant difference between the distributions for positive and negative pairs reflect the inherent differences in their nature. This also suggests that the model requires a more pronounced distinction between these pairs to effectively learn useful representations. Experiments show that the best performance is achieved when au=bu=5, please refer to our answer on W2 for experiment details.