# Improving Molecular Representation Learning with Metric Learning-enhanced Optimal Transport

arXiv: https://arxiv.org/abs/2202.06208 · HTML: https://mlsanity.com/paper/2202.06208/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 84 references that are in the map.
- First posted on arXiv 13 Feb 2022.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Learning Transferable Features with Deep Adaptation Networks](https://mlsanity.com/paper/1502.02791/) (2015)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Generalized Out-of-Distribution Detection: A Survey](https://mlsanity.com/paper/2110.11334/) (2021)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Maximum Classifier Discrepancy for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1712.02560/) (2017)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Generalized ODIN: Detecting Out-of-distribution Image without Learning from Out-of-distribution Data](https://mlsanity.com/paper/2002.11297/) (2020)
- [Domain Generalization via Model-Agnostic Learning of Semantic Features](https://mlsanity.com/paper/1910.13580/) (2019)
- [Return of Frustratingly Easy Domain Adaptation](https://mlsanity.com/paper/1511.05547/) (2015)
- [Analyzing Learned Molecular Representations for Property Prediction](https://mlsanity.com/paper/1904.01561/) (2019)
- [Semi-supervised Domain Adaptation via Minimax Entropy](https://mlsanity.com/paper/1904.06487/) (2019)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [A DIRT-T Approach to Unsupervised Domain Adaptation](https://mlsanity.com/paper/1802.08735/) (2018)
- [DeepJDOT: Deep Joint Distribution Optimal Transport for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1803.10081/) (2018)
- [ChemRL-GEM: Geometry Enhanced Molecular Representation Learning for Property Prediction](https://mlsanity.com/paper/2106.06130/) (2021)
- [Correlation Alignment for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1612.01939/) (2016)
- [ProGen: Language Modeling for Protein Generation](https://mlsanity.com/paper/2004.03497/) (2020)
- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](https://mlsanity.com/paper/2103.03606/) (2021)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [Deep Metric Learning with Hierarchical Triplet Loss](https://mlsanity.com/paper/1810.06951/) (2018)
- [Theoretical Analysis of Domain Adaptation with Optimal Transport](https://mlsanity.com/paper/1610.04420/) (2016)
- [Towards a Theoretical Framework of Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.04496/) (2021)
- [Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking](https://mlsanity.com/paper/2111.07786/) (2021)

## Built on by (most cited first, 9 of 9)

- [Unbalanced Optimal Transport, from Theory to Numerics](https://mlsanity.com/paper/2211.08775/) (2022)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)
- [A Hierarchical Training Paradigm for Antibody Structure-sequence Co-design](https://mlsanity.com/paper/2311.16126/) (2023)
- [Instructor-inspired Machine Learning for Robust Molecular Property Prediction](https://mlsanity.com/paper/2304.03906/) (2023)
- [Graph Data Selection for Domain Adaptation: A Model-Free Approach](https://mlsanity.com/paper/2505.17293/) (2025)
- [Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances](https://mlsanity.com/paper/2509.20508/) (2025)
- [Amortized Optimal Transport from Sliced Potentials](https://mlsanity.com/paper/2604.15114/) (2026)
- [Linear Partial Gromov-Wasserstein Embedding](https://mlsanity.com/paper/2410.16669/) (2024)
- [SurfDesign: Effective Protein Design on Molecular Surfaces](https://mlsanity.com/paper/2606.07567/) (2026)

## Abstract (arXiv)

Training data are usually limited or heterogeneous in many chemical and biological applications. Existing machine learning models for chemistry and materials science fail to consider generalizing beyond training domains. In this article, we develop a novel optimal transport-based algorithm termed MROT to enhance their generalization capability for molecular regression problems. MROT learns a continuous label of the data by measuring a new metric of domain distances and a posterior variance regularization over the transport plan to bridge the chemical domain gap. Among downstream tasks, we consider basic chemical regression tasks in unsupervised and semi-supervised settings, including chemical property prediction and materials adsorption selection. Extensive experiments show that MROT significantly outperforms state-of-the-art models, showing promising potential in accelerating the discovery of new substances with desired properties.