Research map / Label-efficient and robust learning / Recognition and long-tail
A probability contrastive learning framework for 3D molecular representation learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a Bayesian probability-weighted contrastive learning framework for molecular representation learning that infers per-pair weights via stochastic EM to mitigate false positive and negative pairs, integrated into Uni-Mol, I-MolCLR/MolCLR, and Equiformer.
Problem
Conventional graph augmentations in molecular contrastive learning create false positive and false negative pairs, and hard positive/negative labels degrade learned molecular representations.
Approach
Treats per-pair positive and negative weights as latent random variables with Gamma or Bernoulli priors, uses an augmented-variable (Bayesian data augmentation) formulation to get conjugate posteriors, and optimizes the weighted contrastive loss with stochastic EM instead of hard labels or fingerprint similarity thresholds.
Key results
- On MoleculeNet classification (Uni-Mol-based), the Gamma variant reaches 80.1 average ROC-AUC across the nine tasks versus 78.9 for the reported baseline ablation and outperforms Uni-Mol on 13 of 15 benchmarks reported by the authors.
- On MoleculeNet regression, the Bernoulli variant reaches mean RMSE 0.883 versus 0.957 for Uni-Mol, and the Gamma variant reaches mean MAE 12.839 versus 13.940.
- With the MolCLR (non-chirality) setting, the method averages 85.1 ROC-AUC versus 83.1 for i-MolCLR, winning 7 of 9 tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add the probabilistic framework to standard CL (ablation) | MoleculeNet classification · MoleculeNet classification datasets | mean ROC AUC: 75.2 -> 78.6 | Helped |
| add the 3D loss to standard CL (ablation) | MoleculeNet classification · MoleculeNet classification datasets | mean ROC AUC: 75.2 -> 78.1 | Helped |
| add both probabilistic framework and 3D loss to standard CL (ablation) | MoleculeNet classification · MoleculeNet classification datasets | mean ROC AUC: 75.2 -> 80.1 | Helped |
| set a+ to 5 rather than 1 (ablation) | Probability contrastive learning · MoleculeNet classification tasks | average ROC-AUC: 78.8 -> 80.4 | Helped |
| vary a+, a−, b+, and b− across the tested values (ablation) | Probability contrastive learning · MoleculeNet classification tasks | average ROC-AUC: 78.8; 80.4; 79.6; 79.3; 80.0; 79.4; 79.3 | Mixed |
| vary au and bu (ablation) | Probability contrastive learning · MoleculeNet classification tasks | average ROC-AUC: 80.4; 80.2; 80.1; 80.6; 80.7; 80.2; 80.3; 80.1; 80.5 | No clear effect |
| replace Bayesian pair-weight inference with fingerprint-based similarity (ablation) | Probability contrastive learning · MoleculeNet classification tasks | ROC AUC: 77.9; 80.0; 68.9; 64.9; 87.6; 83.6; 78.7; 79.8; 83.8 | Mixed |
| replace Bayesian pair-weight inference with encoder-based similarity (ablation) | Probability contrastive learning · MoleculeNet classification tasks | ROC AUC: 75.2; 79.6; 67.8; 58.5; 90.4; 82.8; 80.5; 81.0; 74.5 | Mixed |
| use Bernoulli rather than Gamma weighting (ablation) | Uni-Mol probability contrastive learning · MoleculeNet classification and regression tasks | performance | Mixed |
Surprising findings
- Fingerprint-based and encoder-based similarity weighting underperform the Bayesian weighting, suggesting direct similarity thresholds are not sufficient.
- The method sometimes assigns lower similarity to positive pairs, which the authors argue helps capture dissimilarity between positives.
Limitations
- Relies on strong assumptions such as feature independence and noisy input data, which may not hold in practice.
- Gains are small on small datasets such as QM9 and results are limited to MoleculeNet and QM9.
- Authors note no error bars reported; experiments averaged over three runs.
Experimental setup
Datasets: MoleculeNet QM9 PCQM4Mv2 pretraining data (19M molecules, 209M conformations as used by Uni-Mol)
Benchmarks: MoleculeNet QM9 PDBbind binding pose prediction
Metrics: ROC AUC RMSE MAE binding pose accuracy (RMSD thresholds)
Baselines: Uni-Mol GEM MolCLR I-MolCLR GraphMVP GROVER Equiformer Transformer-M 3D-MGP GraphCL JOAOv2
Builds on: MolCLR I-MolCLR Uni-Mol Equiformer stochastic EM
Closest papers by meaning
- UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning (ICML 2024 · similarity 0.88)
- SubgDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning (NeurIPS 2024 · similarity 0.87)
- Enhancing Molecular Representation Learning Through the Combination of 3D and 2D Graph Machine Learning (Student Abstract) (AAAI 2025 · similarity 0.86)
- CL-MFAP: A Contrastive Learning-Based Multimodal Foundation Model for Molecular Property Prediction and Antibiotic Screening (ICLR 2025 · similarity 0.86)
- 3D Denoisers Are Good 2D Teachers: Molecular Pretraining via Denoising and Cross-Modal Distillation (AAAI 2025 · similarity 0.86)
- Learning Molecular Representation in a Cell (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.